Users/akataria/improvements (#8)

* code improvements

* code improvements

* code improvements

* Cache push_to_cosmos embeddings; clarify non-turn thread_id

- push_to_cosmos (sync + async): persist generated embeddings back to
  local_memory so a repeat push doesn't re-embed already-embedded
  non-turn records.
- Update test comment: _make_memory auto-generates a UUID thread_id for
  non-turn types when omitted (it does not user-scope them).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Aayush Kataria <aayushkataria@Aayushs-MacBook-Pro-2.local>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 3 people

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -113,7 +113,7 @@ def __init__(
     ) -> None:
         # Local store
         self.local_memory: list[dict[str, Any]] = []
-        self._unflushed_turn_counts: dict[tuple[str, Optional[str]], int] = {}
+        self._unflushed_turn_counts: dict[tuple[str, str], int] = {}
 
         self._background_tasks: set[asyncio.Task[Any]] = set()
         try:
@@ -361,6 +361,11 @@ def add_local(
             ttl=ttl,
             salience=salience,
         )
+        if memory_type == "turn" and not thread_id:
+            raise ValidationError(
+                "thread_id is required for memory_type='turn' so the auto-trigger "
+                "counter can group turns per conversation. Set thread_id explicitly."
+            )
         self.local_memory.append(memory)
         if memory_type == "turn":
             key = (user_id, thread_id)
@@ -491,6 +496,8 @@ async def connect_cosmos(
         try:
             from azure.cosmos.aio import CosmosClient
 
+            await self._drain_cosmos_client()
+
             client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential)
             db = client.get_database_client(self._cosmos_database)
             container_handle = db.get_container_client(self._cosmos_container)
@@ -573,6 +580,8 @@ async def create_memory_store(
             from azure.cosmos import PartitionKey, ThroughputProperties
             from azure.cosmos.aio import CosmosClient
 
+            await self._drain_cosmos_client()
+
             client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential)
 
             db = await client.create_database_if_not_exists(id=self._cosmos_database)
@@ -733,7 +742,34 @@ async def push_to_cosmos(self, batch_size: int = 25) -> None:
 
         for start in range(0, len(records), batch_size):
             batch = records[start : start + batch_size]
-            tasks = [self._container_client.upsert_item(body=r.to_cosmos_dict()) for r in batch]
+            bodies = [r.to_cosmos_dict() for r in batch]
+
+            # Batch-embed non-turn memories that don't already carry a
+            # vector — one /embeddings POST per Cosmos batch instead of
+            # N concurrent ones.
+            to_embed_idx: list[int] = []
+            to_embed_text: list[str] = []
+            for i, body in enumerate(bodies):
+                if body.get("type") != "turn" and body.get("content") and not body.get("embedding"):
+                    to_embed_idx.append(i)
+                    to_embed_text.append(body["content"])
+            if to_embed_text:
+                try:
+                    vectors = await self._embeddings_client.generate_batch(to_embed_text)
+                    for i, vec in zip(to_embed_idx, vectors):
+                        bodies[i]["embedding"] = vec
+                        # Persist the embedding back to local_memory so a
+                        # repeat push_to_cosmos() doesn't re-embed.
+                        self.local_memory[start + i]["embedding"] = vec
+                except Exception as exc:  # noqa: BLE001
+                    logger.warning(
+                        "push_to_cosmos: batch embedding generation failed (%s); "
+                        "proceeding without embeddings for %d records",
+                        exc,
+                        len(to_embed_text),
+                    )
+
+            tasks = [self._container_client.upsert_item(body=b) for b in bodies]
             try:
                 await asyncio.gather(*tasks)
             except Exception as exc:
@@ -1247,7 +1283,7 @@ def _drain_pipeline_resources(self) -> None:
                 try:
                     close()
                 except Exception:
-                    pass
+                    logger.warning("Failed to close prior sync Cosmos client", exc_info=True)
             self._sync_cosmos_client = None
 
         prior_sync_embeddings = getattr(self, "_sync_embeddings_client", None)
@@ -1257,9 +1293,34 @@ def _drain_pipeline_resources(self) -> None:
                 try:
                     close()
                 except Exception:
-                    pass
+                    logger.warning("Failed to close prior sync EmbeddingsClient", exc_info=True)
             self._sync_embeddings_client = None
 
+    async def _drain_cosmos_client(self) -> None:
+        """Close any prior async Cosmos client before reassigning the field.
+
+        Also nulls the cached pipeline, counter container handle, and any
+        lazily-created processor — otherwise they retain references to
+        the drained sync container client and the next op fails with an
+        opaque "Object disposed" error. A user-supplied processor is left
+        intact since the SDK does not own its lifecycle.
+        """
+        prior = self._cosmos_client
+        if prior is None:
+            return
+        close = getattr(prior, "close", None)
+        if callable(close):
+            try:
+                await close()
+            except Exception:
+                logger.warning("Failed to close prior async Cosmos client during reconnect", exc_info=True)
+        self._cosmos_client = None
+        self._container_client = None
+        self._counter_container_client = None
+        self._pipeline = None
+        if not self._processor_explicit:
+            self._processor = None
+
     def _init_pipeline(self) -> None:
         """Initialize the ProcessingPipeline with a sync container client.
 

agent_memory_toolkit/aio/processors/inprocess.py

@@ -115,11 +115,18 @@ async def process_dedup(self, *, user_id: str) -> int:
 
     @staticmethod
     def _extract_dedup_count(dedup: Any) -> int:
-        if isinstance(dedup, dict):
-            for key in ("deduplicated", "merged", "removed", "deduplicated_count"):
-                if key in dedup and isinstance(dedup[key], int):
-                    return dedup[key]
-        return 0
+        """Sum the merged + superseded facts from a ``deduplicate_facts`` result.
+
+        ``ProcessingPipeline.deduplicate_facts`` returns a dict with
+        ``{"kept", "merged", "superseded"}`` — both ``merged`` and
+        ``superseded`` represent facts that were consolidated, so they
+        contribute to the dedup-count metric.
+        """
+        if not isinstance(dedup, dict):
+            return 0
+        merged = dedup.get("merged", 0) if isinstance(dedup.get("merged"), int) else 0
+        superseded = dedup.get("superseded", 0) if isinstance(dedup.get("superseded"), int) else 0
+        return merged + superseded
 
     async def generate_user_summary(
         self,

agent_memory_toolkit/chat.py

@@ -281,6 +281,8 @@ def generate(
         )
 
         attempt = 0
+        unsupported_strips = 0
+        max_unsupported_strips = 5
         while True:
             try:
                 response = client.chat.completions.create(**kwargs)
@@ -315,13 +317,14 @@ def generate(
                 # this does NOT consume a retry slot since it's a request-shape
                 # repair, not a transient failure.
                 bad_param = _unsupported_param(exc) if status == 400 else None
-                if bad_param and bad_param in kwargs:
+                if bad_param and bad_param in kwargs and unsupported_strips < max_unsupported_strips:
                     logger.warning(
                         "LLM model=%s rejected '%s'; retrying without it.",
                         self._model,
                         bad_param,
                     )
                     kwargs.pop(bad_param, None)
+                    unsupported_strips += 1
                     continue
                 if status in _RETRYABLE_STATUS_CODES and attempt < max_retries - 1:
                     delay = base_delay * (2**attempt)
@@ -379,6 +382,8 @@ async def agenerate(
         )
 
         attempt = 0
+        unsupported_strips = 0
+        max_unsupported_strips = 5
         while True:
             try:
                 response = await client.chat.completions.create(**kwargs)
@@ -411,13 +416,14 @@ async def agenerate(
                 # Strip-unsupported-param: request-shape repair, not a transient
                 # failure — does NOT consume a retry slot.
                 bad_param = _unsupported_param(exc) if status == 400 else None
-                if bad_param and bad_param in kwargs:
+                if bad_param and bad_param in kwargs and unsupported_strips < max_unsupported_strips:
                     logger.warning(
                         "LLM model=%s rejected '%s'; retrying without it.",
                         self._model,
                         bad_param,
                     )
                     kwargs.pop(bad_param, None)
+                    unsupported_strips += 1
                     continue
                 if status in _RETRYABLE_STATUS_CODES and attempt < max_retries - 1:
                     delay = base_delay * (2**attempt)

agent_memory_toolkit/cosmos_memory_client.py

@@ -112,7 +112,7 @@ def __init__(
     ) -> None:
         # Local store
         self.local_memory: list[dict[str, Any]] = []
-        self._unflushed_turn_counts: dict[tuple[str, Optional[str]], int] = {}
+        self._unflushed_turn_counts: dict[tuple[str, str], int] = {}
         self._warned_owner_skip: bool = False
         self._warned_counter_unreachable: bool = False
 
@@ -278,6 +278,11 @@ def add_local(
             ttl=ttl,
             salience=salience,
         )
+        if memory_type == "turn" and not thread_id:
+            raise ValidationError(
+                "thread_id is required for memory_type='turn' so the auto-trigger "
+                "counter can group turns per conversation. Set thread_id explicitly."
+            )
         self.local_memory.append(memory)
         if memory_type == "turn":
             key = (user_id, thread_id)
@@ -408,6 +413,8 @@ def connect_cosmos(
         try:
             from azure.cosmos import CosmosClient
 
+            self._drain_cosmos_client()
+
             client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential)
             db = client.get_database_client(self._cosmos_database)
             container_handle = db.get_container_client(self._cosmos_container)
@@ -489,6 +496,8 @@ def create_memory_store(
         try:
             from azure.cosmos import CosmosClient, PartitionKey, ThroughputProperties
 
+            self._drain_cosmos_client()
+
             client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential)
 
             db = client.create_database_if_not_exists(id=self._cosmos_database)
@@ -559,6 +568,28 @@ def _init_pipeline(self) -> None:
         )
         self._warn_on_embedding_dim_mismatch()
 
+    def _drain_cosmos_client(self) -> None:
+        """Close any prior Cosmos client before reassigning the field.
+
+        Repeated ``connect_cosmos`` / ``create_memory_store`` calls on the
+        same instance must not leak the prior client's httpx pool / FDs.
+        """
+        prior = self._cosmos_client
+        if prior is None:
+            return
+        close = getattr(prior, "close", None)
+        if callable(close):
+            try:
+                close()
+            except Exception:
+                logger.warning("Failed to close prior Cosmos client during reconnect", exc_info=True)
+        self._cosmos_client = None
+        self._container_client = None
+        self._counter_container_client = None
+        self._pipeline = None
+        if not self._processor_explicit:
+            self._processor = None
+
     def _warn_on_embedding_dim_mismatch(self) -> None:
         """Log a WARNING if the resolved embedding dim differs from the container's policy.
 
@@ -948,8 +979,34 @@ def push_to_cosmos(self, batch_size: int = 25) -> None:
         records = [MemoryRecord.from_cosmos_dict(dict(m)) for m in self.local_memory]
         for start in range(0, len(records), batch_size):
             batch = records[start : start + batch_size]
-            for record in batch:
-                body = record.to_cosmos_dict()
+            bodies = [r.to_cosmos_dict() for r in batch]
+
+            # Batch-embed non-turn memories that don't already carry a
+            # vector — one /embeddings POST per Cosmos batch instead of
+            # one per record.
+            to_embed_idx: list[int] = []
+            to_embed_text: list[str] = []
+            for i, body in enumerate(bodies):
+                if body.get("type") != "turn" and body.get("content") and not body.get("embedding"):
+                    to_embed_idx.append(i)
+                    to_embed_text.append(body["content"])
+            if to_embed_text:
+                try:
+                    vectors = self._embeddings_client.generate_batch(to_embed_text)
+                    for i, vec in zip(to_embed_idx, vectors):
+                        bodies[i]["embedding"] = vec
+                        # Persist the embedding back to local_memory so a
+                        # repeat push_to_cosmos() doesn't re-embed.
+                        self.local_memory[start + i]["embedding"] = vec
+                except Exception as exc:  # noqa: BLE001
+                    logger.warning(
+                        "push_to_cosmos: batch embedding generation failed (%s); "
+                        "proceeding without embeddings for %d records",
+                        exc,
+                        len(to_embed_text),
+                    )
+
+            for record, body in zip(batch, bodies):
                 try:
                     self._container_client.upsert_item(body=body)
                 except Exception as exc:

agent_memory_toolkit/pipeline.py

@@ -229,13 +229,20 @@ def _load_existing_memories(
         memory_types: list[str],
         limit: int = 100,
     ) -> list[dict[str, Any]]:
-        """Query active (non-superseded) memories for reconciliation context."""
+        """Query active (non-superseded) memories for reconciliation context.
+
+        Results are ordered by ``c._ts DESC`` so the most recently written
+        memories survive the cap — without ORDER BY, Cosmos returns rows
+        in implementation-defined order and the dedup comparison set is
+        non-deterministic.
+        """
         type_placeholders = ", ".join(f"@mtype{i}" for i in range(len(memory_types)))
         query = (
             f"SELECT TOP {limit} * FROM c "
             f"WHERE c.user_id = @user_id "
             f"AND c.type IN ({type_placeholders}) "
-            f"AND (NOT IS_DEFINED(c.superseded_by) OR c.superseded_by = null)"
+            f"AND (NOT IS_DEFINED(c.superseded_by) OR c.superseded_by = null) "
+            f"ORDER BY c._ts DESC"
         )
         parameters: list[dict[str, Any]] = [
             {"name": "@user_id", "value": user_id},
@@ -355,6 +362,7 @@ def extract_memories(
             self._container.query_items(
                 query=query,
                 parameters=parameters,
+                partition_key=[user_id, thread_id],
             )
         )
 
@@ -720,6 +728,7 @@ def generate_thread_summary(
             self._container.query_items(
                 query=query,
                 parameters=parameters,
+                partition_key=[user_id, thread_id],
             )
         )
 

agent_memory_toolkit/processors/inprocess.py

@@ -120,11 +120,18 @@ def process_dedup(self, *, user_id: str) -> int:
 
     @staticmethod
     def _extract_dedup_count(dedup: Any) -> int:
-        if isinstance(dedup, dict):
-            for key in ("deduplicated", "merged", "removed", "deduplicated_count"):
-                if key in dedup and isinstance(dedup[key], int):
-                    return dedup[key]
-        return 0
+        """Sum the merged + superseded facts from a ``deduplicate_facts`` result.
+
+        ``ProcessingPipeline.deduplicate_facts`` returns a dict with
+        ``{"kept", "merged", "superseded"}`` — both ``merged`` and
+        ``superseded`` represent facts that were consolidated, so they
+        contribute to the dedup-count metric.
+        """
+        if not isinstance(dedup, dict):
+            return 0
+        merged = dedup.get("merged", 0) if isinstance(dedup.get("merged"), int) else 0
+        superseded = dedup.get("superseded", 0) if isinstance(dedup.get("superseded"), int) else 0
+        return merged + superseded
 
     def generate_user_summary(
         self,

tests/unit/aio/processors/test_inprocess.py

@@ -14,7 +14,7 @@ async def test_process_thread_calls_summarize_extract_dedup_in_order():
     pipeline = MagicMock()
     pipeline.generate_thread_summary.return_value = {"id": "summary"}
     pipeline.extract_memories.return_value = {"facts": 1}
-    pipeline.deduplicate_facts.return_value = {"deduplicated": 2}
+    pipeline.deduplicate_facts.return_value = {"merged": 2, "superseded": 0, "kept": 3}
 
     proc = AsyncInProcessProcessor(pipeline=pipeline)
     result = await proc.process_thread(user_id="u", thread_id="t", turns=[])