Add provenance field to audit payload to distinguish missing vs. genuinely-zero prompt token counts #1953

Activity

1. hari-kuriakose

   assigned

   johnyrahul

   and

   praveen-formido

   on May 11, 2026
2. chandrasekharan-zipstack commented on May 11, 2026

   

   chandrasekharan-zipstack

   on May 11, 2026

   Contributor

   More actions

   Heads-up after #1929 merged: the problem this issue describes is still real — _record_usage continues to silently record prompt_tokens=0 when the provider omits the field and litellm.token_counter() raises, with no flag distinguishing "estimation failed" from "genuinely zero".

   The Proposed Solution file list is stale, though: _record_usage no longer calls Audit().push_usage_data, TokenCounterCompat is no longer in this path, and usage now flows through self._pending_usage → worker bulk_create_usage. Whoever picks this up should re-anchor against the new shape in unstract/sdk1/src/unstract/sdk1/llm.py and the OSS Usage model rather than audit.py / utils/common.py.
3. ritwik-g commented on May 12, 2026

   

   ritwik-g

   on May 12, 2026

   Contributor

   More actions

   @coderabbitai can we update the description as per #1953 (comment)

   👀React with 👀1Reacted by coderabbitai
4. coderabbitai commented on May 12, 2026

   

   coderabbitaiBot

   on May 12, 2026 – with coderabbitai

   ContributorAuthor

   More actions

   🧩 Analysis chain

   🏁 Script executed:

   #!/bin/bash
   # Get the current issue description
   gh issue view 1953 --repo Zipstack/unstract --json title,body,comments
   
   # Get the specific comment
   gh api repos/Zipstack/unstract/issues/comments/4419838595
   

   Length of output: 6608

   ---

   🏁 Script executed:

   #!/bin/bash
   gh issue edit 1953 --repo Zipstack/unstract --body "## Context
   
   In PR `#1895` (Add a dedicated OpenAI-compatible LLM adapter), the \`_record_usage\` method in \`unstract/sdk1/src/unstract/sdk1/llm.py\` was updated to fall back to \`prompt_tokens=0\` when:
   1. The provider does not return \`usage.prompt_tokens\` in its response, **and**
   2. LiteLLM's \`token_counter()\` raises (e.g., for unmapped custom/OpenAI-compatible models).
   
   A warning is logged in this case, but the recorded usage payload contains \`prompt_llm_token_count=0\` with no flag to distinguish \"estimation failed / data unavailable\" from \"the prompt genuinely consumed zero tokens\".
   
   This was acknowledged as a known limitation and deferred from PR `#1895` because adding a provenance field requires a wider end-to-end contract change across the usage pipeline beyond sdk1.
   
   > **Note (post-#1929):** The usage path has changed — \`_record_usage\` no longer calls \`Audit().push_usage_data\`. Usage now flows through \`self._pending_usage\` → worker \`bulk_create_usage\`. The core problem remains unchanged.
   
   Relevant discussion: https://github.com/Zipstack/unstract/pull/1895#discussion_r3135244964
   
   ## Problem
   
   Downstream consumers of the usage data (cost attribution, analytics, billing) cannot distinguish:
   - **Missing data** — token estimation failed for an unmapped model
   - **Genuinely zero** — the prompt actually consumed zero tokens
   
   This silently understates prompt-token consumption in cost attribution and analytics for long-running workloads against OpenAI-compatible endpoints that do not return \`usage.prompt_tokens\`.
   
   ## Proposed Solution
   
   Add a provenance / sentinel field to the usage payload, for example:
   
   - \`prompt_tokens_source\`: an enum/string such as \`\"provider\"\`, \`\"estimated\"\`, or \`\"unknown\"\`
   - or \`estimation_failed\`: a boolean flag
   
   This would require changes to (re-anchored against the new shape post-#1929):
   1. \`unstract/sdk1/src/unstract/sdk1/llm.py\` — populate the provenance field in \`_record_usage\` before appending to \`self._pending_usage\`
   2. The OSS \`Usage\` model — extend the schema to include the provenance field
   3. The worker \`bulk_create_usage\` path — consume and persist the new field
   4. Downstream platform/usage-record schema — surface and store the new field
   
   Optionally, increment an ops metric/counter when the fallback path is triggered so operations teams can detect silent drift without parsing logs.
   
   ## References
   
   - PR `#1895` — Add a dedicated OpenAI-compatible LLM adapter
   - PR `#1929` — Changed usage flow from \`Audit().push_usage_data\` to \`self._pending_usage\` → \`bulk_create_usage\`
   - Related issues: `#1894`, `#856`, `#1443`
   - Opened by: `@hari-kuriakose`"
   

   Length of output: 49

   ---

   @ritwik-g, the description of #1953 has been updated. Here's a summary of what changed:

   • Context section — Replaced the stale reference to Audit().push_usage_data with a note explaining the post-UN-2946 [FEAT] Prompt Studio lookups bridge, executor hook, and IDE wiring (OSS side) #1929 flow change, and kept the core problem statement intact.
   • Proposed Solution file list — Removed the stale TokenCounterCompat / utils/common.py and audit.py entries, and re-anchored the steps against the new shape:
     1. llm.py → populate provenance in _record_usage before appending to self._pending_usage
     2. OSS Usage model — extend schema
     3. Worker bulk_create_usage path — consume and persist
     4. Downstream platform schema — surface and store
   • References — Added PR #1929 to document the usage-flow change.
5. cschanhniem commented on Jun 19, 2026

   

   cschanhniem

   on Jun 19, 2026

   More actions

   The proposed prompt_tokens_source enum approach looks right. One concrete suggestion on the schema shape:

   class PromptTokensSource(str, Enum):
       provider = "provider"   # usage.prompt_tokens returned by the provider
       estimated = "estimated" # token_counter() succeeded as a fallback
       unknown = "unknown"     # token_counter() also failed; value is 0

   Using a string enum (rather than a bool estimation_failed) lets you add estimated as an intermediate state, which matters if you later want to distinguish "we guessed" from "we have no idea". The unknown case is the one that today silently writes 0, so downstream cost attribution can filter or flag those rows specifically.

   For the bulk_create_usage path post-#1929: the field only needs to survive the worker → DB write. If the DB column is nullable varchar, you can add it with a default of null (mapped to unknown in code) and backfill historical rows as null without a breaking migration.