Consolidate api and message files

Author: abrookins

agent_memory_server/api.py

@@ -3,7 +3,7 @@
 from mcp.server.fastmcp.prompts import base
 from mcp.types import TextContent
 
-from agent_memory_server import long_term_memory, messages, working_memory
+from agent_memory_server import long_term_memory, working_memory
 from agent_memory_server.auth import UserInfo, get_current_user
 from agent_memory_server.config import settings
 from agent_memory_server.dependencies import get_background_tasks
@@ -51,89 +51,6 @@ def _get_effective_window_size(
     return effective_window_size
 
 
-@router.get("/sessions/", response_model=SessionListResponse)
-async def list_sessions(
-    options: GetSessionsQuery = Depends(),
-    current_user: UserInfo = Depends(get_current_user),
-):
-    """
-    Get a list of session IDs, with optional pagination.
-
-    Args:
-        options: Query parameters (page, size, namespace)
-
-    Returns:
-        List of session IDs
-    """
-    redis = await get_redis_conn()
-
-    total, session_ids = await messages.list_sessions(
-        redis=redis,
-        limit=options.limit,
-        offset=options.offset,
-        namespace=options.namespace,
-    )
-
-    return SessionListResponse(
-        sessions=session_ids,
-        total=total,
-    )
-
-
-@router.get("/sessions/{session_id}/memory", response_model=WorkingMemoryResponse)
-async def get_session_memory(
-    session_id: str,
-    namespace: str | None = None,
-    window_size: int = settings.window_size,
-    model_name: ModelNameLiteral | None = None,
-    context_window_max: int | None = None,
-    current_user: UserInfo = Depends(get_current_user),
-):
-    """
-    Get working memory for a session.
-
-    This includes stored conversation messages, context, and structured memory records.
-
-    Args:
-        session_id: The session ID
-        namespace: The namespace to use for the session
-        window_size: The number of messages to include in the response
-        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)
-
-    Returns:
-        Working memory containing messages, context, and structured memory records
-    """
-    redis = await get_redis_conn()
-    effective_window_size = _get_effective_window_size(
-        window_size=window_size,
-        context_window_max=context_window_max,
-        model_name=model_name,
-    )
-
-    # Get unified working memory
-    working_mem = await working_memory.get_working_memory(
-        session_id=session_id,
-        namespace=namespace,
-        redis_client=redis,
-    )
-
-    if not working_mem:
-        # Return empty working memory if none exists
-        working_mem = WorkingMemory(
-            messages=[],
-            memories=[],
-            session_id=session_id,
-            namespace=namespace,
-        )
-
-    # Apply window size to messages if needed
-    if len(working_mem.messages) > effective_window_size:
-        working_mem.messages = working_mem.messages[-effective_window_size:]
-
-    return working_mem
-
-
 async def _summarize_working_memory(
     memory: WorkingMemory,
     window_size: int,
@@ -218,6 +135,89 @@ async def _summarize_working_memory(
     return updated_memory
 
 
+@router.get("/sessions/", response_model=SessionListResponse)
+async def list_sessions(
+    options: GetSessionsQuery = Depends(),
+    current_user: UserInfo = Depends(get_current_user),
+):
+    """
+    Get a list of session IDs, with optional pagination.
+
+    Args:
+        options: Query parameters (page, size, namespace)
+
+    Returns:
+        List of session IDs
+    """
+    redis = await get_redis_conn()
+
+    total, session_ids = await working_memory.list_sessions(
+        redis=redis,
+        limit=options.limit,
+        offset=options.offset,
+        namespace=options.namespace,
+    )
+
+    return SessionListResponse(
+        sessions=session_ids,
+        total=total,
+    )
+
+
+@router.get("/sessions/{session_id}/memory", response_model=WorkingMemoryResponse)
+async def get_session_memory(
+    session_id: str,
+    namespace: str | None = None,
+    window_size: int = settings.window_size,
+    model_name: ModelNameLiteral | None = None,
+    context_window_max: int | None = None,
+    current_user: UserInfo = Depends(get_current_user),
+):
+    """
+    Get working memory for a session.
+
+    This includes stored conversation messages, context, and structured memory records.
+
+    Args:
+        session_id: The session ID
+        namespace: The namespace to use for the session
+        window_size: The number of messages to include in the response
+        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)
+
+    Returns:
+        Working memory containing messages, context, and structured memory records
+    """
+    redis = await get_redis_conn()
+    effective_window_size = _get_effective_window_size(
+        window_size=window_size,
+        context_window_max=context_window_max,
+        model_name=model_name,
+    )
+
+    # Get unified working memory
+    working_mem = await working_memory.get_working_memory(
+        session_id=session_id,
+        namespace=namespace,
+        redis_client=redis,
+    )
+
+    if not working_mem:
+        # Return empty working memory if none exists
+        working_mem = WorkingMemory(
+            messages=[],
+            memories=[],
+            session_id=session_id,
+            namespace=namespace,
+        )
+
+    # Apply window size to messages if needed
+    if len(working_mem.messages) > effective_window_size:
+        working_mem.messages = working_mem.messages[-effective_window_size:]
+
+    return working_mem
+
+
 @router.put("/sessions/{session_id}/memory", response_model=WorkingMemoryResponse)
 async def put_session_memory(
     session_id: str,
@@ -249,7 +249,7 @@ async def put_session_memory(
         if not mem.id:
             raise HTTPException(
                 status_code=400,
-                detail="All memory records in working memory must have an id",
+                detail="All memory records in working memory must have an ID",
             )
 
     # Handle summarization if needed (before storing)

agent_memory_server/long_term_memory.py

@@ -993,13 +993,13 @@ async def search_memories(
                 session_ids_to_search = [session_id.eq]
             else:
                 # Get all sessions for broader search
-                from agent_memory_server import messages
+                from agent_memory_server import working_memory
 
                 namespace_value = None
                 if namespace and hasattr(namespace, "eq"):
                     namespace_value = namespace.eq
 
-                _, session_ids_to_search = await messages.list_sessions(
+                _, session_ids_to_search = await working_memory.list_sessions(
                     redis=redis,
                     limit=1000,  # Get a reasonable number of sessions
                     offset=0,