Resolving comments

Author: Aayush Kataria

azure/cosmos/agent_memory/aio/cosmos_memory_client.py

@@ -27,6 +27,7 @@
 from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore
 from azure.cosmos.agent_memory.exceptions import CosmosNotConnectedError, CosmosOperationError
 from azure.cosmos.agent_memory.logging import get_logger
+from azure.cosmos.agent_memory.services._pipeline_helpers import _normalize_metadata_keys
 from azure.cosmos.agent_memory.thresholds import DEFAULT_TTL_BY_TYPE
 
 if TYPE_CHECKING:  # pragma: no cover - typing-only import
@@ -136,9 +137,7 @@ def __init__(
         self._pipeline: Optional[AsyncPipelineService] = None
         self._processor: Optional[AsyncMemoryProcessor] = processor
         self._processor_explicit = processor is not None
-        self._transcript_metadata_keys: Optional[tuple[str, ...]] = (
-            tuple(transcript_metadata_keys) if transcript_metadata_keys else None
-        )
+        self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys)
         logger.info("AsyncCosmosMemoryClient initialized")
 
     async def __aenter__(self) -> "AsyncCosmosMemoryClient":

azure/cosmos/agent_memory/aio/services/pipeline.py

@@ -16,7 +16,7 @@
 import time
 from collections import defaultdict
 from datetime import datetime, timezone
