resolving comments

Author: Aayush Kataria

agent_memory_toolkit/_utils.py

@@ -378,6 +378,18 @@ def _container_policies(
         ],
         "vectorIndexes": [{"path": "/embedding", "type": "diskANN"}],
         "fullTextIndexes": [{"path": "/content"}],
+        # Procedural synthesis selects TOP N by (salience DESC, created_at ASC, id ASC).
+        # Cosmos requires a composite index for multi-property ORDER BY; without it the
+        # query returns a non-deterministic 50 of N when many docs share the default
+        # salience (0.5), which makes the source-id short-circuit in synthesize_procedural
+        # thrash and burn LLM calls on every reconcile.
+        "compositeIndexes": [
+            [
+                {"path": "/salience", "order": "descending"},
+                {"path": "/created_at", "order": "ascending"},
+                {"path": "/id", "order": "ascending"},
+            ]
+        ],
     }
 
     full_text_policy = {

agent_memory_toolkit/aio/auto_trigger.py

@@ -71,6 +71,7 @@ async def maybe_trigger_steps(
         n_facts=n_facts,
         n_summary=n_summary,
         n_dedup_turns=n_dedup_turns,
+        thresholds=thresholds,
     )
     await _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user)
 
@@ -83,6 +84,7 @@ async def _trigger_thread_steps(
     n_facts: int,
     n_summary: int,
     n_dedup_turns: int,
+    thresholds: Any = None,
 ) -> dict[str, int]:
     user_batch_counts: dict[str, int] = {}
     for (user_id, thread_id), batch_count in turn_counts.items():
@@ -111,6 +113,7 @@ async def _trigger_thread_steps(
             fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts),
             fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary),
             fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns),
+            thresholds=thresholds,
         )
     return user_batch_counts
 
