Support an optional user ID everywhere

Author: abrookins

CLAUDE.md

@@ -16,6 +16,7 @@ uv run ruff check            # Run linting
 uv run ruff format           # Format code
 uv run pytest                # Run tests
 uv run pytest tests/         # Run specific test directory
+uv run pytest --run-api-tests # Run all tests, including API tests
 uv add <dependency>          # Add a dependency to pyproject.toml and update lock file
 uv remove <dependency>       # Remove a dependency from pyproject.toml and update lock file
 

agent-memory-client/agent_memory_client/__init__.py

@@ -5,7 +5,7 @@
 memory management capabilities for AI agents and applications.
 """
 
-__version__ = "0.9.0b3"
+__version__ = "0.9.0b4"
 
 from .client import MemoryAPIClient, MemoryClientConfig, create_memory_client
 from .exceptions import (

agent-memory-client/agent_memory_client/client.py

@@ -166,7 +166,11 @@ async def health_check(self) -> HealthCheckResponse:
             raise
 
     async def list_sessions(
-        self, limit: int = 20, offset: int = 0, namespace: str | None = None
+        self,
+        limit: int = 20,
+        offset: int = 0,
+        namespace: str | None = None,
+        user_id: str | None = None,
     ) -> SessionListResponse:
         """
         List available sessions with optional pagination and namespace filtering.
