Improve trajectory storage and fix schema handling

README.md

@@ -11,6 +11,10 @@ Kaizen is a system designed to help agents improve over time by learning from th
 - **Trajectory Analysis**: Automatically analyzes agent trajectories to generate tips and best practices.
 - **Milvus Integration**: Uses Milvus (or Milvus Lite) for efficient vector storage and retrieval.
 
+## Architecture
+
+<img src="docs/assets/architecture.png" alt="Architecture" width="480">
+
 ## Quick Start
 
 ### Installation

kaizen/backend/filesystem.py

@@ -248,10 +248,14 @@ def _search_entities_internal(
         else:
             # Simple case-insensitive text matching
             query_lower = query.lower()
-            matching = [
-                ent for ent in entities
-                if query_lower in ent.get("content", "").lower()
-            ]
+            matching = []
+            for ent in entities:
+                content = ent.get("content", "")
+                # Convert non-string content to JSON string for searching
+                if not isinstance(content, str):
+                    content = json.dumps(content)
+                if query_lower in content.lower():
+                    matching.append(ent)
             results = matching[:limit]
 
         return [

kaizen/backend/milvus.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 import logging
 import uuid
 
@@ -15,6 +16,22 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("katas-db.milvus")
 
+
+def serialize_content(content) -> str:
+    """Serialize content to string for Milvus storage."""
+    if isinstance(content, str):
+        return content
+    return json.dumps(content)
+
+
+def deserialize_content(content: str):
+    """Deserialize content from Milvus storage."""
+    try:
+        return json.loads(content)
+    except (json.JSONDecodeError, TypeError):
+        return content
+
+
 class MilvusKataBackend(BaseKataBackend):
     milvus = MilvusClient(**milvus_client_settings.model_dump())
     embedding_model = SentenceTransformer(milvus_other_settings.embedding_model)
@@ -93,27 +110,29 @@ def update_entities(
         if enable_conflict_resolution:
             old_entities = []
             for entity in entities:
-                old_entities.extend(self.search_entities(namespace_id=namespace_id, query=entity.content))
+                query_str = serialize_content(entity.content)
+                old_entities.extend(self.search_entities(namespace_id=namespace_id, query=query_str))
 
             updates = resolve_conflicts(old_entities, entities_with_temporary_ids)
             for update in updates:
+                content_str = serialize_content(update.content)
                 match update.event:
                     case 'ADD':
                         entity_id = str(self.milvus.insert(collection_name=namespace_id, data={
                             'type': entity_type,
-                            'content': update.content,
+                            'content': content_str,
                             'created_at': int(now.timestamp()),
-                            'embedding': self.embedding_model.encode(update.content),
+                            'embedding': self.embedding_model.encode(content_str),
                             'metadata': update.metadata,
                         })['ids'][0])
                         update.id = entity_id
                     case 'UPDATE':
                         self.milvus.upsert(collection_name=namespace_id, data={
                             'type': entity_type,
-                            'id': update.id,
-                            'content': update.content,
+                            'id': int(update.id),
+                            'content': content_str,
                             'created_at': int(now.timestamp()),
-                            'embedding': self.embedding_model.encode(update.content),
+                            'embedding': self.embedding_model.encode(content_str),
                             'metadata': update.metadata
                         }, partial_update=True)
                     case 'DELETE':
@@ -123,11 +142,12 @@ def update_entities(
         else:
             updates = []
             for entity in entities:
+                content_str = serialize_content(entity.content)
                 entity_id = str(self.milvus.insert(collection_name=namespace_id, data={
                     'type': entity_type,
-                    'content': entity.content,
+                    'content': content_str,
                     'created_at': int(now.timestamp()),
-                    'embedding': self.embedding_model.encode(entity.content),
+                    'embedding': self.embedding_model.encode(content_str),
                     'metadata': entity.metadata
                 })['ids'][0])
                 updates.append(EntityUpdate(
@@ -175,7 +195,7 @@ def delete_entity_by_id(self, namespace_id: str, entity_id: str):
     # Keep it as an INT64 or else you won't be able to list all entities.
     FieldSchema(name='id', is_primary=True, auto_id=True, dtype=DataType.INT64, max_length=128),
     FieldSchema(name='type', dtype=DataType.VARCHAR, max_length=128),
-    FieldSchema(name='content', dtype=DataType.VARCHAR, max_length=512),
+    FieldSchema(name='content', dtype=DataType.VARCHAR, max_length=65535),
     FieldSchema(name='created_at', dtype=DataType.INT64),
     FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=384),
     FieldSchema(name='metadata', dtype=DataType.JSON),
@@ -185,5 +205,6 @@ def parse_milvus_entity(entity: dict) -> RecordedEntity:
     return RecordedEntity.model_validate({
         **entity,
         'id': str(entity['id']),
+        'content': deserialize_content(entity['content']),
         'created_at': datetime.datetime.fromtimestamp(entity['created_at'], datetime.UTC),
     })

kaizen/db/sqlite_manager.py

@@ -95,8 +95,8 @@ def search_namespaces(
         with self._lock:
             cursor: sqlite3.Cursor = self.connection.cursor()
             cursor.row_factory = Namespace.row_factory
-            cursor.execute(f"""
-                    SELECT id, created_at, user_id, agent_id, app_id
+            cursor.execute("""
+                    SELECT id, created_at
                     FROM namespaces
                     LIMIT ?
                 """,

kaizen/schema/conflict_resolution.py

@@ -6,7 +6,7 @@ class SimpleEntity(BaseModel):
     """Derived from either a `Entity` or `RecordedEntity`. Optimized for LLM-based conflict resolution."""
     id: str = Field(description='The unique ID of an entity.')
     type: str = Field(description='The type of the entity.')
-    content: str = Field(description='The content of the entity.')
+    content: str | list | dict = Field(description='The content of the entity.')
 
     @staticmethod
     def from_recorded_entities(entities: list[RecordedEntity]) -> list['SimpleEntity']:
@@ -17,7 +17,7 @@ class EntityUpdate(BaseModel):
     """Produced by the LLM, to be processed by a kata backend."""
     id: str = Field(description='The unique ID of an entity.')
     type: str = Field(description='The type of the entity.')
-    content: str = Field(description='The content of the entity.')
+    content: str | list | dict = Field(description='The content of the entity.')
     event: Literal['ADD', 'UPDATE', 'DELETE', 'NONE'] = Field(description='The type of update operation to perform.')
     old_entity: str | None = Field(default=None, description='The entity before it was updated.')
     metadata: dict | None = Field(default=None, description='Arbitrary metadata which is related to the entity.')

kaizen/schema/core.py

@@ -16,7 +16,7 @@ def row_factory(cursor: Cursor, row: Row) -> 'Namespace':
 
 class Entity(BaseModel):
     """Basic data stored in the DB"""
-    content: str = Field(description='Some relatively short searchable text.')
+    content: str | list | dict = Field(description='Searchable text or structured data.')
     metadata: dict | None = Field(default=None, description='Arbitrary metadata which is related to the entity.')
     type: str = Field(description='The type of the entity.')
 

kaizen/sync/phoenix_sync.py

@@ -298,29 +298,24 @@ def _process_trajectory(self, trajectory: dict) -> int:
 
         Returns the number of tips generated.
         """
-        # Store trajectory messages
-        entities = []
-        for msg in trajectory.get("messages", []):
-            content = msg.get("content")
-            if isinstance(content, str) and content:
-                entities.append(
-                    Entity(
-                        type="trajectory",
-                        content=content,
-                        metadata={
-                            "trace_id": trajectory["trace_id"],
-                            "span_id": trajectory["span_id"],
-                            "model": trajectory["model"],
-                            "role": msg.get("role"),
-                            "timestamp": trajectory["timestamp"],
-                        },
-                    )
-                )
-
-        if entities:
+        # Store trajectory as a single entity with all messages
+        messages = trajectory.get("messages", [])
+        if messages:
+            entity = Entity(
+                type="trajectory",
+                content=messages,
+                metadata={
+                    "trace_id": trajectory["trace_id"],
+                    "span_id": trajectory["span_id"],
+                    "model": trajectory["model"],
+                    "timestamp": trajectory["timestamp"],
+                    "message_count": len(messages),
+                    "usage": trajectory.get("usage"),
+                },
+            )
             self.client.update_entities(
                 namespace_id=self.namespace_id,
-                entities=entities,
+                entities=[entity],
                 enable_conflict_resolution=False,
             )
 

pyproject.toml

@@ -5,6 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "arize-phoenix>=12.30.0",
     "fastmcp",
     "jinja2",
     "litellm",
@@ -44,4 +45,4 @@ markers = [
     "unit",
     "phoenix"
 ]
-anyio_mode = "auto"
+anyio_mode = "auto"

tests/unit/test_phoenix_sync.py

@@ -638,7 +638,18 @@ def test_sync_processes_valid_spans(self, mock_generate_tips, mock_urlopen, phoe
         mock_urlopen.return_value = mock_response
 
         phoenix_sync.client.search_entities.return_value = []
-        mock_generate_tips.return_value = ["Tip 1", "Tip 2"]
+        # Create mock Tip objects with required attributes
+        mock_tip1 = MagicMock()
+        mock_tip1.content = "Tip 1 content"
+        mock_tip1.category = "strategy"
+        mock_tip1.rationale = "Tip 1 rationale"
+        mock_tip1.trigger = "Tip 1 trigger"
+        mock_tip2 = MagicMock()
+        mock_tip2.content = "Tip 2 content"
+        mock_tip2.category = "optimization"
+        mock_tip2.rationale = "Tip 2 rationale"
+        mock_tip2.trigger = "Tip 2 trigger"
+        mock_generate_tips.return_value = [mock_tip1, mock_tip2]
 
         result = phoenix_sync.sync(limit=10)
 
@@ -684,7 +695,13 @@ def test_sync_returns_correct_counts(self, mock_generate_tips, mock_urlopen, pho
         mock_entity = MagicMock()
         mock_entity.metadata = {"span_id": "old_span"}
         phoenix_sync.client.search_entities.return_value = [mock_entity]
-        mock_generate_tips.return_value = ["Generated tip"]
+        # Create mock Tip object with required attributes
+        mock_tip = MagicMock()
+        mock_tip.content = "Generated tip content"
+        mock_tip.category = "strategy"
+        mock_tip.rationale = "Tip rationale"
+        mock_tip.trigger = "Tip trigger"
+        mock_generate_tips.return_value = [mock_tip]
 
         result = phoenix_sync.sync(limit=10)