Fix mypy issues

abrookins · abrookins · commit 6a5f3ceac8d8 · 2025-06-17T16:17:16.000-07:00
diff --git a/agent-memory-client/agent_memory_client/client.py b/agent-memory-client/agent_memory_client/client.py
@@ -34,6 +34,7 @@
     HealthCheckResponse,
     MemoryRecord,
     MemoryRecordResults,
+    MemoryTypeEnum,
     ModelNameLiteral,
     SessionListResponse,
     WorkingMemory,
@@ -442,7 +443,7 @@ async def add_memories_to_working_memory(
         # Auto-generate IDs for memories that don't have them
         for memory in final_memories:
             if not memory.id:
-                memory.id = str(ulid.new())
+                memory.id = str(ulid.ULID())
 
         # Create new working memory with the memories
         working_memory = WorkingMemory(
@@ -617,136 +618,10 @@ async def search_long_term_memory(
                 exclude_none=True, mode="json"
             )
         if user_id:
-            payload["user_id"] = user_id.model_dump(exclude_none=True)
-        if memory_type:
-            payload["memory_type"] = memory_type.model_dump(exclude_none=True)
-        if distance_threshold is not None:
-            payload["distance_threshold"] = distance_threshold
-
-        try:
-            response = await self._client.post(
-                "/v1/long-term-memory/search",
-                json=payload,
-            )
-            response.raise_for_status()
-            return MemoryRecordResults(**response.json())
-        except httpx.HTTPStatusError as e:
-            self._handle_http_error(e.response)
-            raise
-
-    async def search_memories(
-        self,
-        text: str,
-        session_id: SessionId | dict[str, Any] | None = None,
-        namespace: Namespace | dict[str, Any] | None = None,
-        topics: Topics | dict[str, Any] | None = None,
-        entities: Entities | dict[str, Any] | None = None,
-        created_at: CreatedAt | dict[str, Any] | None = None,
-        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,
-    ) -> MemoryRecordResults:
-        """
-        Search across all memory types (working memory and long-term memory).
-
-        This method searches both working memory (ephemeral, session-scoped) and
-        long-term memory (persistent, indexed) to provide comprehensive results.
-
-        For working memory:
-        - Uses simple text matching
-        - Searches across all sessions (unless session_id filter is provided)
-        - Returns memories that haven't been promoted to long-term storage
-
-        For long-term memory:
-        - Uses semantic vector search
-        - Includes promoted memories from working memory
-        - Supports advanced filtering by topics, entities, etc.
-
-        Args:
-            text: Search query text for semantic similarity
-            session_id: Optional session ID filter
-            namespace: Optional namespace filter
-            topics: Optional topics filter
-            entities: Optional entities filter
-            created_at: Optional creation date filter
-            last_accessed: Optional last accessed date filter
-            user_id: Optional user ID filter
-            distance_threshold: Optional distance threshold for search results
-            memory_type: Optional memory type filter
-            limit: Maximum number of results to return (default: 10)
-            offset: Offset for pagination (default: 0)
-
-        Returns:
-            MemoryRecordResults with matching memories from both memory types
-
-        Raises:
-            MemoryServerError: If the request fails
-
-        Example:
-            ```python
-            # Search for user preferences with topic filtering
-            from .filters import Topics
-
-            results = await client.search_memories(
-                text="user prefers dark mode",
-                topics=Topics(any=["preferences", "ui"]),
-                limit=5
-            )
-
-            for memory in results.memories:
-                print(f"Found: {memory.text}")
-            ```
-        """
-        # Convert dictionary filters to their proper filter objects if needed
-        if isinstance(session_id, dict):
-            session_id = SessionId(**session_id)
-        if isinstance(namespace, dict):
-            namespace = Namespace(**namespace)
-        if isinstance(topics, dict):
-            topics = Topics(**topics)
-        if isinstance(entities, dict):
-            entities = Entities(**entities)
-        if isinstance(created_at, dict):
-            created_at = CreatedAt(**created_at)
-        if isinstance(last_accessed, dict):
-            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:
-            namespace = Namespace(eq=self.config.default_namespace)
-
-        payload = {
-            "text": text,
-            "limit": limit,
-            "offset": offset,
-        }
-
-        # Add filters if provided
-        if session_id:
-            payload["session_id"] = session_id.model_dump(exclude_none=True)
-        if namespace:
-            payload["namespace"] = namespace.model_dump(exclude_none=True)
-        if topics:
-            payload["topics"] = topics.model_dump(exclude_none=True)
-        if entities:
-            payload["entities"] = entities.model_dump(exclude_none=True)
-        if created_at:
-            payload["created_at"] = created_at.model_dump(
-                exclude_none=True, mode="json"
-            )
-        if last_accessed:
-            payload["last_accessed"] = last_accessed.model_dump(
-                exclude_none=True, mode="json"
-            )
-        if user_id:
-            payload["user_id"] = user_id.model_dump(exclude_none=True)
+            if isinstance(user_id, dict):
+                payload["user_id"] = user_id
+            else:
+                payload["user_id"] = user_id.model_dump(exclude_none=True)
         if memory_type:
             payload["memory_type"] = memory_type.model_dump(exclude_none=True)
         if distance_threshold is not None:
@@ -1076,7 +951,7 @@ async def add_memory_tool(
             # Create memory record
             memory = ClientMemoryRecord(
                 text=text,
-                memory_type=memory_type,
+                memory_type=MemoryTypeEnum(memory_type),
                 topics=topics,
                 entities=entities,
                 namespace=namespace or self.config.default_namespace,
@@ -1111,7 +986,7 @@ async def update_memory_data_tool(
         self,
         session_id: str,
         data: dict[str, Any],
-        merge_strategy: str = "merge",
+        merge_strategy: Literal["replace", "merge", "deep_merge"] = "merge",
         namespace: str | None = None,
         user_id: str | None = None,
     ) -> dict[str, Any]:
@@ -1997,70 +1872,6 @@ async def search_all_long_term_memories(
 
             offset += batch_size
 
-    async def search_all_memories(
-        self,
-        text: str,
-        session_id: SessionId | dict[str, Any] | None = None,
-        namespace: Namespace | dict[str, Any] | None = None,
-        topics: Topics | dict[str, Any] | None = None,
-        entities: Entities | dict[str, Any] | None = None,
-        created_at: CreatedAt | dict[str, Any] | None = None,
-        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,
-        batch_size: int = 50,
-    ) -> AsyncIterator[MemoryRecord]:
-        """
-        Auto-paginating version of unified memory search.
-
-        Searches both working memory and long-term memory with automatic pagination.
-
-        Args:
-            text: Search query text
-            session_id: Optional session ID filter
-            namespace: Optional namespace filter
-            topics: Optional topics filter
-            entities: Optional entities filter
-            created_at: Optional creation date filter
-            last_accessed: Optional last accessed date filter
-            user_id: Optional user ID filter
-            distance_threshold: Optional distance threshold
-            memory_type: Optional memory type filter
-            batch_size: Number of results to fetch per API call
-
-        Yields:
-            Individual memory records from all result pages
-        """
-        offset = 0
-        while True:
-            results = await self.search_memories(
-                text=text,
-                session_id=session_id,
-                namespace=namespace,
-                topics=topics,
-                entities=entities,
-                created_at=created_at,
-                last_accessed=last_accessed,
-                user_id=user_id,
-                distance_threshold=distance_threshold,
-                memory_type=memory_type,
-                limit=batch_size,
-                offset=offset,
-            )
-
-            if not results.memories:
-                break
-
-            for memory in results.memories:
-                yield memory
-
-            # If we got fewer results than batch_size, we've reached the end
-            if len(results.memories) < batch_size:
-                break
-
-            offset += batch_size
-
     def validate_memory_record(self, memory: ClientMemoryRecord | MemoryRecord) -> None:
         """
         Validate memory record before sending to server.
@@ -2237,7 +2048,7 @@ async def append_messages_to_working_memory(
                 converted_existing_messages.append(msg)
             else:
                 # Fallback for any other message type - convert to string content
-                converted_existing_messages.append(
+                converted_existing_messages.append(  # type: ignore
                     {"role": "user", "content": str(msg)}
                 )
 
diff --git a/agent-memory-client/agent_memory_client/models.py b/agent-memory-client/agent_memory_client/models.py
@@ -122,7 +122,7 @@ class ClientMemoryRecord(MemoryRecord):
     """A memory record with a client-provided ID"""
 
     id: str = Field(
-        default_factory=lambda: str(ulid.new()),
+        default_factory=lambda: str(ulid.ULID()),
         description="Client-provided ID generated by the client (ULID)",
     )
 
diff --git a/agent-memory-client/tests/test_basic.py b/agent-memory-client/tests/test_basic.py
@@ -83,7 +83,6 @@ def test_enhanced_methods():
 
     # Test pagination
     assert hasattr(client, "search_all_long_term_memories")
-    assert hasattr(client, "search_all_memories")
 
     # Test enhanced convenience methods
     assert hasattr(client, "update_working_memory_data")
diff --git a/agent-memory-client/tests/test_client.py b/agent-memory-client/tests/test_client.py
@@ -297,35 +297,6 @@ async def test_search_all_long_term_memories(self, enhanced_test_client):
             # Should have made 3 API calls
             assert mock_search.call_count == 3
 
-    @pytest.mark.asyncio
-    async def test_search_all_memories(self, enhanced_test_client):
-        """Test auto-paginating unified memory search."""
-        # Similar test for unified search
-        response = MemoryRecordResults(
-            total=25,
-            memories=[
-                MemoryRecordResult(
-                    id=f"memory-{i}",
-                    text=f"Memory text {i}",
-                    dist=0.1,
-                )
-                for i in range(25)
-            ],
-            next_offset=None,
-        )
-
-        with patch.object(enhanced_test_client, "search_memories") as mock_search:
-            mock_search.return_value = response
-
-            all_memories = []
-            async for memory in enhanced_test_client.search_all_memories(
-                text="test query", batch_size=50
-            ):
-                all_memories.append(memory)
-
-            assert len(all_memories) == 25
-            assert mock_search.call_count == 1
-
 
 class TestClientSideValidation:
     """Tests for client-side validation methods."""
diff --git a/agent_memory_server/api.py b/agent_memory_server/api.py
@@ -339,7 +339,7 @@ async def put_working_memory(
 
             memories = [
                 MemoryRecord(
-                    id=str(ulid.new()),
+                    id=str(ulid.ULID()),
                     session_id=session_id,
                     text=f"{msg.role}: {msg.content}",
                     namespace=updated_memory.namespace,
diff --git a/agent_memory_server/extraction.py b/agent_memory_server/extraction.py
@@ -333,7 +333,7 @@ async def extract_discrete_memories(
     if discrete_memories:
         long_term_memories = [
             MemoryRecord(
-                id_=str(ulid.new()),
+                id_=str(ulid.ULID()),
                 text=new_memory["text"],
                 memory_type=new_memory.get("type", "episodic"),
                 topics=new_memory.get("topics", []),
diff --git a/agent_memory_server/long_term_memory.py b/agent_memory_server/long_term_memory.py
@@ -244,7 +244,7 @@ async def merge_memories_with_llm(memories: list[dict], llm_client: Any = None)
     # Create the merged memory
     merged_memory = {
         "text": merged_text.strip(),
-        "id_": str(ulid.new()),
+        "id_": str(ulid.ULID()),
         "user_id": user_id,
         "session_id": session_id,
         "namespace": namespace,
@@ -664,7 +664,7 @@ async def index_long_term_memories(
     async with redis.pipeline(transaction=False) as pipe:
         for idx, vector in enumerate(embeddings):
             memory = processed_memories[idx]
-            id_ = memory.id if memory.id else str(ulid.new())
+            id_ = memory.id if memory.id else str(ulid.ULID())
             key = Keys.memory_key(id_, memory.namespace)
 
             # Generate memory hash for the memory
@@ -1426,7 +1426,7 @@ async def deduplicate_by_semantic_search(
 
             # Convert back to LongTermMemory
             merged_memory_obj = MemoryRecord(
-                id=memory.id or str(ulid.new()),
+                id=memory.id or str(ulid.ULID()),
                 text=merged_memory["text"],
                 user_id=merged_memory["user_id"],
                 session_id=merged_memory["session_id"],
@@ -1646,7 +1646,7 @@ async def extract_memories_from_messages(
 
                     # Create a new memory record from the extraction
                     extracted_memory = MemoryRecord(
-                        id=str(ulid.new()),  # Server-generated ID
+                        id=str(ulid.ULID()),  # Server-generated ID
                         text=memory_data["text"],
                         memory_type=memory_data.get("type", "semantic"),
                         topics=memory_data.get("topics", []),
diff --git a/agent_memory_server/mcp.py b/agent_memory_server/mcp.py
@@ -690,7 +690,7 @@ async def set_working_memory(
             # Handle both MemoryRecord objects and dict inputs
             if isinstance(memory, MemoryRecord):
                 # Already a MemoryRecord object, ensure it has an ID
-                memory_id = memory.id or str(ulid.new())
+                memory_id = memory.id or str(ulid.ULID())
                 processed_memory = memory.model_copy(
                     update={
                         "id": memory_id,
@@ -701,7 +701,7 @@ async def set_working_memory(
                 # Dictionary input, convert to MemoryRecord
                 memory_dict = dict(memory)
                 if not memory_dict.get("id"):
-                    memory_dict["id"] = str(ulid.new())
+                    memory_dict["id"] = str(ulid.ULID())
                 memory_dict["persisted_at"] = None
                 processed_memory = MemoryRecord(**memory_dict)
 
diff --git a/agent_memory_server/migrations.py b/agent_memory_server/migrations.py
@@ -98,7 +98,7 @@ async def migrate_add_discrete_memory_extracted_2(redis: Redis | None = None) ->
         id_ = await redis.hget(name=key, key="id_")  # type: ignore
         if not id_:
             logger.info("Updating memory with no ID to set ID")
-            await redis.hset(name=key, key="id_", value=str(ulid.new()))  # type: ignore
+            await redis.hset(name=key, key="id_", value=str(ulid.ULID()))  # type: ignore
         # extracted: bytes | None = await redis.hget(
         #     name=key, key="discrete_memory_extracted"
         # )  # type: ignore
@@ -126,7 +126,7 @@ async def migrate_add_memory_type_3(redis: Redis | None = None) -> None:
         id_ = await redis.hget(name=key, key="id_")  # type: ignore
         if not id_:
             logger.info("Updating memory with no ID to set ID")
-            await redis.hset(name=key, key="id_", value=str(ulid.new()))  # type: ignore
+            await redis.hset(name=key, key="id_", value=str(ulid.ULID()))  # type: ignore
         memory_type: bytes | None = await redis.hget(name=key, key="memory_type")  # type: ignore
         if not memory_type:
             await redis.hset(name=key, key="memory_type", value="message")  # type: ignore
diff --git a/agent_memory_server/models.py b/agent_memory_server/models.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_client_enhancements.py b/tests/test_client_enhancements.py
diff --git a/tests/test_long_term_memory.py b/tests/test_long_term_memory.py

Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,7 @@ class ClientMemoryRecord(MemoryRecord):`
`122`	`122`	`"""A memory record with a client-provided ID"""`
`123`	`123`
`124`	`124`	`id: str = Field(`
`125`		`- default_factory=lambda: str(ulid.new()),`
	`125`	`+ default_factory=lambda: str(ulid.ULID()),`
`126`	`126`	`description="Client-provided ID generated by the client (ULID)",`
`127`	`127`	`)`
`128`	`128`