feat(chat): report per-call measurements after finalization by BukeLy · Pull Request #581 · VectifyAI/PageIndex · GitHub
Skip to content

feat(chat): report per-call measurements after finalization - #581

Open
BukeLy wants to merge 8 commits into
feat/chat-usage-resultsfrom
feat/chat-measurement-finalization
Open

BukeLy wants to merge 8 commits into
feat/chat-usage-resultsfrom
feat/chat-measurement-finalization

Conversation

@BukeLy

@BukeLy BukeLy commented Oct 13, 2026

Copy link
Copy Markdown
Member

Summary

Applications embedding SDK chat need per-call measurements that distinguish missing usage from zero and remain available after execution fails or is cancelled. Add chat(on_usage=...) to deliver one report after an own-model protocol run finishes or its consumed stream closes.

Depends on #574; this PR targets feat/chat-usage-results and contains only the measurement/finalization changes on top of it. Refs #557.

Goal

Provide the same supported measurement contract for Chat Completions, Responses and Messages, in streaming and non-streaming execution.

What Changed

  • Extend the existing recorder and provider observation points to retain individual call identities, requested/reported models, backend identity and available token quantities before Agents normalizes missing usage.
  • Return ordinary dictionaries containing the execution outcome, call records, known totals and per-field completeness. Preserve cache-write TTL buckets, use a reliably configured Messages lifetime when available, and document token inclusion semantics.
  • Retain in-flight calls with unknown usage on failure or cancellation. Stream closure waits for local cleanup before the synchronous callback runs; a completed run keeps its outcome when the consumer closes after its terminal event.
  • Keep callback failures from replacing an execution failure. Document that local cleanup may wait for a blocking backend and cannot establish remote cancellation or recover undelivered usage.
  • Keep implementation within the existing client and chat modules, using the existing framework hooks and one shared finalization wrapper for the three protocol paths.

Validation: python -m pytest -q -k 'not test_live' --disable-warnings — 909 passed, 5 skipped, 7 live-service tests deselected. Coverage includes multi-turn success, missing usage, genuine zero, later provider failure, cancellation, transport cleanup, and native/LiteLLM cache fields. All provider responses in the added tests are synthetic. git diff --check passes. Requirements and style reviews completed, with regression tests for the identified issues.

Key Files

  • pageindex/client.py: optional callback, argument forwarding and public measurement/lifecycle contract.
  • pageindex/local_chat.py: per-call collection, usage normalization, completeness and finalization across the three protocols.
  • tests/test_local_chat.py: deterministic measurement and lifecycle regression coverage through the SDK entry point and actual framework adapters.

rejojer and others added 5 commits October 13, 2026 05:37
…ment

A local submit made every model call inside the SDK and dropped the
token counts each response carries, so pricing an index meant wrapping
LiteLLM by hand (LiteLLM's own callbacks miss calls: a CustomLogger saw
17 of 18 on one document, success_callback 1 of 18).

submit_document now returns 'usage' in local mode, summed over every
call of that submit, async summaries and expand and the sync description
alike, in the Chat Completions shape LiteLLM returns and chat()'s
default lane speaks: {'prompt_tokens', 'completion_tokens',
'total_tokens', 'prompt_tokens_details': {'cached_tokens',
'cache_write_tokens'}, 'completion_tokens_details': {'reasoning_tokens'}}.
utils.llm_completion and llm_acompletion, the two helpers every indexing
call goes through, add each response into a contextvar dict that
LocalAPI scopes around the indexing operation next to the backend
overrides, so concurrent submits keep separate tallies. A response
without usage, or with fields missing, adds nothing rather than raising
inside the retry loop. Cloud results are unchanged.

Real gpt-6-luna runs on q1-fy25-earnings.pdf matched an independent
count of every LiteLLM call exactly: flash 18 calls, 16,888 prompt
(5,630 cached) and 3,783 completion tokens (31 reasoning); standard 151
calls, 164,174 prompt and 22,262 completion tokens. A fresh flash run
reported 9,073 cache-write tokens: OpenAI writes its prompt cache
without any cache_control.
chat(protocol="messages") reports one usage for a multi-turn tool run,
but only the four top-level token counters were summed; the nested
counts kept the final turn's value. A run whose three turns each thought
500 tokens reported output_tokens 1,500 next to thinking_tokens 500, and
the same went for cache_creation's 5m/1h split and server_tool_use's
web_search/web_fetch requests.

The turns' native usage dicts now merge field by field: counts add up,
inside nested count objects too; the per-iteration breakdown lists
concatenate; and labels (service_tier, fallback_credit, speed) keep the
last value a turn set. Count objects are recognised by shape, so a count
the anthropic SDK adds later is summed without a code change.
chat(protocol="messages", stream=True) forwards each turn's native
events, so a tool run's usage arrived split across turns and, within a
turn, across message_start (input side) and message_delta (output side).
Nothing gave the run's total: a caller reading the last message_delta got
the final turn only.

A finished run now ends with one {"type": "result", "usage": ...} event,
named after the Claude Agent SDK's closing result message. Its usage is
the non-stream envelope's aggregate over each turn's message_stop
snapshot, which the anthropic SDK accumulates from message_start and
message_delta (complete, thinking counts included, from 0.122.0, the
declared floor).

The native events are left untouched; Anthropic's streaming docs ask
clients to handle unknown event types gracefully, and the event is an
anthropic BaseModel, so to_dict() and to_json() work as on the native
ones. A run that raises or is abandoned yields no result event, so a
missing total is never mistaken for a complete one.
…ssages stream transcript

A failed own-model run now raises PageIndexAPIError with `usage`: what
its finished model calls spent, in the lane's usage shape. openai-agents
attaches run data only to its own exceptions, and adds a streamed turn
to raw_responses only after its tools ran, so a provider error on turn 3
lost turns 1-2 and a failing tool lost its own turn. OpenAI runs now
record each response through RunHooks.on_llm_end, streamed or not, and
Messages sums the turns it already holds, tool failures and a failed
transcript read-back included. Any other error keeps usage None.

The Responses stream also counts a call cut short (response.incomplete
or response.failed): openai-agents raises on it before its hooks see
it, yet it was billed.

Every usage lists the run's model calls: `request_usage_entries` on
chat_completions and responses (openai-agents' name; neither wire has a
per-call field), and Anthropic's native `iterations` on Messages, where
a turn the API reports no breakdown for counts as one iteration.

chat_completions usage carries `prompt_tokens_details.cache_write_tokens`,
which the OpenAI SDK type and LiteLLM both define. It was dropped, so
cache writes went uncounted.

The Messages stream's `result` event carries `messages`, the same
continuation transcript the non-stream envelope returns, so a streamed
conversation can continue.

Usage dicts gain keys (request_usage_entries, cache_write_tokens), and
Messages usage.iterations is now always a list rather than None; code
comparing whole usage dicts must update.

Part of #557 and #563.
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 13, 2026 •

Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9626ae2976

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pageindex/local_chat.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8a1ce8eb4b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pageindex/local_chat.py
Comment thread pageindex/local_chat.py Outdated
@BukeLy
BukeLy added this pull request to stack #582 October 13, 2026 19:11
Comment thread pageindex/local_chat.py Fixed

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants