From d7500e420c72936c5ff4e3f2be7e0e9b96b96dec Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Tue, 6 Oct 2026 14:59:12 +0200 Subject: [PATCH 1/3] LCORE-3910: count the summarization calls in the LLM metrics The summarization and fold calls of compaction report their usage since the previous commit, and nobody listened. The calls were missing from ls_llm_token_sent_total, ls_llm_token_received_total and ls_llm_calls_total, although the provider bills them. apply_compaction() now takes the path of the endpoint it runs for and counts every call with SummarizationCalls (new module utils/compaction_usage.py): the tokens and the call are recorded under that endpoint as soon as the call returned, so also when the request fails afterwards. /v1/query, /v1/streaming_query, /v1/responses and A2A pass their path; A2A gets the new constant ENDPOINT_PATH_A2A. A2A has no quota and tells the client no token counts, so it is complete with this commit. Two things ship here without a production caller, because they belong to the same counter and conversation_compaction.py is then edited by one commit only: - the charge callback of apply_compaction(), called once per call; the next commit passes the quota of the user; - CompactionResult.summarization_usage, the usage of the calls added up with the new TokenCounter.__add__; the commit after that reports it to the client. The counter has a module of its own because conversation_compaction.py is at 928 of the 1000 lines pylint allows. Tests: tests/unit/utils/test_compaction_usage.py runs the real summarization and fold calls against a mocked client and checks what is recorded, what charge is called with and what the result holds: for a request that summarizes, one that also folds, one served from an earlier summary, and one that fails after the call because the marker cannot be written. The endpoint tests of /v1/query, /v1/streaming_query and A2A assert the path handed to compaction; the one for /v1/responses comes with the next commit. --- src/app/endpoints/a2a.py | 3 +- src/app/endpoints/query.py | 1 + src/app/endpoints/responses.py | 1 + src/app/endpoints/streaming_query.py | 1 + src/constants.py | 1 + src/utils/compaction_usage.py | 65 ++++++ src/utils/conversation_compaction.py | 50 ++++- src/utils/token_counter.py | 17 ++ tests/unit/app/endpoints/test_a2a.py | 3 + tests/unit/app/endpoints/test_query.py | 4 +- .../app/endpoints/test_streaming_query.py | 7 +- tests/unit/utils/test_compaction_usage.py | 199 ++++++++++++++++++ 12 files changed, 343 insertions(+), 9 deletions(-) create mode 100644 src/utils/compaction_usage.py create mode 100644 tests/unit/utils/test_compaction_usage.py diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index c64718dc2..6d6fbe4e9 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -57,7 +57,7 @@ from authorization.middleware import authorize from client.ogx import AsyncOgxClientHolder from configuration import configuration -from constants import MEDIA_TYPE_EVENT_STREAM +from constants import ENDPOINT_PATH_A2A, MEDIA_TYPE_EVENT_STREAM from log import get_logger from models.api.requests import QueryRequest from models.common.responses.responses_api_params import ResponsesApiParams @@ -241,6 +241,7 @@ async def _compact_a2a_request( responses_params, configuration.inference, configuration.compaction, + endpoint_path=ENDPOINT_PATH_A2A, ) return compaction.params, PendingTurn.for_request( client, compaction.params, compaction.original_input diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 538bd5b1e..1883cf1f4 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -242,6 +242,7 @@ async def _handle_query_with_tracing( cache=configured_conversation_cache(), user_id=user_id, skip_user_id_check=_skip_userid_check, + endpoint_path=endpoint_path, ) responses_params = compaction.params diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 2a39fbe3f..542f70ef3 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -737,6 +737,7 @@ async def handle_responses_with_tracing( # pylint: disable=too-many-locals cache=configured_conversation_cache(), user_id=user_id, skip_user_id_check=skip_userid_check, + endpoint_path=endpoint_path, ) api_params = compaction.params if compaction.compacted: diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 3c8a11de7..0c6ce13f8 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -434,6 +434,7 @@ async def generate_response_with_compaction( cache=configured_conversation_cache(), user_id=context.user_id, skip_user_id_check=context.skip_userid_check, + endpoint_path=endpoint_path, ): if isinstance(item, CompactionStartedEvent): yield stream_compaction_event(context.conversation_id) diff --git a/src/constants.py b/src/constants.py index 041830fba..3ea554a3a 100644 --- a/src/constants.py +++ b/src/constants.py @@ -344,6 +344,7 @@ ENDPOINT_PATH_QUERY: Final[str] = "/v1/query" ENDPOINT_PATH_STREAMING_QUERY: Final[str] = "/v1/streaming_query" ENDPOINT_PATH_RESPONSES: Final[str] = "/v1/responses" +ENDPOINT_PATH_A2A: Final[str] = "/a2a" # Context window (tokens) conversation compaction assumes for a model when # neither its catalog metadata nor automatic resolution gives one. diff --git a/src/utils/compaction_usage.py b/src/utils/compaction_usage.py new file mode 100644 index 000000000..58544983b --- /dev/null +++ b/src/utils/compaction_usage.py @@ -0,0 +1,65 @@ +"""Accounting for the LLM calls conversation compaction makes (LCORE-3910). + +Compaction summarizes older turns with an LLM call of its own, and folds the +summaries with another. The provider bills both, so they are counted: in the +LLM metrics, in the quota of the user, and in the token counts the client is +told. :class:`SummarizationCalls` does the first two and adds the usage up for +the third. +""" + +from dataclasses import dataclass, field +from typing import Optional + +from metrics import recording +from utils.compaction import CallCounter +from utils.query import extract_provider_and_model_from_model_id +from utils.token_counter import TokenCounter + + +@dataclass +class SummarizationCalls: + """The LLM calls compaction makes for one request (LCORE-3910). + + The provider bills them. So each call is recorded in the LLM metrics under + the endpoint that triggered it, and handed to ``charge``. Both happen as + soon as the response of the call arrived, and do not depend on what + becomes of the request afterwards. + + Attributes: + endpoint_path: Path of the endpoint serving the request, the label of + the metrics. No metrics are recorded without it. + charge: Charges one call to whoever pays for the request. Nobody is + charged without it. + usage: The usage of the calls counted so far, added up. + + """ + + endpoint_path: Optional[str] = None + charge: Optional[CallCounter] = None + usage: TokenCounter = field(default_factory=TokenCounter) + + def count(self, model: str, usage: TokenCounter) -> None: + """Record one call in the metrics, charge it, and add it to the total. + + Parameters: + model: Fully-qualified identifier of the model that was called. + usage: The token usage the provider reported for the call. + + Raises: + HTTPException: When charging the call fails. + + """ + self.usage = self.usage + usage + if self.endpoint_path is not None: + provider_id, model_id = extract_provider_and_model_from_model_id(model) + if usage.input_tokens or usage.output_tokens: + recording.record_llm_token_usage( + provider_id, + model_id, + usage.input_tokens, + usage.output_tokens, + self.endpoint_path, + ) + recording.record_llm_call(provider_id, model_id, self.endpoint_path) + if self.charge is not None: + self.charge(model, usage) diff --git a/src/utils/conversation_compaction.py b/src/utils/conversation_compaction.py index af4ba903e..cc1f95fad 100644 --- a/src/utils/conversation_compaction.py +++ b/src/utils/conversation_compaction.py @@ -46,7 +46,7 @@ import re from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Optional, cast from fastapi import HTTPException @@ -64,15 +64,18 @@ from models.compaction import ConversationSummary from models.config import CompactionConfiguration, InferenceConfiguration from utils.compaction import ( + CallCounter, partition_conversation, recursively_resummarize, summarize_chunk, ) +from utils.compaction_usage import SummarizationCalls from utils.context_window import resolve_context_window from utils.conversations import ( build_add_items_request, get_all_conversation_items, ) +from utils.token_counter import TokenCounter from utils.token_estimator import ( DEFAULT_ENCODING_NAME, estimate_conversation_tokens, @@ -190,12 +193,19 @@ class CompactionResult: caller hands it to ``utils.pending_turn.PendingTurn``, which stores this plus the LLM output in the conversation, since the ``conversation`` parameter is no longer passed to OGX. + summarization_usage: Token usage of the LLM calls compaction made for + this request: the summarization of older turns and the fold of + the summaries. Empty when the request made neither. The calls are + recorded and charged by the time the result is returned; the + usage is here for the endpoints that report token counts to the + client (LCORE-3910). """ params: ResponsesApiParams compacted: bool original_input: Optional[ResponseInput] = None + summarization_usage: TokenCounter = field(default_factory=TokenCounter) @property def context_status(self) -> ContextStatus: @@ -649,12 +659,15 @@ async def _maybe_persist_fold( # pylint: disable=too-many-arguments,too-many-po context_window: int, threshold_ratio: float, encoding_name: str, + count_call: Optional[CallCounter] = None, ) -> tuple[list[str], list[ConversationSummary]]: """Run the recursive fold (R3) if persisted summaries crossed the threshold. Requires a persisting cache (marker-only conversations keep additive chunks). Returns the (possibly updated) ``(summaries, cached_summaries)``; on a cache - write failure the unfolded values are returned unchanged. + write failure the unfolded values are returned unchanged. ``count_call`` is + told about the fold call once it returned, also when its result cannot be + stored afterwards (LCORE-3910). """ if cache is None or len(cached_summaries) < 2: return summaries, cached_summaries @@ -668,7 +681,7 @@ async def _maybe_persist_fold( # pylint: disable=too-many-arguments,too-many-po conversation_id, ) folded = await recursively_resummarize( - client, model, cached_summaries, encoding_name + client, model, cached_summaries, encoding_name, count_call=count_call ) try: cache.replace_summaries(user_id, conversation_id, folded, skip_user_id_check) @@ -683,6 +696,7 @@ def _compacted_result( summaries: list[str], recent_items: list[Any], original_input: ResponseInput, + usage: TokenCounter, ) -> CompactionResult: """Build the CompactionResult for compacted mode (explicit input + omit_conversation).""" explicit_input = _build_explicit_input(summaries, recent_items, original_input) @@ -690,7 +704,10 @@ def _compacted_result( update={"input": explicit_input, "omit_conversation": True} ) return CompactionResult( - compacted_params, compacted=True, original_input=original_input + compacted_params, + compacted=True, + original_input=original_input, + summarization_usage=usage, ) @@ -704,6 +721,8 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit cache: Optional[Cache] = None, user_id: str = "", skip_user_id_check: bool = False, + endpoint_path: Optional[str] = None, + charge: Optional[CallCounter] = None, ) -> AsyncIterator[Any]: """Apply conversation compaction to a prepared request, yielding the result. @@ -730,10 +749,19 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit falls back to marker-only summaries with no folding. user_id: User identifier for cache reads/writes. skip_user_id_check: Whether to bypass the cache's user_id validation. + endpoint_path: Path of the endpoint serving the request. The LLM calls + compaction makes are recorded in the metrics under it; without it + they are not recorded. + charge: Called with the model and the token usage of each LLM call + compaction makes, as soon as the call returned. The endpoints + with a quota charge the user here (LCORE-3910). Yields: Zero or more CompactionStartedEvent, then exactly one CompactionResult. + Raises: + HTTPException: When ``charge`` fails. + """ if not compaction_config.enabled: # ``enabled: false`` is a full off-switch: the request passes through @@ -749,6 +777,7 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit conversation_id = params.conversation model = params.model original_input = params.input + calls = SummarizationCalls(endpoint_path, charge) async with _conversation_lock(conversation_id): items = await get_all_conversation_items(client, conversation_id) @@ -781,6 +810,7 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit old_items, summarized_through_turn=_covered_item_count(items, keep_items), encoding_name=encoding_name, + count_call=calls.count, ) await _persist_new_summary_chunk( client, @@ -805,6 +835,7 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit context_window, compaction_config.threshold_ratio, encoding_name, + calls.count, ) if not summaries: @@ -815,7 +846,9 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit # Compacted mode: lightspeed owns the context. Build explicit input and # stop passing the conversation parameter to inference. - yield _compacted_result(params, summaries, recent_items, original_input) + yield _compacted_result( + params, summaries, recent_items, original_input, calls.usage + ) async def apply_compaction_blocking( # pylint: disable=too-many-arguments,too-many-positional-arguments @@ -827,12 +860,15 @@ async def apply_compaction_blocking( # pylint: disable=too-many-arguments,too-m cache: Optional[Cache] = None, user_id: str = "", skip_user_id_check: bool = False, + endpoint_path: Optional[str] = None, + charge: Optional[CallCounter] = None, ) -> CompactionResult: """Non-streaming wrapper around :func:`apply_compaction`. Drains the generator with event emission disabled and returns the final :class:`CompactionResult`. See :func:`apply_compaction` for the ``cache`` / - ``user_id`` / ``skip_user_id_check`` parameters. + ``user_id`` / ``skip_user_id_check`` / ``endpoint_path`` / ``charge`` + parameters. """ result: Optional[CompactionResult] = None async for item in apply_compaction( @@ -845,6 +881,8 @@ async def apply_compaction_blocking( # pylint: disable=too-many-arguments,too-m cache=cache, user_id=user_id, skip_user_id_check=skip_user_id_check, + endpoint_path=endpoint_path, + charge=charge, ): if isinstance(item, CompactionResult): result = item diff --git a/src/utils/token_counter.py b/src/utils/token_counter.py index 2d49c7a03..b7ff659f4 100644 --- a/src/utils/token_counter.py +++ b/src/utils/token_counter.py @@ -24,6 +24,23 @@ class TokenCounter: input_tokens_counted: int = 0 llm_calls: int = 0 + def __add__(self, other: "TokenCounter") -> "TokenCounter": + """Return the usage of two sets of LLM calls taken together. + + Parameters: + other: The token counter to add. + + Returns: + A new counter holding the sums; neither operand is changed. + + """ + return TokenCounter( + input_tokens=self.input_tokens + other.input_tokens, + output_tokens=self.output_tokens + other.output_tokens, + input_tokens_counted=self.input_tokens_counted + other.input_tokens_counted, + llm_calls=self.llm_calls + other.llm_calls, + ) + def __str__(self) -> str: """ Return a human-readable summary of the token usage stored in this TokenCounter. diff --git a/tests/unit/app/endpoints/test_a2a.py b/tests/unit/app/endpoints/test_a2a.py index e87afa25f..2f8da0e53 100644 --- a/tests/unit/app/endpoints/test_a2a.py +++ b/tests/unit/app/endpoints/test_a2a.py @@ -1002,6 +1002,9 @@ async def _event_stream() -> Any: ) apply.assert_awaited_once() + # The summarization calls are recorded under /a2a. There is no quota + # on this endpoint, so no charge is passed (LCORE-3910). + assert apply.await_args.kwargs == {"endpoint_path": "/a2a"} @pytest.mark.asyncio async def test_convert_stream_artifact_carries_conversation_id( diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index cb8fefd32..a9df4a2e9 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -231,7 +231,7 @@ async def test_query_reports_context_status( cast("ResponsesApiParams", mock_responses_params), compacted=compacted, ) - mocker.patch( + apply = mocker.patch( "app.endpoints.query.apply_compaction_blocking", new=mocker.AsyncMock(return_value=compaction_result), ) @@ -259,6 +259,8 @@ async def test_query_reports_context_status( assert isinstance(response, QueryResponse) assert response.context_status == expected_status + # The summarization calls are recorded under this endpoint (LCORE-3910). + assert apply.await_args.kwargs["endpoint_path"] == "/v1/query" @pytest.mark.asyncio async def test_query_merges_inline_and_tool_rag_chunks_and_documents( diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index bb7a6b86c..aaec1878a 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -927,9 +927,12 @@ async def test_threads_context_status_to_agent_response( compaction_result = CompactionResult(responses_params, compacted=compacted) + compaction_kwargs: dict[str, Any] = {} + async def fake_apply_compaction( - *_args: Any, **_kwargs: Any + *_args: Any, **kwargs: Any ) -> AsyncIterator[CompactionResult]: + compaction_kwargs.update(kwargs) yield compaction_result mocker.patch( @@ -983,3 +986,5 @@ async def fake_generate_agent_response( assert events # the start event plus the delegated events assert captured_kwargs["context_status"] == expected_status + # The summarization calls are recorded under this endpoint (LCORE-3910). + assert compaction_kwargs["endpoint_path"] == "/v1/streaming_query" diff --git a/tests/unit/utils/test_compaction_usage.py b/tests/unit/utils/test_compaction_usage.py new file mode 100644 index 000000000..0303d19d3 --- /dev/null +++ b/tests/unit/utils/test_compaction_usage.py @@ -0,0 +1,199 @@ +"""Unit tests for the accounting of the summarization calls (LCORE-3910). + +Compaction makes LLM calls of its own: one to summarize older turns, and one +to fold the summaries when they grow too large. The provider bills them, so +each is counted as soon as its response arrived: recorded in the LLM metrics, +handed to ``charge``, and added to the usage the request is told. +""" + +# pylint: disable=protected-access + +from typing import Any + +import pytest +from ogx_api.openai_responses import OpenAIResponseMessage +from pytest_mock import MockerFixture + +from models.common.responses.responses_api_params import ResponsesApiParams +from models.compaction import ConversationSummary +from models.config import CompactionConfiguration, InferenceConfiguration +from utils import conversation_compaction as cc +from utils.token_counter import TokenCounter + +MODEL = "openai/gpt-4o-mini" +ENDPOINT = "/v1/query" + +SUMMARIZATION = TokenCounter(input_tokens=640, output_tokens=72, llm_calls=1) +FOLD = TokenCounter(input_tokens=310, output_tokens=45, llm_calls=1) + + +def _response(mocker: MockerFixture, text: str, usage: TokenCounter) -> Any: + """Build the result of ``client.responses.create`` with the given usage.""" + return mocker.Mock( + output=[mocker.Mock(content=[mocker.Mock(text=text)])], + usage=mocker.Mock( + input_tokens=usage.input_tokens, output_tokens=usage.output_tokens + ), + ) + + +def _client(mocker: MockerFixture) -> Any: + """Build an OGX client that answers the summarization, then the fold call.""" + client = mocker.AsyncMock() + client.responses.create.side_effect = [ + _response(mocker, "condensed", SUMMARIZATION), + _response(mocker, "folded", FOLD), + ] + return client + + +LONG_CONVERSATION = [ + OpenAIResponseMessage(role="user", content="q1 " * 50), + OpenAIResponseMessage(role="assistant", content="a1 " * 50), +] +"""One stored turn, long enough for the next request to summarize it.""" + + +@pytest.fixture(name="recorded") +def recorded_fixture(mocker: MockerFixture) -> Any: + """Replace the metric recorders and return them.""" + return mocker.patch("utils.compaction_usage.recording") + + +async def _compact( # pylint: disable=too-many-arguments + mocker: MockerFixture, + items: list[Any], + *, + charge: Any = None, + cache: Any = None, + write_marker: Any = None, + threshold_ratio: float = 0.1, +) -> cc.CompactionResult: + """Apply compaction to a follow-up request on the stored items.""" + mocker.patch.object( + cc, "get_all_conversation_items", mocker.AsyncMock(return_value=items) + ) + mocker.patch.object(cc, "_write_summary_marker", write_marker or mocker.AsyncMock()) + return await cc.apply_compaction_blocking( + client=_client(mocker), + params=ResponsesApiParams( + input="follow-up", + model=MODEL, + conversation="conv_abc123", + instructions="system prompt", + store=True, + stream=False, + ), + inference_config=InferenceConfiguration(context_windows={MODEL: 50}), + compaction_config=CompactionConfiguration( + enabled=True, + threshold_ratio=threshold_ratio, + token_floor=0, + buffer_turns=0, + buffer_max_ratio=0.3, + ), + cache=cache, + user_id="u1", + endpoint_path=ENDPOINT, + charge=charge, + ) + + +def test_token_counters_add_up() -> None: + """The sum counts the tokens and the calls of both, and changes neither.""" + turn = TokenCounter(100, 50, input_tokens_counted=90, llm_calls=1) + + assert turn + SUMMARIZATION == TokenCounter( + input_tokens=740, output_tokens=122, input_tokens_counted=90, llm_calls=2 + ) + assert SUMMARIZATION + turn == turn + SUMMARIZATION + assert turn == TokenCounter(100, 50, input_tokens_counted=90, llm_calls=1) + assert SUMMARIZATION == TokenCounter(640, 72, llm_calls=1) + + +@pytest.mark.asyncio +async def test_request_that_summarizes(mocker: MockerFixture, recorded: Any) -> None: + """The call is recorded in the metrics, charged, and the request is told.""" + charge = mocker.Mock() + + result = await _compact(mocker, LONG_CONVERSATION, charge=charge) + + assert result.compacted is True + assert result.summarization_usage == SUMMARIZATION + charge.assert_called_once_with(MODEL, SUMMARIZATION) + recorded.record_llm_token_usage.assert_called_once_with( + "openai", "gpt-4o-mini", 640, 72, ENDPOINT + ) + recorded.record_llm_call.assert_called_once_with("openai", "gpt-4o-mini", ENDPOINT) + + +@pytest.mark.asyncio +async def test_request_served_from_an_earlier_summary( + mocker: MockerFixture, recorded: Any +) -> None: + """A request that makes no summarization call counts nothing.""" + charge = mocker.Mock() + marker = OpenAIResponseMessage( + role="user", content=f"{cc.MARKER_SENTINEL} an earlier summary" + ) + + result = await _compact(mocker, [marker], charge=charge, threshold_ratio=0.9) + + assert result.compacted is True + assert result.summarization_usage == TokenCounter() + charge.assert_not_called() + recorded.record_llm_token_usage.assert_not_called() + recorded.record_llm_call.assert_not_called() + + +@pytest.mark.asyncio +async def test_call_is_counted_when_the_marker_cannot_be_written( + mocker: MockerFixture, recorded: Any +) -> None: + """The request fails after the call was made; the call is counted all the same.""" + charge = mocker.Mock() + + with pytest.raises(RuntimeError, match="the store is down"): + await _compact( + mocker, + LONG_CONVERSATION, + charge=charge, + write_marker=mocker.AsyncMock( + side_effect=RuntimeError("the store is down") + ), + ) + + charge.assert_called_once_with(MODEL, SUMMARIZATION) + recorded.record_llm_call.assert_called_once_with("openai", "gpt-4o-mini", ENDPOINT) + + +@pytest.mark.asyncio +async def test_summarization_and_fold_are_both_counted( + mocker: MockerFixture, recorded: Any +) -> None: + """A request that summarizes and then folds counts each of the two calls.""" + earlier = ConversationSummary( + summary_text="earlier", + summarized_through_turn=2, + token_count=20, + created_at="2026-09-28T00:00:00Z", + model_used=MODEL, + ) + cache = mocker.Mock() + cache.get_summaries.return_value = [earlier, earlier] + charge = mocker.Mock() + + result = await _compact(mocker, LONG_CONVERSATION, charge=charge, cache=cache) + + assert result.summarization_usage == TokenCounter( + input_tokens=950, output_tokens=117, llm_calls=2 + ) + assert charge.call_args_list == [ + mocker.call(MODEL, SUMMARIZATION), + mocker.call(MODEL, FOLD), + ] + assert recorded.record_llm_token_usage.call_args_list == [ + mocker.call("openai", "gpt-4o-mini", 640, 72, ENDPOINT), + mocker.call("openai", "gpt-4o-mini", 310, 45, ENDPOINT), + ] + assert recorded.record_llm_call.call_count == 2 From a5d5be03a94f3e206b66cf9be1f8ce7ac8176234 Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Tue, 6 Oct 2026 15:30:31 +0200 Subject: [PATCH 2/3] LCORE-3910: charge the summarization calls to the user's quota The summarization and fold calls of compaction consumed no quota. Quota is consumed from the usage of the turn, and these calls are made before the turn and outside of it. A user could go over the quota without it showing, by more the longer the conversation was. /v1/query, /v1/streaming_query and /v1/responses now pass consume_summarization_tokens(), bound to the user of the request, as the charge callback of compaction. A call is therefore charged when it returned and not together with the turn: the summary is stored 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 then fails. A later request is served from the stored summary and makes no call. Nothing is consumed for a call the provider reported no usage for. /v1/responses keeps the usage object of the response as OGX reported it, to stay a drop-in for clients of the OpenAI Responses API. Its available_quotas is read from the limiters after the answer, so it includes the charge. A2A has no quota and passes no callback. Until the next commit the input_tokens and output_tokens of /v1/query and /v1/streaming_query still cover the answer alone, so available_quotas falls by more than they add up to. Tests: - tests/integration/endpoints/test_compaction_token_usage.py runs /v1/query against a real quota limiter on SQLite and the real summarization call, answered by the mocked OGX client: the balance falls by the turn and the summarization, available_quotas says so, and a turn that fails after the summarization pays for the summarization alone. Both cases fail without the callback. - The endpoint tests of /v1/query, /v1/streaming_query and /v1/responses assert the callback and the user it is bound to; the one for /v1/responses is new and asserts the endpoint path as well. The other two endpoints are covered at the call site only. - tests/unit/utils/test_query.py covers consume_summarization_tokens(). --- src/app/endpoints/query.py | 3 + src/app/endpoints/responses.py | 6 + src/app/endpoints/streaming_query.py | 3 + src/utils/query.py | 28 +++ .../endpoints/test_compaction_token_usage.py | 181 ++++++++++++++++++ tests/unit/app/endpoints/test_query.py | 7 +- tests/unit/app/endpoints/test_responses.py | 50 +++++ .../app/endpoints/test_streaming_query.py | 7 +- tests/unit/utils/test_query.py | 36 ++++ 9 files changed, 319 insertions(+), 2 deletions(-) create mode 100644 tests/integration/endpoints/test_compaction_token_usage.py diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 1883cf1f4..3bb457316 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -1,6 +1,7 @@ """Handler for REST API call to provide answer to query using Response API.""" import datetime +from functools import partial from typing import Annotated from fastapi import APIRouter, Depends, Request @@ -50,6 +51,7 @@ from utils.pending_turn import pending_turn from utils.query import ( consume_query_tokens, + consume_summarization_tokens, store_query_results, validate_attachments_metadata, validate_model_provider_override, @@ -243,6 +245,7 @@ async def _handle_query_with_tracing( user_id=user_id, skip_user_id_check=_skip_userid_check, endpoint_path=endpoint_path, + charge=partial(consume_summarization_tokens, user_id), ) responses_params = compaction.params diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 542f70ef3..45aba0800 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -6,6 +6,7 @@ import time from collections.abc import AsyncIterator, Sequence from datetime import UTC, datetime +from functools import partial from typing import TYPE_CHECKING, Annotated, Any, Final, NoReturn, Optional, cast from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request @@ -90,6 +91,7 @@ from utils.prompts import get_system_prompt from utils.query import ( consume_query_tokens, + consume_summarization_tokens, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, is_context_length_error, @@ -738,6 +740,10 @@ async def handle_responses_with_tracing( # pylint: disable=too-many-locals user_id=user_id, skip_user_id_check=skip_userid_check, endpoint_path=endpoint_path, + # The summarization calls are charged to the user. The usage + # object of the response stays as OGX reported it, so it covers + # the answer alone. + charge=partial(consume_summarization_tokens, user_id), ) api_params = compaction.params if compaction.compacted: diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 0c6ce13f8..1ab629f5d 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -3,6 +3,7 @@ import asyncio import datetime from collections.abc import AsyncIterator +from functools import partial from typing import TYPE_CHECKING, Annotated, Optional from fastapi import APIRouter, Depends, HTTPException, Request @@ -69,6 +70,7 @@ ) from utils.pending_turn import PendingTurn from utils.query import ( + consume_summarization_tokens, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, is_context_length_error, @@ -435,6 +437,7 @@ async def generate_response_with_compaction( user_id=context.user_id, skip_user_id_check=context.skip_userid_check, endpoint_path=endpoint_path, + charge=partial(consume_summarization_tokens, context.user_id), ): if isinstance(item, CompactionStartedEvent): yield stream_compaction_event(context.conversation_id) diff --git a/src/utils/query.py b/src/utils/query.py index b6df0f12f..dabab8cbc 100644 --- a/src/utils/query.py +++ b/src/utils/query.py @@ -317,6 +317,34 @@ def consume_query_tokens( raise HTTPException(**response.model_dump()) from e +def consume_summarization_tokens( + user_id: str, + model_id: str, + summarization_usage: TokenCounter, +) -> None: + """Consume the tokens of an LLM call compaction made for a request. + + Compaction summarizes older turns with LLM calls of its own, and the + provider bills them. Each is charged as soon as it returned and not with + the turn, so that it is charged also when the turn is blocked, fails or + is interrupted. Nothing is consumed for a call the provider reported no + usage for. + + Parameters: + user_id: The authenticated user ID. + model_id: The full model identifier in "provider/model" format. + summarization_usage: The token usage the provider reported. + + Raises: + HTTPException: On database errors during token consumption. + + """ + if summarization_usage.input_tokens or summarization_usage.output_tokens: + consume_query_tokens( + user_id=user_id, model_id=model_id, token_usage=summarization_usage + ) + + def is_transcripts_enabled() -> bool: """Check if transcripts is enabled. diff --git a/tests/integration/endpoints/test_compaction_token_usage.py b/tests/integration/endpoints/test_compaction_token_usage.py new file mode 100644 index 000000000..7634079da --- /dev/null +++ b/tests/integration/endpoints/test_compaction_token_usage.py @@ -0,0 +1,181 @@ +"""Integration tests for the quota charge of the summarization calls (LCORE-3910). + +Compaction makes an LLM call of its own to summarize older turns. The provider +bills it, so it is charged to the user's quota. It is charged when it is made, +so also when the turn that triggered it fails afterward. + +The quota here is a real limiter on a SQLite database and the summarization +call is the real one, answered by the mocked OGX client. Every test reads what +the limiter holds after the request. +""" + +from collections.abc import Generator +from dataclasses import dataclass +from pathlib import Path + +import pytest +from fastapi import HTTPException, Request +from pytest_mock import AsyncMockType, MockerFixture +from sqlalchemy.orm import Session + +from app.endpoints.query import query_endpoint_handler +from authentication.interface import AuthTuple +from configuration import AppConfig +from models.api.requests import QueryRequest +from models.api.responses.successful.query import QueryResponse +from models.config import ( + QuotaHandlersConfiguration, + QuotaLimiterConfiguration, + SQLiteDatabaseConfiguration, +) +from tests.integration.conftest import ( + InMemoryConversationStore, + make_openai_response_object, + set_query_agent_run, +) +from tests.integration.endpoints._compaction_helpers import ( + CONV_ID_LLAMA, + EXISTING_CONV_ID, + assert_marker_count, + create_existing_conversation, + enable_compaction, + msg, +) + +INITIAL_QUOTA = 100_000 + +# what the provider reports for the turn itself and for the summarization call +TURN_INPUT, TURN_OUTPUT = 100, 50 +SUMMARY_INPUT, SUMMARY_OUTPUT = 640, 72 +TURN = TURN_INPUT + TURN_OUTPUT +SUMMARY = SUMMARY_INPUT + SUMMARY_OUTPUT + + +@pytest.fixture(name="quota") +def quota_fixture( + test_config: AppConfig, tmp_path: Path +) -> Generator[AppConfig, None, None]: + """Give every user a quota, kept by a real limiter on a SQLite database.""" + # pylint: disable=protected-access + assert test_config._configuration is not None + test_config._configuration.quota_handlers = QuotaHandlersConfiguration( + sqlite=SQLiteDatabaseConfiguration(db_path=str(tmp_path / "quota.db")), + limiters=[ + QuotaLimiterConfiguration( + type="user_limiter", + name="user quota", + initial_quota=INITIAL_QUOTA, + quota_increase=0, + period="1 day", + ) + ], + ) + test_config._quota_limiters = [] + yield test_config + for limiter in test_config._quota_limiters: + if limiter.connection is not None: + limiter.connection.close() + test_config._quota_limiters = [] + + +@dataclass +class Scene: + """A user with a quota, on a conversation the next request summarizes. + + Attributes: + config: The configuration, with the quota limiter. + store: The conversation store behind the mocked OGX client. + agent: The mocked agent that answers the turn. + request: The request object the handler gets. + auth: The authenticated user. + + """ + + config: AppConfig + store: InMemoryConversationStore + agent: AsyncMockType + request: Request + auth: AuthTuple + + async def query(self) -> QueryResponse: + """Send a new query on the stored conversation.""" + return await query_endpoint_handler( + request=self.request, + query_request=QueryRequest( + query="What else can you help with?", conversation_id=EXISTING_CONV_ID + ), + auth=self.auth, + mcp_headers={}, + ) + + def quota_left(self) -> int: + """Read the quota the user has left, from the limiter.""" + (limiter,) = self.config.quota_limiters + return limiter.available_quota(self.auth[0]) + + +@pytest.fixture(name="scene") +async def scene_fixture( # pylint: disable=too-many-arguments,too-many-positional-arguments + quota: AppConfig, + mock_ogx_client: AsyncMockType, + mock_conversation_store: InMemoryConversationStore, + mock_query_agent: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, + patch_db_session: Session, + mocker: MockerFixture, +) -> Scene: + """Set the scene: the quota, a long conversation, and what the two LLM calls report.""" + create_existing_conversation(patch_db_session, test_auth[0]) + enable_compaction(quota, context_window=200) + await mock_conversation_store.create( + conversation_id=CONV_ID_LLAMA, + items=[ + msg("user", "question one " * 20), + msg("assistant", "answer one " * 20), + msg("user", "question two " * 20), + msg("assistant", "answer two " * 20), + ], + ) + # The summarization call goes through client.responses.create, the turn + # itself is answered by the agent. + mock_ogx_client.responses.create = mocker.AsyncMock( + return_value=make_openai_response_object( + content="condensed earlier turns", + input_tokens=SUMMARY_INPUT, + output_tokens=SUMMARY_OUTPUT, + ) + ) + set_query_agent_run( + mock_query_agent, mocker, input_tokens=TURN_INPUT, output_tokens=TURN_OUTPUT + ) + return Scene( + quota, mock_conversation_store, mock_query_agent, test_request, test_auth + ) + + +@pytest.mark.asyncio +async def test_turn_that_summarizes_pays_for_the_summarization(scene: Scene) -> None: + """The quota falls by the turn and by the summarization call, and the client is told.""" + response = await scene.query() + + assert_marker_count(scene.store, CONV_ID_LLAMA, 1) + assert response.context_status == "summarized" + assert response.available_quotas == { + "UserQuotaLimiter": INITIAL_QUOTA - TURN - SUMMARY + } + assert scene.quota_left() == INITIAL_QUOTA - TURN - SUMMARY + + +@pytest.mark.asyncio +async def test_failed_turn_that_summarized_pays_for_the_summarization( + scene: Scene, +) -> None: + """The summary is made and kept before the model call that then fails.""" + scene.agent.run.side_effect = RuntimeError("the model is down") + + with pytest.raises(HTTPException): + await scene.query() + + assert_marker_count(scene.store, CONV_ID_LLAMA, 1) + assert scene.quota_left() == INITIAL_QUOTA - SUMMARY diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index a9df4a2e9..89935d10e 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -22,6 +22,7 @@ ) from models.database.conversations import UserConversation from utils.conversation_compaction import CompactionResult +from utils.query import consume_summarization_tokens # User ID must be proper UUID MOCK_AUTH = ( @@ -259,8 +260,12 @@ async def test_query_reports_context_status( assert isinstance(response, QueryResponse) assert response.context_status == expected_status - # The summarization calls are recorded under this endpoint (LCORE-3910). + # The summarization calls are recorded under this endpoint and charged + # to the user of the request (LCORE-3910). assert apply.await_args.kwargs["endpoint_path"] == "/v1/query" + charge = apply.await_args.kwargs["charge"] + assert charge.func is consume_summarization_tokens + assert charge.args == (MOCK_AUTH[0],) @pytest.mark.asyncio async def test_query_merges_inline_and_tool_rag_chunks_and_documents( diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index 934be9f5e..df9c1006c 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -47,6 +47,8 @@ from models.config import Action, ModelContextProtocolServer from models.database.conversations import UserConversation from tests.unit.conftest import mock_async_ogx_client +from utils.conversation_compaction import CompactionResult +from utils.query import consume_summarization_tokens MOCK_AUTH = ( "00000001-0001-0001-0001-000000000001", @@ -495,6 +497,54 @@ async def test_responses_streaming_returns_streaming_response( ) assert response is mock_streaming + @pytest.mark.asyncio + async def test_summarization_calls_are_recorded_and_charged( + self, + dummy_request: Request, + minimal_config: AppConfig, + mocker: MockerFixture, + ) -> None: + """Compaction is told the endpoint and whom to charge (LCORE-3910).""" + minimal_config.compaction.enabled = True + _patch_base(mocker, minimal_config) + _patch_client(mocker) + _patch_resolve_response_context(mocker, conversation=VALID_CONV_ID) + mocker.patch( + f"{MODULE}.select_model_for_responses", + new=mocker.AsyncMock(return_value="provider/model1"), + ) + mocker.patch( + f"{MODULE}.check_model_configured", + new=mocker.AsyncMock(return_value=True), + ) + _patch_rag(mocker) + _patch_moderation(mocker, decision="passed") + mocker.patch(f"{MODULE}.configured_conversation_cache", return_value=None) + mocker.patch( + f"{MODULE}.handle_non_streaming_response", + new=mocker.AsyncMock(return_value=_make_responses_response()), + ) + apply = mocker.patch( + f"{MODULE}.apply_compaction_blocking", + new=mocker.AsyncMock( + side_effect=lambda _client, params, *_args, **_kwargs: ( + CompactionResult(params, compacted=False) + ) + ), + ) + + await responses_endpoint_handler( + request=dummy_request, + responses_request=ResponsesRequest(input="Hello"), + auth=MOCK_AUTH, + mcp_headers={}, + ) + + assert apply.await_args.kwargs["endpoint_path"] == "/v1/responses" + charge = apply.await_args.kwargs["charge"] + assert charge.func is consume_summarization_tokens + assert charge.args == (MOCK_AUTH[0],) + @pytest.mark.asyncio async def test_responses_azure_token_refresh( self, diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index aaec1878a..40b752bed 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -35,6 +35,7 @@ from models.config import Action from utils.conversation_compaction import CompactionResult from utils.otel_tracing import SpanAttributes, SpanEvents +from utils.query import consume_summarization_tokens INTERRUPTED_INDICATOR = f"\n\n*{INTERRUPTED_RESPONSE_MESSAGE}*" @@ -986,5 +987,9 @@ async def fake_generate_agent_response( assert events # the start event plus the delegated events assert captured_kwargs["context_status"] == expected_status - # The summarization calls are recorded under this endpoint (LCORE-3910). + # The summarization calls are recorded under this endpoint and charged + # to the user of the request (LCORE-3910). assert compaction_kwargs["endpoint_path"] == "/v1/streaming_query" + charge = compaction_kwargs["charge"] + assert charge.func is consume_summarization_tokens + assert charge.args == ("user_123",) diff --git a/tests/unit/utils/test_query.py b/tests/unit/utils/test_query.py index c311f0620..d5db16e8e 100644 --- a/tests/unit/utils/test_query.py +++ b/tests/unit/utils/test_query.py @@ -31,6 +31,7 @@ from utils.query import ( build_multimodal_input, consume_query_tokens, + consume_summarization_tokens, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, is_transcripts_enabled, @@ -615,6 +616,41 @@ def test_consume_tokens_database_error(self, mocker: MockerFixture) -> None: assert exc_info.value.status_code == 500 +class TestConsumeSummarizationTokens: + """Tests for consume_summarization_tokens function.""" + + def test_summarization_call_is_charged(self, mocker: MockerFixture) -> None: + """A call is charged for its input also when it produced no output tokens.""" + mock_consume = mocker.patch("utils.query.consume_tokens") + + consume_summarization_tokens( + "user1", + "provider1/model1", + TokenCounter(input_tokens=640, output_tokens=0, llm_calls=1), + ) + + mock_consume.assert_called_once() + charged = mock_consume.call_args.kwargs + assert charged["user_id"] == "user1" + assert (charged["input_tokens"], charged["output_tokens"]) == (640, 0) + assert (charged["provider_id"], charged["model_id"]) == ( + "provider1", + "model1", + ) + + def test_call_without_reported_usage_consumes_nothing( + self, mocker: MockerFixture + ) -> None: + """A call the provider reported no usage for leaves nothing to consume.""" + mock_consume = mocker.patch("utils.query.consume_tokens") + + consume_summarization_tokens( + "user1", "provider1/model1", TokenCounter(llm_calls=1) + ) + + mock_consume.assert_not_called() + + class TestStoreQueryResults: """Tests for store_query_results function.""" From f1f0652897aa80efff480a70677f8a9033b751de Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Tue, 6 Oct 2026 16:10:22 +0200 Subject: [PATCH 3/3] LCORE-3910: include the summarization calls in the reported token counts The summarization and fold calls of compaction are charged to the user since the previous commit, and the input_tokens / output_tokens the client is told still covered the answer alone. So a request that summarized consumed more quota than it reported. /v1/query now reports the usage of the turn plus CompactionResult.summarization_usage. /v1/streaming_query hands that usage to generate_agent_response(), which reports the sum in the end event. No field is added, and the model text "Number of tokens sent to LLM" still holds, so the OpenAPI document does not change. The turn summary is left as it is. What is stored with the turn, what is charged with the turn and what the llm.inference span carries stay the usage of the answer alone; the summarization was charged when its call returned and is not charged a second time here. /v1/responses is not touched: its usage object stays as the provider reported it. Documentation, for the three commits together: the user guide gets a section on token usage and quota per endpoint and two FAQ answers; the design document describes who pays for the calls, that the quota is checked once before compaction, and that the trace keeps the answer's usage; ARCHITECTURE.md and query_endpoint.md name the exception to "a failed or interrupted request consumes nothing"; the open question in the guardrails design document is updated. Tests: - tests/unit/utils/agents/test_streaming.py: the end event reports the sum, and consume_query_tokens() gets the usage of the turn alone. - The handler test of /v1/query asserts the sum in the response; the one of /v1/streaming_query asserts the usage is handed on. - The integration case on /v1/query now also reads the counts the client is told, next to the quota the limiter holds. --- .../conversation-compaction.md | 55 +++++++++++++++++-- .../prompt-guardrails/prompt-guardrails.md | 3 +- docs/devel_doc/ARCHITECTURE.md | 6 +- docs/devel_doc/query_endpoint.md | 7 ++- docs/user_doc/conversation_compaction.md | 37 ++++++++++++- src/app/endpoints/query.py | 7 ++- src/app/endpoints/streaming_query.py | 4 ++ src/utils/agents/streaming.py | 12 +++- .../endpoints/test_compaction_token_usage.py | 4 +- tests/unit/app/endpoints/test_query.py | 5 ++ .../app/endpoints/test_streaming_query.py | 11 +++- tests/unit/utils/agents/test_streaming.py | 34 ++++++++++-- 12 files changed, 161 insertions(+), 24 deletions(-) diff --git a/docs/design/conversation-compaction/conversation-compaction.md b/docs/design/conversation-compaction/conversation-compaction.md index eb1a06214..17b5e36a9 100644 --- a/docs/design/conversation-compaction/conversation-compaction.md +++ b/docs/design/conversation-compaction/conversation-compaction.md @@ -311,7 +311,7 @@ Add `compaction` field to the root `Configuration` class. |----------------------------------------|----------------------------------------------------------------------------| | `pyproject.toml` | Add `tiktoken` dependency | | `src/utils/token_estimator.py` | New module: `estimate_tokens()`, `estimate_conversation_tokens()` | -| `src/utils/compaction.py` | New module: summarization logic, partitioning, additive summary management | +| `src/utils/compaction.py` | New module: summarization logic, partitioning, additive summary management; reports the usage of each LLM call it makes — LCORE-3910 | | `src/models/config.py` | Add `CompactionConfiguration` (near `ConversationHistoryConfiguration`) | | `src/configuration.py` | Add `compaction_configuration` property to `AppConfig` singleton | | `src/utils/conversation_compaction.py` | New module: `apply_compaction()` / `apply_compaction_blocking()`, `needs_compaction_path()`, marker helpers, per-conversation lock | @@ -333,9 +333,11 @@ reusable unit in `src/utils/conversation_compaction.py` that each endpoint calls after its params are prepared: - `apply_compaction_blocking(client, params, inference_config, compaction_config)` - returns a `CompactionResult` (possibly-rewritten params, a `summarized` - flag, and the `original_input`). Non-streaming `/v1/query`, A2A, and - `/v1/responses` use this. + returns a `CompactionResult` (possibly-rewritten params, a `compacted` + flag, the `original_input`, and the `summarization_usage`). Non-streaming + `/v1/query`, A2A, and `/v1/responses` use this. The endpoints also pass + `endpoint_path` and, where there is a quota, `charge`; see + [Who pays for the summarization calls](#who-pays-for-the-summarization-calls). - `apply_compaction(..., emit_events=True)` is the async-generator variant that yields a `CompactionStartedEvent` before the summarization LLM call; the native `/v1/streaming_query` SSE path uses it to satisfy R12. @@ -441,6 +443,51 @@ Example config files go in `examples/`. Compaction adds latency only on the trigger turn. In PoC testing, compaction turns took 14-40 seconds vs 9-20 seconds for normal turns (gpt-5.6-luna). +## Who pays for the summarization calls + +The provider bills the summarization call and the fold call like any other, so +lightspeed-stack counts them (LCORE-3910). Before that, their usage was +discarded: a user could go over the quota without it ever showing, and by more +the longer the conversation was. + +`summarize_chunk()` and `recursively_resummarize()` report each call they make +to a `count_call` callback, with the model and the usage the provider reported. +They do it as soon as the response arrived and before they look at it, so a +call that returned no text is reported too. `apply_compaction()` takes the +path of the endpoint and a `charge` callback, and counts the calls with +`SummarizationCalls` (`src/utils/compaction_usage.py`): for each call it +records the token and call metrics under that endpoint, calls `charge`, and +adds the usage up. The sum is handed to the endpoint as +`CompactionResult.summarization_usage`. + +| Endpoint | Quota | Counts the client sees | +|------------------------|---------|----------------------------------------------------------| +| `/v1/query` | charged | `input_tokens` / `output_tokens` include summarization | +| `/v1/streaming_query` | charged | the same, in the `end` event | +| `/v1/responses` | charged | `usage` as the provider reported it, the answer alone | +| `/a2a` | none | none; the endpoint has no quota, the metrics are recorded | + +The endpoints with a quota pass `consume_summarization_tokens()` +(`src/utils/query.py`), 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 a cost of the conversation. It is charged also +when the turn is blocked, fails or is interrupted, and when compaction itself +fails after the call (the marker cannot be written, the fold fails). `/a2a` +passes no `charge`. The quota is checked once, before compaction: a request +that started with quota left runs to its end, as before. + +`/v1/query` and `/v1/streaming_query` add `summarization_usage` to the usage of +the turn in what they report to the client. The `llm.inference` span keeps the +usage of the answer alone. The sum is not stored with the turn; the token +usage history receives the charges one by one. +`/v1/responses` passes the `usage` object of the response through unchanged, to +stay a drop-in for clients of the OpenAI Responses API. + +Two limits. A call that fails, or is cancelled before its response arrived, +has no usage to count, whatever the provider bills for it. And a call the +provider reports no usage for is counted as a call and charges nothing. + # Open Questions for Future Work - **Compaction-proof instructions**: Allow "pinned" messages that always survive compaction (inspired by Claude Code's CLAUDE.md pattern). Not needed for v1. diff --git a/docs/design/prompt-guardrails/prompt-guardrails.md b/docs/design/prompt-guardrails/prompt-guardrails.md index 3cfee9181..886c5fc0e 100644 --- a/docs/design/prompt-guardrails/prompt-guardrails.md +++ b/docs/design/prompt-guardrails/prompt-guardrails.md @@ -442,7 +442,8 @@ product need justifies it. version-specific prompt. - **Guardian token usage:** whether guardian calls should count against user quota or be tracked as service overhead. Compaction's summarization calls - raise the same question. + are charged to the user's quota (LCORE-3910); the same choice is open for + guardian calls. - **Streaming checkpoint sizing:** defaults for LCORE-3391 (spike Decision T4, 70% confidence); tune with real latency data. - **Cheap classifier tier for `tool`:** Prompt Guard 2-class, and its diff --git a/docs/devel_doc/ARCHITECTURE.md b/docs/devel_doc/ARCHITECTURE.md index c6ae2c43d..8ec5a0275 100644 --- a/docs/devel_doc/ARCHITECTURE.md +++ b/docs/devel_doc/ARCHITECTURE.md @@ -274,9 +274,9 @@ The system defines 30+ actions that can be authorized. Examples (see `docs/user_ - Update quota counters 3. **On Error:** - - If LLM call fails, no tokens are consumed - - Quota remains unchanged - - User can retry the request + - If LLM call fails, no tokens are consumed for that call + - Quota remains unchanged, with one exception: the calls conversation compaction made to summarize older turns for the request were charged when they were made + - User can retry the request; the summary is kept, so the retry does not pay for it again --- diff --git a/docs/devel_doc/query_endpoint.md b/docs/devel_doc/query_endpoint.md index bc431c969..ed9658919 100644 --- a/docs/devel_doc/query_endpoint.md +++ b/docs/devel_doc/query_endpoint.md @@ -308,7 +308,7 @@ Cancels an in-progress streaming query. | `interrupted` | boolean | Whether an active stream was interrupted (`false` if already completed) | | `message` | string | Human-readable status message | -When a stream is interrupted, any partial response is persisted to conversation history and token consumption is skipped. +When a stream is interrupted, any partial response is persisted to conversation history and token consumption for the answer is skipped. The summarization calls conversation compaction made for the request were charged when they were made (see [Quota and Token Counting](#quota-and-token-counting)). --- @@ -360,9 +360,10 @@ If the server configuration sets `disable_query_system_prompt` to `true`, reques ## Quota and Token Counting - **Pre-request:** `check_tokens_available()` verifies the user/cluster has available quota (429 if not) -- **Post-response:** `consume_query_tokens()` deducts `input_tokens` and `output_tokens` from configured quota limiters +- **Post-response:** `consume_query_tokens()` deducts the input and output tokens of the answer from configured quota limiters +- **Compaction:** each LLM call that summarizes older turns is charged by `consume_summarization_tokens()` as soon as it returned, so also when the turn is then blocked, fails or is interrupted. The `input_tokens` and `output_tokens` the response reports include these calls - **Available quotas:** Remaining balances per limiter are included in the response (`available_quotas` field in sync, `end` event in streaming) -- **Stream interruption:** Token consumption is skipped for interrupted streams +- **Stream interruption:** Token consumption is skipped for interrupted streams, except for the summarization calls, which were charged before the answer started --- diff --git a/docs/user_doc/conversation_compaction.md b/docs/user_doc/conversation_compaction.md index 148bfd32a..b1206ceb3 100644 --- a/docs/user_doc/conversation_compaction.md +++ b/docs/user_doc/conversation_compaction.md @@ -187,6 +187,37 @@ Compaction acquires a per-conversation lock to prevent concurrent requests on th Over very long conversations, multiple compaction summaries may accumulate. When the total size of cached summaries approaches the context window threshold, they are recursively folded into a single summary using a dedicated re-summarization prompt. This prevents summaries from themselves exceeding the context window. +### Token usage and quota + +Summarizing older turns takes an LLM call of its own, and so does folding the +summaries. The provider bills these calls, so they are counted: + +| Endpoint | Quota | Reported token counts | Metrics | +|---|---|---|---| +| `POST /v1/query` | Charged to the user. | `input_tokens` and `output_tokens` include the summarization calls. | Counted under `/v1/query`. | +| `POST /v1/streaming_query` | Charged to the user. | `input_tokens` and `output_tokens` of the `end` event include the summarization calls. | Counted under `/v1/streaming_query`. | +| `POST /v1/responses` | Charged to the user. | `usage` is what the provider reported for the response, so it covers the answer alone. `available_quotas` shows the quota left after both. | Counted under `/v1/responses`. | +| `POST /a2a` | Not charged. The endpoint does not use quotas. | None. | The summarization calls are counted under `/a2a`. The call that answers the request is not recorded in these metrics on this endpoint. | + +The metrics are `ls_llm_calls_total`, `ls_llm_token_sent_total` and +`ls_llm_token_received_total`, labelled with the endpoint. + +Only a request that makes a summarization call pays for it: one whose +estimated input crosses the threshold. The other requests are served from the +stored summary and cost what they cost without compaction. With the default +settings a summarization happens once every several turns; with a small +context window or a low `threshold_ratio` it can happen on every request. + +Each summarization call is charged as soon as it returned, before the model is +asked for the answer. It is therefore charged also when the request is blocked +by a shield, fails afterward, or is interrupted by the client. The summary is +kept in these cases, so the next request on the conversation does not pay for +it again. A summarization call that itself fails, or is cut off before it +returned, reports no usage and is not charged. + +On `/v1/responses`, the difference between the quota consumed and the `usage` +of the response is the summarization. + ## When compaction is disabled When compaction is disabled (the default), requests that cause the conversation history to exceed the model's context window will fail with HTTP 413 (Prompt Too Long). Clients must manage conversation length themselves, for example by starting new conversations or deleting old ones. @@ -201,7 +232,11 @@ Compaction summarizes older turns, so fine-grained details from early in the con **Does compaction use extra tokens?** -Yes. The summarization step requires an additional LLM call, which consumes tokens. These tokens are counted against the user's quota. The trade-off is that the conversation can continue instead of failing with HTTP 413. +Yes. The summarization step requires an additional LLM call, which consumes tokens. On `/v1/query` and `/v1/streaming_query` these tokens are counted against the user's quota and are included in the `input_tokens` and `output_tokens` of the request that triggered the summarization. `/v1/responses` charges the quota and leaves `usage` as the provider reported it; `/a2a` has no quota. See [Token usage and quota](#token-usage-and-quota). The trade-off is that the conversation can continue instead of failing with HTTP 413. + +**Why did one request use many more tokens than the ones before it?** + +On `/v1/query` and `/v1/streaming_query`, that request triggered a summarization. Its `input_tokens` include the older turns that were sent to the LLM to be summarized, and its `output_tokens` include the summary. `context_status` is `"summarized"` on that request, and also on the requests after it that are served from the stored summary and do not pay for it again. **Can I use compaction with all LLM providers?** diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 3bb457316..7985a6e31 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -310,6 +310,9 @@ async def _handle_query_with_tracing( model_id=responses_params.model, token_usage=turn_summary.token_usage, ) + # The turn reports the summarization calls made for it as part of its + # usage. They were charged when they were made. + reported_usage = turn_summary.token_usage + compaction.summarization_usage logger.info("Getting available quotas") available_quotas = get_available_quotas( @@ -355,7 +358,7 @@ async def _handle_query_with_tracing( referenced_documents=turn_summary.referenced_documents, truncated=False, context_status=compaction.context_status, - input_tokens=turn_summary.token_usage.input_tokens, - output_tokens=turn_summary.token_usage.output_tokens, + input_tokens=reported_usage.input_tokens, + output_tokens=reported_usage.output_tokens, available_quotas=available_quotas, ) diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 1ab629f5d..7a1b926f6 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -91,6 +91,7 @@ stream_start_event, ) from utils.suid import get_suid, normalize_conversation_id +from utils.token_counter import TokenCounter from utils.types import Responses from utils.vector_search import build_rag_context @@ -426,6 +427,7 @@ async def generate_response_with_compaction( turn: Optional[PendingTurn] = None context_status: ContextStatus = "full" + summarization_usage = TokenCounter() try: async for item in apply_compaction( context.client, @@ -447,6 +449,7 @@ async def generate_response_with_compaction( context.client, item.params, item.original_input ) context_status = item.context_status + summarization_usage = item.summarization_usage generator, turn_summary = await retrieve_agent_response_generator( responses_params=responses_params, @@ -501,6 +504,7 @@ async def generate_response_with_compaction( emit_start=False, turn=turn, context_status=context_status, + summarization_usage=summarization_usage, ): yield event finally: diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 3c3aa36a6..6192487f1 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -88,6 +88,7 @@ persist_interrupted_turn, register_interrupt_callback, ) +from utils.token_counter import TokenCounter if TYPE_CHECKING: from models.common.query import Attachment @@ -195,6 +196,7 @@ async def generate_agent_response( # pylint: disable=too-many-statements emit_start: bool = True, turn: Optional[PendingTurn] = None, context_status: ContextStatus = "full", + summarization_usage: Optional[TokenCounter] = None, ) -> AsyncIterator[str]: """Wrap an agent SSE generator with cleanup logic. @@ -218,6 +220,9 @@ async def generate_agent_response( # pylint: disable=too-many-statements context_status: Whether the conversation context was sent in full ("full") or older turns were replaced by a summary ("summarized"). Reported to the client in the SSE end event. + summarization_usage: Usage of the summarization calls compaction made + for this request (LCORE-3910). They were charged when they were + made; here they are added to the usage the turn reports. Yields: SSE-formatted strings from the wrapped generator. @@ -337,6 +342,9 @@ async def generate_agent_response( # pylint: disable=too-many-statements model_id=responses_params.model, token_usage=turn_summary.token_usage, ) + # The turn reports the summarization calls made for it as part of its + # usage. They were charged when they were made. + reported_usage = turn_summary.token_usage + (summarization_usage or TokenCounter()) logger.info("Getting available quotas") available_quotas = get_available_quotas( quota_limiters=configuration.quota_limiters, @@ -345,8 +353,8 @@ async def generate_agent_response( # pylint: disable=too-many-statements end_payload = EndStreamPayload.create( referenced_documents=turn_summary.referenced_documents, context_status=context_status, - input_tokens=turn_summary.token_usage.input_tokens, - output_tokens=turn_summary.token_usage.output_tokens, + input_tokens=reported_usage.input_tokens, + output_tokens=reported_usage.output_tokens, available_quotas=available_quotas, ) yield serialize_event(end_payload, media_type) diff --git a/tests/integration/endpoints/test_compaction_token_usage.py b/tests/integration/endpoints/test_compaction_token_usage.py index 7634079da..ffdbb8aed 100644 --- a/tests/integration/endpoints/test_compaction_token_usage.py +++ b/tests/integration/endpoints/test_compaction_token_usage.py @@ -156,11 +156,13 @@ async def scene_fixture( # pylint: disable=too-many-arguments,too-many-position @pytest.mark.asyncio async def test_turn_that_summarizes_pays_for_the_summarization(scene: Scene) -> None: - """The quota falls by the turn and by the summarization call, and the client is told.""" + """The turn and the summarization call are charged, and the client is told both.""" response = await scene.query() assert_marker_count(scene.store, CONV_ID_LLAMA, 1) assert response.context_status == "summarized" + assert response.input_tokens == TURN_INPUT + SUMMARY_INPUT + assert response.output_tokens == TURN_OUTPUT + SUMMARY_OUTPUT assert response.available_quotas == { "UserQuotaLimiter": INITIAL_QUOTA - TURN - SUMMARY } diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 89935d10e..2610e4684 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -23,6 +23,7 @@ from models.database.conversations import UserConversation from utils.conversation_compaction import CompactionResult from utils.query import consume_summarization_tokens +from utils.token_counter import TokenCounter # User ID must be proper UUID MOCK_AUTH = ( @@ -231,6 +232,7 @@ async def test_query_reports_context_status( compaction_result = CompactionResult( cast("ResponsesApiParams", mock_responses_params), compacted=compacted, + summarization_usage=TokenCounter(input_tokens=640, output_tokens=72), ) apply = mocker.patch( "app.endpoints.query.apply_compaction_blocking", @@ -239,6 +241,7 @@ async def test_query_reports_context_status( mock_turn_summary = TurnSummary() mock_turn_summary.llm_response = "An answer" + mock_turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) mocker.patch( "app.endpoints.query.retrieve_agent_response", new=mocker.AsyncMock(return_value=mock_turn_summary), @@ -266,6 +269,8 @@ async def test_query_reports_context_status( charge = apply.await_args.kwargs["charge"] assert charge.func is consume_summarization_tokens assert charge.args == (MOCK_AUTH[0],) + # The counts the client is told include the summarization calls. + assert (response.input_tokens, response.output_tokens) == (643, 79) @pytest.mark.asyncio async def test_query_merges_inline_and_tool_rag_chunks_and_documents( diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index 40b752bed..d1e10b37e 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -36,6 +36,7 @@ from utils.conversation_compaction import CompactionResult from utils.otel_tracing import SpanAttributes, SpanEvents from utils.query import consume_summarization_tokens +from utils.token_counter import TokenCounter INTERRUPTED_INDICATOR = f"\n\n*{INTERRUPTED_RESPONSE_MESSAGE}*" @@ -926,7 +927,11 @@ async def test_threads_context_status_to_agent_response( query="What is OpenShift?" ) # pyright: ignore[reportCallIssue] - compaction_result = CompactionResult(responses_params, compacted=compacted) + compaction_result = CompactionResult( + responses_params, + compacted=compacted, + summarization_usage=TokenCounter(input_tokens=640, output_tokens=72), + ) compaction_kwargs: dict[str, Any] = {} @@ -993,3 +998,7 @@ async def fake_generate_agent_response( charge = compaction_kwargs["charge"] assert charge.func is consume_summarization_tokens assert charge.args == ("user_123",) + # Their usage is handed on to be reported in the end event. + assert captured_kwargs["summarization_usage"] == TokenCounter( + input_tokens=640, output_tokens=72 + ) diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index 2dfd7301c..ed08037f0 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -631,21 +631,36 @@ async def inner() -> AsyncIterator[str]: @pytest.mark.asyncio @pytest.mark.parametrize( - ("generate_kwargs", "expected_status"), + ("generate_kwargs", "expected_status", "expected_tokens"), [ - ({}, "full"), - ({"context_status": "summarized"}, "summarized"), + ({}, "full", (3, 7)), + ( + { + "context_status": "summarized", + "summarization_usage": TokenCounter( + input_tokens=640, output_tokens=72, llm_calls=1 + ), + }, + "summarized", + (643, 79), + ), ], ) - async def test_end_event_reports_context_status( + async def test_end_event_reports_context_status( # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals self, mocker: MockerFixture, make_generator_context: Callable[..., ResponseGeneratorContext], responses_params: ResponsesApiParams, generate_kwargs: dict[str, Any], expected_status: str, + expected_tokens: tuple[int, int], ) -> None: - """Test the end event carries context_status ("full" by default).""" + """Test the end event carries context_status and the token counts. + + The counts include the summarization calls compaction made for the + request (LCORE-3910). Those were charged when they were made, so the + turn alone is charged here. + """ context = make_generator_context() turn_summary = TurnSummary() turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) @@ -657,7 +672,7 @@ async def inner() -> AsyncIterator[str]: MEDIA_TYPE_JSON, ) - mocker.patch("utils.agents.streaming.consume_query_tokens") + consume_mock = mocker.patch("utils.agents.streaming.consume_query_tokens") mocker.patch( "utils.agents.streaming.get_available_quotas", return_value={"daily": 100}, @@ -693,6 +708,13 @@ async def inner() -> AsyncIterator[str]: ] assert len(end_events) == 1 assert end_events[0]["data"]["context_status"] == expected_status + assert ( + end_events[0]["data"]["input_tokens"], + end_events[0]["data"]["output_tokens"], + ) == expected_tokens + assert consume_mock.call_args.kwargs["token_usage"] == TokenCounter( + input_tokens=3, output_tokens=7 + ) @pytest.mark.asyncio async def test_cancelled_persists_interrupted_turn(