Resolving comments

Author: Aayush Kataria

.env.template

@@ -1,9 +1,22 @@
 # This is a template for the .env file. Copy this file to .env and fill in the values for your accounts.
+
+# ---- Cosmos DB ----
 COSMOS_DB_ENDPOINT=https://<your-account>.documents.azure.com:443/
+# COSMOS_DB__accountEndpoint is required for the Azure Functions change feed trigger
+# (identity-based connection). Set it to the same value as COSMOS_DB_ENDPOINT.
+COSMOS_DB__accountEndpoint=https://<your-account>.documents.azure.com:443/
 COSMOS_DB_DATABASE=ai_memory
 COSMOS_DB_CONTAINER=memories
+COSMOS_DB_COUNTERS_CONTAINER=counter
+COSMOS_DB_LEASE_CONTAINER=leases
 COSMOS_DB_AUTOSCALE_MAX_RU=1000
 
+# ---- Change Feed Thresholds (set to 0 to disable) ----
+THREAD_SUMMARY_EVERY_N=0
+FACT_EXTRACTION_EVERY_N=0
+USER_SUMMARY_EVERY_N=0
+
+# ---- AI Foundry / Azure OpenAI ----
 AI_FOUNDRY_ENDPOINT=https://<your-account>.openai.azure.com/
 AI_FOUNDRY_API_KEY=
 EMBEDDING_MODEL=text-embedding-3-large
@@ -14,5 +27,6 @@ FULL_TEXT_LANGUAGE=en-US
 
 LLM_MODEL=<your-model-deployment>
 
+# ---- Azure Durable Functions ----
 ADF_ENDPOINT=http://localhost:7071/api
 ADF_KEY=

Samples/Demo.ipynb

@@ -36,7 +36,6 @@ class MemoryType(str, Enum):
     summary = "summary"
     fact = "fact"
     user_summary = "user_summary"
-    counter = "counter"
 
 
 # ---------------------------------------------------------------------------

agent_memory_toolkit/models.py

@@ -126,7 +126,13 @@ async def _get_cosmos_database_client():
     """Return the Cosmos DB database client, connecting on first call."""
     global _cosmos_client, _cosmos_database_client
     if _cosmos_database_client is None:
