feat(retain): expose processed_content_tokens on RetainResult

Delta retain already knows, at chunk-level granularity, which content
was new vs unchanged on an upsert to an existing document_id. Surface
that signal to post-retain hooks so extensions can reason about "how
much content actually went through the extraction pipeline" without
re-implementing the dedup logic.

New field `RetainResult.processed_content_tokens: int | None`:
  * None — the retain went through the full (non-delta) path or has
           no dedup signal. Consumers should treat this as "the full
           submitted payload was processed."
  * 0    — the submission matched prior content exactly; no chunks
           went through extraction (metadata-only update).
  * N>0  — only N tokens of content+context were actually re-extracted.
           The remainder matched existing chunks by content_hash and
           was skipped.

Populated in three places:
  * Streaming / full retain path → None
  * `_try_delta_retain` no-changes fallthrough (`_delta_metadata_only`)
    → 0
  * `_try_delta_retain` partial-delta success → sum of
    count_tokens(content) + count_tokens(context) across the chunks
    built for extraction (delta_contents)

Sub-batch aggregation propagates None if any sub-batch bypassed dedup,
so callers never accidentally undercount when only part of a large
batch was eligible for delta processing.

Tests exercise the full path, unchanged-resubmit, appended-content,
and no-document-id cases plus a unit check on the aggregation helper.

Author: cdbartholomew

hindsight-api-slim/hindsight_api/engine/memory_engine.py