@@ -175,6 +179,7 @@ async def list_sessions(
             limit: Maximum number of sessions to return (default: 20)
             offset: Offset for pagination (default: 0)
             namespace: Optional namespace filter
+            user_id: Optional user ID filter
 
         Returns:
             SessionListResponse containing session IDs and total count
@@ -188,6 +193,9 @@ async def list_sessions(
         elif self.config.default_namespace is not None:
             params["namespace"] = self.config.default_namespace
 
+        if user_id is not None:
+            params["user_id"] = user_id
+
         try:
             response = await self._client.get("/v1/working-memory/", params=params)
             response.raise_for_status()
@@ -199,6 +207,7 @@ async def list_sessions(
     async def get_working_memory(
         self,
         session_id: str,
+        user_id: str | None = None,
         namespace: str | None = None,
         window_size: int | None = None,
         model_name: ModelNameLiteral | None = None,
@@ -209,6 +218,7 @@ async def get_working_memory(
 
         Args:
             session_id: The session ID to retrieve working memory for
+            user_id: The user ID to retrieve working memory for
             namespace: Optional namespace for the session
             window_size: Optional number of messages to include
             model_name: Optional model name to determine context window size
@@ -304,14 +314,15 @@ async def put_working_memory(
             raise
 
     async def delete_working_memory(
-        self, session_id: str, namespace: str | None = None
+        self, session_id: str, namespace: str | None = None, user_id: str | None = None
     ) -> AckResponse:
         """
         Delete working memory for a session.
 
         Args:
             session_id: The session ID to delete memory for
             namespace: Optional namespace for the session
+            user_id: Optional user ID for the session
 
         Returns:
             AckResponse indicating success
@@ -322,6 +333,9 @@ async def delete_working_memory(
         elif self.config.default_namespace is not None:
             params["namespace"] = self.config.default_namespace
 
+        if user_id is not None:
+            params["user_id"] = user_id
+
         try:
             response = await self._client.delete(
                 f"/v1/working-memory/{session_id}", params=params
@@ -571,7 +585,7 @@ async def search_long_term_memory(
 
             print(f"Found {results.total} memories")
             for memory in results.memories:
-                print(f"- {memory.text[:100]}... (distance: {memory.distance})")
+                print(f"- {memory.text[:100]}... (distance: {memory.dist})")
             ```
         """
         # Convert dictionary filters to their proper filter objects if needed
@@ -651,11 +665,12 @@ async def search_memory_tool(
         user_id: str | None = None,
     ) -> dict[str, Any]:
         """
-        Simplified memory search designed for LLM tool use.
+        Simplified long-term memory search designed for LLM tool use.
 
         This method provides a streamlined interface for LLMs to search
         long-term memory with common parameters and user-friendly output.
-        Perfect for exposing as a tool to LLM frameworks.
+        Perfect for exposing as a tool to LLM frameworks. Note: This only
+        searches long-term memory, not working memory.
 
         Args:
             query: The search query text
@@ -664,6 +679,7 @@ async def search_memory_tool(
             memory_type: Optional memory type ("episodic", "semantic", "message")
             max_results: Maximum results to return (default: 5)
             min_relevance: Optional minimum relevance score (0.0-1.0)
+            user_id: Optional user ID to filter memories by
 
         Returns:
             Dict with 'memories' list and 'summary' for LLM consumption
@@ -729,8 +745,8 @@ async def search_memory_tool(
                     "created_at": memory.created_at.isoformat()
                     if memory.created_at
                     else None,
-                    "relevance_score": 1.0 - memory.distance
-                    if hasattr(memory, "distance") and memory.distance is not None
+                    "relevance_score": 1.0 - memory.dist
+                    if hasattr(memory, "dist") and memory.dist is not None
                     else None,
                 }
             )
@@ -784,7 +800,7 @@ async def handle_tool_calls(client, tool_calls):
             "type": "function",
             "function": {
                 "name": "search_memory",
-                "description": "Search long-term memory for relevant information based on a query. Use this when you need to recall past conversations, user preferences, or previously stored information.",
+                "description": "Search long-term memory for relevant information based on a query. Use this when you need to recall past conversations, user preferences, or previously stored information. Note: This searches only long-term memory, not current working memory.",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -820,6 +836,10 @@ async def handle_tool_calls(client, tool_calls):
                             "maximum": 1.0,
                             "description": "Optional minimum relevance score (0.0-1.0, higher = more relevant)",
                         },
+                        "user_id": {
+                            "type": "string",
+                            "description": "Optional user ID to filter memories by (e.g., 'user123')",
+                        },
                     },
                     "required": ["query"],
                 },
@@ -864,6 +884,7 @@ async def get_working_memory_tool(
             result = await self.get_working_memory(
                 session_id=session_id,
                 namespace=namespace or self.config.default_namespace,
+                user_id=user_id,
             )
 
             # Format for LLM consumption
@@ -2037,18 +2058,24 @@ async def append_messages_to_working_memory(
         # Convert existing messages to dict format if they're objects
         converted_existing_messages = []
         for msg in existing_messages:
-            if hasattr(msg, "model_dump"):
-                converted_existing_messages.append(msg.model_dump())
-            elif hasattr(msg, "role") and hasattr(msg, "content"):
+            if hasattr(msg, "model_dump") and callable(
+                getattr(msg, "model_dump", None)
+            ):
+                converted_existing_messages.append(msg.model_dump())  # type: ignore
+            elif (
+                hasattr(msg, "role")
+                and hasattr(msg, "content")
+                and not isinstance(msg, dict)
+            ):
                 converted_existing_messages.append(
-                    {"role": msg.role, "content": msg.content}
+                    {"role": msg.role, "content": msg.content}  # type: ignore
                 )
             elif isinstance(msg, dict):
                 # Message is already a dictionary, use as-is
                 converted_existing_messages.append(msg)
             else:
                 # Fallback for any other message type - convert to string content
-                converted_existing_messages.append(  # type: ignore
+                converted_existing_messages.append(
                     {"role": "user", "content": str(msg)}
                 )
 

agent_memory_server/__init__.py

@@ -1,3 +1,3 @@
 """Redis Agent Memory Server - A memory system for conversational AI."""
 
-__version__ = "0.9.0b3"
+__version__ = "0.9.0b4"

agent_memory_server/api.py

@@ -188,7 +188,7 @@ async def list_sessions(
     Get a list of session IDs, with optional pagination.
 
     Args:
-        options: Query parameters (page, size, namespace)
+        options: Query parameters (limit, offset, namespace, user_id)
 
     Returns:
         List of session IDs
@@ -200,6 +200,7 @@ async def list_sessions(
         limit=options.limit,
         offset=options.offset,
         namespace=options.namespace,
+        user_id=options.user_id,
     )
 
     return SessionListResponse(
@@ -211,8 +212,8 @@ async def list_sessions(
 @router.get("/v1/working-memory/{session_id}", response_model=WorkingMemoryResponse)
 async def get_working_memory(
     session_id: str,
+    user_id: str | None = None,
     namespace: str | None = None,
-    window_size: int = settings.window_size,  # Deprecated: kept for backward compatibility
     model_name: ModelNameLiteral | None = None,
     context_window_max: int | None = None,
     current_user: UserInfo = Depends(get_current_user),
@@ -225,8 +226,8 @@ async def get_working_memory(
 
     Args:
         session_id: The session ID
+        user_id: The user ID to retrieve working memory for
         namespace: The namespace to use for the session
-        window_size: DEPRECATED - The number of messages to include (kept for backward compatibility)
         model_name: The client's LLM model name (will determine context window size if provided)
         context_window_max: Direct specification of the context window max tokens (overrides model_name)
 
@@ -240,6 +241,7 @@ async def get_working_memory(
         session_id=session_id,
         namespace=namespace,
         redis_client=redis,
+        user_id=user_id,
     )
 
     if not working_mem:
@@ -249,6 +251,7 @@ async def get_working_memory(
             memories=[],
             session_id=session_id,
             namespace=namespace,
+            user_id=user_id,
         )
 
     # Apply token-based truncation if we have messages and model info
@@ -266,17 +269,14 @@ async def get_working_memory(
                     break
             working_mem.messages = truncated_messages
 
-    # Fallback to message-count truncation for backward compatibility
-    elif len(working_mem.messages) > window_size:
-        working_mem.messages = working_mem.messages[-window_size:]
-
     return working_mem
 
 
 @router.put("/v1/working-memory/{session_id}", response_model=WorkingMemoryResponse)
 async def put_working_memory(
     session_id: str,
     memory: WorkingMemory,
+    user_id: str | None = None,
     model_name: ModelNameLiteral | None = None,
     context_window_max: int | None = None,
     background_tasks=Depends(get_background_tasks),
@@ -291,6 +291,7 @@ async def put_working_memory(
     Args:
         session_id: The session ID
         memory: Working memory to save
+        user_id: Optional user ID for the session (overrides user_id in memory object)
         model_name: The client's LLM model name for context window determination
         context_window_max: Direct specification of context window max tokens
         background_tasks: DocketBackgroundTasks instance (injected automatically)
@@ -303,6 +304,10 @@ async def put_working_memory(
     # Ensure session_id matches
     memory.session_id = session_id
 
+    # Override user_id if provided as query parameter
+    if user_id is not None:
+        memory.user_id = user_id
+
     # Validate that all structured memories have id (if any)
     for mem in memory.memories:
         if not mem.id:
@@ -359,6 +364,7 @@ async def put_working_memory(
 @router.delete("/v1/working-memory/{session_id}", response_model=AckResponse)
 async def delete_working_memory(
     session_id: str,
+    user_id: str | None = None,
     namespace: str | None = None,
     current_user: UserInfo = Depends(get_current_user),
 ):
@@ -369,6 +375,7 @@ async def delete_working_memory(
 
     Args:
         session_id: The session ID
+        user_id: Optional user ID for the session
         namespace: Optional namespace for the session
 
     Returns:
@@ -379,6 +386,7 @@ async def delete_working_memory(
     # Delete unified working memory
     await working_memory.delete_working_memory(
         session_id=session_id,
+        user_id=user_id,
         namespace=namespace,
         redis_client=redis,
     )
@@ -558,6 +566,7 @@ async def memory_prompt(
         working_mem = await working_memory.get_working_memory(
             session_id=params.session.session_id,
             namespace=params.session.namespace,
+            user_id=params.session.user_id,
             redis_client=redis,
         )
 

agent_memory_server/config.py

@@ -20,7 +20,6 @@ def load_yaml_settings():
 class Settings(BaseSettings):
     redis_url: str = "redis://localhost:6379"
     long_term_memory: bool = True
-    window_size: int = 20
     openai_api_key: str | None = None
     anthropic_api_key: str | None = None
     generation_model: str = "gpt-4o-mini"
@@ -66,6 +65,9 @@ class Settings(BaseSettings):
     auth0_client_id: str | None = None
     auth0_client_secret: str | None = None
 
+    # Working memory settings
+    window_size: int = 20  # Default number of recent messages to return
+
     # Other Application settings
     log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"