Merge pull request #62 from redis/fix/change-get-or-create-working-memory

Add new_session field to WorkingMemoryResponse to indicate if a session was created, new client method get_or_create_working_memory_session()

Author: abrookins

agent-memory-client/agent_memory_client/__init__.py

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

agent-memory-client/agent_memory_client/client.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import logging  # noqa: F401
 import re
 from collections.abc import AsyncIterator, Sequence
 from typing import TYPE_CHECKING, Any, Literal, TypedDict
@@ -39,7 +40,6 @@
     RecencyConfig,
     SessionListResponse,
     WorkingMemory,
-    WorkingMemoryGetOrCreateResponse,
     WorkingMemoryResponse,
 )
 
@@ -120,10 +120,16 @@ def __init__(self, config: MemoryClientConfig):
         Args:
             config: MemoryClientConfig instance with server connection details
         """
+        from . import __version__
+
         self.config = config
         self._client = httpx.AsyncClient(
             base_url=config.base_url,
             timeout=config.timeout,
+            headers={
+                "User-Agent": f"agent-memory-client/{__version__}",
+                "X-Client-Version": __version__,
+            },
         )
 
     async def close(self) -> None:
@@ -289,11 +295,11 @@ async def get_or_create_working_memory(
         namespace: str | None = None,
         model_name: ModelNameLiteral | None = None,
         context_window_max: int | None = None,
-    ) -> WorkingMemoryGetOrCreateResponse:
+    ) -> tuple[bool, WorkingMemory]:
         """
         Get working memory for a session, creating it if it doesn't exist.
 
-        This method returns both the working memory and whether it was created or found.
+        This method returns a tuple with the creation status and the working memory.
         This is important for applications that need to know if they're working with
         a new session or an existing one.
 
