Clean up some client handling of user IDs

Author: abrookins

agent-memory-client/README.md

@@ -206,8 +206,8 @@ await client.update_working_memory_data(
 
 # Append messages
 new_messages = [
-    MemoryMessage(role="user", content="What's the weather?"),
-    MemoryMessage(role="assistant", content="It's sunny today!")
+    {"role": "user", "content": "What's the weather?"},
+    {"role": "assistant", "content": "It's sunny today!"}
 ]
 
 await client.append_messages_to_working_memory(

agent-memory-client/agent_memory_client/client.py

@@ -233,6 +233,9 @@ async def get_working_memory(
         """
         params = {}
 
+        if user_id is not None:
+            params["user_id"] = user_id
+
         if namespace is not None:
             params["namespace"] = namespace
         elif self.config.default_namespace is not None:
@@ -258,7 +261,12 @@ async def get_working_memory(
                 f"/v1/working-memory/{session_id}", params=params
             )
             response.raise_for_status()
-            return WorkingMemoryResponse(**response.json())
+
+            # Get the raw JSON response
+            response_data = response.json()
+
+            # Messages from JSON parsing are already in the correct dict format
+            return WorkingMemoryResponse(**response_data)
         except httpx.HTTPStatusError as e:
             self._handle_http_error(e.response)
             raise
@@ -267,6 +275,7 @@ async def put_working_memory(
         self,
         session_id: str,
         memory: WorkingMemory,
+        user_id: str | None = None,
         model_name: str | None = None,
         context_window_max: int | None = None,
     ) -> WorkingMemoryResponse:
@@ -276,6 +285,7 @@ async def put_working_memory(
         Args:
             session_id: The session ID to store memory for
             memory: WorkingMemory object with messages and optional context
+            user_id: Optional user ID for the session (overrides user_id in memory object)
             model_name: Optional model name for context window management
             context_window_max: Optional direct specification of context window max tokens
 
@@ -289,6 +299,9 @@ async def put_working_memory(
         # Build query parameters for model-aware summarization
         params = {}
 
+        if user_id is not None:
+            params["user_id"] = user_id
+
         # Use provided model_name or fall back to config default
         effective_model_name = model_name or self.config.default_model_name
         if effective_model_name is not None:
@@ -2023,9 +2036,8 @@ async def update_working_memory_data(
     async def append_messages_to_working_memory(
         self,
         session_id: str,
-        messages: list[Any],  # Using Any since MemoryMessage isn't imported
+        messages: list[dict[str, Any]],  # Expect proper message dicts
         namespace: str | None = None,
-        auto_summarize: bool = True,
         model_name: str | None = None,
         context_window_max: int | None = None,
     ) -> WorkingMemoryResponse:
@@ -2036,9 +2048,8 @@ async def append_messages_to_working_memory(
 
         Args:
             session_id: Target session
-            messages: List of messages to append
+            messages: List of message dictionaries with 'role' and 'content' keys
             namespace: Optional namespace
-            auto_summarize: Whether to allow automatic summarization
             model_name: Optional model name for token-based summarization
             context_window_max: Optional direct specification of context window max tokens
 
@@ -2052,48 +2063,19 @@ async def append_messages_to_working_memory(
                 session_id=session_id, namespace=namespace
             )
 
-        # Combine messages - convert MemoryMessage objects to dicts if needed
-        existing_messages = existing_memory.messages if existing_memory else []
-
-        # Convert existing messages to dict format if they're objects
-        converted_existing_messages = []
-        for msg in existing_messages:
-            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}  # 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(
-                    {"role": "user", "content": str(msg)}
+        # Validate new messages have required structure
+        for msg in messages:
+            if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
+                raise ValueError(
+                    "All messages must be dictionaries with 'role' and 'content' keys"
                 )
 
-        # Convert new messages to dict format if they're objects
-        new_messages = []
-        for msg in messages:
-            if hasattr(msg, "model_dump"):
-                new_messages.append(msg.model_dump())
-            elif hasattr(msg, "role") and hasattr(msg, "content"):
-                new_messages.append({"role": msg.role, "content": msg.content})
-            elif isinstance(msg, dict):
-                # Message is already a dictionary, use as-is
-                new_messages.append(msg)
-            else:
-                # Fallback - assume it's already in the right format
-                new_messages.append(msg)
+        # Get existing messages (already in proper dict format from get_working_memory)
+        existing_messages = []
+        if existing_memory and existing_memory.messages:
+            existing_messages = existing_memory.messages
 
-        final_messages = converted_existing_messages + new_messages
+        final_messages = existing_messages + messages
 
         # Create updated working memory
         working_memory = WorkingMemory(

agent-memory-client/tests/test_client.py

@@ -527,8 +527,8 @@ async def test_append_messages_to_working_memory(self, enhanced_test_client):
         )
 
         new_messages = [
-            MemoryMessage(role="assistant", content="Second message"),
-            MemoryMessage(role="user", content="Third message"),
+            {"role": "assistant", "content": "Second message"},
+            {"role": "user", "content": "Third message"},
         ]
 
         with (

agent_memory_server/dependencies.py

@@ -27,10 +27,19 @@ async def add_task(
             logger.info("Scheduling task through Docket")
             # Get the Redis connection that's already configured (will use testcontainer in tests)
             redis_conn = await get_redis_conn()
-            # Use the connection's URL instead of settings.redis_url directly
-            redis_url = redis_conn.connection_pool.connection_kwargs.get(
-                "url", settings.redis_url
-            )
+
+            # Extract Redis URL from the connection pool
+            connection_kwargs = redis_conn.connection_pool.connection_kwargs
+            if "host" in connection_kwargs and "port" in connection_kwargs:
+                redis_url = (
+                    f"redis://{connection_kwargs['host']}:{connection_kwargs['port']}"
+                )
+                if "db" in connection_kwargs:
+                    redis_url += f"/{connection_kwargs['db']}"
+            else:
+                # Fallback to settings if we can't extract from connection
+                redis_url = settings.redis_url
+
             logger.info("redis_url: %s", redis_url)
             logger.info("docket_name: %s", settings.docket_name)
             async with Docket(

tests/conftest.py

@@ -275,11 +275,11 @@ def redis_url(redis_container):
 
 
 @pytest.fixture()
-def async_redis_client(redis_url):
+def async_redis_client(use_test_redis_connection):
     """
-    An async Redis client that uses the dynamic `redis_url`.
+    An async Redis client that uses the same connection as other test fixtures.
     """
-    return AsyncRedis.from_url(redis_url)
+    return use_test_redis_connection
 
 
 @pytest.fixture()

tests/test_cli.py

@@ -229,11 +229,11 @@ class TestTaskWorker:
 
     @patch("docket.Worker.run")
     @patch("agent_memory_server.cli.settings")
-    def test_task_worker_success(self, mock_settings, mock_worker_run):
+    def test_task_worker_success(self, mock_settings, mock_worker_run, redis_url):
         """Test successful task worker start."""
         mock_settings.use_docket = True
         mock_settings.docket_name = "test-docket"
-        mock_settings.redis_url = "redis://localhost:6379/0"
+        mock_settings.redis_url = redis_url
 
         mock_worker_run.return_value = None
 
@@ -258,11 +258,13 @@ def test_task_worker_docket_disabled(self, mock_settings):
 
     @patch("docket.Worker.run")
     @patch("agent_memory_server.cli.settings")
-    def test_task_worker_default_params(self, mock_settings, mock_worker_run):
+    def test_task_worker_default_params(
+        self, mock_settings, mock_worker_run, redis_url
+    ):
         """Test task worker with default parameters."""
         mock_settings.use_docket = True
         mock_settings.docket_name = "test-docket"
-        mock_settings.redis_url = "redis://localhost:6379/0"
+        mock_settings.redis_url = redis_url
 
         mock_worker_run.return_value = None
 

tests/test_client_enhancements.py

@@ -533,8 +533,8 @@ async def test_append_messages_to_working_memory(self, enhanced_test_client):
         )
 
         new_messages = [
-            MemoryMessage(role="assistant", content="Second message"),
-            MemoryMessage(role="user", content="Third message"),
+            {"role": "assistant", "content": "Second message"},
+            {"role": "user", "content": "Third message"},
         ]
 
         with (