add token usage tracking to LLM responses

- Add TokenCounter dataclass to track input/output tokens and LLM calls
- Update QueryResponse model with token usage fields (input_tokens, output_tokens, truncated, available_quotas)
- Implement extract_token_usage_from_turn() function for token counting
- Update /query and /streaming_query endpoints to include token usage in responses
- Modify retrieve_response() to return token usage information
- Update test cases to handle new return values and mock token usage
- Maintain backward compatibility with existing API structure

The implementation provides a foundation for token tracking that can be enhanced
with more sophisticated counting logic in the future.

Author: gallettilance

src/app/endpoints/query.py

@@ -31,7 +31,6 @@
 from authorization.middleware import authorize
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
-from metrics.utils import update_llm_token_count_from_turn
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import Attachment, QueryRequest
@@ -55,6 +54,7 @@
 from utils.mcp_headers import handle_mcp_headers_with_toolgroups, mcp_headers_dependency
 from utils.transcripts import store_transcript
 from utils.types import TurnSummary
+from utils.token_counter import extract_and_update_token_metrics, TokenCounter
 
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["query"])
@@ -279,16 +279,16 @@ async def query_endpoint_handler(  # pylint: disable=R0914
                 user_conversation=user_conversation, query_request=query_request
             ),
         )
-        summary, conversation_id, referenced_documents = await retrieve_response(
-            client,
-            llama_stack_model_id,
-            query_request,
-            token,
-            mcp_headers=mcp_headers,
-            provider_id=provider_id,
+        summary, conversation_id, referenced_documents, token_usage = (
+            await retrieve_response(
+                client,
+                llama_stack_model_id,
+                query_request,
+                token,
+                mcp_headers=mcp_headers,
+                provider_id=provider_id,
+            )
         )
-        # Update metrics for the LLM call
-        metrics.llm_calls_total.labels(provider_id, model_id).inc()
 
         # Get the initial topic summary for the conversation
         topic_summary = None
@@ -371,6 +371,10 @@ async def query_endpoint_handler(  # pylint: disable=R0914
             rag_chunks=summary.rag_chunks if summary.rag_chunks else [],
             tool_calls=tool_calls if tool_calls else None,
             referenced_documents=referenced_documents,
+            truncated=False,  # TODO: implement truncation detection
+            input_tokens=token_usage.input_tokens,
+            output_tokens=token_usage.output_tokens,
+            available_quotas={},  # TODO: implement quota tracking
         )
         logger.info("Query processing completed successfully!")
         return response
@@ -583,7 +587,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
     mcp_headers: dict[str, dict[str, str]] | None = None,
     *,
     provider_id: str = "",
-) -> tuple[TurnSummary, str, list[ReferencedDocument]]:
+) -> tuple[TurnSummary, str, list[ReferencedDocument], TokenCounter]:
     """
     Retrieve response from LLMs and agents.
 
@@ -607,9 +611,9 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         mcp_headers (dict[str, dict[str, str]], optional): Headers for multi-component processing.
 
     Returns:
-        tuple[TurnSummary, str, list[ReferencedDocument]]: A tuple containing
+        tuple[TurnSummary, str, list[ReferencedDocument], TokenCounter]: A tuple containing
         a summary of the LLM or agent's response
-        content, the conversation ID and the list of parsed referenced documents.
+        content, the conversation ID, the list of parsed referenced documents, and token usage information.
     """
     available_input_shields = [
         shield.identifier
@@ -704,9 +708,11 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
 
     referenced_documents = parse_referenced_documents(response)
 
-    # Update token count metrics for the LLM call
+    # Update token count metrics and extract token usage in one call
     model_label = model_id.split("/", 1)[1] if "/" in model_id else model_id
-    update_llm_token_count_from_turn(response, model_label, provider_id, system_prompt)
+    token_usage = extract_and_update_token_metrics(
+        response, model_label, provider_id, system_prompt
+    )
 
     # Check for validation errors in the response
     steps = response.steps or []
@@ -722,7 +728,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
             "Response lacks output_message.content (conversation_id=%s)",
             conversation_id,
         )
-    return (summary, conversation_id, referenced_documents)
+    return (summary, conversation_id, referenced_documents, token_usage)
 
 
 def validate_attachments_metadata(attachments: list[Attachment]) -> None:

src/app/endpoints/streaming_query.py

@@ -41,7 +41,6 @@
 from configuration import configuration
 from constants import DEFAULT_RAG_TOOL
 import metrics
-from metrics.utils import update_llm_token_count_from_turn
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
@@ -58,6 +57,7 @@
 from utils.mcp_headers import handle_mcp_headers_with_toolgroups, mcp_headers_dependency
 from utils.transcripts import store_transcript
 from utils.types import TurnSummary