-        endpoint = os.environ["COSMOS_DB__accountEndpoint"]
+        endpoint = os.environ.get("COSMOS_DB__accountEndpoint") or os.environ.get("COSMOS_DB_ENDPOINT")
+        if not endpoint:
+            raise ValueError(
+                "Cosmos DB endpoint not configured. "
+                "Set COSMOS_DB__accountEndpoint (required for change feed trigger) "
+                "or COSMOS_DB_ENDPOINT."
+            )
         database = os.environ["COSMOS_DB_DATABASE"]
         logger.info(
             "Connecting to Cosmos DB endpoint=%s database=%s",

azure_functions/activities.py

@@ -49,6 +49,8 @@
 Required Cosmos DB containers:
 
   - ``memories`` – existing container for memory documents
+  - ``counter``  – dedicated container for change feed counter documents
+                   (configurable via ``COSMOS_DB_COUNTERS_CONTAINER``)
   - ``leases``   – auto-created by the trigger for change feed checkpointing
 """
 
@@ -59,6 +61,7 @@
 
 import azure.functions as func
 import azure.durable_functions as df
+from azure.core import MatchConditions
 from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceNotFoundError
 
 from activities import _get_cosmos_counter_container, bp as activities_bp
@@ -83,11 +86,49 @@ def _utc_now_iso() -> str:
     return datetime.now(timezone.utc).isoformat()
 
 
-async def increment_counter_by(counter_id: str, user_id: str, thread_id: str, count: int) -> tuple[int, int]:
+def _parse_threshold(name: str) -> int:
+    """Parse an integer threshold from an environment variable.
+
+    Returns 0 (disabled) if the variable is missing, empty, or not a valid
+    integer, and logs a warning so misconfigurations are visible.
+    """
+    raw = os.environ.get(name, "0")
+    try:
+        return int(raw)
+    except (ValueError, TypeError):
+        logger.warning(
+            "Invalid value for %s=%r, defaulting to 0 (disabled)", name, raw,
+        )
+        return 0
+
+
+async def increment_counter_by(
+    counter_id: str,
+    user_id: str,
+    thread_id: str,
+    count: int,
+    *,
+    batch_max_lsn: int | None = None,
+) -> tuple[int, int]:
     """Atomically increment a counter document by *count* using ETag concurrency.
 
     Returns ``(old_count, new_count)``.  Creates the counter document if it
     does not exist.  Retries up to 3 times on ETag conflicts (HTTP 412).
+
+    If *batch_max_lsn* is provided, the counter stores it alongside the
+    pre-increment count.  On a change-feed retry (same batch replayed),
+    the function detects the duplicate via LSN comparison and returns the
+    cached ``(pre_batch_count, current_count)`` **without** writing,
+    preserving threshold-crossing semantics for the caller.
+
+    .. note::
+
+       LSN-based replay detection is perfect for **thread-scoped** counters
+       (single logical partition → monotonic LSNs).  For **user-scoped**
+       counters that aggregate across partitions, it handles the common
+       single-partition-range retry but may not detect cross-partition
+       interleaving.  Deterministic orchestration instance IDs provide
+       an additional safety net against duplicate orchestration starts.
     """
     container = await _get_cosmos_counter_container()
     max_retries = 3
@@ -108,12 +149,27 @@ async def increment_counter_by(counter_id: str, user_id: str, thread_id: str, co
         except CosmosHttpResponseError:
             raise
 
+        # ---- Replay detection via LSN ----
+        if (
+            batch_max_lsn is not None
+            and existing_doc is not None
+            and existing_doc.get("last_batch_lsn") == batch_max_lsn
+        ):
+            replay_old = existing_doc.get("last_batch_old_count", old_count)
+            logger.info(
+                "Counter replay detected counter_id=%s lsn=%s, returning cached result",
+                counter_id, batch_max_lsn,
+            )
+            return (replay_old, old_count)
+
         new_count = old_count + count
         new_doc = {
             "id": counter_id,
             "user_id": user_id,
             "thread_id": thread_id,
             "count": new_count,
+            "last_batch_lsn": batch_max_lsn,
+            "last_batch_old_count": old_count,
             "created_at": existing_doc.get("created_at", _utc_now_iso()) if existing_doc else _utc_now_iso(),
             "updated_at": _utc_now_iso(),
         }
@@ -124,10 +180,28 @@ async def increment_counter_by(counter_id: str, user_id: str, thread_id: str, co
                 await container.upsert_item(
                     body=new_doc,
                     etag=etag,
-                    match_condition="IfMatch",
+                    match_condition=MatchConditions.IfNotModified,
                 )
             else:
-                await container.upsert_item(body=new_doc)
+                # First-time creation: use create_item to avoid last-writer-wins
+                # race when multiple Function instances see 404 concurrently.
+                try:
+                    await container.create_item(body=new_doc)
+                except CosmosHttpResponseError as create_exc:
+                    if create_exc.status_code == 409 and attempt < max_retries - 1:
+                        # Another instance created it first — retry with read-modify-write
+                        logger.warning(
+                            "Counter create conflict counter_id=%s attempt=%d/%d, retrying",
+                            counter_id, attempt + 1, max_retries,
+                        )
+                        continue
+                    if create_exc.status_code == 409:
+                        logger.warning(
+                            "Counter create conflict exhausted retries counter_id=%s, skipping",
+                            counter_id,
+                        )
+                        return (old_count, old_count)
+                    raise
             return (old_count, new_count)
         except CosmosHttpResponseError as exc:
             if exc.status_code == 412 and attempt < max_retries - 1:
@@ -330,16 +404,18 @@ async def process_changefeed_batch(documents: list[dict], starter) -> None:
     Extracted from the trigger function so it can be tested without the
     Durable Functions middleware.
     """
-    n_thread = int(os.environ.get("THREAD_SUMMARY_EVERY_N", "0"))
-    n_facts = int(os.environ.get("FACT_EXTRACTION_EVERY_N", "0"))
-    n_user = int(os.environ.get("USER_SUMMARY_EVERY_N", "0"))
+    n_thread = _parse_threshold("THREAD_SUMMARY_EVERY_N")
+    n_facts = _parse_threshold("FACT_EXTRACTION_EVERY_N")
+    n_user = _parse_threshold("USER_SUMMARY_EVERY_N")
 
     if n_thread == 0 and n_facts == 0 and n_user == 0:
         return  # all processing disabled
 
     # ---- Step 1: Filter to turns, group by scope ----
     thread_counts: dict[tuple[str, str], int] = defaultdict(int)
     user_counts: dict[str, int] = defaultdict(int)
+    thread_max_lsn: dict[tuple[str, str], int] = {}
+    user_max_lsn: dict[str, int] = {}
 
     for doc in documents:
         # Counter writes land in the separate counter container, so only raw
@@ -356,6 +432,13 @@ async def process_changefeed_batch(documents: list[dict], starter) -> None:
         thread_counts[(user_id, thread_id)] += 1
         user_counts[user_id] += 1
 
+        # Track max _lsn per scope for replay detection
+        lsn = doc.get("_lsn")
+        if lsn is not None:
+            key = (user_id, thread_id)
+            thread_max_lsn[key] = max(thread_max_lsn.get(key, 0), lsn)
+            user_max_lsn[user_id] = max(user_max_lsn.get(user_id, 0), lsn)
+
     thread_counters_enabled = (n_thread > 0 or n_facts > 0)
     user_counters_enabled = (n_user > 0)
     enabled_thread_groups = len(thread_counts) if thread_counters_enabled else 0
@@ -370,59 +453,102 @@ async def process_changefeed_batch(documents: list[dict], starter) -> None:
     )
 
     # ---- Step 2: Thread-scoped counters and threshold checks ----
+    orchestration_errors: list[Exception] = []
+
     if thread_counters_enabled:
         for (user_id, thread_id), batch_count in thread_counts.items():
