code improvements

Author: Aayush Kataria

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -113,6 +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._background_tasks: set[asyncio.Task[Any]] = set()
         try:
@@ -155,6 +156,8 @@ def __init__(
         # accidentally close a user-supplied credential or leak one we created.
         self._owns_cosmos_credential = False
         self._owns_ai_foundry_credential = False
+        self._owns_sync_cosmos_credential = False
+        self._owns_sync_ai_foundry_credential = False
         if self._cosmos_credential is None and self._cosmos_key:
             self._cosmos_credential = self._cosmos_key
         # Keep a sync credential for the pipeline's sync Cosmos container
@@ -188,6 +191,7 @@ def __init__(
                     from azure.identity import DefaultAzureCredential as SyncDefaultAzureCredential
 
                     self._sync_cosmos_credential = SyncDefaultAzureCredential()
+                    self._owns_sync_cosmos_credential = True
                 except ImportError:
                     self._sync_cosmos_credential = None
             # Sync credential for AI Foundry (used by sync ChatClient inside
@@ -197,6 +201,7 @@ def __init__(
                     from azure.identity import DefaultAzureCredential as SyncDefaultAzureCredential
 
                     self._sync_ai_foundry_credential = SyncDefaultAzureCredential()
+                    self._owns_sync_ai_foundry_credential = True
                 except ImportError:
                     self._sync_ai_foundry_credential = None
 
@@ -300,9 +305,11 @@ async def close(self) -> None:
                     await close()
                 except Exception:
                     pass
-        # Sync credentials (Cosmos + AI Foundry fallback) close synchronously.
-        for sync_cred in (self._sync_cosmos_credential, self._sync_ai_foundry_credential):
-            if sync_cred is None:
+        for owns, sync_cred in (
+            (self._owns_sync_cosmos_credential, self._sync_cosmos_credential),
+            (self._owns_sync_ai_foundry_credential, self._sync_ai_foundry_credential),
+        ):
+            if not owns or sync_cred is None:
                 continue
             close = getattr(sync_cred, "close", None)
             if callable(close):
@@ -343,6 +350,9 @@ def add_local(
             salience=salience,
         )
         self.local_memory.append(memory)
+        if memory_type == "turn":
+            key = (user_id, thread_id)
+            self._unflushed_turn_counts[key] = self._unflushed_turn_counts.get(key, 0) + 1
         logger.debug("add_local id=%s role=%s type=%s", memory["id"], role, memory_type)
 
     def get_local(
@@ -701,11 +711,6 @@ async def push_to_cosmos(self, batch_size: int = 25) -> None:
             batch_size,
         )
         records = [MemoryRecord.from_cosmos_dict(dict(m)) for m in self.local_memory]
-        turn_counts: dict[tuple[str, str], int] = {}
-        for r in records:
-            if str(getattr(r, "memory_type", "")) == "turn":
-                key = (r.user_id, r.thread_id)
-                turn_counts[key] = turn_counts.get(key, 0) + 1
 
         for start in range(0, len(records), batch_size):
             batch = records[start : start + batch_size]
@@ -717,6 +722,9 @@ async def push_to_cosmos(self, batch_size: int = 25) -> None:
 
         logger.info("Async upserted batch of %d records", len(records))
 
+        turn_counts = self._unflushed_turn_counts
+        self._unflushed_turn_counts = {}
+
         if turn_counts:
             try:
                 task = asyncio.create_task(self._maybe_auto_trigger(turn_counts))

agent_memory_toolkit/cosmos_memory_client.py

@@ -112,6 +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._warned_owner_skip: bool = False
         self._warned_counter_unreachable: bool = False
 
@@ -278,6 +279,9 @@ def add_local(
             salience=salience,
         )
         self.local_memory.append(memory)
+        if memory_type == "turn":
+            key = (user_id, thread_id)
+            self._unflushed_turn_counts[key] = self._unflushed_turn_counts.get(key, 0) + 1
         logger.debug("add_local id=%s role=%s type=%s", memory["id"], role, memory_type)
 
     def get_local(
@@ -935,7 +939,6 @@ def push_to_cosmos(self, batch_size: int = 25) -> None:
             batch_size,
         )
         records = [MemoryRecord.from_cosmos_dict(dict(m)) for m in self.local_memory]
-        turn_counts: dict[tuple[str, str], int] = {}
         for start in range(0, len(records), batch_size):
             batch = records[start : start + batch_size]
             for record in batch:
@@ -944,11 +947,11 @@ def push_to_cosmos(self, batch_size: int = 25) -> None:
                     self._container_client.upsert_item(body=body)
                 except Exception as exc:
                     raise CosmosOperationError(f"Upsert failed for record {record.id}: {exc}") from exc
-                if str(getattr(record, "memory_type", "")) == "turn":
-                    key = (record.user_id, record.thread_id)
-                    turn_counts[key] = turn_counts.get(key, 0) + 1
         logger.info("Upserted batch of %d records", len(records))
 
+        turn_counts = self._unflushed_turn_counts
+        self._unflushed_turn_counts = {}
+
         try:
             self._maybe_auto_trigger(turn_counts)
         except Exception as exc:  # pragma: no cover - defensive

agent_memory_toolkit/pipeline.py

@@ -16,7 +16,7 @@
 from typing import Any, Optional
 
 from ._utils import DEFAULT_TTL_BY_TYPE, compute_content_hash
-from .exceptions import ValidationError
+from .exceptions import LLMError, ValidationError
 
 logger = logging.getLogger(__name__)
 
@@ -301,16 +301,24 @@ def _mark_superseded(self, old_doc: dict[str, Any], superseder_id: str) -> bool:
             return False
 
     @staticmethod
-    def _parse_llm_json(text: str) -> dict[str, Any]:
+    def _parse_llm_json(text: str | None) -> dict[str, Any]:
         """Parse JSON from an LLM response, stripping markdown fences."""
+        if text is None:
+            raise LLMError("LLM returned no content (None response body)")
         cleaned = text.strip()
         if cleaned.startswith("```"):
-            # Remove opening fence (possibly with language hint)
-            first_newline = cleaned.index("\n")
-            cleaned = cleaned[first_newline + 1 :]
+            first_newline = cleaned.find("\n")
+            if first_newline >= 0:
+                cleaned = cleaned[first_newline + 1 :]
+            else:
+                cleaned = cleaned.lstrip("`").lstrip()
         if cleaned.endswith("```"):
             cleaned = cleaned[:-3]
-        return json.loads(cleaned.strip())
+        try:
+            return json.loads(cleaned.strip())
+        except json.JSONDecodeError as exc:
+            preview = (text or "")[:200].replace("\n", " ")
+            raise LLMError(f"LLM returned invalid JSON (preview={preview!r}): {exc}") from exc
 
     # -----------------------------------------------------------------------
     # Public API
@@ -651,7 +659,9 @@ def extract_memories(
             self._upsert_memory(doc)
 
         result = {
-            "facts_count": sum(1 for d in docs_to_embed if d["type"] == "fact"),
+            "facts_count": sum(
+                1 for d in docs_to_embed if d["type"] == "fact" and "sys:unclassified" not in d.get("tags", [])
+            ),
             "procedural_count": sum(1 for d in docs_to_embed if d["type"] == "procedural"),
             "episodic_count": sum(1 for d in docs_to_embed if d["type"] == "episodic"),
             "unclassified_count": sum(1 for d in docs_to_embed if "sys:unclassified" in d.get("tags", [])),
@@ -704,6 +714,8 @@ def generate_thread_summary(
             query += " AND c.created_at > @since"
             parameters.append({"name": "@since", "value": since})
 
+        query_started_at = datetime.now(timezone.utc).isoformat()
+
         items = list(
             self._container.query_items(
                 query=query,
@@ -762,7 +774,6 @@ def generate_thread_summary(
         topic_tags = [f"topic:{t}" for t in topics]
         tags = ["sys:summary"] + topic_tags
 
-        now = datetime.now(timezone.utc).isoformat()
         summary_doc: dict[str, Any] = {
             "id": summary_id,
             "user_id": user_id,
@@ -779,8 +790,8 @@ def generate_thread_summary(
                 "recent_k": recent_k,
                 "incremental_update": existing_summary is not None,
             },
-            "created_at": existing_summary["created_at"] if existing_summary else now,
-            "updated_at": now,
+            "created_at": existing_summary["created_at"] if existing_summary else query_started_at,
+            "updated_at": query_started_at,
         }
 
         self._upsert_memory(summary_doc)
@@ -834,6 +845,8 @@ def generate_user_summary(
             for i, tid in enumerate(thread_ids):
                 parameters.append({"name": f"@tid{i}", "value": tid})
 
+        query_started_at = datetime.now(timezone.utc).isoformat()
+
         items = list(
             self._container.query_items(
                 query=query,
@@ -904,7 +917,6 @@ def generate_user_summary(
             all_thread_ids = sorted(new_thread_ids)
             total_memory_count = len(items)
 
-        now = datetime.now(timezone.utc).isoformat()
         summary_doc: dict[str, Any] = {
             "id": user_summary_id,
             "user_id": user_id,
@@ -923,8 +935,8 @@ def generate_user_summary(
                 "recent_k": recent_k,
                 "incremental_update": existing_summary is not None,
             },
-            "created_at": existing_summary["created_at"] if existing_summary else now,
-            "updated_at": now,
+            "created_at": existing_summary["created_at"] if existing_summary else query_started_at,
+            "updated_at": query_started_at,
         }
 
         self._upsert_memory(summary_doc)
@@ -1040,6 +1052,18 @@ def deduplicate_facts(
                         cluster_facts[0],
                     )
 
+                    merged_tags: list[str] = []
+                    seen_tags: set[str] = set()
+                    for f in cluster_facts:
+                        if f["id"] not in source_ids:
+                            continue
+                        for t in f.get("tags", []):
+                            if t not in seen_tags:
+                                seen_tags.add(t)
+                                merged_tags.append(t)
+                    if not merged_tags:
+                        merged_tags = ["sys:fact"]
+
                     source_confidences = [
                         c for f in cluster_facts if f["id"] in source_ids and (c := f.get("confidence")) is not None
                     ]
@@ -1059,7 +1083,7 @@ def deduplicate_facts(
                         },
                         "salience": act.get("salience"),
                         "supersedes_ids": source_ids,
-                        "tags": source_fact.get("tags", ["sys:fact"]),
+                        "tags": merged_tags,
                         "created_at": now,
                     }
                     if merged_confidence is not None:

function_app/_vendor/agent_memory_toolkit-0.1.0-py3-none-any.whl

@@ -90,10 +90,14 @@ async def increment_counter_by(
       times on HTTP 412.
     * Uses ``create_item`` for the first-write path, retrying on HTTP 409 in
       case multiple Function workers raced to seed the counter.
-    * If ``batch_max_lsn`` matches the LSN persisted on the existing doc, this
-      is treated as a change-feed replay and we return the cached
-      ``(pre_batch_count, current_count)`` **without writing** so the
-      threshold-crossing semantics are preserved without double-counting.
+    * If ``batch_max_lsn`` is *less than or equal to* the LSN persisted on the
+      existing doc, this is treated as a change-feed replay (immediate or
+      after a lease re-balance / host crash where checkpoints regressed) and
+      we return without writing. For the equal case we return the cached
+      ``(pre_batch_count, current_count)`` so threshold-crossing semantics
+      hold; for the strict-less case we return ``(current, current)`` (no
+      crossing) because some other batch already advanced the counter past
+      this one.
     * Preserves SDK-written failure breadcrumbs (``last_failure_at`` /
       ``last_failure_reason``) so monitors don't flap when the FA writes
       after an SDK failure stamp.
@@ -123,18 +127,35 @@ async def increment_counter_by(
             )
 
         # ---- Replay detection via LSN ----
+        # Use ``>=`` not ``==`` so out-of-order redeliveries (lease
+        # re-balance, host crash → checkpoint regression where another
+        # batch landed in between) also short-circuit. For the exact
+        # match we replay the cached result; for the strict-greater case
+        # we return (current, current) — no threshold crossing — because
+        # the batch's effect is already absorbed in a later state we have
+        # no cached pre-batch value for.
         if (
             batch_max_lsn is not None
             and existing_doc is not None
-            and existing_doc.get("last_batch_lsn") == batch_max_lsn
+            and existing_doc.get("last_batch_lsn") is not None
+            and existing_doc["last_batch_lsn"] >= batch_max_lsn
         ):
-            replay_old = existing_doc.get("last_batch_old_count", old_count)
+            stored_lsn = existing_doc["last_batch_lsn"]
+            if stored_lsn == batch_max_lsn:
+                replay_old = existing_doc.get("last_batch_old_count", old_count)
+                logger.info(
+                    "Counter replay detected counter_id=%s lsn=%s, returning cached result",
+                    counter_id,
+                    batch_max_lsn,
+                )
+                return (replay_old, old_count)
             logger.info(
-                "Counter replay detected counter_id=%s lsn=%s, returning cached result",
+                "Counter out-of-order replay counter_id=%s redelivered_lsn=%s stored_lsn=%s; no-op",
                 counter_id,
                 batch_max_lsn,
+                stored_lsn,
             )
-            return (replay_old, old_count)
+            return (old_count, old_count)
 
         new_count = old_count + count
         new_doc = {