code improvements

Author: Aayush Kataria

README.md

@@ -211,11 +211,12 @@ By default, the **InProcess processor** runs each pipeline step independently as
 
 | Env var | Default | Step that fires | Async behavior |
 |---|---|---|---|
-| `FACT_EXTRACTION_EVERY_N` | `1` (every turn) | `process_extract_memories` (extract + dedup) | scheduled via `asyncio.create_task` |
+| `FACT_EXTRACTION_EVERY_N` | `1` (every turn) | `process_extract_memories` | scheduled via `asyncio.create_task` |
+| `DEDUP_EVERY_N` | `5` | `process_dedup` (fires every Nth extract → effectively every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns) | scheduled via `asyncio.create_task` |
 | `THREAD_SUMMARY_EVERY_N` | `10` | `process_thread_summary` | scheduled via `asyncio.create_task` |
 | `USER_SUMMARY_EVERY_N` | `20` | `process_user_summary` | scheduled via `asyncio.create_task` |
 
-Each `*_EVERY_N=0` disables only that step. The Durable backend uses the same defaults via the change-feed function app, so the in-process and durable backends fire on the same turn boundaries — only the *where* differs. Calling `process_now()` is normally redundant — it remains as an explicit "process now" hook for tests, manual workflows, and operators who set every threshold to `0`.
+Each `*_EVERY_N=0` disables only that step. Dedup is gated independently of extract because cross-thread dedup is dramatically more expensive than per-thread extract (it reads every active fact for the user) — running it on every extract slammed AI Foundry. The Durable backend uses the same defaults via the change-feed function app (the function-app `azd` deploy bumps `FACT_EXTRACTION_EVERY_N` to `5` since the FA path is intended for higher-volume workloads). Calling `process_now()` is normally redundant — it remains as an explicit "process now" hook for tests, manual workflows, and operators who set every threshold to `0`.
 
 The async client (`AsyncCosmosMemoryClient.push_to_cosmos`) does **not** await the auto-trigger; it schedules it as a background `asyncio.Task` so the write call returns as soon as the Cosmos upserts complete. Background failures are surfaced via `logger.warning` (search for `"Background auto-trigger task failed"`).
 

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -154,7 +154,11 @@ def __init__(
         #   2. Explicit cosmos_key / ai_foundry_api_key (relief for environments
         #      where Cosmos control-plane RBAC is in private preview, etc.)
         #   3. DefaultAzureCredential() when use_default_credential=True
-        self._owns_credential = False
+        #
+        # Track ownership separately for each credential so close() doesn't
+        # accidentally close a user-supplied credential or leak one we created.
+        self._owns_cosmos_credential = False
+        self._owns_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
@@ -171,14 +175,17 @@ def __init__(
                 try:
                     from azure.identity.aio import DefaultAzureCredential
 
-                    _default = DefaultAzureCredential()
-                    self._owns_credential = True
+                    # Use independent instances per consumer so close()
+                    # ordering is unambiguous and one consumer's lifecycle
+                    # can't tear down a credential another is still using.
+                    if needs_cosmos:
+                        self._cosmos_credential = DefaultAzureCredential()
+                        self._owns_cosmos_credential = True
+                    if needs_embed:
+                        self._ai_foundry_credential = DefaultAzureCredential()
+                        self._owns_ai_foundry_credential = True
                 except ImportError:
-                    _default = None
-                if needs_cosmos and _default is not None:
-                    self._cosmos_credential = _default
-                if needs_embed and _default is not None:
-                    self._ai_foundry_credential = _default
+                    pass
             # Sync credential for pipeline (only when we don't already have a key)
             if self._sync_cosmos_credential is None:
                 try:
@@ -266,10 +273,20 @@ async def close(self) -> None:
                     pass
             self._sync_cosmos_client = None
         await self._embeddings_client.close()
-        if self._owns_credential and self._cosmos_credential is not None:
+        if self._owns_cosmos_credential and self._cosmos_credential is not None:
             close = getattr(self._cosmos_credential, "close", None)
             if close is not None:
-                await close()
+                try:
+                    await close()
+                except Exception:
+                    pass
+        if self._owns_ai_foundry_credential and self._ai_foundry_credential is not None:
+            close = getattr(self._ai_foundry_credential, "close", None)
+            if close is not None:
+                try:
+                    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:
@@ -1297,6 +1314,7 @@ async def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) ->
         from ..thresholds import (
             PROCESSOR_OWNER_DURABLE,
             PROCESSOR_OWNER_INPROCESS,
+            get_dedup_every_n,
             get_fact_extraction_every_n,
             get_processor_owner,
             get_thread_summary_every_n,
@@ -1330,9 +1348,14 @@ async def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) ->
         n_facts = get_fact_extraction_every_n()
         n_summary = get_thread_summary_every_n()
         n_user = get_user_summary_every_n()
+        n_dedup = get_dedup_every_n()
         if n_facts == 0 and n_summary == 0 and n_user == 0:
             return
 
+        # Dedup fires every Nth EXTRACT, not every Nth turn — see the
+        # sync-client mirror for the same combined-threshold pattern.
+        n_dedup_turns = n_facts * n_dedup if (n_facts > 0 and n_dedup > 0) else 0
+
         counter_container = self._get_counter_container()
         if counter_container is None:
             return
@@ -1349,6 +1372,7 @@ async def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) ->
                 n_facts=n_facts,
                 n_summary=n_summary,
                 n_user=n_user,
