Repository navigation
LCORE-3910: count the summarization calls in quota, token counts and metrics - #2797
max-svistunov wants to merge 2 commits into
Conversation
OGX appends a turn to the conversation only when it is handed the conversation parameter and runs the inference. In compacted mode that parameter is dropped, so lightspeed-stack appends the turn itself. Every endpoint made that write on its own, and nothing enforced it: two unrelated cleanups removed the calls, and for six weeks conversations stopped growing after a compaction without anything failing (LCORE-3883). The write now has one owner, PendingTurn in utils/pending_turn.py. An endpoint creates it from the parameters the request is sent with and the original input of the compaction result, and reports how the turn ended: store_completed, store_blocked, store_interrupted or drop. The owner decides whether the turn is ours, which input is stored, and it stores a turn once: the first report settles the turn, every later one does nothing. It covers all nine places that stored a turn, not only the compacted ones, because they are the same duty: a turn OGX does not store. That is a conversation in compacted mode, a request a shield blocked, a stream the client interrupted, a continuation from previous_response_id. store_compacted_turn is removed; no endpoint appends a turn except through the owner. The shield capabilities still store the turn they rejected themselves, from inside the agent run, and only when the model was handed the conversation, so never in compacted mode. A compacted request that loses its turn through a change in the code fails: - creating the owner for compacted parameters without the original input raises ValueError, which is how a caller that lost the hand-over looks; - ensure_settled raises TurnNotStoredError when a compacted request reaches the end of its handler and nobody tried to store its turn or dropped it on purpose. /v1/query runs in a scope that makes the check when it is left; the streaming paths, /v1/responses and A2A make it after their write, in the function that calls the write, not in the one that performs it. The error is not a RuntimeError, because the endpoints report that one as an inference failure. The check does not cover a stream the client stops reading, a /v1/responses stream without a final response, and a failed write that is logged. They store nothing, as before. Every condition of the old call sites is kept: where the write happens, what is stored for a failed or incomplete turn, what a failed write does to the request (it fails /v1/query and /v1/responses, it is logged on /v1/streaming_query, A2A and the interrupt path), and the order relative to quota consumption. /v1/query still stores nothing for a request blocked on a compacted conversation; that is now an explicit drop, and LCORE-3788 settles what all endpoints should do. The guard that lets only one of stream end, cancellation handler and interrupt callback finish a turn is unchanged. One behaviour changes, on purpose. A blocked request on a conversation that is not compacted has its refusal turn stored before the stream starts. An interrupt while the refusal was streamed stored the turn a second time, with the interruption notice for an answer. It is stored once now; the interrupt still records the turn in the database. The path cannot be reached today, because run_shield_moderation always passes. Tests: - integration, tests/integration/endpoints/test_turn_persistence.py: for /v1/query, /v1/streaming_query, /v1/responses (blocking and streaming) and A2A, every way a turn can end (completed, blocked, interrupted, failed, cut short by the model, abandoned by the client, continued from a previous response, store off), compacted and not; what a failed write does to the request on each endpoint; and what happens when the step that stores the turn is taken out of an endpoint. Each test runs the real handler and compares the conversation item by item afterwards, so a missing write, a second write and a write of the wrong input all fail it. 44 of the 50 tests pass on main unchanged. The six that fail there are the double write above and the five that take the write out of an endpoint, which main does not notice. - unit: the owner, the hand-over from the compaction seam (omit_conversation and the original input are set together), and a scan of src that fails when a module other than the known writers mentions one of the functions that append to a conversation. - existing unit tests follow the new parameter. Those that asserted on the arguments of a patched helper in the agent and interrupt paths now assert on what is stored. Mocked request parameters got the fields the owner reads. The A2A helpers of the compaction tests moved to _compaction_helpers.py, where both test modules import them. The tests were checked by mutation: with the write removed from any of the four endpoints, a check removed, the write doubled, the settled state ignored, the interrupt guard ignored, the explicit rewrite stored in place of the input, a turn OGX stores stored by us as well, or the store flag ignored, tests fail. Docs: the design document describes the owner and has a table of what is stored per endpoint and per way a turn can end. The endpoint guide names the exception for an interrupted stream.
…metrics Compaction summarizes older turns with an LLM call of its own, and folds the summaries with another. The provider bills both. Their usage was discarded: summarize_chunk() and recursively_resummarize() read the text of the response and nothing else. So the calls consumed no quota, were missing from the input_tokens / output_tokens the client is told, and were missing from the token and call metrics. A user could go over the quota without it showing, by more the longer the conversation was. The user guide said the opposite. What is decided (by the owner of the ticket): - the summarization tokens are charged to the user's quota; - they are part of the existing input_tokens / output_tokens of /v1/query and of the end event of /v1/streaming_query, no new fields; - /v1/responses passes the usage object of the response through unchanged, to stay a drop-in for clients of the OpenAI Responses API; the quota is charged and available_quotas shows it; - /a2a has no quota, so the calls show in the metrics only. How it works: - summarize_chunk() and recursively_resummarize() take a count_call callback and report each LLM call to it, with the model and the usage the provider reported. They do so as soon as the response arrived and before they look at it, so a call that returned no text is reported as well. - apply_compaction() takes the path of the endpoint and a charge callback, and counts the calls with SummarizationCalls (new module utils/compaction_usage.py). For each call it records ls_llm_token_sent_total, ls_llm_token_received_total and ls_llm_calls_total under that endpoint, calls charge, and adds the usage up. The sum is returned as CompactionResult.summarization_usage. Without an endpoint path no metrics are recorded, without charge nobody is charged. The counter has a module of its own because conversation_compaction.py is close to the 1000 lines pylint allows once the open pull requests on it are merged. - The endpoints with a quota pass consume_summarization_tokens(), bound to the user, as charge. A call is therefore charged when it returned, and not together with the turn. The summary is written before the model is asked for the answer and is kept whatever becomes of the turn, so the call is charged also when the turn is blocked by a shield, fails or is interrupted, and when compaction itself fails after the call (the marker cannot be written, the fold fails). The next request is served from the stored summary and does not pay again. - /v1/query and /v1/streaming_query add the summarization usage to the usage of the turn in what they report to the client and set on the request span. The turn summary itself is left as it is. TokenCounter gained __add__ for the sum; it returns a new counter. - /a2a passes its endpoint path (new constant ENDPOINT_PATH_A2A) and no charge. Known limits: - A call that fails, or is cancelled before its response arrived, has no usage to count, whatever the provider bills for it. - A call the provider reports no usage for is counted as a call and charges nothing. Tests: - tests/integration/endpoints/test_compaction_token_usage.py runs the four endpoints against a real quota limiter on SQLite and the real summarization path, answered by the mocked OGX client. It reads what the limiter holds after the request, what the client is told, and the Prometheus registry. It covers a turn that summarizes, one that also folds (against a real SQLite conversation cache), a turn served from a stored summary, a short conversation, and a turn that is blocked, fails or is interrupted after the summarization. 14 of the 19 tests fail without the endpoint changes; the other five are the controls that nothing extra is charged. - tests/unit/utils/test_summarization_usage.py covers what the two calls report, and what a request records, charges and is told, including a call that returned no text, a marker that cannot be written, a fold that fails and a fold that cannot be stored. - tests/unit/utils/test_query.py covers consume_summarization_tokens(). Documentation: the user guide gets a section on token usage and quota and the FAQ answers are made exact per endpoint; the design document describes who pays for the calls; the statements about failed and interrupted requests in ARCHITECTURE.md and query_endpoint.md now name the exception; the open question in the guardrails design document is updated.
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: lightspeed-core/lightspeed-stack/.coderabbit.yaml Review profile: ASSERTIVE Plan: Advanced Run ID: 📒 Files selected for processing (35)
💤 Files with no reviewable changes (1)
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review. 📜 Recent review details⏰ Context from checks skipped due to timeout. (22)
🧰 Additional context used📓 Path-based instructions (2)Source excerpt: **Required**: Use pytest for all unit and integration tests Source excerpt: **Forbidden**: Do not use unittest - pytest is the standard for this project📄 CodeRabbit inference engine (AGENTS.md) Files:
Source excerpt: Check `constants.py` for shared constants before defining new ones Source excerpt: Central `constants.py` for shared constants with descriptive comments📄 CodeRabbit inference engine (AGENTS.md) Files:
🧠 Learnings (1)📚 Learning: 2026-06-24T13:45:37.249ZApplied to files:
🪛 ast-grep (0.45.3)tests/integration/endpoints/_compaction_helpers.py[warning] 274-274: Do not make http calls without encryption (requests-http) [info] 305-305: use jsonify instead of json.dumps for JSON output (use-jsonify) 🪛 LanguageTooldocs/user_doc/conversation_compaction.md[locale-violation] ~170-~170: In American English, ‘afterward’ is the preferred variant. ‘Afterwards’ is more commonly used in British English and other dialects. (AFTERWARDS_US) docs/design/conversation-compaction/conversation-compaction.md[style] ~388-~388: ‘with success’ might be wordy. Consider a shorter alternative. (EN_WORDINESS_PREMIUM_WITH_SUCCESS) [style] ~389-~389: ‘with success’ might be wordy. Consider a shorter alternative. (EN_WORDINESS_PREMIUM_WITH_SUCCESS) 🔇 Additional comments (34)
WalkthroughCompaction now reports and charges summarization and fold usage. Endpoint handlers add or preserve that usage according to endpoint behavior. ChangesCompaction accounting
Turn persistence
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant Endpoint
participant apply_compaction
participant summarize_chunk
participant SummarizationCalls
participant consume_summarization_tokens
Endpoint->>apply_compaction: request compaction with endpoint and charge callback
apply_compaction->>summarize_chunk: summarize or fold items
summarize_chunk->>SummarizationCalls: report provider usage after response
SummarizationCalls->>consume_summarization_tokens: charge when callback is configured
SummarizationCalls-->>apply_compaction: accumulated usage
apply_compaction-->>Endpoint: CompactionResult with summarization_usage
Suggested reviewers: Merge Risk: ⚪ Minimal · up to Conversation-summarization calls are now counted in quotas, token totals, and metrics. Conversation turns are stored once per request across endpoints. No concrete defects were identified, so the change appears ready to merge. Security Architecture ReviewSecurity architecture risk: 🟡 Moderate · up to Summarization calls are now charged to the authenticated user, improving accountability. A request checks quota before those calls, however, and no second check is visible before it starts the answer after summarization has consumed the remaining balance. The practical limit depends on how the quota backend handles deductions. Retained concerns
Security review detailsSecurity Blast Radius
Security Findings and Attack Paths
Trust Boundaries and Controls
Resilience and Maintainability Implications
Hardening Proposals
🚥 Pre-merge checks | ✅ 7✅ Passed checks (7 passed)
✨ Finishing Touches 💡 1🧪 Generate unit tests (beta)
✨ Simplify code
🛠️ Fix failing CI checks 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Description
LCORE-3910. Compaction summarizes older turns with an LLM call of its own, and folds the summaries with another. The provider bills both. Their usage was discarded, so the calls consumed no quota, were missing from the
input_tokens/output_tokensthe client is told, and were missing from the token and call metrics. A user could go over the quota without it showing, by more the longer the conversation was. The user guide said the opposite.What was decided on the ticket.
/v1/queryinput_tokens/output_tokensinclude the summarization calls/v1/query/v1/streaming_queryendevent/v1/streaming_query/v1/responsesusageas the provider reported it, the answer alone;available_quotasshows the quota left after both/v1/responses/a2a/a2aNo new response fields.
/v1/responseskeeps theusageobject untouched to stay a drop-in for clients of the OpenAI Responses API.How it works.
summarize_chunk()andrecursively_resummarize()take acount_callcallback and report each LLM call to it, with the model and the usage the provider reported. They do so as soon as the response arrived and before they look at it, so a call that returned no text is reported as well.apply_compaction()takes the path of the endpoint and achargecallback. For each call it recordsls_llm_token_sent_total,ls_llm_token_received_totalandls_llm_calls_totalunder that endpoint, callscharge, and adds the usage up. The sum is returned asCompactionResult.summarization_usage.consume_summarization_tokens(), bound to the user, ascharge. A call is therefore charged when it returned, and not together with the turn. The summary is written before the model is asked for the answer and is kept whatever becomes of the turn. So the call is charged also when the turn is blocked by a shield, fails or is interrupted, and when compaction itself fails after the call (the marker cannot be written, the fold fails). The next request does not pay for the summary again./v1/queryand/v1/streaming_queryadd the summarization usage to the usage of the turn in what they report to the client and set on the request span.TokenCountergained__add__for the sum; it returns a new counter.ConversationSummaryand the cache are not changed.Known limits.
Seen while testing, not changed here (both are on
mainwithout this PR, see the "before" run below):/v1/streaming_querycounts two LLM calls per request inls_llm_calls_total, and non-streaming/v1/responsescounts the tokens of the answer twice in the token metrics.src/utils/compaction.pycount_callsrc/utils/compaction_usage.pysrc/utils/conversation_compaction.pysummarization_usagesrc/utils/query.pyconsume_summarization_tokens()src/utils/token_counter.pyTokenCounter.__add__src/app/endpoints/query.py,streaming_query.py,responses.py,src/utils/agents/streaming.pysrc/app/endpoints/a2a.py,src/constants.pytests/integration/endpoints/test_compaction_token_usage.pytests/unit/utils/test_summarization_usage.py,tests/unit/utils/test_query.pydocs/user_doc/conversation_compaction.md,docs/design/...,docs/devel_doc/query_endpoint.md,docs/devel_doc/ARCHITECTURE.mdType of change
pyproject.toml+uv.lock]requirements.*.txtfor Konflux]Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
1. Manual verification
Start the service in library mode with compaction enabled (
context_windows: openai/gpt-4o-mini: 2000,threshold_ratio: 0.1,token_floor: 100,buffer_turns: 1) and a user quota limiter on SQLite (initial_quota: 1000000). Send the five queries of the e2e compaction scenario through each endpoint. After every request read the quota from the limiter's database and compare it with what the client was told. Expected: on a request that summarizes, the quota goes down by the answer plus the summarization;/v1/queryand/v1/streaming_queryreport both,/v1/responsesreports the answer.Before, on the branch this one starts from (requests 3 and 5 summarize; nothing but the answer is counted):
After, on this branch:
The service log of request 3 on
/v1/queryshows the two charges, the summarization first:148 + 398 = 546 and 95 + 10 = 105, which is what the client was told. The two runs are separate model calls, so the numbers before and after are not comparable token for token. The token usage history holds 7722 input and 1406 output tokens at the end, 9128 together, and the limiter holds 1000000 - 9128 = 990872.
Two things in the output are on
mainas well and are not changed here:/v1/streaming_querycounts two LLM calls per request (12 = 2 x 5 + 2), and/v1/responseswithout a stream counts the tokens of the answer twice (3800 = 2 x 1465 + 870).2. Tests of the change
Expected: all pass. Without the endpoint changes (compaction layer kept) 14 of the 19 integration tests fail; the five that pass are the controls that nothing extra is charged.
3. Full suites
The one failure is
tests/integration/container_lifecycle/test_container_lifecycle.py::TestContainerLifecycle::test_container_lifecycle, which needs a container runtime and fails onmainthe same way.Linters: black, ruff, pylint (10.00/10), pyright (0 errors), pydocstyle pass. mypy on
srcreports the one error that is onmain(src/models/config.py,CustomProfile).Summary by CodeRabbit