-from typing import Any, Literal, Optional
+from typing import Any, Iterable, Literal, Optional
 
 from azure.cosmos.exceptions import (
     CosmosHttpResponseError,
@@ -48,6 +48,7 @@
 from azure.cosmos.agent_memory.services._pipeline_helpers import (
     VALID_VALENCES,
     PromptyLoader,
+    _normalize_metadata_keys,
     build_topic_tags,
     build_transcript,
     cap_structured_summary,
@@ -176,7 +177,7 @@ def __init__(
         prompts_dir: str | None = None,
         *,
         containers: dict[ContainerKey, Any],
-        transcript_metadata_keys: Optional[tuple[str, ...]] = None,
+        transcript_metadata_keys: Optional[Iterable[str]] = None,
     ) -> None:
         self._store = store
         self._containers = containers
@@ -187,9 +188,7 @@ def __init__(
         self._chat_client = chat_client
         self._embeddings = embeddings_client
         self._prompty = PromptyLoader(prompts_dir)
-        self._transcript_metadata_keys: Optional[tuple[str, ...]] = (
-            tuple(transcript_metadata_keys) if transcript_metadata_keys else None
-        )
+        self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys)
 
     async def _query_items(self, container: Any, **kwargs: Any) -> list[dict[str, Any]]:
         result = container.query_items(**kwargs)
@@ -271,6 +270,8 @@ def _build_transcript(
         *,
         group_by_thread: bool = False,
     ) -> str:
+        # getattr fallback covers unit tests that build AsyncPipelineService
+        # via __new__ to bypass __init__ (and therefore the metadata-keys stash).
         return build_transcript(
             items,
             group_by_thread=group_by_thread,

azure/cosmos/agent_memory/cosmos_memory_client.py

@@ -25,6 +25,7 @@
 from .embeddings import EmbeddingsClient
 from .exceptions import CosmosOperationError
 from .processors import InProcessProcessor, MemoryProcessor
+from .services._pipeline_helpers import _normalize_metadata_keys
 from .services.pipeline import PipelineService
 from .store import MemoryStore
 from .thresholds import DEFAULT_TTL_BY_TYPE
@@ -128,9 +129,7 @@ def __init__(
         self._pipeline: Optional[PipelineService] = None
         self._processor: Optional[MemoryProcessor] = processor
         self._processor_explicit = processor is not None
-        self._transcript_metadata_keys: Optional[tuple[str, ...]] = (
-            tuple(transcript_metadata_keys) if transcript_metadata_keys else None
-        )
+        self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys)
         if self._cosmos_endpoint:
             self.create_memory_store()
         logger.info("CosmosMemoryClient initialized")

azure/cosmos/agent_memory/services/_pipeline_helpers.py

@@ -181,13 +181,34 @@ def extract_prompty_params(p: Any) -> dict[str, Any]:
     return params
 
 
+def _normalize_metadata_keys(
+    value: Optional[Iterable[str]],
+) -> Optional[tuple[str, ...]]:
+    """Validate + coerce a ``transcript_metadata_keys`` argument to a tuple.
+
+    Rejects ``str`` outright (a bare string is iterable char-by-char, which
+    would silently produce a one-letter allow-list). Returns ``None`` for
+    empty or missing input.
+    """
+    if value is None:
+        return None
+    if isinstance(value, str):
+        raise TypeError(
+            "transcript_metadata_keys must be a sequence of keys "
+            "(list/tuple/set), not a single str. "
+            f"Got: {value!r}. Did you mean [{value!r}]?"
+        )
+    keys = tuple(str(k) for k in value if str(k))
+    return keys or None
+
+
 def _format_metadata_segment(
     metadata: Any,
-    metadata_keys: Optional[Iterable[str]],
+    metadata_keys: Optional[tuple[str, ...]],
 ) -> str:
     """Render the trailing ``[metadata: {...}]`` segment for a transcript line.
 
-    Returns an empty string unless ``metadata_keys`` is a non-empty iterable
+    Returns an empty string unless ``metadata_keys`` is a non-empty tuple
     AND at least one of those keys is present in ``metadata``. Only the
     explicitly allow-listed keys are serialized, in the iteration order of
     ``metadata_keys``.
@@ -197,7 +218,8 @@ def _format_metadata_segment(
     filtered = {k: metadata[k] for k in metadata_keys if k in metadata}
     if not filtered:
         return ""
-    return f" [metadata: {json.dumps(filtered)}]"
+    payload = json.dumps(filtered, separators=(",", ":"), ensure_ascii=False, default=str)
+    return f" [metadata: {payload}]"
 
 
 def build_transcript(
@@ -227,13 +249,18 @@ def build_transcript(
         metadata blobs (raw tool calls, IDE schema, etc.) out of every
         prompt — they're often 10-100x larger than the dialog itself and
         dilute extraction quality.
+
+        Accepts any iterable of strings except ``str`` itself (which would
+        be interpreted char-by-char). Generators are coerced to a tuple so
+        the allow-list is reusable across turns.
     """
+    keys = _normalize_metadata_keys(metadata_keys)
     if not group_by_thread:
         lines: list[str] = []
         for m in items:
             role = m.get("role", "unknown")
             content = m.get("content", "")
-            meta_str = _format_metadata_segment(m.get("metadata", {}), metadata_keys)
+            meta_str = _format_metadata_segment(m.get("metadata", {}), keys)
             lines.append(f"[{role}]: {content}{meta_str}")
         return "\n".join(lines)
 
@@ -247,7 +274,7 @@ def build_transcript(
         for m in thread_items:
             role = m.get("role", "unknown")
             content = m.get("content", "")
-            meta_str = _format_metadata_segment(m.get("metadata", {}), metadata_keys)
+            meta_str = _format_metadata_segment(m.get("metadata", {}), keys)
             parts.append(f"[{role}]: {content}{meta_str}")
         parts.append("")
     return "\n".join(parts)

azure/cosmos/agent_memory/services/pipeline.py

@@ -14,7 +14,7 @@
 import time
 from collections import defaultdict
 from datetime import datetime, timezone
-from typing import Any, Literal, Optional
+from typing import Any, Iterable, Literal, Optional
 
 from azure.cosmos.exceptions import (
     CosmosHttpResponseError,
@@ -46,6 +46,7 @@
 from azure.cosmos.agent_memory.services._pipeline_helpers import (
     VALID_VALENCES,
     PromptyLoader,
+    _normalize_metadata_keys,
     build_topic_tags,
     build_transcript,
     cap_structured_summary,
@@ -161,7 +162,7 @@ def __init__(
         prompts_dir: str | None = None,
         *,
         containers: dict[ContainerKey, Any],
-        transcript_metadata_keys: Optional[tuple[str, ...]] = None,
+        transcript_metadata_keys: Optional[Iterable[str]] = None,
     ) -> None:
         self._store = store
         self._containers = containers
@@ -172,9 +173,7 @@ def __init__(
         self._chat_client = chat_client
         self._embeddings = embeddings_client
         self._prompty = PromptyLoader(prompts_dir)
-        self._transcript_metadata_keys: Optional[tuple[str, ...]] = (
-            tuple(transcript_metadata_keys) if transcript_metadata_keys else None
-        )
+        self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys)
 
     def _run_prompty(
         self,
@@ -224,6 +223,8 @@ def _build_transcript(
         *,
         group_by_thread: bool = False,
     ) -> str:
+        # getattr fallback covers unit tests that build PipelineService via
+        # __new__ to bypass __init__ (and therefore the metadata-keys stash).
         return build_transcript(
             items,
             group_by_thread=group_by_thread,

function_app/orchestrators/extract_memories.py

@@ -56,9 +56,7 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext):
         )
         if config.get_procedural_synthesis_auto():
             count = payload.get("count")
-            instance_id = (
-                f"procedural:{user_id}:{thread_id}:{count}" if count is not None else None
-            )
+            instance_id = f"procedural:{user_id}:{thread_id}:{count}" if count is not None else None
             try:
                 procedural = yield context.call_sub_orchestrator_with_retry(
                     "SynthesizeProceduralOrchestrator",

function_app/shared/pipeline_factory.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import os
 from typing import Any
 
 from . import config
@@ -18,6 +19,20 @@
 _pipeline: Any | None = None
 
 
+def _read_transcript_metadata_keys() -> tuple[str, ...] | None:
+    """Parse ``AGENT_MEMORY_TRANSCRIPT_METADATA_KEYS`` (comma-separated allow-list).
+
+    Mirrors the ``transcript_metadata_keys`` ctor kwarg on
+    :class:`CosmosMemoryClient` so the Durable-Functions backend produces
+    the same prompt content as the in-process backend.
+    """
+    raw = os.environ.get("AGENT_MEMORY_TRANSCRIPT_METADATA_KEYS", "").strip()
+    if not raw:
+        return None
+    keys = tuple(part.strip() for part in raw.split(",") if part.strip())
+    return keys or None
+
+
 def get_pipeline():
     """Return the cached :class:`PipelineService` for this worker."""
     global _pipeline
@@ -59,5 +74,11 @@ def get_pipeline():
         ContainerKey.SUMMARIES: summaries_container,
     }
     store = MemoryStore(containers=containers, embeddings_client=embeddings)
-    _pipeline = PipelineService(store, chat, embeddings, containers=containers)
+    _pipeline = PipelineService(
+        store,
+        chat,
+        embeddings,
+        containers=containers,
+        transcript_metadata_keys=_read_transcript_metadata_keys(),
+    )
     return _pipeline

tests/unit/function_app/test_pipeline_factory.py

@@ -107,6 +107,7 @@ def test_builds_pipeline_from_complete_env(mocks):
         mocks.chat_instance,
         mocks.embed_instance,
         containers=expected_containers,
+        transcript_metadata_keys=None,
     )
 
 
@@ -144,6 +145,31 @@ def test_passes_credential_to_both_clients(mocks):
     assert mocks.embed_ctor.call_args.kwargs["credential"] is mocks.credential
 
 
+def test_transcript_metadata_keys_env_threads_into_pipeline(monkeypatch, mocks):
+    monkeypatch.setenv(
+        "AGENT_MEMORY_TRANSCRIPT_METADATA_KEYS",
+        " agent_id , timestamp ,  , model_id",
+    )
+    pipeline_factory.get_pipeline()
+
+    kwargs = mocks.pipeline_ctor.call_args.kwargs
+    assert kwargs["transcript_metadata_keys"] == ("agent_id", "timestamp", "model_id")
+
+
+def test_transcript_metadata_keys_env_unset_yields_none(monkeypatch, mocks):
+    monkeypatch.delenv("AGENT_MEMORY_TRANSCRIPT_METADATA_KEYS", raising=False)
+    pipeline_factory.get_pipeline()
+
+    assert mocks.pipeline_ctor.call_args.kwargs["transcript_metadata_keys"] is None
+
+
+def test_transcript_metadata_keys_env_empty_string_yields_none(monkeypatch, mocks):
+    monkeypatch.setenv("AGENT_MEMORY_TRANSCRIPT_METADATA_KEYS", "   ")
+    pipeline_factory.get_pipeline()
+
+    assert mocks.pipeline_ctor.call_args.kwargs["transcript_metadata_keys"] is None
+
+
 # ---------------------------------------------------------------------------
 # Caching / idempotence
 # ---------------------------------------------------------------------------