+            lsn = thread_max_lsn.get((user_id, thread_id))
             old_count, new_count = await increment_counter_by(
                 f"thread_counter_{user_id}_{thread_id}", user_id, thread_id, batch_count,
+                batch_max_lsn=lsn,
             )
 
             if n_thread > 0 and crosses_threshold(old_count, new_count, n_thread):
+                bucket = new_count // n_thread
+                instance_id = f"ts_{user_id}_{thread_id}_{bucket}"
                 logger.info(
-                    "on_memory_change: triggering thread_summary user_id=%s thread_id=%s count=%d",
-                    user_id, thread_id, new_count,
-                )
-                await starter.start_new(
-                    "memory_orchestrator",
-                    client_input={
-                        "thread_summary_only": True,
-                        "user_id": user_id,
-                        "thread_id": thread_id,
-                    },
+                    "on_memory_change: triggering thread_summary user_id=%s thread_id=%s count=%d instance=%s",
+                    user_id, thread_id, new_count, instance_id,
                 )
+                try:
+                    await starter.start_new(
+                        "memory_orchestrator",
+                        instance_id=instance_id,
+                        client_input={
+                            "thread_summary_only": True,
+                            "user_id": user_id,
+                            "thread_id": thread_id,
+                        },
+                    )
+                except Exception as exc:
+                    logger.exception(
+                        "Failed to start thread_summary orchestration user_id=%s thread_id=%s",
+                        user_id, thread_id,
+                    )
+                    orchestration_errors.append(exc)
 
             if n_facts > 0 and crosses_threshold(old_count, new_count, n_facts):
+                bucket = new_count // n_facts
+                instance_id = f"ef_{user_id}_{thread_id}_{bucket}"
                 logger.info(
-                    "on_memory_change: triggering extract_facts user_id=%s thread_id=%s count=%d",
-                    user_id, thread_id, new_count,
-                )
-                await starter.start_new(
-                    "memory_orchestrator",
-                    client_input={
-                        "extract_facts_only": True,
-                        "user_id": user_id,
-                        "thread_id": thread_id,
-                    },
+                    "on_memory_change: triggering extract_facts user_id=%s thread_id=%s count=%d instance=%s",
+                    user_id, thread_id, new_count, instance_id,
                 )
+                try:
+                    await starter.start_new(
+                        "memory_orchestrator",
+                        instance_id=instance_id,
+                        client_input={
+                            "extract_facts_only": True,
+                            "user_id": user_id,
+                            "thread_id": thread_id,
+                        },
+                    )
+                except Exception as exc:
+                    logger.exception(
+                        "Failed to start extract_facts orchestration user_id=%s thread_id=%s",
+                        user_id, thread_id,
+                    )
+                    orchestration_errors.append(exc)
 
     # ---- Step 3: User-scoped counters and threshold checks ----
     if user_counters_enabled:
         for user_id, batch_count in user_counts.items():
+            lsn = user_max_lsn.get(user_id)
             old_count, new_count = await increment_counter_by(
                 f"user_counter_{user_id}", user_id, USER_COUNTER_THREAD_ID, batch_count,
+                batch_max_lsn=lsn,
             )
 
             if crosses_threshold(old_count, new_count, n_user):
+                bucket = new_count // n_user
+                instance_id = f"us_{user_id}_{bucket}"
                 logger.info(
-                    "on_memory_change: triggering user_summary user_id=%s count=%d",
-                    user_id, new_count,
-                )
-                await starter.start_new(
-                    "memory_orchestrator",
-                    client_input={
-                        "user_summary_only": True,
-                        "user_id": user_id,
-                    },
+                    "on_memory_change: triggering user_summary user_id=%s count=%d instance=%s",
+                    user_id, new_count, instance_id,
                 )
+                try:
+                    await starter.start_new(
+                        "memory_orchestrator",
+                        instance_id=instance_id,
+                        client_input={
+                            "user_summary_only": True,
+                            "user_id": user_id,
+                        },
+                    )
+                except Exception as exc:
+                    logger.exception(
+                        "Failed to start user_summary orchestration user_id=%s",
+                        user_id,
+                    )
+                    orchestration_errors.append(exc)
+
+    # Re-raise so the change feed batch retries and thresholds re-fire
+    if orchestration_errors:
+        raise RuntimeError(
+            f"Failed to start {len(orchestration_errors)} orchestration(s); "
+            "raising to retry the change feed batch"
+        ) from orchestration_errors[0]
 @df_app.cosmos_db_trigger(
     arg_name="documents",
     connection="COSMOS_DB",

azure_functions/function_app.py

@@ -38,7 +38,7 @@ def cosmos_clients():
     from azure.cosmos import CosmosClient
     from azure.identity import DefaultAzureCredential
 
-    endpoint = os.environ["COSMOS_DB__accountEndpoint"]
+    endpoint = os.environ.get("COSMOS_DB__accountEndpoint") or os.environ.get("COSMOS_DB_ENDPOINT")
     database_name = os.environ.get("COSMOS_DB_DATABASE", "ai_memory")
     memories_container_name = os.environ.get("COSMOS_DB_CONTAINER", "memories")
     counter_container_name = os.environ.get("COSMOS_DB_COUNTERS_CONTAINER", "counter")