Merge pull request #1379 from Kiln-AI/leonard/kil-606-fix-usage-cost-summing-in-chat-multi-turn-conversations

feat: include usage in each message

Author: leonardmq

app/web_ui/src/lib/api_schema.d.ts

@@ -3732,6 +3732,7 @@ export interface components {
             tool_calls?: components["schemas"]["ChatCompletionMessageFunctionToolCallParam"][];
             /** Latency Ms */
             latency_ms?: number | null;
+            usage?: components["schemas"]["MessageUsage"] | null;
         };
         /**
          * ChatCompletionAssistantMessageParamWrapper
@@ -3762,6 +3763,7 @@ export interface components {
             tool_calls?: components["schemas"]["ChatCompletionMessageFunctionToolCallParam"][];
             /** Latency Ms */
             latency_ms?: number | null;
+            usage?: components["schemas"]["MessageUsage"] | null;
         };
         /** ChatCompletionContentPartImageParam */
         ChatCompletionContentPartImageParam: {
@@ -7226,6 +7228,47 @@ export interface components {
              */
             mean_total_llm_latency_ms?: number | null;
         };
+        /**
+         * MessageUsage
+         * @description Token usage and cost for a single LLM call or a multi-message sum.
+         *
+         *     Carries only the fields that are meaningfully aggregatable across
+         *     messages: token counts and cost. Per-call latency lives on the
+         *     individual message's ``latency_ms`` field; aggregating it across the
+         *     full trace would mix latencies from different points in time, so
+         *     ``MessageUsage`` does NOT carry ``total_llm_latency_ms``.
+         *
+         *     The :class:`Usage` subclass adds ``total_llm_latency_ms`` for the
+         *     in-flight per-run accumulator that tracks how long this run spent
+         *     waiting on LLM calls.
+         */
+        MessageUsage: {
+            /**
+             * Input Tokens
+             * @description The number of input tokens used.
+             */
+            input_tokens?: number | null;
+            /**
+             * Output Tokens
+             * @description The number of output tokens used.
+             */
+            output_tokens?: number | null;
+            /**
+             * Total Tokens
+             * @description The total number of tokens used.
+             */
+            total_tokens?: number | null;
+            /**
+             * Cost
+             * @description The cost in US dollars, saved at runtime (prices can change over time).
+             */
+            cost?: number | null;
+            /**
+             * Cached Tokens
+             * @description Number of tokens served from prompt cache. None if not reported.
+             */
+            cached_tokens?: number | null;
+        };
         /** ModelDetails */
         ModelDetails: {
             /** Id */
@@ -9599,6 +9642,8 @@ export interface components {
             tags: string[];
             /** @description Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used. */
             usage?: components["schemas"]["Usage"] | null;
+            /** @description Sum of per-message token usage and cost across the entire trace, including any seeded prior trace. None on records created before this field existed. For a fresh (non-seeded) run, the token / cost fields equal those of `usage`. */
+            cumulative_usage?: components["schemas"]["MessageUsage"] | null;
             /**
              * Trace
              * @description The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.
@@ -9676,6 +9721,8 @@ export interface components {
             tags: string[];
             /** @description Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used. */
             usage?: components["schemas"]["Usage"] | null;
+            /** @description Sum of per-message token usage and cost across the entire trace, including any seeded prior trace. None on records created before this field existed. For a fresh (non-seeded) run, the token / cost fields equal those of `usage`. */
+            cumulative_usage?: components["schemas"]["MessageUsage"] | null;
             /**
              * Trace
              * @description The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.
@@ -10184,27 +10231,33 @@ export interface components {
         };
         /**
          * Usage
-         * @description Token usage and cost information for a task run.
+         * @description Token usage, cost, and aggregate LLM latency for a per-run accumulator.
+         *
+         *     Extends :class:`MessageUsage` with ``total_llm_latency_ms``, which is
+         *     only meaningful while a single run is in flight (its model calls run
+         *     sequentially in real time). For per-message records and full-trace
+         *     sums use :class:`MessageUsage` — those values would mix latencies
+         *     from different points in time, so the field doesn't apply.
          */
         Usage: {
             /**
              * Input Tokens
-             * @description The number of input tokens used in the task run.
+             * @description The number of input tokens used.
              */
             input_tokens?: number | null;
             /**
              * Output Tokens
-             * @description The number of output tokens used in the task run.
+             * @description The number of output tokens used.
              */
             output_tokens?: number | null;
             /**
              * Total Tokens
-             * @description The total number of tokens used in the task run.
+             * @description The total number of tokens used.
              */
             total_tokens?: number | null;
             /**
              * Cost
-             * @description The cost of the task run in US dollars, saved at runtime (prices can change over time).
+             * @description The cost in US dollars, saved at runtime (prices can change over time).
              */
             cost?: number | null;
             /**

libs/core/kiln_ai/adapters/litellm_utils/litellm_streaming.py

@@ -30,6 +30,17 @@ class StreamingCompletion:
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         kwargs = dict(kwargs)
         kwargs.pop("stream", None)
+        # LiteLLM's streaming responses don't include a usage block by
+        # default — ``stream_options={"include_usage": True}`` is required
+        # for the final assembled ModelResponse to carry token counts (and
+        # downstream cost). Force it on; merge with caller-provided
+        # ``stream_options`` without clobbering unrelated keys, but always
+        # override ``include_usage`` since usage tracking is mandatory.
+        caller_stream_options = kwargs.get("stream_options") or {}
+        kwargs["stream_options"] = {
+            **caller_stream_options,
+            "include_usage": True,
+        }
         self._args = args
         self._kwargs = kwargs
         self._response: Optional[Union[ModelResponse, TextCompletionResponse]] = None

libs/core/kiln_ai/adapters/litellm_utils/test_litellm_streaming.py

@@ -137,3 +137,70 @@ async def test_empty_stream(self, mock_acompletion, mock_chunk_builder):
 
         assert received == []
         assert stream.response is None
+
+    async def test_stream_options_include_usage_added_by_default(
+        self, mock_acompletion, mock_chunk_builder
+    ):
+        """Default ``stream_options`` must request usage so the assembled response carries it."""
+        mock_acompletion.return_value = _async_iter([])
+        mock_chunk_builder.return_value = None
+
+        stream = StreamingCompletion(model="test", messages=[])
+        async for _ in stream:
+            pass
+
+        _, call_kwargs = mock_acompletion.call_args
+        assert call_kwargs["stream_options"] == {"include_usage": True}
+
+    async def test_stream_options_include_usage_merged_with_caller_options(
+        self, mock_acompletion, mock_chunk_builder
+    ):
+        """Caller-provided ``stream_options`` keys must survive; ``include_usage`` is forced on."""
+        mock_acompletion.return_value = _async_iter([])
+        mock_chunk_builder.return_value = None
+
+        stream = StreamingCompletion(
+            model="test",
+            messages=[],
+            stream_options={"some_other_flag": True},
+        )
+        async for _ in stream:
+            pass
+
+        _, call_kwargs = mock_acompletion.call_args
+        assert call_kwargs["stream_options"] == {
+            "some_other_flag": True,
+            "include_usage": True,
+        }
+
+    async def test_caller_provided_include_usage_false_is_overridden(
+        self, mock_acompletion, mock_chunk_builder
+    ):
+        """Streaming usage tracking is mandatory; an explicit ``False`` is overridden."""
+        mock_acompletion.return_value = _async_iter([])
+        mock_chunk_builder.return_value = None
+
+        stream = StreamingCompletion(
+            model="test",
+            messages=[],
+            stream_options={"include_usage": False},
+        )
+        async for _ in stream:
+            pass
+
+        _, call_kwargs = mock_acompletion.call_args
+        assert call_kwargs["stream_options"]["include_usage"] is True
+
+    async def test_stream_options_none_treated_as_empty(
+        self, mock_acompletion, mock_chunk_builder
+    ):
+        """Passing ``stream_options=None`` must not crash — treated as empty."""
+        mock_acompletion.return_value = _async_iter([])
+        mock_chunk_builder.return_value = None
+
+        stream = StreamingCompletion(model="test", messages=[], stream_options=None)
+        async for _ in stream:
+            pass
+
+        _, call_kwargs = mock_acompletion.call_args
+        assert call_kwargs["stream_options"] == {"include_usage": True}

libs/core/kiln_ai/adapters/model_adapters/adapter_stream.py

@@ -19,7 +19,7 @@
     ToolCallEventType,
 )
 from kiln_ai.adapters.run_output import RunOutput
-from kiln_ai.datamodel import Usage
+from kiln_ai.datamodel import MessageUsage, Usage
 
 if TYPE_CHECKING:
     from kiln_ai.adapters.model_adapters.litellm_adapter import LiteLlmAdapter
@@ -63,7 +63,12 @@ def __init__(
         self._top_logprobs = top_logprobs
         self._result: AdapterStreamResult | None = None
         self._iterated = False
+        # Per-LLM-call latency / usage, keyed by index in the messages list.
+        # Mirrors the non-streaming adapter — we don't own the LiteLLM
+        # message objects, so we accumulate side-channel state and attach
+        # it during ``all_messages_to_trace`` at finalization.
         self._message_latency: dict[int, int] = {}
+        self._message_usage: dict[int, MessageUsage] = {}
 
     @property
     def result(self) -> AdapterStreamResult:
@@ -134,7 +139,7 @@ async def __aiter__(self) -> AsyncIterator[AdapterStreamEvent]:
             raise RuntimeError(f"assistant message is not a string: {prior_output}")
 
         trace = self._adapter.all_messages_to_trace(
-            self._messages, self._message_latency
+            self._messages, self._message_latency, self._message_usage
         )
         self._result = AdapterStreamResult(
             run_output=RunOutput(
@@ -170,7 +175,8 @@ async def _stream_model_turn(
             call_latency_ms = int((time.monotonic() - start) * 1000)
 
             response, response_choice = _validate_response(stream.response)
-            usage += self._adapter.usage_from_response(response)
+            call_usage = self._adapter.usage_from_response(response)
+            usage += call_usage
             usage.total_llm_latency_ms = (
                 usage.total_llm_latency_ms or 0
             ) + call_latency_ms
@@ -184,6 +190,7 @@ async def _stream_model_turn(
 
             self._messages.append(response_choice.message)
             self._message_latency[len(self._messages) - 1] = call_latency_ms
+            self._message_usage[len(self._messages) - 1] = call_usage
 
             if tool_calls and len(tool_calls) > 0:
                 # Check for return_on_tool_call BEFORE processing

libs/core/kiln_ai/adapters/model_adapters/base_adapter.py

@@ -38,6 +38,7 @@
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
+    MessageUsage,
     Task,
     TaskOutput,
     TaskRun,
@@ -683,6 +684,7 @@ def generate_run(
             tags=self.base_adapter_config.default_tags or [],
             usage=usage,
             trace=trace,
+            cumulative_usage=MessageUsage.from_trace(trace),
         )
 
     def _properties_for_task_output(self) -> Dict[str, str | int | float]: