# Implementation Plan: Consolidate Nested Dictionary Operations

## Issue

- **Issue Number**: #604
- **Title**: Consolidate nested dictionary helper functions into
  `lionagi.utils.nested`
- **Link**: [Issue #604](https://github.com/khive-ai/lionagi/issues/604)

## Overview

Custom implementations for nested dictionary operations (e.g., `flatten`,
`unflatten`, `nset`, `nget`, `npop`) exist in `lionagi/libs/nested` and
potentially other utility modules. These replicate functionality available in
mature third-party libraries and introduce potential inconsistencies or
edge-case bugs. This implementation plan outlines the approach to consolidate
these operations into a single, well-defined module.

## Research Summary

### Current Implementation

The current implementation in `lionagi/libs/nested/` includes:

- `flatten.py` / `unflatten.py`: Convert between nested and flat dictionary
  structures
- `nget.py`: Retrieve values from nested structures using path indices
- `nset.py`: Set values in nested structures at specified paths
- `npop.py`: Remove and return values from nested structures
- `ninsert.py`: Insert values into nested structures
- `nmerge.py`: Merge nested structures
- `nfilter.py`: Filter nested structures
- `utils.py`: Helper functions for nested operations

### External Libraries

1. **boltons**:
   - Provides `remap()` for traversing and transforming nested structures
   - Includes `get_path()` for retrieving values from nested structures
   - Has `research()` for finding values in nested structures

2. **flatten-dict**:
   - Specialized in flattening/unflattening dictionaries
   - Supports various key joining methods
   - Handles empty structures and custom delimiters

## Approach

After evaluating the options, I recommend consolidating our custom
implementation rather than adopting external libraries for the following
reasons:

1. The current implementation is already well-structured and covers
   functionality not fully available in external libraries
2. Avoiding external dependencies simplifies deployment and maintenance
3. We can ensure backward compatibility with existing code
4. We can optimize for the specific use cases in lionagi

## Implementation Steps

### 1. Create New Module Structure

Create a consolidated module at `lionagi/utils/nested.py` with all nested
dictionary operations:

```python
# lionagi/utils/nested.py
"""
Utility functions for working with nested dictionary/list structures.

This module provides functions for manipulating nested data structures:
- flatten/unflatten: Convert between nested and flat dictionary structures
- nget/nset/npop: Access and modify values in nested structures
- ninsert: Insert values into nested structures
- nmerge: Merge nested structures
- nfilter: Filter nested structures
"""

from typing import Any, Callable, Iterable, Literal, Mapping, Sequence, TypeVar, Union, overload
from collections.abc import Mapping, Sequence

from lionagi.utils import UNDEFINED, to_list

# Type definitions
T = TypeVar("T")
IndicesType = Union[str, int, Sequence[Union[str, int]]]

# Re-implement core functions
def flatten(...): ...
def unflatten(...): ...
def nget(...): ...
def nset(...): ...
def npop(...): ...
def ninsert(...): ...
def nmerge(...): ...
def nfilter(...): ...

# Helper functions
def get_target_container(...): ...
def ensure_list_index(...): ...
def is_homogeneous(...): ...
def is_same_dtype(...): ...
def is_structure_homogeneous(...): ...
def deep_update(...): ...
```

### 2. Implement Core Functions

Transfer and optimize the implementations from the existing files to the new
module, ensuring:

- Consistent parameter naming and ordering
- Comprehensive docstrings with examples
- Type hints for better IDE support
- Optimized implementations where possible

### 3. Create Backward Compatibility Layer

Create a backward compatibility layer in the original module locations:

```python
# lionagi/libs/nested/flatten.py
from lionagi.utils.nested import flatten

# Re-export for backward compatibility
__all__ = ['flatten']
```

Repeat for all existing modules to maintain backward compatibility.

### 4. Update Import References

Update imports in the codebase to use the new consolidated module:

```python
# Before
from lionagi.libs.nested.flatten import flatten
from lionagi.libs.nested.nget import nget

# After
from lionagi.utils.nested import flatten, nget
```

### 5. Add Comprehensive Tests

Create comprehensive tests for the consolidated module:

```python
# tests/utils/test_nested.py
import pytest
from lionagi.utils.nested import flatten, unflatten, nget, nset, npop, ninsert, nmerge, nfilter

# Test cases for flatten/unflatten
def test_flatten_basic(): ...
def test_flatten_with_max_depth(): ...
def test_unflatten_basic(): ...

# Test cases for nget/nset/npop
def test_nget_basic(): ...
def test_nset_basic(): ...
def test_npop_basic(): ...

# Test cases for ninsert/nmerge/nfilter
def test_ninsert_basic(): ...
def test_nmerge_basic(): ...
def test_nfilter_basic(): ...

# Edge cases
def test_empty_structures(): ...
def test_nested_lists(): ...
def test_mixed_types(): ...
```

## Risks and Mitigations

| Risk                            | Mitigation                                                             |
| ------------------------------- | ---------------------------------------------------------------------- |
| Breaking existing functionality | Comprehensive test coverage; backward compatibility layer              |
| Performance regression          | Benchmark against current implementation; optimize critical paths      |
| Missing edge cases              | Review existing tests; add tests for all identified edge cases         |
| Inconsistent API                | Maintain consistent parameter naming and ordering; clear documentation |

## Future Considerations

- Consider adding more advanced operations like deep comparison, path-based
  filtering
- Add support for additional data structures (e.g., namedtuples, custom classes)
- Optimize for memory usage with large nested structures

## Timeline

1. Create new module structure (1 day)
2. Implement core functions (2 days)
3. Create backward compatibility layer (0.5 day)
4. Update import references (0.5 day)
5. Add comprehensive tests (1 day)
6. Review and finalize (1 day)

Total: 6 days
