code improvements

Author: Aayush Kataria

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -115,18 +115,14 @@ def __init__(
         self.local_memory: list[dict[str, Any]] = []
 
         self._background_tasks: set[asyncio.Task[Any]] = set()
-        # Backpressure for fire-and-forget auto-trigger tasks. Without a
-        # cap, FACT_EXTRACTION_EVERY_N=1 + a burst of pushes spawns one
-        # extract pipeline per push and slams AI Foundry. Configurable via
-        # MEMORY_AUTO_TRIGGER_CONCURRENCY (default 4).
         try:
             _max = int(os.environ.get("MEMORY_AUTO_TRIGGER_CONCURRENCY", "4"))
         except ValueError:
             _max = 4
         self._auto_trigger_semaphore: asyncio.Semaphore = asyncio.Semaphore(max(1, _max))
-        # One-shot WARN guard: skip-due-to-MEMORY_PROCESSOR_OWNER=durable
-        # would otherwise log on every push.
         self._warned_owner_skip: bool = False
+        self._warned_counter_unreachable: bool = False
+        self._pipeline_init_error: Exception | None = None
 
         # Store kwargs directly
         self._cosmos_endpoint = cosmos_endpoint
@@ -249,16 +245,23 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
     async def close(self) -> None:
         """Close all underlying async clients.
 
-        Drains any in-flight fire-and-forget auto-trigger tasks first so
-        they don't continue running against torn-down Cosmos/embeddings
-        handles.
+        Cancels any in-flight fire-and-forget auto-trigger tasks first so
+        ``close()`` / ``__aexit__`` doesn't block for tens of seconds while
+        a mid-LLM extract runs to completion. Auto-trigger work is
+        recoverable on the next push, so cancellation is safe.
         """
-        # Drain fire-and-forget auto-trigger tasks before tearing down the
-        # underlying clients, otherwise an in-flight extract+dedup will
-        # crash with CosmosOperationError mid-shutdown.
+        # Cancel + bounded-wait. Without the cancel, an in-flight extract+dedup
+        # can hold close() for tens of seconds. With cancel + 5s wait we
+        # guarantee bounded shutdown latency; remaining tasks raise
+        # CancelledError and will be retried on the next push_to_cosmos.
         pending = list(self._background_tasks)
         if pending:
-            await asyncio.gather(*pending, return_exceptions=True)
+            for task in pending:
+                task.cancel()
+            try:
+                await asyncio.wait(pending, timeout=5.0)
+            except Exception:
+                pass
         if self._cosmos_client is not None:
             await self._cosmos_client.close()
             self._cosmos_client = None
@@ -273,6 +276,16 @@ async def close(self) -> None:
                     pass
             self._sync_cosmos_client = None
         await self._embeddings_client.close()
+        # Async client's ChatClient is shared with the sync pipeline; close
+        # its sync httpx pool too so we don't leak across __aexit__.
+        try:
+            self._chat_client.close_sync()
+        except Exception:
+            pass
+        try:
+            await self._chat_client.close()
+        except Exception:
+            pass
         if self._owns_cosmos_credential and self._cosmos_credential is not None:
             close = getattr(self._cosmos_credential, "close", None)
             if close is not None:
@@ -1213,7 +1226,14 @@ def _init_pipeline(self) -> None:
             sync_db = sync_client.get_database_client(self._cosmos_database)
             sync_container = sync_db.get_container_client(self._cosmos_container)
             self._sync_cosmos_client = sync_client
+            self._pipeline_init_error = None
         except Exception as exc:
+            # Capture the real failure so _require_pipeline can surface it
+            # later — the alternative is the user calls connect_cosmos()
+            # successfully, then sees CosmosNotConnectedError("call
+            # connect_cosmos() first") on the first extract, with the real
+            # cause buried in an old WARN.
+            self._pipeline_init_error = exc
             logger.warning("Failed to create sync Cosmos client for pipeline: %s", exc)
             sync_container = None
 
@@ -1268,6 +1288,14 @@ def _warn_on_embedding_dim_mismatch(self, sync_container: Any) -> None:
     def _require_pipeline(self) -> None:
         """Raise if the processing pipeline is not available."""
         if self._pipeline is None:
+            if self._pipeline_init_error is not None:
+                raise CosmosNotConnectedError(
+                    "Processing pipeline failed to initialize "
+                    f"({type(self._pipeline_init_error).__name__}: "
+                    f"{self._pipeline_init_error}). The async client connected "
+                    "but the sync Cosmos client used by the pipeline could "
+                    "not be built. Check credentials and endpoint."
+                ) from self._pipeline_init_error
             raise CosmosNotConnectedError(
                 "Processing pipeline requires Cosmos DB connection. "
                 "Call connect_cosmos() or create_memory_store() first."
@@ -1281,7 +1309,11 @@ def _get_processor(self) -> AsyncMemoryProcessor:
         return self._processor
 
     def _get_counter_container(self) -> Any:
-        """Lazy handle to the counter container (best-effort, returns None on failure)."""
+        """Lazy handle to the counter container (best-effort, returns None on failure).
+
+        Logs a one-shot WARN on first failure so operators don't silently
+        lose all memory processing because of a missing counter container.
+        """
         if self._counter_container_client is not None:
             return self._counter_container_client
         if self._cosmos_client is None:
@@ -1291,12 +1323,18 @@ def _get_counter_container(self) -> Any:
             self._counter_container_client = db.get_container_client(self._cosmos_counter_container)
             return self._counter_container_client
         except Exception as exc:  # pragma: no cover - defensive
-            logger.warning(
-                "Counter container %s/%s unreachable; auto-trigger disabled: %s",
-                self._cosmos_database,
-                self._cosmos_counter_container,
-                exc,
-            )
+            if not self._warned_counter_unreachable:
+                self._warned_counter_unreachable = True
+                logger.warning(
+                    "Counter container %s/%s unreachable (%s: %s); "
+                    "auto-trigger DISABLED for the lifetime of this client. "
+                    "Provision the container or set MEMORY_PROCESSOR_OWNER=durable "
+                    "if processing runs in the Function App.",
+                    self._cosmos_database,
+                    self._cosmos_counter_container,
+                    type(exc).__name__,
+                    exc,
+                )
             return None
 
     async def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
@@ -1421,6 +1459,14 @@ async def _run_auto_trigger_steps(
             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)
 
+            # Order matters: extract → dedup → summary.
+            # Each step gates on its OWN threshold (independent crossings),
+            # but when multiple thresholds cross in the same batch we still
+            # want the data flow to be: write fresh facts, deduplicate them,
+            # then fold the deduplicated set into the summary. Reordering
+            # would risk a summary that includes since-removed duplicates
+            # or omits just-extracted facts. Keep in sync with the sync
+            # client and function_app/orchestrators/.
             if fire_extract:
                 try:
                     await processor.process_extract_memories(user_id=user_id, thread_id=thread_id)

agent_memory_toolkit/chat.py

@@ -10,6 +10,7 @@
 
 import asyncio
 import logging
+import re
 import time
 from typing import Any
 
@@ -62,10 +63,18 @@ async def _provider() -> str:
 def _unsupported_param(exc: Exception) -> str | None:
     """If *exc* is a 400 about an unsupported sampling param, return its name."""
     msg = str(exc).lower()
-    if "400" not in msg and "unsupported" not in msg and "does not support" not in msg:
+    if "400" not in msg:
+        return None
+    if not (
+        "does not support" in msg
+        or "is not supported" in msg
+        or "unsupported parameter" in msg
+        or "unsupported value" in msg
+    ):
         return None
     for p in _SAMPLING_PARAMS:
-        if p in msg:
+        pattern = rf"(?<![a-z_]){re.escape(p)}(?![a-z_])"
+        if re.search(pattern, msg):
             return p
     return None
 
@@ -434,3 +443,19 @@ async def close(self) -> None:
         if self._async_client is not None:
             await self._async_client.close()
             self._async_client = None
+
+    def close_sync(self) -> None:
+        """Close the underlying sync HTTP client, if one has been created.
+
+        ``openai.AzureOpenAI`` owns an httpx connection pool that leaks
+        across ``with`` blocks unless closed explicitly. Sync callers should
+        invoke this from their own ``close()`` to drain the pool.
+        """
+        if self._client is not None:
+            close = getattr(self._client, "close", None)
+            if callable(close):
+                try:
+                    close()
+                except Exception:
+                    pass
+            self._client = None

agent_memory_toolkit/cosmos_memory_client.py

@@ -112,9 +112,8 @@ def __init__(
     ) -> None:
         # Local store
         self.local_memory: list[dict[str, Any]] = []
-        # One-shot WARN guard for MEMORY_PROCESSOR_OWNER=durable skip path
-        # (see _maybe_auto_trigger).
         self._warned_owner_skip: bool = False
+        self._warned_counter_unreachable: bool = False
 
         # Store kwargs directly
         self._cosmos_endpoint = cosmos_endpoint
@@ -219,6 +218,16 @@ def close(self) -> None:
             self._container_client = None
             self._counter_container_client = None
             logger.info("Cosmos client closed")
+        # Drain LLM/embeddings httpx pools — openai.AzureOpenAI keeps them
+        # open across `with` blocks otherwise.
+        try:
+            self._chat_client.close_sync()
+        except Exception:
+            pass
+        try:
+            self._embeddings_client.close()
+        except Exception:
+            pass
         # Close credentials we created ourselves (sync DefaultAzureCredential
         # holds an underlying token cache + HTTP transport).
         for owns, cred in (
@@ -595,6 +604,9 @@ def _get_counter_container(self) -> Any:
         operator brought their own memory container without provisioning a
         counter container). Auto-trigger callers must tolerate ``None`` and
         skip the increment.
+
+        Logs a one-shot WARN on first failure so operators don't silently
+        lose all memory processing because of a missing counter container.
         """
         if self._counter_container_client is not None:
             return self._counter_container_client
@@ -605,12 +617,18 @@ def _get_counter_container(self) -> Any:
             self._counter_container_client = db.get_container_client(self._cosmos_counter_container)
             return self._counter_container_client
         except Exception as exc:  # pragma: no cover - defensive
-            logger.warning(
-                "Counter container %s/%s unreachable; auto-trigger disabled: %s",
-                self._cosmos_database,
-                self._cosmos_counter_container,
-                exc,
-            )
+            if not self._warned_counter_unreachable:
+                self._warned_counter_unreachable = True
+                logger.warning(
+                    "Counter container %s/%s unreachable (%s: %s); "
+                    "auto-trigger DISABLED for the lifetime of this client. "
+                    "Provision the container or set MEMORY_PROCESSOR_OWNER=durable "
+                    "if processing runs in the Function App.",
+                    self._cosmos_database,
+                    self._cosmos_counter_container,
+                    type(exc).__name__,
+                    exc,
+                )
             return None
 
     def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
@@ -720,6 +738,15 @@ def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None:
             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)
 
+            # Order matters: extract → dedup → summary.
+            # Each step gates on its OWN threshold (independent crossings),
+            # but when multiple thresholds cross in the same batch we still
+            # want the data flow to be: write fresh facts, deduplicate them,
+            # then fold the deduplicated set into the summary. Reordering
+            # would risk a summary that includes since-removed duplicates
+            # or omits just-extracted facts. DO NOT reorder without
+            # updating the equivalent block in aio/cosmos_memory_client.py
+            # and function_app/orchestrators/.
             if fire_extract:
                 try:
                     processor.process_extract_memories(user_id=user_id, thread_id=thread_id)

agent_memory_toolkit/embeddings.py

@@ -153,3 +153,18 @@ def generate_batch(self, texts: list[str]) -> list[list[float]]:
         # the caller receives embeddings in the same order as the input.
         sorted_data = sorted(response.data, key=lambda d: d.index)
         return [item.embedding for item in sorted_data]
+
+    def close(self) -> None:
+        """Close the underlying sync HTTP client, if one has been created.
+
+        ``openai.AzureOpenAI`` owns an httpx connection pool that leaks
+        across ``with`` blocks unless closed explicitly.
+        """
+        if self._client is not None:
+            close = getattr(self._client, "close", None)
+            if callable(close):
+                try:
+                    close()
+                except Exception:
+                    pass
+            self._client = None

function_app/host.json

@@ -4,7 +4,7 @@
     "applicationInsights": {
       "samplingSettings": {
         "isEnabled": true,
-        "excludedTypes": "Request"
+        "excludedTypes": "Request;Trace"
       }
     }
   },

function_app/shared/pipeline_factory.py

@@ -22,6 +22,7 @@ def get_pipeline():
 
     from azure.identity import DefaultAzureCredential
 
+    from agent_memory_toolkit._utils import _resolve_embedding_dimensions
     from agent_memory_toolkit.chat import ChatClient
     from agent_memory_toolkit.embeddings import EmbeddingsClient
     from agent_memory_toolkit.pipeline import ProcessingPipeline
@@ -30,6 +31,8 @@ def get_pipeline():
     container = get_memories_container()
     ai_endpoint = config.get_ai_foundry_endpoint()
 
+    embedding_dimensions = _resolve_embedding_dimensions(None)
+
     llm = ChatClient(
         endpoint=ai_endpoint,
         credential=credential,
@@ -39,6 +42,7 @@ def get_pipeline():
         endpoint=ai_endpoint,
         credential=credential,
         model=config.get_embedding_deployment_name(),
+        dimensions=embedding_dimensions,
     )
 
     _pipeline = ProcessingPipeline(container, llm, embeddings)

infra/modules/cosmos.bicep

@@ -184,6 +184,7 @@ resource counterContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/co
           '/thread_id'
         ]
       }
+      defaultTtl: 7776000
     }
   }
 }

infra/modules/functions.bicep

@@ -47,6 +47,9 @@ param aiFoundryEndpoint string
 @description('Embedding model deployment name.')
 param embeddingDeploymentName string = 'text-embedding-3-large'
 
+@description('Embedding output dimensions. MUST match the dimensions configured in the Cosmos memories container vectorEmbeddingPolicy (default 1536).')
+param embeddingDimensions int = 1536
+
 @description('LLM model deployment name.')
 param chatDeploymentName string = 'gpt-4o-mini'
 
@@ -251,6 +254,15 @@ resource functionApp 'Microsoft.Web/sites@2023-12-01' = {
           name: 'AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME'
           value: embeddingDeploymentName
         }
+        {
+          // Pins the embedding output dim. Without this, text-embedding-3-large
+          // returns its native 3072-dim vectors and Cosmos accepts them silently —
+          // but DiskANN (configured for 1536 in cosmos.bicep) cannot match
+          // them, so every FA-written memory becomes invisible to vector /
+          // hybrid search. Must equal cosmos.bicep vectorEmbeddingPolicy dimensions.
+          name: 'AI_FOUNDRY_EMBEDDING_DIMENSIONS'
+          value: string(embeddingDimensions)
+        }
         {
           name: 'AI_FOUNDRY_CHAT_DEPLOYMENT_NAME'
           value: chatDeploymentName