@@ -125,10 +128,11 @@ async def _fire_thread_steps(
     fire_extract: bool,
     fire_summary: bool,
     fire_dedup: bool,
+    thresholds: Any = None,
 ) -> None:
     fire_procedural = fire_dedup and bool(
         _threshold_value(
-            default_thresholds,
+            thresholds,
             "get_procedural_synthesis_auto",
             "PROCEDURAL_SYNTHESIS_AUTO",
             default=True,

agent_memory_toolkit/aio/chat.py

@@ -136,7 +136,6 @@ def _build_kwargs(
         self,
         messages: list[dict[str, str]],
         *,
-        temperature: float | None = None,
         response_format: dict | None = None,
         **extra: Any,
     ) -> dict[str, Any]:
@@ -156,7 +155,6 @@ async def generate(
         self,
         messages: list[dict[str, str]],
         *,
-        temperature: float | None = None,
         response_format: dict | None = None,
         max_retries: int = 3,
         base_delay: float = 2.0,
@@ -185,7 +183,6 @@ async def generate(
         client = self._ensure_client()
         kwargs = self._build_kwargs(
             messages,
-            temperature=temperature,
             response_format=response_format,
             **extra,
         )

agent_memory_toolkit/aio/services/pipeline.py

@@ -18,7 +18,7 @@
 from datetime import datetime, timezone
 from typing import Any, Literal, Optional
 
-from azure.cosmos.exceptions import CosmosResourceExistsError
+from azure.cosmos.exceptions import CosmosResourceExistsError, CosmosResourceNotFoundError
 
 from agent_memory_toolkit._utils import DEFAULT_TTL_BY_TYPE, compute_content_hash
 from agent_memory_toolkit.aio.store import AsyncMemoryStore
@@ -672,7 +672,7 @@ async def synthesize_procedural(
                 "AND ((IS_DEFINED(c.metadata.category) "
                 "AND c.metadata.category IN ('preference', 'requirement')) "
                 "OR (IS_DEFINED(c.salience) AND c.salience >= @min_salience)) "
-                "ORDER BY c.salience DESC"
+                "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC"
             ),
             parameters=[
                 {"name": "@uid", "value": user_id},
@@ -681,14 +681,6 @@ async def synthesize_procedural(
             ],
             enable_cross_partition_query=True,
         )
-        behavioral_fact_docs.sort(
-            key=lambda doc: (
-                float(doc.get("salience") or 0.0),
-                doc.get("created_at", ""),
-                doc.get("id", ""),
-            ),
-            reverse=True,
-        )
         behavioral_fact_docs = [
             doc
             for doc in behavioral_fact_docs
@@ -703,22 +695,14 @@ async def synthesize_procedural(
                 f"AND {active_filter} "
                 "AND IS_DEFINED(c.metadata.lesson) "
                 "AND c.metadata.lesson != null "
-                "ORDER BY c.salience DESC"
+                "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC"
             ),
             parameters=[
                 {"name": "@uid", "value": user_id},
                 {"name": "@type", "value": "episodic"},
             ],
             enable_cross_partition_query=True,
         )
-        episodic_docs.sort(
-            key=lambda doc: (
-                float(doc.get("salience") or 0.0),
-                doc.get("created_at", ""),
-                doc.get("id", ""),
-            ),
-            reverse=True,
-        )
         episodic_with_lessons = [
             doc
             for doc in episodic_docs
@@ -742,6 +726,13 @@ async def synthesize_procedural(
             )
             return {"status": "unchanged", "procedural": prior_doc}
 
+        if not current_source_ids:
+            logger.info(
+                "synthesize_procedural skipping LLM user_id=%s — no behavioral facts or episodic lessons",
+                user_id,
+            )
+            return {"status": "unchanged", "procedural": prior_doc}
+
         name_docs = await self._container.query_items(
             query=(
                 "SELECT TOP 1 * FROM c WHERE c.user_id = @uid "
@@ -854,7 +845,7 @@ async def generate_thread_summary_dry(
         existing_summary: Optional[dict[str, Any]] = None
         try:
             existing_summary = await self._container.read_item(item=summary_id, partition_key=[user_id, thread_id])
-        except Exception:
+        except CosmosResourceNotFoundError:
             pass
 
         query = "SELECT * FROM c WHERE c.user_id = @user_id AND c.thread_id = @thread_id AND c.type != 'summary'"
@@ -995,7 +986,7 @@ async def generate_user_summary_dry(
                 item=user_summary_id,
                 partition_key=[user_id, "__user_summary__"],
             )
-        except Exception:
+        except CosmosResourceNotFoundError:
             pass
 
         query = "SELECT * FROM c WHERE c.user_id = @user_id AND c.type != 'user_summary'"

agent_memory_toolkit/aio/store/memory_store.py

@@ -339,14 +339,22 @@ async def get_memories(
             query = f"SELECT * FROM c{where}"
 
         logger.debug("async get_memories query: %s", query)
-        results = await self._query_items(
-            query=query,
-            parameters=parameters or None,
-            operation="async get_memories query",
-            container=self._container_for_query(memory_types),
-        )
+        # Iterate over all containers that may hold the requested memory types
+        # so mixed turn + non-turn queries do not silently drop one container.
+        results: list[dict] = []
+        for container in self._containers_for_query(memory_types):
+            results.extend(
+                await self._query_items(
+                    query=query,
+                    parameters=parameters or None,
+                    operation="async get_memories query",
+                    container=container,
+                )
+            )
 
         if recent_k is not None:
+            results.sort(key=lambda i: i.get("_ts") or 0, reverse=True)
+            results = results[:recent_k]
             results.reverse()
         if min_salience is not None:
             results = [i for i in results if (i.get("salience") or 0.0) >= min_salience]
@@ -501,13 +509,16 @@ async def list_tags(
         thread_id: Optional[str] = None,
         prefix: Optional[str] = None,
         include_sys: bool = False,
+        include_superseded: bool = False,
     ) -> list[str]:
         """Return sorted distinct tags for a user, optionally scoped to one thread."""
         query = "SELECT VALUE c.tags FROM c WHERE c.user_id = @user_id AND ARRAY_LENGTH(c.tags) > 0"
         parameters = [{"name": "@user_id", "value": user_id}]
         if thread_id is not None:
             query += " AND c.thread_id = @thread_id"
             parameters.append({"name": "@thread_id", "value": thread_id})
+        if not include_superseded:
+            query += " AND (NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))"
 
         prefix_norm = prefix.strip().lower() if prefix else None
         partition_key, cross_partition = query_scope(user_id, thread_id)

agent_memory_toolkit/auto_trigger.py

@@ -115,6 +115,7 @@ def maybe_trigger_steps(
         n_facts=n_facts,
         n_summary=n_summary,
         n_dedup_turns=n_dedup_turns,
+        thresholds=thresholds,
     )
     _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user)
 
@@ -127,6 +128,7 @@ def _trigger_thread_steps(
     n_facts: int,
     n_summary: int,
     n_dedup_turns: int,
+    thresholds: Any = None,
 ) -> dict[str, int]:
     user_batch_counts: dict[str, int] = {}
     for (user_id, thread_id), batch_count in turn_counts.items():
@@ -155,6 +157,7 @@ def _trigger_thread_steps(
             fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts),
             fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary),
             fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns),
+            thresholds=thresholds,
         )
     return user_batch_counts
 
@@ -169,10 +172,11 @@ def _fire_thread_steps(
     fire_extract: bool,
     fire_summary: bool,
     fire_dedup: bool,
+    thresholds: Any = None,
 ) -> None:
     fire_procedural = fire_dedup and bool(
         _threshold_value(
-            default_thresholds,
+            thresholds,
             "get_procedural_synthesis_auto",
             "PROCEDURAL_SYNTHESIS_AUTO",
             default=True,

agent_memory_toolkit/chat.py

@@ -144,7 +144,6 @@ def _build_kwargs(
         self,
         messages: list[dict[str, str]],
         *,
-        temperature: float | None = None,
         response_format: dict | None = None,
         **extra: Any,
     ) -> dict[str, Any]:
@@ -170,7 +169,6 @@ def generate(
         self,
         messages: list[dict[str, str]],
         *,
-        temperature: float | None = None,
         response_format: dict | None = None,
         max_retries: int = 3,
         base_delay: float = 2.0,
@@ -201,7 +199,6 @@ def generate(
         client = self._ensure_client()
         kwargs = self._build_kwargs(
             messages,
-            temperature=temperature,
             response_format=response_format,
             **extra,
         )

agent_memory_toolkit/prompts/dedup.prompty

@@ -100,9 +100,9 @@ N. ID: <uuid> | Content: "<text>" | Confidence: 0.85 | Salience: 0.7 | Created:
 
 Some facts may show `Confidence: N/A`, `Salience: N/A`, or `Created: N/A`.
 
-- Treat `N/A` confidence/salience as **unknown** — omit those fields from the
-  `duplicate_groups` output; the pipeline will fall back to
-  `max(source.confidence)` / `max(source.salience)`.
+- Treat `N/A` confidence/salience as **unknown** — set those fields to `null`
+  in the `duplicate_groups` output (do NOT omit the keys); the pipeline will
+  fall back to `max(source.confidence)` / `max(source.salience)`.
 - Treat `N/A` created_at as **older than any real timestamp** for
   winner-selection purposes (a fact with a real date beats one with N/A).
 
@@ -121,9 +121,9 @@ You must output ONLY valid JSON matching this exact schema. No preamble, no expl
   "kept_ids": ["<id>", "<id>"]
 }
 ```
-<!-- confidence/salience: float 0..1, or omit the field if unknown -->
+<!-- confidence/salience: float 0..1, or null if you cannot infer -->
 
-The `confidence` and `salience` values **must** be real numbers between 0 and 1 in the actual output — the `0.0` above is a structural placeholder, not a literal value to echo. Compute them as the maximum across source facts in the group. If you cannot determine confidence or salience, omit the field entirely (the pipeline will fall back to `max(source.*)` from the source records).
+The `confidence` and `salience` values **must** be real numbers between 0 and 1 in the actual output — the `0.0` above is a structural placeholder, not a literal value to echo. Compute them as the maximum across source facts in the group. If you cannot determine confidence or salience, set the field to `null` (the pipeline will fall back to `max(source.*)` from the source records). **Do not omit the keys** — the response schema requires both fields to appear on every group, with either a number or `null` as the value.
 
 If a bucket is empty, emit it as an empty array (`[]`) rather than omitting the key.
 

agent_memory_toolkit/services/pipeline.py

@@ -16,7 +16,7 @@
 from datetime import datetime, timezone
 from typing import Any, Literal, Optional
 
-from azure.cosmos.exceptions import CosmosResourceExistsError
+from azure.cosmos.exceptions import CosmosResourceExistsError, CosmosResourceNotFoundError
 
 from agent_memory_toolkit._utils import DEFAULT_TTL_BY_TYPE, compute_content_hash
 from agent_memory_toolkit.exceptions import (
@@ -716,7 +716,7 @@ def synthesize_procedural(
                     "AND ((IS_DEFINED(c.metadata.category) "
                     "AND c.metadata.category IN ('preference', 'requirement')) "
                     "OR (IS_DEFINED(c.salience) AND c.salience >= @min_salience)) "
-                    "ORDER BY c.salience DESC"
+                    "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC"
                 ),
                 parameters=[
                     {"name": "@uid", "value": user_id},
@@ -726,14 +726,6 @@ def synthesize_procedural(
                 enable_cross_partition_query=True,
             )
         )
-        behavioral_fact_docs.sort(
-            key=lambda doc: (
-                float(doc.get("salience") or 0.0),
-                doc.get("created_at", ""),
-                doc.get("id", ""),
-            ),
-            reverse=True,
-        )
         behavioral_fact_docs = [
             doc
             for doc in behavioral_fact_docs
@@ -749,7 +741,7 @@ def synthesize_procedural(
                     f"AND {active_filter} "
                     "AND IS_DEFINED(c.metadata.lesson) "
                     "AND c.metadata.lesson != null "
-                    "ORDER BY c.salience DESC"
+                    "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC"
                 ),
                 parameters=[
                     {"name": "@uid", "value": user_id},
@@ -758,14 +750,6 @@ def synthesize_procedural(
                 enable_cross_partition_query=True,
             )
         )
-        episodic_docs.sort(
-            key=lambda doc: (
-                float(doc.get("salience") or 0.0),
-                doc.get("created_at", ""),
-                doc.get("id", ""),
-            ),
-            reverse=True,
-        )
         episodic_with_lessons = [
             doc
             for doc in episodic_docs
@@ -789,6 +773,13 @@ def synthesize_procedural(
             )
             return {"status": "unchanged", "procedural": prior_doc}
 
+        if not current_source_ids:
+            logger.info(
+                "synthesize_procedural skipping LLM user_id=%s — no behavioral facts or episodic lessons",
+                user_id,
+            )
+            return {"status": "unchanged", "procedural": prior_doc}
+
         name_docs = list(
             self._container.query_items(
                 query=(
@@ -906,7 +897,7 @@ def generate_thread_summary_dry(
         existing_summary: Optional[dict[str, Any]] = None
         try:
             existing_summary = self._container.read_item(item=summary_id, partition_key=[user_id, thread_id])
-        except Exception:
+        except CosmosResourceNotFoundError:
             pass
 
         query = "SELECT * FROM c WHERE c.user_id = @user_id AND c.thread_id = @thread_id AND c.type != 'summary'"
@@ -1052,7 +1043,7 @@ def generate_user_summary_dry(
                 item=user_summary_id,
                 partition_key=[user_id, "__user_summary__"],
             )
-        except Exception:
+        except CosmosResourceNotFoundError:
             pass
 
         query = "SELECT * FROM c WHERE c.user_id = @user_id AND c.type != 'user_summary'"