Better summarization, memory extraction, etc.

Author: abrookins

.gitignore

@@ -221,3 +221,4 @@ libs/redis/docs/.Trash*
 .python-version
 .idea/*
 .vscode/settings.json
+.cursor

agent_memory_server/cli.py

@@ -79,11 +79,11 @@ async def run_migrations():
 @click.option("--reload", is_flag=True, help="Enable auto-reload")
 def api(port: int, host: str, reload: bool):
     """Run the REST API server."""
-    from agent_memory_server.main import app, on_start_logger
+    from agent_memory_server.main import on_start_logger
 
     on_start_logger(port)
     uvicorn.run(
-        app,
+        "agent_memory_server.main:app",
         host=host,
         port=port,
         reload=reload,

agent_memory_server/client/api.py

@@ -13,6 +13,7 @@
     CreatedAt,
     Entities,
     LastAccessed,
+    MemoryType,
     Namespace,
     SessionId,
     Topics,
@@ -273,6 +274,7 @@ async def search_long_term_memory(
         last_accessed: LastAccessed | dict[str, Any] | None = None,
         user_id: UserId | dict[str, Any] | None = None,
         distance_threshold: float | None = None,
+        memory_type: MemoryType | dict[str, Any] | None = None,
         limit: int = 10,
         offset: int = 0,
     ) -> LongTermMemoryResults:
@@ -313,6 +315,8 @@ async def search_long_term_memory(
             last_accessed = LastAccessed(**last_accessed)
         if isinstance(user_id, dict):
             user_id = UserId(**user_id)
+        if isinstance(memory_type, dict):
+            memory_type = MemoryType(**memory_type)
 
         # Apply default namespace if needed and no namespace filter specified
         if namespace is None and self.config.default_namespace is not None:
@@ -328,6 +332,7 @@ async def search_long_term_memory(
             last_accessed=last_accessed,
             user_id=user_id,
             distance_threshold=distance_threshold,
+            memory_type=memory_type,
             limit=limit,
             offset=offset,
         )

agent_memory_server/extraction.py

@@ -343,4 +343,7 @@ async def extract_discrete_memories(redis: Redis | None = None):
             for new_memory in discrete_memories
         ]
 
-        await index_long_term_memories(long_term_memories)
+        await index_long_term_memories(
+            long_term_memories,
+            deduplicate=True,
+        )

agent_memory_server/long_term_memory.py

@@ -66,28 +66,24 @@ async def get_session_memory(
     if not session_exists:
         return None
 
-    # Retrieve messages and metadata
     async with redis.pipeline() as pipe:
         pipe.lrange(messages_key, -window_size, -1)  # Get the most recent messages
         pipe.hgetall(metadata_key)
         messages_data, metadata = await pipe.execute()
 
-    # Parse messages
     messages = []
     for msg_data in messages_data:
         if isinstance(msg_data, bytes):
             msg_data = msg_data.decode("utf-8")
         msg = json.loads(msg_data)
         messages.append(MemoryMessage(**msg))
 
-    # Parse metadata
     metadata_dict = {}
     for k, v in metadata.items():
         key = k.decode("utf-8") if isinstance(k, bytes) else k
         value = v.decode("utf-8") if isinstance(v, bytes) else v
         metadata_dict[key] = value
 
-    # Create SessionMemory object
     return SessionMemory(messages=messages, **metadata_dict)
 
 
@@ -112,6 +108,7 @@ async def set_session_memory(
     messages_json = [json.dumps(msg.model_dump()) for msg in memory.messages]
     metadata = memory.model_dump(
         exclude_none=True,
+        exclude_unset=True,
         exclude={"messages"},
     )
 
@@ -138,11 +135,7 @@ async def set_session_memory(
 
     # Check if window size is exceeded
     current_size = await redis.llen(messages_key)  # type: ignore
-    print(
-        f"Current size: {current_size}", "Current window size: ", settings.window_size
-    )
     if current_size > settings.window_size:
-        print("Queuing summarizing session task")
         # Add summarization task
         await background_tasks.add_task(
             summarize_session,
@@ -152,7 +145,6 @@ async def set_session_memory(
         )
 
     # If long-term memory is enabled, index messages
-    print("Long-term memory is enabled: ", settings.long_term_memory)
     if settings.long_term_memory:
         memories = [
             LongTermMemory(
@@ -164,7 +156,6 @@ async def set_session_memory(
             for msg in memory.messages
         ]
 
-        print("Adding a task")
         await background_tasks.add_task(
             index_long_term_memories,
             memories,

agent_memory_server/messages.py

@@ -8,6 +8,7 @@
     CreatedAt,
     Entities,
     LastAccessed,
+    MemoryType,
     Namespace,
     SessionId,
     Topics,
@@ -213,6 +214,10 @@ class SearchPayload(BaseModel):
         default=None,
         description="Optional distance threshold to filter by",
     )
+    memory_type: MemoryType | None = Field(
+        default=None,
+        description="Optional memory type to filter by",
+    )
     limit: int = Field(
         default=10,
         ge=1,
@@ -250,4 +255,7 @@ def get_filters(self):
         if self.last_accessed is not None:
             filters["last_accessed"] = self.last_accessed
 
+        if self.memory_type is not None:
+            filters["memory_type"] = self.memory_type
+
         return filters

agent_memory_server/models.py

@@ -98,7 +98,7 @@ async def _incremental_summary(
         response = await client.create_chat_completion(model, progressive_prompt)
 
         # Extract completion text
-        completion = response.choices[0]["message"]["content"]
+        completion = response.choices[0].message.content
 
         # Get token usage
         tokens_used = response.total_tokens
@@ -219,6 +219,7 @@ async def summarize_session(
                 metadata["tokens"] = str(total_tokens)
 
                 pipe.hmset(metadata_key, mapping=metadata)
+                print("Metadata: ", metadata_key, metadata)
 
                 # Messages that were summarized
                 num_summarized = len(messages_to_summarize)

agent_memory_server/summarization.py

@@ -147,12 +147,11 @@ def __init__(self, docs):
 
         assert mock_index.query.call_count == 1
 
-        assert len(results.memories) == 2
+        assert len(results.memories) == 1
         assert isinstance(results.memories[0], LongTermMemoryResult)
         assert results.memories[0].text == "Hello, world!"
         assert results.memories[0].dist == 0.25
-        assert results.memories[1].text == "Hi there!"
-        assert results.memories[1].dist == 0.75
+        assert results.memories[0].memory_type == "message"
 
 
 @pytest.mark.requires_api_keys

tests/test_long_term_memory.py

@@ -70,9 +70,7 @@ async def test_merge_memories_with_llm(mock_openai_client, monkeypatch):
         },
     ]
 
-    merged = await merge_memories_with_llm(
-        memories, "hash", llm_client=mock_openai_client
-    )
+    merged = await merge_memories_with_llm(memories, llm_client=mock_openai_client)
     assert merged["text"] == "Merged content"
     assert merged["created_at"] == memories[1]["created_at"]
     assert merged["last_accessed"] == memories[1]["last_accessed"]