Allow nested JSON structures in metadata

app/models.py

@@ -10,8 +10,7 @@
 from pydantic import field_validator
 from pydantic import model_validator
 
-# Define MetadataDict directly to avoid import issues
-type MetadataDict = dict[str, str | int | float | bool | None]
+from app.types import MetadataDict
 
 
 class SourceType(StrEnum):

app/types.py

@@ -12,8 +12,8 @@
 # Metadata value types - simpler non-recursive type for metadata fields
 type MetadataValue = str | int | float | bool | None
 
-# Metadata dictionary type for use in models
-type MetadataDict = dict[str, MetadataValue]
+# Metadata dictionary type for use in models - supports nested JSON structures
+type MetadataDict = dict[str, JsonValue]
 
 
 # API Response TypedDicts for proper return type annotations

tests/test_metadata_filtering.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import math
 import time
 from typing import TYPE_CHECKING
 
@@ -571,3 +572,202 @@ async def test_all_operators(
 
     assert 'entries' in result
     assert len(result['entries']) == expected_count
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('initialized_server')
+class TestNestedJSONMetadata:
+    """Test nested JSON structures in metadata."""
+
+    @pytest.mark.asyncio
+    async def test_store_nested_objects(self) -> None:
+        """Test storing nested JSON objects in metadata."""
+        complex_metadata = {
+            'status': 'active',
+            'config': {
+                'database': {
+                    'connection': {
+                        'pool': {'size': 10, 'timeout': 30},
+                        'retry': {'max_attempts': 3, 'backoff': 2.5},
+                    },
+                },
+                'cache': {'enabled': True, 'ttl': 300},
+            },
+            'user': {'id': 123, 'name': 'Alice Johnson', 'preferences': {'theme': 'dark', 'language': 'en'}},
+        }
+
+        result = await store_context.fn(
+            thread_id='test_nested_json',
+            source='agent',
+            text='Test nested metadata storage',
+            metadata=complex_metadata,
+            ctx=None,
+        )
+
+        assert result['success'] is True
+        assert 'context_id' in result
+
+        # Retrieve and verify the metadata is preserved
+        search_result = await search_context.fn(thread_id='test_nested_json', ctx=None)
+        assert len(search_result['entries']) == 1
+
+        stored_metadata = search_result['entries'][0]['metadata']
+        assert stored_metadata['status'] == 'active'
+        assert stored_metadata['config']['database']['connection']['pool']['size'] == 10
+        assert stored_metadata['config']['database']['connection']['pool']['timeout'] == 30
+        assert stored_metadata['config']['database']['connection']['retry']['max_attempts'] == 3
+        assert stored_metadata['config']['database']['connection']['retry']['backoff'] == 2.5
+        assert stored_metadata['config']['cache']['enabled'] is True
+        assert stored_metadata['user']['preferences']['theme'] == 'dark'
+        assert stored_metadata['user']['preferences']['language'] == 'en'
+
+    @pytest.mark.asyncio
+    async def test_store_arrays_in_metadata(self) -> None:
+        """Test storing arrays in metadata."""
+        metadata_with_arrays = {
+            'tags': ['urgent', 'backend', 'production'],
+            'priority_levels': [1, 2, 3, 4, 5],
+            'mixed_array': ['string', 42, math.pi, True, None],
+            'nested_arrays': [[1, 2], [3, 4], [5, 6]],
+        }
+
+        result = await store_context.fn(
+            thread_id='test_arrays',
+            source='agent',
+            text='Test array metadata',
+            metadata=metadata_with_arrays,
+            ctx=None,
+        )
+
+        assert result['success'] is True
+
+        # Retrieve and verify arrays are preserved
+        search_result = await search_context.fn(thread_id='test_arrays', ctx=None)
+        stored_metadata = search_result['entries'][0]['metadata']
+
+        assert stored_metadata['tags'] == ['urgent', 'backend', 'production']
+        assert stored_metadata['priority_levels'] == [1, 2, 3, 4, 5]
+        assert stored_metadata['mixed_array'] == ['string', 42, math.pi, True, None]
+        assert stored_metadata['nested_arrays'] == [[1, 2], [3, 4], [5, 6]]
+
+    @pytest.mark.asyncio
+    async def test_query_nested_paths(self) -> None:
+        """Test querying nested JSON paths."""
+        # Store multiple entries with nested metadata
+        await store_context.fn(
+            thread_id='test_nested_paths',
+            source='agent',
+            text='Entry 1',
+            metadata={'user': {'preferences': {'theme': 'dark', 'notifications': {'email': True}}}},
+            ctx=None,
+        )
+
+        await store_context.fn(
+            thread_id='test_nested_paths',
+            source='agent',
+            text='Entry 2',
+            metadata={'user': {'preferences': {'theme': 'light', 'notifications': {'email': False}}}},
+            ctx=None,
+        )
+
+        # Query using nested path
+        result = await search_context.fn(
+            thread_id='test_nested_paths',
+            metadata={'user.preferences.theme': 'dark'},
+            ctx=None,
+        )
+
+        assert len(result['entries']) == 1
+        assert result['entries'][0]['text_content'] == 'Entry 1'
+        assert result['entries'][0]['metadata']['user']['preferences']['theme'] == 'dark'
+
+    @pytest.mark.asyncio
+    async def test_complex_nested_structure(self) -> None:
+        """Test very complex nested structure with multiple levels."""
+        complex_structure = {
+            'level1': {
+                'level2': {
+                    'level3': {
+                        'level4': {
+                            'value': 'deeply_nested',
+                            'number': 42,
+                            'array': [1, 2, 3],
+                            'object': {'key': 'value'},
+                        },
+                    },
+                },
+            },
+            'metrics': {
+                'cpu': 45.5,
+                'memory': 512,
+                'disk': {'used': 80.5, 'total': 100.0, 'partitions': ['/dev/sda1', '/dev/sda2']},
+            },
+            'features': {
+                'enabled': ['feature_a', 'feature_b', 'feature_c'],
+                'disabled': [],
+                'experimental': {'count': 3, 'names': ['exp_1', 'exp_2', 'exp_3']},
+            },
+        }
+
+        result = await store_context.fn(
+            thread_id='test_complex',
+            source='agent',
+            text='Complex nested structure test',
+            metadata=complex_structure,
+            ctx=None,
+        )
+
+        assert result['success'] is True
+
+        # Verify structure is preserved
+        search_result = await search_context.fn(thread_id='test_complex', ctx=None)
+        stored_metadata = search_result['entries'][0]['metadata']
+
+        # Verify deep nesting
+        assert stored_metadata['level1']['level2']['level3']['level4']['value'] == 'deeply_nested'
+        assert stored_metadata['level1']['level2']['level3']['level4']['number'] == 42
+        assert stored_metadata['level1']['level2']['level3']['level4']['array'] == [1, 2, 3]
+        assert stored_metadata['level1']['level2']['level3']['level4']['object']['key'] == 'value'
+
+        # Verify metrics
+        assert stored_metadata['metrics']['cpu'] == 45.5
+        assert stored_metadata['metrics']['disk']['used'] == 80.5
+        assert stored_metadata['metrics']['disk']['partitions'] == ['/dev/sda1', '/dev/sda2']
+
+        # Verify features
+        assert stored_metadata['features']['enabled'] == ['feature_a', 'feature_b', 'feature_c']
+        assert stored_metadata['features']['disabled'] == []
+        assert stored_metadata['features']['experimental']['count'] == 3
+
+    @pytest.mark.asyncio
+    async def test_mixed_flat_and_nested(self) -> None:
+        """Test mixing flat and nested metadata structures."""
+        mixed_metadata = {
+            'simple_string': 'value',
+            'simple_int': 42,
+            'simple_bool': True,
+            'nested': {'level1': {'level2': 'deep_value'}},
+            'array': [1, 2, 3],
+        }
+
+        result = await store_context.fn(
+            thread_id='test_mixed',
+            source='agent',
+            text='Mixed flat and nested',
+            metadata=mixed_metadata,
+            ctx=None,
+        )
+
+        assert result['success'] is True
+
+        # Query using both flat and nested paths
+        search_result = await search_context.fn(thread_id='test_mixed', metadata={'simple_string': 'value'}, ctx=None)
+        assert len(search_result['entries']) == 1
+
+        # Verify all types are preserved
+        stored_metadata = search_result['entries'][0]['metadata']
+        assert stored_metadata['simple_string'] == 'value'
+        assert stored_metadata['simple_int'] == 42
+        assert stored_metadata['simple_bool'] is True
+        assert stored_metadata['nested']['level1']['level2'] == 'deep_value'
+        assert stored_metadata['array'] == [1, 2, 3]