@@ -2267,6 +2267,11 @@ async def retain_batch_async(
         # Calculate total token count
         total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
         total_usage = TokenUsage()
+        # Aggregate "content tokens that actually went through extraction after
+        # chunk-level dedup" across sub-batches. ``None`` in any sub-batch
+        # means that sub-batch bypassed dedup, so the aggregate is None
+        # (see RetainResult.processed_content_tokens).
+        total_processed_content_tokens: int | None = 0
 
         # Get batch size threshold from config
         config = get_config()
@@ -2318,7 +2323,7 @@ async def retain_batch_async(
                     f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
                 )
 
-                sub_results, sub_usage = await self._retain_batch_async_internal(
+                sub_results, sub_usage, sub_processed = await self._retain_batch_async_internal(
                     bank_id=bank_id,
                     contents=sub_batch,
                     request_context=request_context,
@@ -2334,6 +2339,10 @@ async def retain_batch_async(
                 )
                 all_results.extend(sub_results)
                 total_usage = total_usage + sub_usage
+                if total_processed_content_tokens is None or sub_processed is None:
+                    total_processed_content_tokens = None
+                else:
+                    total_processed_content_tokens = total_processed_content_tokens + sub_processed
 
             total_time = time.time() - start_time
             logger.info(
@@ -2342,7 +2351,7 @@ async def retain_batch_async(
             result = all_results
         else:
             # Small batch - use internal method directly
-            result, total_usage = await self._retain_batch_async_internal(
+            result, total_usage, total_processed_content_tokens = await self._retain_batch_async_internal(
                 bank_id=bank_id,
                 contents=contents,
                 request_context=request_context,
@@ -2371,6 +2380,7 @@ async def retain_batch_async(
                 llm_input_tokens=total_usage.input_tokens,
                 llm_output_tokens=total_usage.output_tokens,
                 llm_total_tokens=total_usage.total_tokens,
+                processed_content_tokens=total_processed_content_tokens,
             )
             try:
                 await self._operation_validator.on_retain_complete(result_ctx)
@@ -2403,7 +2413,7 @@ async def _retain_batch_async_internal(
         operation_id: str | None = None,
         outbox_callback: "Callable[[asyncpg.Connection], Awaitable[None]] | None" = None,
         strategy: str | None = None,
-    ) -> tuple[list[list[str]], "TokenUsage"]:
+    ) -> tuple[list[list[str]], "TokenUsage", int | None]:
         """
         Internal method for batch processing without chunking logic.
 
@@ -2422,7 +2432,9 @@ async def _retain_batch_async_internal(
             document_tags: Tags applied to all items in this batch
 
         Returns:
-            Tuple of (unit ID lists, token usage for fact extraction)
+            Tuple of (unit ID lists, LLM token usage, processed_content_tokens).
+            See ``RetainResult.processed_content_tokens`` for the semantics of
+            the third element.
         """
         # Use the new modular orchestrator
         from .retain import orchestrator

hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py

@@ -16,7 +16,7 @@
 
 from ...worker.stage import set_stage
 from ..db_utils import acquire_with_retry
-from ..memory_engine import fq_table
+from ..memory_engine import count_tokens, fq_table
 from . import bank_utils
 
 
@@ -25,6 +25,32 @@ def utcnow():
     return datetime.now(UTC)
 
 
+def _merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
+    """Combine the processed-content-tokens signal across sub-results.
+
+    Semantics (see RetainResult.processed_content_tokens):
+      * None means "this part of the retain did not go through chunk-level
+        dedup" — i.e. the entire submitted payload was processed. If any
+        sub-result is None, the aggregate is None so callers conservatively
+        bill the full content.
+      * Otherwise, accumulate the int values.
+    """
+    if a is None or b is None:
+        return None
+    return a + b
+
+
+def _count_delta_content_tokens(delta_contents: list["RetainContent"]) -> int:
+    """Sum content + context tokens across the chunk items that were
+    actually fed into the extraction pipeline on a partial-delta retain.
+    """
+    total = 0
+    for c in delta_contents:
+        total += count_tokens(c.content or "")
+        total += count_tokens(c.context or "")
+    return total
+
+
 def parse_datetime_flexible(value: Any) -> datetime:
     """
     Parse a datetime value that could be either a datetime object or an ISO string.
@@ -417,13 +443,21 @@ async def retain_batch(
     schema: str | None = None,
     outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
     db_semaphore: "asyncio.Semaphore | None" = None,
-) -> tuple[list[list[str]], TokenUsage]:
+) -> tuple[list[list[str]], TokenUsage, int | None]:
     """
     Process a batch of content through the retain pipeline.
 
     Supports delta retain: when upserting a document that already has chunks,
     only re-processes chunks whose content has changed. Unchanged chunks keep
     their existing facts, entities, and links.
+
+    Returns a three-tuple of:
+      * per-content-item unit ID lists
+      * aggregate LLM token usage
+      * processed_content_tokens — content+context tokens that actually went
+        through extraction after chunk-level dedup, or ``None`` if this path
+        didn't dedup (caller should treat as "bill full submitted content").
+        See ``RetainResult.processed_content_tokens`` for details.
     """
     start_time = time.time()
     total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
@@ -463,8 +497,9 @@ async def retain_batch(
             # Process each group and merge results back in original order
             result_unit_ids: list[list[str]] = [[] for _ in contents_dicts]
             total_usage = TokenUsage()
+            total_processed_tokens: int | None = 0
             for doc_key, (group_dicts, group_contents) in groups.items():
-                group_ids, group_usage = await retain_batch(
+                group_ids, group_usage, group_processed = await retain_batch(
                     pool=pool,
                     embeddings_model=embeddings_model,
                     llm_config=llm_config,
@@ -486,7 +521,8 @@ async def retain_batch(
                     if group_idx < len(group_ids):
                         result_unit_ids[orig_idx] = group_ids[group_idx]
                 total_usage = total_usage + group_usage
-            return result_unit_ids, total_usage
+                total_processed_tokens = _merge_processed_content_tokens(total_processed_tokens, group_processed)
+            return result_unit_ids, total_usage, total_processed_tokens
 
     # Resolve effective document_id early so both delta and streaming paths
     # can find existing chunks from a prior attempt. On retry, a generated
@@ -595,7 +631,9 @@ async def retain_batch(
                 f"{datetime.fromtimestamp(start_time, tz=UTC).isoformat()})"
             )
             logger.info("\n" + "\n".join(log_buffer) + "\n")
-            return [[] for _ in contents], TokenUsage()
+            # No new content was processed — report 0 so callers can skip
+            # billing cleanly instead of falling back to full-content billing.
+            return [[] for _ in contents], TokenUsage(), 0
 
     # --- Delta retain: check if we can skip unchanged chunks ---
     if is_first_batch:
@@ -1388,7 +1426,10 @@ async def _run_mini_batch_db_work() -> None:
     # Map all unit_ids back to the original content items.
     # For streaming mode with a single document, all units belong to content 0.
     result_unit_ids = [all_unit_ids] + [[] for _ in contents[1:]]
-    return result_unit_ids, total_usage
+    # The streaming path doesn't compute per-chunk content-hash dedup in
+    # a way that lets us report a partial-processed tokens count — signal
+    # ``None`` so callers bill against the full submitted payload.
+    return result_unit_ids, total_usage, None
 
 
 # ---------------------------------------------------------------------------
@@ -1416,10 +1457,15 @@ async def _try_delta_retain(
     schema,
     outbox_callback,
     db_semaphore: "asyncio.Semaphore | None" = None,
-):
+) -> tuple[list[list[str]], TokenUsage, int | None] | None:
     """
     Attempt delta retain for a document upsert. Returns result tuple if delta
     was performed, or None to fall back to full retain.
+
+    When a result tuple is returned, the third element is the content+context
+    token count for the chunks that actually went through extraction
+    (``0`` if the submission matched prior content exactly and nothing was
+    re-extracted).
     """
     # Need a single document_id
     effective_doc_id = document_id
@@ -1678,7 +1724,12 @@ async def _run_delta_db_work() -> None:
             await _run_delta_db_work()
     else:
         await _run_delta_db_work()
-    return result_unit_ids, usage
+    # Count content + context tokens that actually went through extraction.
+    # ``delta_contents`` holds the per-chunk RetainContent items for the
+    # changed/new chunks (see ``_build_delta_contents``) — i.e. exactly what
+    # the LLM pipeline saw this call. Unchanged chunks contribute zero.
+    processed_tokens = _count_delta_content_tokens(delta_contents)
+    return result_unit_ids, usage, processed_tokens
 
 
 async def _delta_metadata_only(
@@ -1718,7 +1769,11 @@ async def _delta_metadata_only(
     total_time = time.time() - start_time
     log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s")
     logger.info("\n" + "\n".join(log_buffer) + "\n")
-    return [[] for _ in contents], TokenUsage()
+    # Nothing went through the extraction pipeline — report 0 processed
+    # content tokens so callers can bill accordingly (a caller that's been
+    # told ``0`` knows the retain was a pure metadata update and should
+    # charge nothing for content).
+    return [[] for _ in contents], TokenUsage(), 0
 
 
 # ---------------------------------------------------------------------------

hindsight-api-slim/hindsight_api/extensions/operation_validator.py

@@ -176,6 +176,22 @@ class RetainResult:
     llm_input_tokens: int | None = None
     llm_output_tokens: int | None = None
     llm_total_tokens: int | None = None
+    # Content tokens the retain pipeline actually processed, after
+    # chunk-level content-hash deduplication. Semantics:
+    #   None — no dedup signal available (e.g. a first-time retain or a
+    #          path that doesn't compute it). Callers that care about
+    #          "what was actually new on this retain" should treat None
+    #          as "the full submitted content was processed."
+    #   0    — the entire submission was a duplicate of prior content
+    #          (all chunks matched by content_hash); nothing went
+    #          through LLM extraction.
+    #   N>0  — only N tokens of content + context went through the
+    #          extraction pipeline. The remainder was dedup'd against
+    #          existing chunks.
+    # This is the basis most billing/metering extensions want to use
+    # when the customer's client resubmits growing payloads to the same
+    # document_id (e.g. a session transcript appended to on each turn).
+    processed_content_tokens: int | None = None
 
 
 @dataclass