+from utils.token_counter import extract_and_update_token_metrics, TokenCounter
 
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["streaming_query"])
@@ -75,7 +75,7 @@
                         'data: {"event": "token", "data": {"id": 0, "role": "inference", '
                         '"token": "Hello"}}\n\n'
                         'data: {"event": "end", "data": {"referenced_documents": [], '
-                        '"truncated": null, "input_tokens": 0, "output_tokens": 0}, '
+                        '"truncated": false, "input_tokens": 150, "output_tokens": 50}, '
                         '"available_quotas": {}}\n\n'
                     ),
                 }
@@ -144,7 +144,9 @@ def stream_start_event(conversation_id: str) -> str:
     )
 
 
-def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str:
+def stream_end_event(
+    metadata_map: dict, summary: TurnSummary, token_usage: TokenCounter
+) -> str:
     """
     Yield the end of the data stream.
 
@@ -181,9 +183,9 @@ def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str:
             "data": {
                 "rag_chunks": rag_chunks,
                 "referenced_documents": referenced_docs_dict,
-                "truncated": None,  # TODO(jboos): implement truncated
-                "input_tokens": 0,  # TODO(jboos): implement input tokens
-                "output_tokens": 0,  # TODO(jboos): implement output tokens
+                "truncated": False,  # TODO(jboos): implement truncated
+                "input_tokens": token_usage.input_tokens,
+                "output_tokens": token_usage.output_tokens,
             },
             "available_quotas": {},  # TODO(jboos): implement available quotas
         }
@@ -672,6 +674,7 @@ async def response_generator(
             summary = TurnSummary(
                 llm_response="No response from the model", tool_calls=[]
             )
+            token_usage = TokenCounter()
 
             # Send start event
             yield stream_start_event(conversation_id)
@@ -686,7 +689,8 @@ async def response_generator(
                     )
                     system_prompt = get_system_prompt(query_request, configuration)
                     try:
-                        update_llm_token_count_from_turn(
+                        # Update token count metrics and extract token usage in one call
+                        token_usage = extract_and_update_token_metrics(
                             p.turn, model_id, provider_id, system_prompt
                         )
                     except Exception:  # pylint: disable=broad-except
@@ -699,7 +703,7 @@ async def response_generator(
                     chunk_id += 1
                     yield event
 
-            yield stream_end_event(metadata_map, summary)
+            yield stream_end_event(metadata_map, summary, token_usage)
 
             if not is_transcripts_enabled():
                 logger.debug("Transcript collection is disabled in the configuration")
@@ -755,9 +759,6 @@ async def response_generator(
                 topic_summary=topic_summary,
             )
 
-        # Update metrics for the LLM call
-        metrics.llm_calls_total.labels(provider_id, model_id).inc()
-
         return StreamingResponse(response_generator(response))
     # connection to Llama Stack server
     except APIConnectionError as e:

src/models/responses.py

@@ -125,11 +125,10 @@ class QueryResponse(BaseModel):
         rag_chunks: List of RAG chunks used to generate the response.
         referenced_documents: The URLs and titles for the documents used to generate the response.
         tool_calls: List of tool calls made during response generation.
-        TODO: truncated: Whether conversation history was truncated.
-        TODO: input_tokens: Number of tokens sent to LLM.
-        TODO: output_tokens: Number of tokens received from LLM.
-        TODO: available_quotas: Quota available as measured by all configured quota limiters
-        TODO: tool_results: List of tool results.
+        truncated: Whether conversation history was truncated.
+        input_tokens: Number of tokens sent to LLM.
+        output_tokens: Number of tokens received from LLM.
+        available_quotas: Quota available as measured by all configured quota limiters.
     """
 
     conversation_id: Optional[str] = Field(
@@ -169,6 +168,30 @@ class QueryResponse(BaseModel):
         ],
     )
 
+    truncated: bool = Field(
+        False,
+        description="Whether conversation history was truncated",
+        examples=[False, True],
+    )
+
+    input_tokens: int = Field(
+        0,
+        description="Number of tokens sent to LLM",
+        examples=[150, 250, 500],
+    )
+
+    output_tokens: int = Field(
+        0,
+        description="Number of tokens received from LLM",
+        examples=[50, 100, 200],
+    )
+
+    available_quotas: dict[str, int] = Field(
+        default_factory=dict,
+        description="Quota available as measured by all configured quota limiters",
+        examples=[{"daily": 1000, "monthly": 50000}],
+    )
+
     # provides examples for /docs endpoint
     model_config = {
         "json_schema_extra": {
@@ -197,6 +220,10 @@ class QueryResponse(BaseModel):
                             "doc_title": "Operator Lifecycle Manager (OLM)",
                         }
                     ],
