feat: support separate turns container for short-term/long-term memory separation

Add optional cosmos_turns_container parameter to CosmosMemoryClient and
AsyncCosmosMemoryClient that routes turn-type documents to a dedicated
container while keeping derived memories (facts, summaries, episodic,
procedural) in the main container.

- New _container_for_type() routes writes by memory_type
- New _container_for_query() returns single container for simple cases
- New _containers_for_query() returns list of containers for multi-container merging
- get_thread uses _containers_for_query and merges with chronological sort
- ProcessingPipeline now accepts cosmos_turns_container and queries both
  containers for thread-wide operations (extract_memories, generate_user_summary)
- create_memory_store provisions turns container when configured
- update_cosmos/delete_cosmos check both containers (main first, then turns)
- close/drain methods clear _turns_container_client
- Fully backward compatible: omit the param for single-container behavior

Author: TheovanKraay

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -98,6 +98,7 @@ def __init__(
         cosmos_key: Optional[str] = None,
         cosmos_database: Optional[str] = None,
         cosmos_container: Optional[str] = None,
+        cosmos_turns_container: Optional[str] = None,
         cosmos_counter_container: Optional[str] = None,
         cosmos_lease_container: Optional[str] = None,
         cosmos_throughput_mode: Optional[str] = None,
@@ -131,6 +132,7 @@ def __init__(
         self._cosmos_key = cosmos_key
         self._cosmos_database = cosmos_database or "ai_memory"
         self._cosmos_container = cosmos_container or "memories"
+        self._cosmos_turns_container = cosmos_turns_container  # None means use main container
         self._cosmos_counter_container = cosmos_counter_container or "counter"
         self._cosmos_lease_container = cosmos_lease_container or "leases"
         self._cosmos_throughput_mode = _resolve_cosmos_throughput_mode(cosmos_throughput_mode)
@@ -208,6 +210,7 @@ def __init__(
         # Internal Cosmos SDK handles
         self._cosmos_client: Any = None
         self._container_client: Any = None
+        self._turns_container_client: Any = None  # Separate container for turns (optional)
         self._counter_container_client: Any = None
 
         # Composed sub-clients
@@ -274,6 +277,7 @@ async def close(self) -> None:
             await self._cosmos_client.close()
             self._cosmos_client = None
             self._container_client = None
+            self._turns_container_client = None
             self._counter_container_client = None
         if self._sync_cosmos_client is not None:
             close = getattr(self._sync_cosmos_client, "close", None)
@@ -505,6 +509,19 @@ async def connect_cosmos(
 
             self._cosmos_client = client
             self._container_client = container_handle
+
+            # Connect turns container if configured separately
+            if self._cosmos_turns_container:
+                turns_handle = db.get_container_client(self._cosmos_turns_container)
+                self._turns_container_client = turns_handle
+                logger.info(
+                    "Async connected turns container: %s/%s",
+                    self._cosmos_database,
+                    self._cosmos_turns_container,
+                )
+            else:
+                self._turns_container_client = None
+
             self._init_pipeline()
         except Exception as exc:
             raise CosmosOperationError(f"Failed to connect to Cosmos DB (async): {exc}") from exc
@@ -519,6 +536,7 @@ async def create_memory_store(
         self,
         database: Optional[str] = None,
         container: Optional[str] = None,
+        turns_container: Optional[str] = None,
         counter_container: Optional[str] = None,
         lease_container: Optional[str] = None,
         endpoint: Optional[str] = None,
@@ -560,6 +578,8 @@ async def create_memory_store(
             self._cosmos_key = key
         self._cosmos_database = database or self._cosmos_database
         self._cosmos_container = container or self._cosmos_container
+        if turns_container is not None:
+            self._cosmos_turns_container = turns_container
         self._cosmos_counter_container = counter_container or self._cosmos_counter_container
         self._cosmos_lease_container = lease_container or self._cosmos_lease_container
         self._cosmos_throughput_mode = _resolve_cosmos_throughput_mode(
@@ -630,6 +650,29 @@ async def create_memory_store(
             )
             self._cosmos_client = client
             self._container_client = container_handle
+
+            # Create and connect separate turns container if configured
+            if self._cosmos_turns_container:
+                turns_handle = await db.create_container_if_not_exists(
+                    **_build_container_kwargs(
+                        container_id=self._cosmos_turns_container,
+                        partition_key=partition_key,
+                        offer_throughput=offer_throughput,
+                        default_ttl=-1,
+                        indexing_policy=idx_policy,
+                        vector_embedding_policy=vec_policy,
+                        full_text_policy=ft_policy,
+                    )
+                )
+                self._turns_container_client = turns_handle
+                logger.info(
+                    "Created turns container: %s/%s",
+                    self._cosmos_database,
+                    self._cosmos_turns_container,
+                )
+            else:
+                self._turns_container_client = None
+
             self._init_pipeline()
         except Exception as exc:
             raise CosmosOperationError(f"Failed to create memory store (async): {exc}") from exc
@@ -647,6 +690,60 @@ async def _require_cosmos(self) -> None:
         if self._container_client is None:
             raise CosmosNotConnectedError()
 
+    def _container_for_type(self, memory_type: str) -> Any:
+        """Return the appropriate container client based on memory type."""
+        if memory_type == "turn" and self._turns_container_client is not None:
+            return self._turns_container_client
+        return self._container_client
+
+    def _container_for_query(self, memory_types: Optional[list[str]] = None) -> Any:
+        """Return the appropriate container for a read query.
+
+        In the default single-container configuration this returns only the
+        main memories container. When a dedicated turns container is
+        configured:
+
+        * turn-only queries should search only the turns container
+        * non-turn-only queries should search only the main memories container
+        * mixed or unspecified queries may need to inspect both containers
+        """
+        if not memory_types:
+            return self._container_client
+
+        has_turn = any(t == "turn" for t in memory_types)
+        has_not_turn = any(t != "turn" for t in memory_types)
+
+        if self._turns_container_client is not None:
+            if has_turn and not has_not_turn:
+                return self._turns_container_client
+            if not has_turn:
+                return self._container_client
+            return self._container_client
+        return self._container_client
+
+    def _containers_for_query(self, memory_types: Optional[list[str]] = None) -> list[Any]:
+        """Return candidate containers for a read query.
+
+        Unlike :meth:`_container_for_query` which returns a single container,
+        this returns all containers that should be queried to get complete
+        results, enabling callers to merge across containers.
+        """
+        if self._turns_container_client is None:
+            return [self._container_client]
+
+        if not memory_types:
+            return [self._container_client, self._turns_container_client]
+
+        has_turn = any(t == "turn" for t in memory_types)
+        has_not_turn = any(t != "turn" for t in memory_types)
+
+        if has_turn and has_not_turn:
+            return [self._container_client, self._turns_container_client]
+
+        if has_turn:
+            return [self._turns_container_client]
+        return [self._container_client]
+
     # ------------------------------------------------------------------
     # Cosmos DB CRUD operations (async)
     # ------------------------------------------------------------------
@@ -705,7 +802,8 @@ async def add_cosmos(
                 )
 
         try:
-            await self._container_client.upsert_item(body=body)
+            container = self._container_for_type(memory_type)
+            await container.upsert_item(body=body)
         except Exception as exc:
             raise CosmosOperationError(f"Async upsert failed for record {record.id}: {exc}") from exc
         logger.info("add_cosmos id=%s role=%s type=%s", record.id, role, memory_type)
@@ -770,7 +868,7 @@ async def push_to_cosmos(self, batch_size: int = 25) -> None:
                         len(to_embed_text),
                     )
 
-            tasks = [self._container_client.upsert_item(body=b) for b in bodies]
+            tasks = [self._container_for_type(b.get("type", "turn")).upsert_item(body=b) for b in bodies]
             try:
                 await asyncio.gather(*tasks)
             except Exception as exc:
@@ -870,7 +968,8 @@ async def get_memories(
         logger.debug("async get_memories query: %s", query)
 
         try:
-            items_iter = self._container_client.query_items(
+            container = self._container_for_query(memory_types)
+            items_iter = container.query_items(
                 query=query,
                 parameters=parameters or None,
             )
@@ -914,6 +1013,14 @@ async def update_cosmos(
                 parameters=[{"name": "@id", "value": memory_id}],
             )
             docs = [item async for item in items_iter]
+            target_container = self._container_client
+            if not docs and self._turns_container_client is not None:
+                items_iter = self._turns_container_client.query_items(
+                    query="SELECT * FROM c WHERE c.id = @id",
+                    parameters=[{"name": "@id", "value": memory_id}],
+                )
+                docs = [item async for item in items_iter]
+                target_container = self._turns_container_client
         except Exception as exc:
             raise CosmosOperationError(f"async update query failed: {exc}") from exc
 
@@ -932,7 +1039,7 @@ async def update_cosmos(
         doc["updated_at"] = datetime.now(timezone.utc).isoformat()
 
         try:
-            await self._container_client.replace_item(item=doc["id"], body=doc)
+            await target_container.replace_item(item=doc["id"], body=doc)
         except Exception as exc:
             raise CosmosOperationError(f"async update replace failed for {memory_id}: {exc}") from exc
 
@@ -962,14 +1069,29 @@ async def delete_cosmos(self, memory_id: str, thread_id: str, user_id: str) -> N
                 ],
             )
             docs = [item async for item in items_iter]
+            target_container = self._container_client
+            if not docs and self._turns_container_client is not None:
+                items_iter = self._turns_container_client.query_items(
+                    query=(
+                        "SELECT TOP 1 c.id FROM c WHERE c.id = @id"
+                        " AND c.thread_id = @thread_id AND c.user_id = @user_id"
+                    ),
+                    parameters=[
+                        {"name": "@id", "value": memory_id},
+                        {"name": "@thread_id", "value": thread_id},
+                        {"name": "@user_id", "value": user_id},
+                    ],
+                )
+                docs = [item async for item in items_iter]
+                target_container = self._turns_container_client
         except Exception as exc:
             raise CosmosOperationError(f"async delete lookup failed: {exc}") from exc
 
         if not docs:
             raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id)
 
         try:
-            await self._container_client.delete_item(item=memory_id, partition_key=[user_id, thread_id])
+            await target_container.delete_item(item=memory_id, partition_key=[user_id, thread_id])
         except Exception as exc:
             raise CosmosOperationError(f"async delete failed for {memory_id}: {exc}") from exc
 
@@ -1126,8 +1248,14 @@ async def get_thread(
         logger.debug("async get_thread query: %s", query)
 
         try:
-            items_iter = self._container_client.query_items(query=query, parameters=parameters)
-            items = [item async for item in items_iter]
+            containers = self._containers_for_query(memory_types)
+            items = []
+            for container in containers:
+                items_iter = container.query_items(query=query, parameters=parameters)
+                items.extend([item async for item in items_iter])
+            # Re-sort merged results to preserve chronological ordering
+            if len(containers) > 1:
+                items.sort(key=lambda x: x.get("created_at", ""), reverse=True)
         except Exception as exc:
             raise CosmosOperationError(f"async get_thread query failed: {exc}") from exc
 
@@ -1317,6 +1445,7 @@ async def _drain_cosmos_client(self) -> None:
                 logger.warning("Failed to close prior async Cosmos client during reconnect", exc_info=True)
         self._cosmos_client = None
         self._container_client = None
+        self._turns_container_client = None
         self._counter_container_client = None
         self._pipeline = None
         if not self._processor_explicit:
@@ -1368,10 +1497,18 @@ def _init_pipeline(self) -> None:
         self._sync_embeddings_client = sync_embeddings
 
         if sync_container is not None:
+            # Build sync turns container if configured
+            sync_turns_container = None
+            if self._cosmos_turns_container:
+                try:
+                    sync_turns_container = sync_db.get_container_client(self._cosmos_turns_container)
+                except Exception:
+                    logger.warning("Failed to create sync turns container for pipeline")
             self._pipeline = ProcessingPipeline(
                 cosmos_container=sync_container,
                 chat_client=self._chat_client,
                 embeddings_client=sync_embeddings,
+                cosmos_turns_container=sync_turns_container,
             )
             self._warn_on_embedding_dim_mismatch(sync_container)