Fixing builds

Author: Aayush Kataria

agent_memory_toolkit/_utils.py

@@ -96,17 +96,36 @@ def _make_memory(
     return memory
 
 
-def _resolve_embedding_dimensions(val: Optional[int]) -> Optional[int]:
+def _resolve_embedding_dimensions(val: Optional[int]) -> int:
     """Resolve embedding dimensions from explicit value or ``AI_FOUNDRY_EMBEDDING_DIMENSIONS`` env var.
 
     Defaults to 1536 (the dimension we ship with for ``text-embedding-3-large``
     truncated to 1536, which is the size DiskANN is tuned for in our containers).
+
+    Raises :class:`ConfigurationError` if the env var is set but cannot be
+    parsed as a positive integer.
     """
     if val is not None:
         return val
-    raw = os.environ.get("AI_FOUNDRY_EMBEDDING_DIMENSIONS", "1536") or "1536"
-    parsed = int(raw)
-    return parsed if parsed else 1536
+    raw = os.environ.get("AI_FOUNDRY_EMBEDDING_DIMENSIONS")
+    if raw is None or raw == "":
+        return 1536
+    try:
+        parsed = int(raw)
+    except (ValueError, TypeError) as exc:
+        raise ConfigurationError(
+            message=(
+                f"Invalid configuration for embedding_dimensions: AI_FOUNDRY_EMBEDDING_DIMENSIONS"
+                f" must be a positive integer, got {raw!r}"
+            ),
+            parameter="embedding_dimensions",
+        ) from exc
+    if parsed <= 0:
+        raise ConfigurationError(
+            message=(f"Invalid configuration for embedding_dimensions: must be a positive integer, got {parsed}"),
+            parameter="embedding_dimensions",
+        )
+    return parsed
 
 
 def _resolve_cosmos_throughput_mode(val: Optional[str]) -> str:

agent_memory_toolkit/aio/processors/durable.py

@@ -33,7 +33,7 @@ async def process_thread(
             thread_id,
             len(turns) if turns else 0,
         )
-        return ProcessThreadResult(thread_summary=None, extracted=[], deduplicated_count=0, elapsed_ms=0)
+        return ProcessThreadResult(thread_summary=None, extracted_counts={}, deduplicated_count=0, elapsed_ms=0)
 
     async def generate_user_summary(
         self,

agent_memory_toolkit/aio/processors/inprocess.py

@@ -75,18 +75,14 @@ async def process_thread(
                     deduped_count = dedup[key]
                     break
 
-        extracted_list: list[dict[str, Any]]
-        if isinstance(extracted, list):
-            extracted_list = extracted
-        elif isinstance(extracted, dict):
-            extracted_list = [extracted]
-        else:
-            extracted_list = []
+        extracted_counts: dict[str, int] = (
+            {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {}
+        )
 
         elapsed_ms = int((time.monotonic() - start) * 1000)
         return ProcessThreadResult(
             thread_summary=thread_summary if isinstance(thread_summary, dict) else None,
-            extracted=extracted_list,
+            extracted_counts=extracted_counts,
             deduplicated_count=deduped_count,
             elapsed_ms=elapsed_ms,
         )

agent_memory_toolkit/processors/base.py

@@ -16,10 +16,18 @@
 
 @dataclass
 class ProcessThreadResult:
-    """Outcome of a single ``process_thread`` invocation."""
+    """Outcome of a single ``process_thread`` invocation.
+
+    ``extracted_counts`` carries the per-type extraction telemetry returned by
+    :meth:`ProcessingPipeline.extract_memories` (e.g.
+    ``{"facts_count": 2, "procedural_count": 1, "episodic_count": 0, "updated_count": 0}``).
+    The actual extracted memory documents are persisted to Cosmos DB by the
+    pipeline; query them back via the SDK's ``get_memories`` if you need the
+    raw docs.
+    """
 
     thread_summary: Optional[dict[str, Any]] = None
-    extracted: list[dict[str, Any]] = field(default_factory=list)
+    extracted_counts: dict[str, int] = field(default_factory=dict)
     deduplicated_count: int = 0
     elapsed_ms: int = 0
 

agent_memory_toolkit/processors/durable.py

@@ -38,7 +38,7 @@ def process_thread(
             thread_id,
             len(turns) if turns else 0,
         )
-        return ProcessThreadResult(thread_summary=None, extracted=[], deduplicated_count=0, elapsed_ms=0)
+        return ProcessThreadResult(thread_summary=None, extracted_counts={}, deduplicated_count=0, elapsed_ms=0)
 
     def generate_user_summary(
         self,

agent_memory_toolkit/processors/inprocess.py

@@ -71,18 +71,14 @@ def process_thread(
                     deduped_count = dedup[key]
                     break
 
-        extracted_list: list[dict[str, Any]]
-        if isinstance(extracted, list):
-            extracted_list = extracted
-        elif isinstance(extracted, dict):
-            extracted_list = [extracted]
-        else:
-            extracted_list = []
+        extracted_counts: dict[str, int] = (
+            {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {}
+        )
 
         elapsed_ms = int((time.monotonic() - start) * 1000)
         return ProcessThreadResult(
             thread_summary=thread_summary if isinstance(thread_summary, dict) else None,
-            extracted=extracted_list,
+            extracted_counts=extracted_counts,
             deduplicated_count=deduped_count,
             elapsed_ms=elapsed_ms,
         )

function_app/orchestrators/extract_memories.py

@@ -64,9 +64,11 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext):
 def em_ExtractMemories(payload: dict) -> dict:
     """Run the LLM extraction step.
 
-    Returns the per-type counts produced by ``pipeline.extract_memories``
-    (``{"facts": N, "procedural": N, "episodic": N}``-shaped). Salience-based
-    filtering is delegated to the pipeline since it owns the schema.
+    Returns the per-type counts produced by ``pipeline.extract_memories``,
+    shaped like
+    ``{"facts_count": N, "procedural_count": N, "episodic_count": N, "updated_count": N}``.
+    Salience-based filtering is delegated to the pipeline since it owns the
+    schema.
 
     The pipeline loads recent turns internally, so we do NOT pre-load them in
     a separate activity (which would duplicate the query, waste RUs, and open

function_app/shared/config.py

@@ -49,21 +49,32 @@ def _parse_threshold(name: str, default: int) -> int:
     """Parse an integer threshold env var.
 
     Returns ``default`` when the env var is unset, empty, or invalid (with a
-    warning logged on invalid input). Use this for the ``*_EVERY_N`` knobs
-    where ``"0"`` is a valid explicit-disable value but a missing setting
-    should fall back to the documented default rather than silently disabling
-    the orchestrator.
+    warning logged on invalid input). Negative values are also rejected and
+    fall back to ``default`` so misconfiguration is explicit instead of
+    silently disabling the orchestrator. Use this for the ``*_EVERY_N`` knobs
+    where ``"0"`` is a valid explicit-disable value but anything else
+    nonsensical should fall back.
     """
     raw = os.environ.get(name)
     if raw is None or raw == "":
         return default
     try:
-        return int(raw)
+        parsed = int(raw)
     except (ValueError, TypeError):
         logger.warning(
             "Invalid value for %s=%r, using default %d", name, raw, default,
         )
         return default
+    if parsed < 0:
+        logger.warning(
+            "Negative value for %s=%r is not allowed; using default %d "
+            "(set to 0 to explicitly disable)",
+            name,
+            raw,
+            default,
+        )
+        return default
+    return parsed
 
 
 def _parse_int(name: str, default: int) -> int:

tests/integration/test_processor_integration.py

@@ -38,9 +38,12 @@ def test_flush_invokes_pipeline_with_correct_args(self):
             "type": "summary",
             "content": "Conversation about Paris.",
         }
-        pipeline.extract_memories.return_value = [
-            {"id": "fact-1", "type": "fact", "content": "User likes Paris."},
-        ]
+        pipeline.extract_memories.return_value = {
+            "facts_count": 1,
+            "procedural_count": 0,
+            "episodic_count": 0,
+            "updated_count": 0,
+        }
         pipeline.deduplicate_facts.return_value = {"deduplicated": 0}
 
         processor = InProcessProcessor(pipeline=pipeline)
@@ -60,9 +63,12 @@ def test_flush_invokes_pipeline_with_correct_args(self):
             "type": "summary",
             "content": "Conversation about Paris.",
         }
-        assert result.extracted == [
-            {"id": "fact-1", "type": "fact", "content": "User likes Paris."},
-        ]
+        assert result.extracted_counts == {
+            "facts_count": 1,
+            "procedural_count": 0,
+            "episodic_count": 0,
+            "updated_count": 0,
+        }
         client.get_thread.assert_called_once_with(thread_id="thread-paris", user_id="u-paris", memory_type="turn")
         pipeline.generate_thread_summary.assert_called_once_with("u-paris", "thread-paris")
         pipeline.extract_memories.assert_called_once_with("u-paris", "thread-paris")
@@ -90,7 +96,7 @@ def test_flush_is_a_noop(self, caplog):
 
         assert isinstance(result, ProcessThreadResult)
         assert result.thread_summary is None
-        assert result.extracted == []
+        assert result.extracted_counts == {}
         assert result.deduplicated_count == 0
 
         pipeline.generate_thread_summary.assert_not_called()

tests/integration/test_processor_integration_async.py

@@ -34,9 +34,12 @@ async def test_flush_invokes_pipeline_with_correct_args(self):
             "type": "summary",
             "content": "Conversation about Paris.",
         }
-        pipeline.extract_memories.return_value = [
-            {"id": "fact-1", "type": "fact", "content": "User likes Paris."},
-        ]
+        pipeline.extract_memories.return_value = {
+            "facts_count": 1,
+            "procedural_count": 0,
+            "episodic_count": 0,
+            "updated_count": 0,
+        }
         pipeline.deduplicate_facts.return_value = {"deduplicated": 0}
 
         processor = AsyncInProcessProcessor(pipeline=pipeline)
@@ -56,9 +59,12 @@ async def test_flush_invokes_pipeline_with_correct_args(self):
             "type": "summary",
             "content": "Conversation about Paris.",
         }
-        assert result.extracted == [
-            {"id": "fact-1", "type": "fact", "content": "User likes Paris."},
-        ]
+        assert result.extracted_counts == {
+            "facts_count": 1,
+            "procedural_count": 0,
+            "episodic_count": 0,
+            "updated_count": 0,
+        }
         client.get_thread.assert_awaited_once_with(thread_id="thread-paris", user_id="u-paris", memory_type="turn")
         pipeline.generate_thread_summary.assert_called_once_with("u-paris", "thread-paris")
         pipeline.extract_memories.assert_called_once_with("u-paris", "thread-paris")
@@ -85,7 +91,7 @@ async def test_flush_is_a_noop(self, caplog):
 
         assert isinstance(result, ProcessThreadResult)
         assert result.thread_summary is None
-        assert result.extracted == []
+        assert result.extracted_counts == {}
         assert result.deduplicated_count == 0
 
         pipeline.generate_thread_summary.assert_not_called()