@@ -305,24 +311,24 @@ async def get_or_create_working_memory(
             context_window_max: Optional direct specification of context window tokens
 
         Returns:
-            WorkingMemoryGetOrCreateResponse containing the memory and creation status
+            Tuple of (created: bool, memory: WorkingMemory)
+            - created: True if the session was created, False if it already existed
+            - memory: The WorkingMemory object
 
         Example:
             ```python
             # Get or create session memory
-            result = await client.get_or_create_working_memory(
+            created, memory = await client.get_or_create_working_memory(
                 session_id="chat_session_123",
                 user_id="user_456"
             )
 
-            if result.created:
-                print("Created new session")
+            if created:
+                logging.info("Created new session")
             else:
-                print("Found existing session")
+                logging.info("Found existing session")
 
-            # Access the memory
-            memory = result.memory
-            print(f"Session has {len(memory.messages)} messages")
+            logging.info(f"Session has {len(memory.messages)} messages")
             ```
         """
         try:
@@ -334,29 +340,54 @@ async def get_or_create_working_memory(
                 model_name=model_name,
                 context_window_max=context_window_max,
             )
-            return WorkingMemoryGetOrCreateResponse(
-                memory=existing_memory, created=False
-            )
-        except Exception:
-            # Session doesn't exist, create it
-            empty_memory = WorkingMemory(
-                session_id=session_id,
-                namespace=namespace or self.config.default_namespace,
-                messages=[],
-                memories=[],
-                data={},
-                user_id=user_id,
-            )
 
-            created_memory = await self.put_working_memory(
-                session_id=session_id,
-                memory=empty_memory,
-                user_id=user_id,
-                model_name=model_name,
-                context_window_max=context_window_max,
-            )
+            # Check if this is an unsaved session (deprecated behavior for old clients)
+            if getattr(existing_memory, "unsaved", None) is True:
+                # This is an unsaved session - we need to create it properly
+                empty_memory = WorkingMemory(
+                    session_id=session_id,
+                    namespace=namespace or self.config.default_namespace,
+                    messages=[],
+                    memories=[],
+                    data={},
+                    user_id=user_id,
+                )
+
+                created_memory = await self.put_working_memory(
+                    session_id=session_id,
+                    memory=empty_memory,
+                    user_id=user_id,
+                    model_name=model_name,
+                    context_window_max=context_window_max,
+                )
 
-            return WorkingMemoryGetOrCreateResponse(memory=created_memory, created=True)
+                return (True, created_memory)
+
+            return (False, existing_memory)
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                # Session doesn't exist, create it
+                empty_memory = WorkingMemory(
+                    session_id=session_id,
+                    namespace=namespace or self.config.default_namespace,
+                    messages=[],
+                    memories=[],
+                    data={},
+                    user_id=user_id,
+                )
+
+                created_memory = await self.put_working_memory(
+                    session_id=session_id,
+                    memory=empty_memory,
+                    user_id=user_id,
+                    model_name=model_name,
+                    context_window_max=context_window_max,
+                )
+
+                return (True, created_memory)
+            else:
+                # Re-raise other HTTP errors
+                raise
 
     async def put_working_memory(
         self,
@@ -484,11 +515,10 @@ async def set_working_memory_data(
         existing_memory = None
         if preserve_existing:
             try:
-                result_obj = await self.get_or_create_working_memory(
+                created, existing_memory = await self.get_or_create_working_memory(
                     session_id=session_id,
                     namespace=namespace,
                 )
-                existing_memory = result_obj.memory
             except Exception:
                 existing_memory = None
 
@@ -544,11 +574,10 @@ async def add_memories_to_working_memory(
             ```
         """
         # Get existing memory
-        result_obj = await self.get_or_create_working_memory(
+        created, existing_memory = await self.get_or_create_working_memory(
             session_id=session_id,
             namespace=namespace,
         )
-        existing_memory = result_obj.memory
 
         # Determine final memories list
         if replace or not existing_memory:
@@ -610,7 +639,7 @@ async def create_long_term_memory(
             ]
 
             response = await client.create_long_term_memory(memories)
-            print(f"Stored memories: {response.status}")
+            logging.info(f"Stored memories: {response.status}")
             ```
         """
         # Apply default namespace and ensure IDs are present
@@ -764,9 +793,9 @@ async def search_long_term_memory(
                 distance_threshold=0.3
             )
 
-            print(f"Found {results.total} memories")
+            logging.info(f"Found {results.total} memories")
             for memory in results.memories:
-                print(f"- {memory.text[:100]}... (distance: {memory.dist})")
+                logging.info(f"- {memory.text[:100]}... (distance: {memory.dist})")
             ```
         """
         # Convert dictionary filters to their proper filter objects if needed
@@ -916,9 +945,9 @@ async def search_memory_tool(
                 min_relevance=0.7
             )
 
-            print(result["summary"])  # "Found 2 relevant memories for: user preferences about UI themes"
+            logging.info(result["summary"])  # "Found 2 relevant memories for: user preferences about UI themes"
             for memory in result["memories"]:
-                print(f"- {memory['text']} (score: {memory['relevance_score']})")
+                logging.info(f"- {memory['text']} (score: {memory['relevance_score']})")
             ```
 
         LLM Framework Integration:
@@ -1119,18 +1148,17 @@ async def get_working_memory_tool(
                 session_id="current_session"
             )
 
-            print(memory_state["summary"])  # Human-readable summary
-            print(f"Messages: {memory_state['message_count']}")
-            print(f"Memories: {len(memory_state['memories'])}")
+            logging.info(memory_state["summary"])  # Human-readable summary
+            logging.info(f"Messages: {memory_state['message_count']}")
+            logging.info(f"Memories: {len(memory_state['memories'])}")
             ```
         """
         try:
-            result_obj = await self.get_or_create_working_memory(
+            created, result = await self.get_or_create_working_memory(
                 session_id=session_id,
                 namespace=namespace or self.config.default_namespace,
                 user_id=user_id,
             )
-            result = result_obj.memory
 
             # Format for LLM consumption
             message_count = len(result.messages) if result.messages else 0
@@ -1200,24 +1228,23 @@ async def get_or_create_working_memory_tool(
             )
 
             if memory_state["created"]:
-                print("Created new session")
+                logging.info("Created new session")
             else:
-                print("Found existing session")
+                logging.info("Found existing session")
 
-            print(memory_state["summary"])  # Human-readable summary
-            print(f"Messages: {memory_state['message_count']}")
-            print(f"Memories: {len(memory_state['memories'])}")
+            logging.info(memory_state["summary"])  # Human-readable summary
+            logging.info(f"Messages: {memory_state['message_count']}")
+            logging.info(f"Memories: {len(memory_state['memories'])}")
             ```
         """
         try:
-            result_obj = await self.get_or_create_working_memory(
+            created, result = await self.get_or_create_working_memory(
                 session_id=session_id,
                 namespace=namespace or self.config.default_namespace,
                 user_id=user_id,
             )
 
             # Format for LLM consumption
-            result = result_obj.memory
             message_count = len(result.messages) if result.messages else 0
             memory_count = len(result.memories) if result.memories else 0
             data_keys = list(result.data.keys()) if result.data else []
@@ -1238,11 +1265,11 @@ async def get_or_create_working_memory_tool(
                         }
                     )
 
-            status_text = "new session" if result_obj.created else "existing session"
+            status_text = "new session" if created else "existing session"
 
             return {
                 "session_id": session_id,
-                "created": result_obj.created,
+                "created": created,
                 "message_count": message_count,
                 "memory_count": memory_count,
                 "memories": formatted_memories,
@@ -1299,7 +1326,7 @@ async def add_memory_tool(
                 entities=["vegetarian", "restaurants"]
             )
 
-            print(result["summary"])  # "Successfully stored semantic memory"
+            logging.info(result["summary"])  # "Successfully stored semantic memory"
             ```
         """
         try:
@@ -1373,7 +1400,7 @@ async def update_memory_data_tool(
                 }
             )
 
-            print(result["summary"])  # "Successfully updated 3 data entries"
+            logging.info(result["summary"])  # "Successfully updated 3 data entries"
             ```
         """
         try:
@@ -1948,9 +1975,9 @@ async def resolve_tool_call(
             )
 
             if result["success"]:
-                print(result["formatted_response"])
+                logging.info(result["formatted_response"])
             else:
-                print(f"Error: {result['error']}")
+                logging.error(f"Error: {result['error']}")
             ```
         """
         try:
@@ -2004,7 +2031,7 @@ async def resolve_tool_calls(
 
             for result in results:
                 if result["success"]:
-                    print(f"{result['function_name']}: {result['formatted_response']}")
+                    logging.info(f"{result['function_name']}: {result['formatted_response']}")
             ```
         """
         results = []
@@ -2062,9 +2089,9 @@ async def resolve_function_call(
                     )
 
                     if result["success"]:
-                        print(result["formatted_response"])
+                        logging.info(result["formatted_response"])
                     else:
-                        print(f"Error: {result['error']}")
+                        logging.error(f"Error: {result['error']}")
             ```
         """
         import json
@@ -2352,7 +2379,7 @@ async def resolve_function_calls(
             results = await client.resolve_function_calls(calls, "session123")
             for result in results:
                 if result["success"]:
-                    print(f"{result['function_name']}: {result['formatted_response']}")
+                    logging.info(f"{result['function_name']}: {result['formatted_response']}")
             ```
         """
         results = []
@@ -2395,10 +2422,9 @@ async def promote_working_memories_to_long_term(
             Acknowledgement of promotion operation
         """
         # Get current working memory
-        result_obj = await self.get_or_create_working_memory(
+        created, working_memory = await self.get_or_create_working_memory(
             session_id=session_id, namespace=namespace
         )
-        working_memory = result_obj.memory
 
         # Filter memories if specific IDs are requested
         memories_to_promote = working_memory.memories
@@ -2611,10 +2637,9 @@ async def update_working_memory_data(
             WorkingMemoryResponse with updated memory
         """
         # Get existing memory
-        result_obj = await self.get_or_create_working_memory(
+        created, existing_memory = await self.get_or_create_working_memory(
             session_id=session_id, namespace=namespace, user_id=user_id
         )
-        existing_memory = result_obj.memory
 
         # Determine final data based on merge strategy
         if existing_memory and existing_memory.data:
@@ -2667,10 +2692,9 @@ async def append_messages_to_working_memory(
             WorkingMemoryResponse with updated memory (potentially summarized if token limit exceeded)
         """
         # Get existing memory
-        result_obj = await self.get_or_create_working_memory(
+        created, existing_memory = await self.get_or_create_working_memory(
             session_id=session_id, namespace=namespace, user_id=user_id
         )
-        existing_memory = result_obj.memory
 
         # Convert messages to MemoryMessage objects
         converted_messages = []