+                n_dedup_turns=n_dedup_turns,
             )
 
     async def _run_auto_trigger_steps(
@@ -1361,6 +1385,7 @@ async def _run_auto_trigger_steps(
         n_facts: int,
         n_summary: int,
         n_user: int,
+        n_dedup_turns: int = 0,
     ) -> None:
         from .._counters import (
             USER_COUNTER_THREAD_ID,
@@ -1394,6 +1419,7 @@ async def _run_auto_trigger_steps(
 
             fire_extract = n_facts > 0 and crosses_threshold(old_count, new_count, n_facts)
             fire_summary = n_summary > 0 and crosses_threshold(old_count, new_count, n_summary)
+            fire_dedup = n_dedup_turns > 0 and crosses_threshold(old_count, new_count, n_dedup_turns)
 
             if fire_extract:
                 try:
@@ -1413,6 +1439,23 @@ async def _run_auto_trigger_steps(
                         f"process_extract_memories: {exc!r}",
                     )
 
+            if fire_dedup:
+                try:
+                    await processor.process_dedup(user_id=user_id)
+                except Exception as exc:
+                    logger.warning(
+                        "Auto-trigger process_dedup failed for %s: %s",
+                        user_id,
+                        exc,
+                    )
+                    await stamp_failure_async(
+                        counter_container,
+                        thread_counter_id(user_id, thread_id),
+                        user_id,
+                        thread_id,
+                        f"process_dedup: {exc!r}",
+                    )
+
             if fire_summary:
                 try:
                     await processor.process_thread_summary(user_id=user_id, thread_id=thread_id)

agent_memory_toolkit/aio/processors/inprocess.py

@@ -89,10 +89,6 @@ async def process_extract_memories(
         thread_id: str,
     ) -> dict[str, int]:
         extracted = await asyncio.to_thread(self._pipeline.extract_memories, user_id, thread_id)
-        try:
-            await asyncio.to_thread(self._pipeline.deduplicate_facts, user_id)
-        except Exception as exc:  # pragma: no cover - defensive
-            logger.warning("post-extract dedup failed user_id=%s: %s", user_id, exc)
         return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {}
 
     async def process_thread_summary(

agent_memory_toolkit/chat.py

@@ -75,6 +75,25 @@ def _unsupported_param(exc: Exception) -> str | None:
 # ---------------------------------------------------------------------------
 
 
+def _extract_content(response: Any, model: str) -> str:
+    """Pull the assistant content out of a chat-completions response.
+
+    Raises ``LLMError`` when the model returned no content (content filter,
+    max-tokens-no-output, or certain reasoning-model paths). Without this
+    guard, downstream JSON parsing crashes with ``AttributeError: 'NoneType'
+    object has no attribute 'strip'`` and the actual root cause is invisible
+    in App Insights.
+    """
+    if not response.choices:
+        raise LLMError(f"LLM returned no choices (model={model})")
+    choice = response.choices[0]
+    content = getattr(choice.message, "content", None)
+    if content is None:
+        finish_reason = getattr(choice, "finish_reason", "unknown")
+        raise LLMError(f"LLM returned no content (model={model}, finish_reason={finish_reason})")
+    return content
+
+
 class ChatClient:
     """Synchronous LLM chat completion client backed by Azure OpenAI.
 
@@ -265,7 +284,7 @@ def generate(
                         usage.completion_tokens,
                         usage.total_tokens,
                     )
-                return response.choices[0].message.content
+                return _extract_content(response, self._model)
             except openai.RateLimitError as exc:
                 if attempt < max_retries - 1:
                     delay = base_delay * (2**attempt)
@@ -363,7 +382,7 @@ async def agenerate(
                         usage.completion_tokens,
                         usage.total_tokens,
                     )
-                return response.choices[0].message.content
+                return _extract_content(response, self._model)
             except openai.RateLimitError as exc:
                 if attempt < max_retries - 1:
                     delay = base_delay * (2**attempt)

agent_memory_toolkit/cosmos_memory_client.py

@@ -142,6 +142,11 @@ def __init__(
         #   2. Explicit cosmos_key / ai_foundry_api_key (relief for environments
         #      where Cosmos control-plane RBAC is in private preview, etc.)
         #   3. DefaultAzureCredential() when use_default_credential=True
+        #
+        # Track ownership separately so close() doesn't accidentally close a
+        # user-supplied credential or leak one we created.
+        self._owns_cosmos_credential = False
+        self._owns_ai_foundry_credential = False
         if self._cosmos_credential is None and self._cosmos_key:
             self._cosmos_credential = self._cosmos_key
 
@@ -152,13 +157,14 @@ def __init__(
                 try:
                     from azure.identity import DefaultAzureCredential
 
-                    _default = DefaultAzureCredential()
+                    if needs_cosmos:
+                        self._cosmos_credential = DefaultAzureCredential()
+                        self._owns_cosmos_credential = True
+                    if needs_embed:
+                        self._ai_foundry_credential = DefaultAzureCredential()
+                        self._owns_ai_foundry_credential = True
                 except ImportError:
-                    _default = None
-                if needs_cosmos and _default is not None:
-                    self._cosmos_credential = _default
-                if needs_embed and _default is not None:
-                    self._ai_foundry_credential = _default
+                    pass
 
         # Internal Cosmos SDK handles
         self._cosmos_client: Any = None
@@ -213,6 +219,20 @@ def close(self) -> None:
             self._container_client = None
             self._counter_container_client = None
             logger.info("Cosmos client closed")
+        # Close credentials we created ourselves (sync DefaultAzureCredential
+        # holds an underlying token cache + HTTP transport).
+        for owns, cred in (
+            (self._owns_cosmos_credential, self._cosmos_credential),
+            (self._owns_ai_foundry_credential, self._ai_foundry_credential),
+        ):
+            if not owns or cred is None:
+                continue
+            close = getattr(cred, "close", None)
+            if callable(close):
+                try:
+                    close()
+                except Exception:
+                    pass
 
     # ------------------------------------------------------------------
     # Read-only properties
@@ -621,6 +641,7 @@ def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
         from .thresholds import (
             PROCESSOR_OWNER_DURABLE,
             PROCESSOR_OWNER_INPROCESS,
+            get_dedup_every_n,
             get_fact_extraction_every_n,
             get_processor_owner,
             get_thread_summary_every_n,
@@ -660,9 +681,15 @@ def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
         n_facts = get_fact_extraction_every_n()
         n_summary = get_thread_summary_every_n()
         n_user = get_user_summary_every_n()
+        n_dedup = get_dedup_every_n()
         if n_facts == 0 and n_summary == 0 and n_user == 0:
             return
 
+        # Dedup fires every Nth EXTRACT, not every Nth turn. Express that
+        # via the same thread counter using the combined threshold
+        # n_facts * n_dedup. Disabled when either knob is 0.
+        n_dedup_turns = n_facts * n_dedup if (n_facts > 0 and n_dedup > 0) else 0
+
         counter_container = self._get_counter_container()
         if counter_container is None:
             return
@@ -691,6 +718,7 @@ def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
 
             fire_extract = n_facts > 0 and crosses_threshold(old_count, new_count, n_facts)
             fire_summary = n_summary > 0 and crosses_threshold(old_count, new_count, n_summary)
+            fire_dedup = n_dedup_turns > 0 and crosses_threshold(old_count, new_count, n_dedup_turns)
 
             if fire_extract:
                 try:
@@ -710,6 +738,23 @@ def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
                         f"process_extract_memories: {exc!r}",
                     )
 
+            if fire_dedup:
+                try:
+                    processor.process_dedup(user_id=user_id)
+                except Exception as exc:
+                    logger.warning(
+                        "Auto-trigger process_dedup failed for %s: %s",
+                        user_id,
+                        exc,
+                    )
+                    stamp_failure_sync(
+                        counter_container,
+                        thread_counter_id(user_id, thread_id),
+                        user_id,
+                        thread_id,
+                        f"process_dedup: {exc!r}",
+                    )
+
             if fire_summary:
                 try:
                     processor.process_thread_summary(user_id=user_id, thread_id=thread_id)