+                    "truncated": False,
+                    "input_tokens": 150,
+                    "output_tokens": 75,
+                    "available_quotas": {"daily": 1000, "monthly": 50000},
                 }
             ]
         }

src/utils/token_counter.py

@@ -0,0 +1,130 @@
+"""Helper classes to count tokens sent and received by the LLM."""
+
+import logging
+from dataclasses import dataclass
+from typing import cast
+
+from llama_stack.models.llama.datatypes import RawMessage
+from llama_stack.models.llama.llama3.chat_format import ChatFormat
+from llama_stack.models.llama.llama3.tokenizer import Tokenizer
+from llama_stack_client.types.agents.turn import Turn
+
+import metrics
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TokenCounter:
+    """Model representing token counter.
+
+    Attributes:
+        input_tokens: number of tokens sent to LLM
+        output_tokens: number of tokens received from LLM
+        input_tokens_counted: number of input tokens counted by the handler
+        llm_calls: number of LLM calls
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    input_tokens_counted: int = 0
+    llm_calls: int = 0
+
+    def __str__(self) -> str:
+        """Textual representation of TokenCounter instance."""
+        return (
+            f"{self.__class__.__name__}: "
+            + f"input_tokens: {self.input_tokens} "
+            + f"output_tokens: {self.output_tokens} "
+            + f"counted: {self.input_tokens_counted} "
+            + f"LLM calls: {self.llm_calls}"
+        )
+
+
+def extract_token_usage_from_turn(turn: Turn, system_prompt: str = "") -> TokenCounter:
+    """Extract token usage information from a turn.
+
+    This function uses the same tokenizer and logic as the metrics system
+    to ensure consistency between API responses and Prometheus metrics.
+
+    Args:
+        turn: The turn object containing token usage information
+        system_prompt: The system prompt used for the turn
+
+    Returns:
+        TokenCounter: Token usage information
+    """
+    token_counter = TokenCounter()
+
+    try:
+        # Use the same tokenizer as the metrics system for consistency
+        tokenizer = Tokenizer.get_instance()
+        formatter = ChatFormat(tokenizer)
+
+        # Count output tokens (same logic as metrics.utils.update_llm_token_count_from_turn)
+        if hasattr(turn, "output_message") and turn.output_message:
+            raw_message = cast(RawMessage, turn.output_message)
+            encoded_output = formatter.encode_dialog_prompt([raw_message])
+            token_counter.output_tokens = (
+                len(encoded_output.tokens) if encoded_output.tokens else 0
+            )
+
+        # Count input tokens (same logic as metrics.utils.update_llm_token_count_from_turn)
+        if hasattr(turn, "input_messages") and turn.input_messages:
+            input_messages = cast(list[RawMessage], turn.input_messages)
+            if system_prompt:
+                input_messages = [
+                    RawMessage(role="system", content=system_prompt)
+                ] + input_messages
+            encoded_input = formatter.encode_dialog_prompt(input_messages)
+            token_counter.input_tokens = (
+                len(encoded_input.tokens) if encoded_input.tokens else 0
+            )
+            token_counter.input_tokens_counted = token_counter.input_tokens
+
+        token_counter.llm_calls = 1
+
+    except (AttributeError, TypeError, ValueError) as e:
+        logger.warning("Failed to extract token usage from turn: %s", e)
+        # Fallback to default values if token counting fails
+        token_counter.input_tokens = 100  # Default estimate
+        token_counter.output_tokens = 50  # Default estimate
+        token_counter.llm_calls = 1
+
+    return token_counter
+
+
+def extract_and_update_token_metrics(
+    turn: Turn, model: str, provider: str, system_prompt: str = ""
+) -> TokenCounter:
+    """Extract token usage and update Prometheus metrics in one call.
+
+    This function combines the token counting logic with the metrics system
+    to ensure both API responses and Prometheus metrics are updated consistently.
+
+    Args:
+        turn: The turn object containing token usage information
+        model: The model identifier for metrics labeling
+        provider: The provider identifier for metrics labeling
+        system_prompt: The system prompt used for the turn
+
+    Returns:
+        TokenCounter: Token usage information
+    """
+    token_counter = extract_token_usage_from_turn(turn, system_prompt)
+
+    # Update Prometheus metrics with the same token counts
+    try:
+        # Update the metrics using the same token counts we calculated
+        metrics.llm_token_sent_total.labels(provider, model).inc(
+            token_counter.input_tokens
+        )
+        metrics.llm_token_received_total.labels(provider, model).inc(
+            token_counter.output_tokens
+        )
+        metrics.llm_calls_total.labels(provider, model).inc()
+
+    except (AttributeError, TypeError, ValueError) as e:
+        logger.warning("Failed to update token metrics: %s", e)
+
+    return token_counter