diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql new file mode 100644 index 00000000..3011c33b --- /dev/null +++ b/infra/dashboards/token-usage-queries.kql @@ -0,0 +1,199 @@ +// ============================================================ +// KQL Queries for LLM Token Usage Monitoring +// Run these in Application Insights > Logs +// ============================================================ + +// 1. Overall token usage summary (last 7 days) +customEvents +| where name == 'LLM_Token_Usage_Summary' +| where timestamp > ago(7d) +| extend input_tokens = toint(customDimensions['total_input_tokens']) +| extend output_tokens = toint(customDimensions['total_output_tokens']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + TotalRequests = count(), + TotalInputTokens = sum(input_tokens), + TotalOutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + AvgTokensPerRequest = round(avg(total_tokens), 0); + +// 2. Token usage by agent +customEvents +| where name == 'LLM_Agent_Token_Usage' +| where timestamp > ago(7d) +| extend agent = tostring(customDimensions['agent_name']) +| extend input_tokens = toint(customDimensions['input_tokens']) +| extend output_tokens = toint(customDimensions['output_tokens']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + Invocations = count() + by Agent = agent +| order by TotalTokens desc; + +// 3. Token usage over time (hourly) +customEvents +| where name == 'LLM_Token_Usage_Summary' +| where timestamp > ago(7d) +| extend input_tokens = toint(customDimensions['total_input_tokens']) +| extend output_tokens = toint(customDimensions['total_output_tokens']) +| summarize InputTokens = sum(input_tokens), OutputTokens = sum(output_tokens) by bin(timestamp, 1h) +| order by timestamp asc +| render areachart; + +// 4. Estimated cost (GPT-4o pricing: $2.50/1M input, $10.00/1M output) +let input_price_per_million = 2.50; +let output_price_per_million = 10.00; +customEvents +| where name == 'LLM_Token_Usage_Summary' +| where timestamp > ago(30d) +| extend input_tokens = toint(customDimensions['total_input_tokens']) +| extend output_tokens = toint(customDimensions['total_output_tokens']) +| summarize TotalInput = sum(input_tokens), TotalOutput = sum(output_tokens) by bin(timestamp, 1d) +| extend InputCost = round(TotalInput * input_price_per_million / 1000000.0, 4) +| extend OutputCost = round(TotalOutput * output_price_per_million / 1000000.0, 4) +| extend TotalCost = InputCost + OutputCost +| project Day = timestamp, TotalInput, TotalOutput, InputCost, OutputCost, TotalCost +| order by Day desc; + +// 5. Top token consumers by user +customEvents +| where name == 'LLM_Token_Usage_Summary' +| where timestamp > ago(7d) +| extend total_tokens = toint(customDimensions['total_tokens']) +| extend user_id = tostring(customDimensions['user_id']) +| summarize TotalTokens = sum(total_tokens), Requests = count() by user_id +| order by TotalTokens desc +| take 20; + +// 6. Agent token distribution (pie chart) +customEvents +| where name == 'LLM_Agent_Token_Usage' +| where timestamp > ago(7d) +| extend agent = tostring(customDimensions['agent_name']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize TotalTokens = sum(total_tokens) by agent +| render piechart; + +// 7. Token usage percentiles per task +customEvents +| where name == 'LLM_Token_Usage_Summary' +| where timestamp > ago(7d) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + p50 = percentile(total_tokens, 50), + p90 = percentile(total_tokens, 90), + p95 = percentile(total_tokens, 95), + p99 = percentile(total_tokens, 99), + Max = max(total_tokens); + +// 8. Token usage by model deployment +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(7d) +| extend model = tostring(customDimensions['model_deployment_name']) +| extend input_tokens = toint(customDimensions['input_tokens']) +| extend output_tokens = toint(customDimensions['output_tokens']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + Invocations = count() + by Model = model +| order by TotalTokens desc; + +// 9. Token usage by model over time (hourly) +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(7d) +| extend model = tostring(customDimensions['model_deployment_name']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize TotalTokens = sum(total_tokens) by bin(timestamp, 1h), model +| order by timestamp asc +| render areachart; + +// 10. Model token distribution (pie chart) +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(7d) +| extend model = tostring(customDimensions['model_deployment_name']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize TotalTokens = sum(total_tokens) by model +| render piechart; + +// 11. Estimated cost by model (adjust pricing per model) +let gpt41_input = 2.00; +let gpt41_output = 8.00; +let gpt41_mini_input = 0.40; +let gpt41_mini_output = 1.60; +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(30d) +| extend model = tolower(tostring(customDimensions['model_deployment_name'])) +| extend input_tokens = toint(customDimensions['input_tokens']) +| extend output_tokens = toint(customDimensions['output_tokens']) +| summarize TotalInput = sum(input_tokens), TotalOutput = sum(output_tokens) by model +| extend InputPrice = case( + model has "4.1-mini", gpt41_mini_input, + model has "4.1", gpt41_input, + gpt41_mini_input) +| extend OutputPrice = case( + model has "4.1-mini", gpt41_mini_output, + model has "4.1", gpt41_output, + gpt41_mini_output) +| extend InputCost = round(TotalInput * InputPrice / 1000000.0, 4) +| extend OutputCost = round(TotalOutput * OutputPrice / 1000000.0, 4) +| extend TotalCost = InputCost + OutputCost +| project Model = model, TotalInput, TotalOutput, InputCost, OutputCost, TotalCost +| order by TotalCost desc; + +// 12. Agent-to-model mapping with token usage +customEvents +| where name == 'LLM_Agent_Token_Usage' +| where timestamp > ago(7d) +| extend agent = tostring(customDimensions['agent_name']) +| extend model = tostring(customDimensions['model_deployment_name']) +| extend input_tokens = toint(customDimensions['input_tokens']) +| extend output_tokens = toint(customDimensions['output_tokens']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + Invocations = count() + by Agent = agent, Model = model +| order by TotalTokens desc; + +// 13. OpenTelemetry auto-instrumented OpenAI calls (if available) +dependencies +| where name has "openai" or target has "openai" +| where timestamp > ago(7d) +| extend input_tokens = tolong(customDimensions["gen_ai.usage.input_tokens"]) +| extend output_tokens = tolong(customDimensions["gen_ai.usage.output_tokens"]) +| extend model = tostring(customDimensions["gen_ai.request.model"]) +| where isnotnull(input_tokens) +| summarize + Calls = count(), + TotalInput = sum(input_tokens), + TotalOutput = sum(output_tokens) + by model +| order by TotalInput desc; + +// 14. Total time taken per invocation (conversation) +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(7d) +| extend conversation_id = tostring(customDimensions['conversation_id']) +| extend model = tostring(customDimensions['model_deployment_name']) +| extend total_tokens = toint(customDimensions['total_tokens']) +| summarize + StartTime = min(timestamp), + EndTime = max(timestamp), + TotalTokens = sum(total_tokens), + Calls = count() + by conversation_id +| extend TotalDuration_sec = round(datetime_diff('second', EndTime, StartTime), 0) +| order by EndTime desc; diff --git a/src/api/python/chat.py b/src/api/python/chat.py index c47d9a03..e3539c55 100644 --- a/src/api/python/chat.py +++ b/src/api/python/chat.py @@ -18,7 +18,6 @@ # Azure SDK from azure.core.exceptions import HttpResponseError -from azure.monitor.events.extension import track_event from azure.ai.projects.aio import AIProjectClient # Agent Framework @@ -28,6 +27,9 @@ from auth.auth_utils import get_authenticated_user_details from auth.azure_credential_utils import get_azure_credential_async +# Token usage telemetry helpers (modular, reusable) +from token_usage import track_event_if_configured, UsageAccumulator + load_dotenv() # Constants @@ -108,15 +110,6 @@ async def _delete_thread_async(self, thread_conversation_id: str): await credential.close() -def track_event_if_configured(event_name: str, event_data: dict): - """Track event to Application Insights if configured.""" - instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING") - if instrumentation_key: - track_event(event_name, event_data) - else: - logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name) - - # Global thread cache thread_cache = None @@ -141,6 +134,8 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" complete_response = "" credential = None db_connection = None + # Accumulator for LLM token usage across all iterations of this request + usage = UsageAccumulator() try: if not query: @@ -184,6 +179,9 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" extra_body={"agent_reference": {"name": os.getenv("AGENT_NAME_CHAT"), "type": "agent_reference"}} ) + # Accumulate token usage from initial response + usage.add_from_response(response) + # Process response - handle function calls iteratively max_iterations = 10 iteration = 0 @@ -258,6 +256,9 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" extra_body={"agent_reference": {"name": os.getenv("AGENT_NAME_CHAT"), "type": "agent_reference"}} ) + # Accumulate token usage from follow-up response + usage.add_from_response(response) + if iteration >= max_iterations: logger.warning("Max iterations reached for conversation %s", conversation_id) yield "\n\n(Response processing reached maximum iterations)" @@ -270,6 +271,14 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" "response_length": str(len(complete_response)), }) + # Emit LLM token usage telemetry + usage.emit( + agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + user_id=user_id, + conversation_id=conversation_id, + ) + except HttpResponseError as e: complete_response = str(e) if "Rate limit is exceeded" in str(e) or e.status_code == 429: @@ -348,6 +357,8 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: complete_response = "" credential = None db_connection = None + # Accumulator for LLM token usage across all streaming updates + usage = UsageAccumulator() try: if not query: @@ -420,6 +431,9 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: if raw_repr: _extract_mcp_from_raw(raw_repr, mcp_docs) + # Accumulate token usage from this streaming update + usage.add_from_update(chunk) + chunk_text = str(chunk.text) if chunk.text else "" if not chunk_text: continue @@ -473,6 +487,14 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: "citation_count": str(len(mcp_docs)), }) + # Emit LLM token usage telemetry + usage.emit( + agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + user_id=user_id, + conversation_id=conversation_id, + ) + # Yield citations as a tool message — deduplicated by source citation_list = [] if original_markers: diff --git a/src/api/python/history.py b/src/api/python/history.py index 590c10ad..2daf1523 100644 --- a/src/api/python/history.py +++ b/src/api/python/history.py @@ -8,7 +8,6 @@ from azure.core.exceptions import HttpResponseError from azure.cosmos import exceptions from azure.cosmos.aio import CosmosClient -from azure.monitor.events.extension import track_event from fastapi import APIRouter, HTTPException, Query, Request, status from fastapi.responses import JSONResponse from opentelemetry import trace @@ -18,6 +17,9 @@ from auth.auth_utils import get_authenticated_user_details from auth.azure_credential_utils import get_azure_credential_async +# Token usage telemetry helpers (modular, reusable) +from token_usage import track_event_if_configured, UsageAccumulator + router = APIRouter() logger = logging.getLogger(__name__) @@ -39,15 +41,6 @@ AGENT_NAME_TITLE = os.getenv("AGENT_NAME_TITLE") -def track_event_if_configured(event_name: str, event_data: dict): - """Track events to Application Insights if configured.""" - instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING") - if instrumentation_key: - track_event(event_name, event_data) - else: - logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name) - - class CosmosConversationClient: """Client for managing conversations and messages in CosmosDB.""" @@ -279,7 +272,7 @@ async def init_cosmosdb_client(): raise -async def generate_title(conversation_messages): +async def generate_title(conversation_messages, user_id: str = "", conversation_id: str = ""): """Generate a title for a conversation using Azure AI Foundry agent.""" try: user_messages = [msg for msg in conversation_messages if msg["role"] == "user"] @@ -314,6 +307,15 @@ async def generate_title(conversation_messages): extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} ) + _title_usage = UsageAccumulator() + _title_usage.add_from_response(response) + _title_usage.emit( + agent_name=AGENT_NAME_TITLE or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + user_id=user_id or "", + conversation_id=conversation_id or "", + ) + result_text = "" for item in response.output: if getattr(item, 'type', None) == 'message': @@ -363,7 +365,7 @@ async def add_conversation(user_id: str, request_json: dict): raise ValueError("CosmosDB is not configured or unavailable") if not conversation_id: - title = await generate_title(messages) + title = await generate_title(messages, user_id=user_id) conversation_dict = await cosmos_conversation_client.create_conversation(user_id, title) conversation_id = conversation_dict["id"] history_metadata["title"] = title @@ -397,7 +399,7 @@ async def update_conversation(user_id: str, request_json: dict): # Retrieve or create conversation conversation = await cosmos_conversation_client.get_conversation(user_id, conversation_id) if not conversation: - title = await generate_title(messages) + title = await generate_title(messages, user_id=user_id, conversation_id=conversation_id) conversation = await cosmos_conversation_client.create_conversation( user_id=user_id, conversation_id=conversation_id, title=title ) diff --git a/src/api/python/history_sql.py b/src/api/python/history_sql.py index 1f775c6e..bb22ddff 100644 --- a/src/api/python/history_sql.py +++ b/src/api/python/history_sql.py @@ -10,7 +10,6 @@ from pydantic import BaseModel, ConfigDict import pyodbc from azure.identity.aio import AzureCliCredential -from azure.monitor.events.extension import track_event from fastapi import APIRouter, HTTPException, Query, Request, status from fastapi.responses import JSONResponse from opentelemetry import trace @@ -21,6 +20,9 @@ from azure.core.exceptions import HttpResponseError +# Token usage telemetry helpers (modular, reusable) +from token_usage import track_event_if_configured, UsageAccumulator + router = APIRouter() logger = logging.getLogger(__name__) @@ -32,21 +34,6 @@ # Database configuration -def track_event_if_configured(event_name: str, event_data: dict): - """ - Track an event with Application Insights if configured. - - Args: - event_name (str): The name of the event to track. - event_data (dict): The data to associate with the event. - """ - instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING") - if instrumentation_key: - track_event(event_name, event_data) - else: - logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name) - - async def get_azure_sql_connection(): """ Get a connection to Azure SQL Server using DefaultAzureCredential. @@ -523,12 +510,14 @@ async def rename_conversation(user_id: str, conversation_id, title) -> bool: return False -async def generate_title(conversation_messages): +async def generate_title(conversation_messages, user_id: str = "", conversation_id: str = ""): """ Generate a concise title for a conversation using Azure AI Foundry agent. Args: conversation_messages (list): List of messages in the conversation. + user_id (str): Optional user id propagated to telemetry events. + conversation_id (str): Optional conversation id propagated to telemetry events. Returns: str: A 4-word or less title summarizing the conversation. @@ -549,29 +538,42 @@ async def generate_title(conversation_messages): logger.warning("Azure AI Agent endpoint not configured, using fallback title generation") return generate_fallback_title(conversation_messages) - async with AIProjectClient( - endpoint=AZURE_AI_AGENT_ENDPOINT, - credential=await get_azure_credential_async() - ) as project_client: - openai_client = project_client.get_openai_client() - conversation = await openai_client.conversations.create() - - response = await openai_client.responses.create( - conversation=conversation.id, - input=final_prompt, - extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} - ) + credential = await get_azure_credential_async() + try: + async with AIProjectClient( + endpoint=AZURE_AI_AGENT_ENDPOINT, + credential=credential + ) as project_client: + openai_client = project_client.get_openai_client() + conversation = await openai_client.conversations.create() + + response = await openai_client.responses.create( + conversation=conversation.id, + input=final_prompt, + extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} + ) - # Extract text from response output - result_text = "" - for item in response.output: - if getattr(item, 'type', None) == 'message': - if hasattr(item, 'content') and item.content is not None: - for content in item.content: - if hasattr(content, 'text'): - result_text += content.text + _title_usage = UsageAccumulator() + _title_usage.add_from_response(response) + _title_usage.emit( + agent_name=AGENT_NAME_TITLE or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + user_id=user_id or "", + conversation_id=conversation_id or "", + ) + + # Extract text from response output + result_text = "" + for item in response.output: + if getattr(item, 'type', None) == 'message': + if hasattr(item, 'content') and item.content is not None: + for content in item.content: + if hasattr(content, 'text'): + result_text += content.text - return result_text.strip() if result_text else generate_fallback_title(conversation_messages) + return result_text.strip() if result_text else generate_fallback_title(conversation_messages) + finally: + await credential.close() except HttpResponseError as sre: logger.warning("HttpResponseError generating title with Azure AI Foundry agent: %s", sre) @@ -756,7 +758,7 @@ async def update_conversation(user_id: str, request_json: dict): conversation = await run_query_params(query, (conversation_id,)) if not conversation or len(conversation) == 0: - title = await generate_title(messages) + title = await generate_title(messages, user_id=user_id, conversation_id=conversation_id) await create_conversation(user_id=user_id, conversation_id=conversation_id, title=title) messages = request_json["messages"] diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py new file mode 100644 index 00000000..6e1de992 --- /dev/null +++ b/src/api/python/token_usage.py @@ -0,0 +1,302 @@ +""" +Reusable LLM token usage tracking helpers. + +This module centralizes: + * Application Insights event emission (guarded by configuration). + * Token-usage extraction from various LLM SDK shapes + (dict payloads, agent_framework streaming updates, OpenAI Responses). + * Emission of standardized telemetry events: + - LLM_Token_Usage_Summary + - LLM_Agent_Token_Usage + - LLM_Model_Token_Usage + +Import these helpers from any module that needs to record LLM usage, +instead of duplicating the logic. +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional, Tuple + +from azure.monitor.events.extension import track_event + +logger = logging.getLogger(__name__) + +UsageTuple = Tuple[int, int, int] # (input_tokens, output_tokens, total_tokens) + +# Module-level flag so we only emit the "App Insights not configured" warning +# once per process. Without this, every track_event_if_configured() call (3+ per +# request from token usage tracking) would flood logs in dev/misconfigured envs. +_app_insights_warning_emitted = False + + +def track_event_if_configured(event_name: str, event_data: dict) -> None: + """Track event to Application Insights if a connection string is configured.""" + global _app_insights_warning_emitted + instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING") + if instrumentation_key: + track_event(event_name, event_data) + else: + if not _app_insights_warning_emitted: + logger.warning( + "Skipping track_event for %s as Application Insights is not " + "configured (further occurrences will be logged at DEBUG)", + event_name, + ) + _app_insights_warning_emitted = True + else: + logger.debug( + "Skipping track_event for %s as Application Insights is not configured", + event_name, + ) + + +def _first_non_none(*values, default=0): + """Return the first value that is not None; otherwise the default. + + Unlike ``a or b``, this preserves explicit zero values (treating ``0`` as + a legitimate token count rather than as "missing"). + """ + for v in values: + if v is not None: + return v + return default + + +def extract_usage_from_dict(d: dict) -> Optional[UsageTuple]: + """Extract (input, output, total) token counts from a usage-like dict.""" + if not isinstance(d, dict) or not d: + return None + try: + inp = _first_non_none( + d.get("input_token_count"), + d.get("prompt_tokens"), + d.get("input_tokens"), + default=0, + ) + out = _first_non_none( + d.get("output_token_count"), + d.get("completion_tokens"), + d.get("output_tokens"), + default=0, + ) + tot = _first_non_none( + d.get("total_token_count"), + d.get("total_tokens"), + default=None, + ) + inp_i, out_i = int(inp), int(out) + tot_i = int(tot) if tot is not None else inp_i + out_i + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) + return None + + +def _extract_usage_obj(usage_obj) -> Optional[UsageTuple]: + """Extract a UsageTuple from a dict or object exposing token attributes.""" + if usage_obj is None: + return None + if isinstance(usage_obj, dict): + return extract_usage_from_dict(usage_obj) + try: + inp = _first_non_none( + getattr(usage_obj, "prompt_tokens", None), + getattr(usage_obj, "input_tokens", None), + default=0, + ) + out = _first_non_none( + getattr(usage_obj, "completion_tokens", None), + getattr(usage_obj, "output_tokens", None), + default=0, + ) + tot = getattr(usage_obj, "total_tokens", None) + inp_i, out_i = int(inp), int(out) + tot_i = int(tot) if tot is not None else inp_i + out_i + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) + return None + + +def extract_usage_from_update(update) -> Optional[UsageTuple]: + """Extract (input, output, total) token counts from an agent_framework streaming update. + + Checks, in order: + 1. update.contents[*].usage_details (dict) + 2. update.contents[*].raw_representation.usage (workshop streaming) + 3. update.contents[*].raw_representation.response.usage + (OpenAI Responses completion event nested inside a content item) + 4. update.raw_representation.usage + 5. update.raw_representation.response.usage + (OpenAI Responses completion event surfaced at the update level) + """ + contents = getattr(update, "contents", None) or [] + for item in contents: + usage_details = getattr(item, "usage_details", None) + if isinstance(usage_details, dict): + result = extract_usage_from_dict(usage_details) + if result: + return result + item_raw = getattr(item, "raw_representation", None) + if item_raw is not None: + result = _extract_usage_obj(getattr(item_raw, "usage", None)) + if result: + return result + item_response = getattr(item_raw, "response", None) + if item_response is not None: + result = _extract_usage_obj(getattr(item_response, "usage", None)) + if result: + return result + + raw = getattr(update, "raw_representation", None) + if raw is not None: + result = _extract_usage_obj(getattr(raw, "usage", None)) + if result: + return result + response = getattr(raw, "response", None) + if response is not None: + result = _extract_usage_obj(getattr(response, "usage", None)) + if result: + return result + return None + + +def extract_usage_from_response(response) -> Optional[UsageTuple]: + """Extract (input, output, total) tokens from an OpenAI Responses API response object.""" + if response is None: + return None + usage_obj = getattr(response, "usage", None) + if usage_obj is None: + return None + if isinstance(usage_obj, dict): + return extract_usage_from_dict(usage_obj) + try: + inp = _first_non_none( + getattr(usage_obj, "input_tokens", None), + getattr(usage_obj, "prompt_tokens", None), + default=0, + ) + out = _first_non_none( + getattr(usage_obj, "output_tokens", None), + getattr(usage_obj, "completion_tokens", None), + default=0, + ) + tot = getattr(usage_obj, "total_tokens", None) + inp_i, out_i = int(inp), int(out) + tot_i = int(tot) if tot is not None else inp_i + out_i + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) + return None + + +def track_token_usage( + agent_name: str, + model_deployment_name: str, + input_tokens: int, + output_tokens: int, + total_tokens: int, + user_id: str = "", + conversation_id: str = "", +) -> None: + """Emit LLM token usage events to Application Insights. + + Emits three events: + - LLM_Token_Usage_Summary : overall totals per request + - LLM_Agent_Token_Usage : usage attributed to the agent + - LLM_Model_Token_Usage : usage attributed to the model deployment + """ + if total_tokens <= 0: + return + try: + track_event_if_configured("LLM_Token_Usage_Summary", { + "total_input_tokens": str(input_tokens), + "total_output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + "agent_count": "1", + "model_count": "1", + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + track_event_if_configured("LLM_Agent_Token_Usage", { + "agent_name": agent_name or "", + "input_tokens": str(input_tokens), + "output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + "model_deployment_name": model_deployment_name or "", + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + track_event_if_configured("LLM_Model_Token_Usage", { + "model_deployment_name": model_deployment_name or "", + "input_tokens": str(input_tokens), + "output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + logger.debug( + "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d", + agent_name, model_deployment_name, input_tokens, output_tokens, total_tokens, + ) + except Exception as e: + logger.warning("Failed to emit token usage telemetry: %s", e) + + +class UsageAccumulator: + """Accumulates LLM token usage across multiple model calls / streaming chunks. + + Typical use: + usage = UsageAccumulator() + usage.add_from_response(response) # OpenAI Responses object + async for chunk in agent.run(...): + usage.add_from_update(chunk) # agent_framework streaming update + usage.emit(agent_name, model_deployment_name, + user_id=user_id, conversation_id=conversation_id) + """ + + __slots__ = ("input", "output", "total") + + def __init__(self) -> None: + self.input = 0 + self.output = 0 + self.total = 0 + + def add(self, usage: Optional[UsageTuple]) -> None: + """Add a (input, output, total) usage tuple, ignoring None.""" + if usage: + self.input += usage[0] + self.output += usage[1] + self.total += usage[2] + + def add_from_response(self, response) -> None: + """Extract and add usage from an OpenAI Responses API response object.""" + self.add(extract_usage_from_response(response)) + + def add_from_update(self, update) -> None: + """Extract and add usage from an agent_framework streaming update.""" + self.add(extract_usage_from_update(update)) + + def emit( + self, + agent_name: str, + model_deployment_name: str, + user_id: str = "", + conversation_id: str = "", + ) -> None: + """Emit accumulated token usage telemetry to Application Insights.""" + track_token_usage( + agent_name=agent_name, + model_deployment_name=model_deployment_name, + input_tokens=self.input, + output_tokens=self.output, + total_tokens=self.total, + user_id=user_id, + conversation_id=conversation_id, + ) diff --git a/src/test/api/python/test_chat.py b/src/test/api/python/test_chat.py index a46c774c..3c04508a 100644 --- a/src/test/api/python/test_chat.py +++ b/src/test/api/python/test_chat.py @@ -141,7 +141,7 @@ def test_track_event_if_configured_without_key(self): """Test track_event when no instrumentation key is set.""" from chat import track_event_if_configured - with patch('chat.track_event') as mock_track: + with patch('token_usage.track_event') as mock_track: track_event_if_configured("TestEvent", {"key": "value"}) # Should not call track_event when no instrumentation key mock_track.assert_not_called() @@ -1067,11 +1067,14 @@ def test_app_insights_not_configured_on_import(self, monkeypatch): def test_track_event_if_configured_without_instrumentation_key(self, monkeypatch): """Test line 125: track_event_if_configured when APPLICATIONINSIGHTS_CONNECTION_STRING is not set.""" from chat import track_event_if_configured + import token_usage # Ensure no instrumentation key monkeypatch.delenv("APPLICATIONINSIGHTS_CONNECTION_STRING", raising=False) + # Reset the warn-once flag so we can deterministically observe the warning. + monkeypatch.setattr(token_usage, "_app_insights_warning_emitted", False) - with patch('chat.logging.warning') as mock_warning: + with patch('token_usage.logger.warning') as mock_warning: track_event_if_configured("test_event", {"data": "value"}) mock_warning.assert_called_once() assert "Skipping track_event" in str(mock_warning.call_args) @@ -1127,6 +1130,6 @@ def test_track_event_with_instrumentation_key(self, monkeypatch): monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "InstrumentationKey=test-key") - with patch('chat.track_event') as mock_track: + with patch('token_usage.track_event') as mock_track: track_event_if_configured("test_event", {"key": "value"}) mock_track.assert_called_once_with("test_event", {"key": "value"}) diff --git a/src/test/api/python/test_history.py b/src/test/api/python/test_history.py index 7b9622f0..2ad45f05 100644 --- a/src/test/api/python/test_history.py +++ b/src/test/api/python/test_history.py @@ -33,7 +33,7 @@ def test_configuration_loaded(self): def test_track_event_configured(self, monkeypatch): from history import track_event_if_configured monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "test") - with patch('history.track_event'): + with patch('token_usage.track_event'): track_event_if_configured("event", {}) def test_track_event_not_configured(self, monkeypatch): @@ -715,7 +715,77 @@ async def test_generate_title_success(self, monkeypatch): with patch('history.AIProjectClient', return_value=mock_project): result = await generate_title([{"role": "user", "content": "Hello"}]) assert result == "Generated Title" - + + @pytest.mark.asyncio + async def test_generate_title_emits_token_telemetry(self, monkeypatch): + """generate_title should emit the three LLM_* telemetry events with + the supplied user_id/conversation_id when response.usage is present.""" + from history import generate_title + + monkeypatch.setattr('history.AZURE_AI_AGENT_ENDPOINT', 'https://test.ai.azure.com') + monkeypatch.setattr('history.AGENT_NAME_TITLE', 'title-agent') + monkeypatch.setenv('APPLICATIONINSIGHTS_CONNECTION_STRING', 'test-conn-str') + monkeypatch.setenv('AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME', 'gpt-4.1-mini') + + mock_content = MagicMock() + mock_content.text = "Generated Title" + mock_item = MagicMock() + mock_item.type = 'message' + mock_item.content = [mock_content] + mock_response = MagicMock() + mock_response.output = [mock_item] + mock_response.usage = { + "input_tokens": 12, + "output_tokens": 7, + "total_tokens": 19, + } + + mock_conversation = MagicMock() + mock_conversation.id = "conv-id" + + mock_openai = AsyncMock() + mock_openai.conversations.create = AsyncMock(return_value=mock_conversation) + mock_openai.responses.create = AsyncMock(return_value=mock_response) + + mock_project = AsyncMock() + mock_project.get_openai_client = MagicMock(return_value=mock_openai) + mock_project.__aenter__ = AsyncMock(return_value=mock_project) + mock_project.__aexit__ = AsyncMock(return_value=False) + + mock_credential = AsyncMock() + mock_credential.close = AsyncMock() + + with patch('history.get_azure_credential_async', return_value=mock_credential): + with patch('history.AIProjectClient', return_value=mock_project): + with patch('token_usage.track_event') as mock_track: + result = await generate_title( + [{"role": "user", "content": "Hello"}], + user_id="user-42", + conversation_id="conv-99", + ) + + assert result == "Generated Title" + event_names = [call.args[0] for call in mock_track.call_args_list] + assert "LLM_Token_Usage_Summary" in event_names + assert "LLM_Agent_Token_Usage" in event_names + assert "LLM_Model_Token_Usage" in event_names + + events_by_name = {call.args[0]: call.args[1] for call in mock_track.call_args_list} + summary = events_by_name["LLM_Token_Usage_Summary"] + assert summary.get("user_id") == "user-42" + assert summary.get("conversation_id") == "conv-99" + assert int(summary.get("total_tokens") or summary.get("total_token_count") or 0) == 19 + + agent_evt = events_by_name["LLM_Agent_Token_Usage"] + assert agent_evt.get("agent_name") == "title-agent" + assert agent_evt.get("user_id") == "user-42" + assert agent_evt.get("conversation_id") == "conv-99" + + model_evt = events_by_name["LLM_Model_Token_Usage"] + assert model_evt.get("model_deployment_name") == "gpt-4.1-mini" + assert model_evt.get("user_id") == "user-42" + assert model_evt.get("conversation_id") == "conv-99" + @pytest.mark.asyncio async def test_generate_title_fallback(self, monkeypatch): from history import generate_title diff --git a/src/test/api/python/test_history_sql.py b/src/test/api/python/test_history_sql.py index 8e7d12a4..8e2dbec3 100644 --- a/src/test/api/python/test_history_sql.py +++ b/src/test/api/python/test_history_sql.py @@ -266,7 +266,7 @@ def test_track_event_with_instrumentation_key(self, monkeypatch): monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "InstrumentationKey=test") - with patch('history_sql.track_event') as mock_track: + with patch('token_usage.track_event') as mock_track: track_event_if_configured("TestEvent", {"key": "value"}) mock_track.assert_called_once_with("TestEvent", {"key": "value"}) @@ -276,7 +276,7 @@ def test_track_event_without_instrumentation_key(self, monkeypatch): monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "") - with patch('history_sql.track_event') as mock_track: + with patch('token_usage.track_event') as mock_track: track_event_if_configured("TestEvent", {"key": "value"}) mock_track.assert_not_called() @@ -710,6 +710,73 @@ async def test_generate_title_with_agent(self): assert isinstance(result, str) assert result == "AI Generated Title" + @pytest.mark.asyncio + async def test_generate_title_emits_token_telemetry(self, monkeypatch): + """generate_title should emit the three LLM_* telemetry events with the + supplied user_id/conversation_id when response.usage is present.""" + from history_sql import generate_title + + messages = [{"role": "user", "content": "Hello"}] + monkeypatch.setenv('APPLICATIONINSIGHTS_CONNECTION_STRING', 'test-conn-str') + monkeypatch.setenv('AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME', 'gpt-4.1-mini') + + with patch('history_sql.AZURE_AI_AGENT_ENDPOINT', 'http://test'), \ + patch('history_sql.AGENT_NAME_TITLE', 'title-agent'), \ + patch('history_sql.AIProjectClient') as mock_client, \ + patch('history_sql.get_azure_credential_async') as mock_cred, \ + patch('token_usage.track_event') as mock_track: + + mock_cred.return_value = AsyncMock() + + mock_project = AsyncMock() + mock_openai = AsyncMock() + mock_conv = Mock(id="title_thread") + mock_openai.conversations.create = AsyncMock(return_value=mock_conv) + + mock_response = Mock() + mock_message_item = Mock() + mock_message_item.type = 'message' + mock_content = Mock() + mock_content.text = "AI Generated Title" + mock_message_item.content = [mock_content] + mock_response.output = [mock_message_item] + mock_response.usage = { + "input_tokens": 12, + "output_tokens": 7, + "total_tokens": 19, + } + mock_openai.responses.create = AsyncMock(return_value=mock_response) + + mock_project.get_openai_client = Mock(return_value=mock_openai) + mock_client.return_value.__aenter__ = AsyncMock(return_value=mock_project) + mock_client.return_value.__aexit__ = AsyncMock() + + result = await generate_title( + messages, user_id="user-42", conversation_id="conv-99" + ) + + assert result == "AI Generated Title" + event_names = [call.args[0] for call in mock_track.call_args_list] + assert "LLM_Token_Usage_Summary" in event_names + assert "LLM_Agent_Token_Usage" in event_names + assert "LLM_Model_Token_Usage" in event_names + + events_by_name = {call.args[0]: call.args[1] for call in mock_track.call_args_list} + summary = events_by_name["LLM_Token_Usage_Summary"] + assert summary.get("user_id") == "user-42" + assert summary.get("conversation_id") == "conv-99" + assert int(summary.get("total_tokens") or 0) == 19 + + agent_evt = events_by_name["LLM_Agent_Token_Usage"] + assert agent_evt.get("agent_name") == "title-agent" + assert agent_evt.get("user_id") == "user-42" + assert agent_evt.get("conversation_id") == "conv-99" + + model_evt = events_by_name["LLM_Model_Token_Usage"] + assert model_evt.get("model_deployment_name") == "gpt-4.1-mini" + assert model_evt.get("user_id") == "user-42" + assert model_evt.get("conversation_id") == "conv-99" + class TestGenerateFallbackTitleFunction: """Tests for generate_fallback_title function.""" diff --git a/src/test/api/python/test_token_usage.py b/src/test/api/python/test_token_usage.py new file mode 100644 index 00000000..a75a2b09 --- /dev/null +++ b/src/test/api/python/test_token_usage.py @@ -0,0 +1,221 @@ +"""Unit tests for token_usage helpers.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from token_usage import ( + UsageAccumulator, + extract_usage_from_dict, + extract_usage_from_response, + extract_usage_from_update, + track_event_if_configured, + track_token_usage, +) + + +class TestExtractUsageFromDict: + def test_openai_style_keys(self): + assert extract_usage_from_dict( + {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + ) == (10, 5, 15) + + def test_agent_framework_style_keys(self): + assert extract_usage_from_dict( + {"input_token_count": 7, "output_token_count": 3, "total_token_count": 10} + ) == (7, 3, 10) + + def test_alternate_input_output_keys(self): + assert extract_usage_from_dict( + {"input_tokens": 4, "output_tokens": 6, "total_tokens": 10} + ) == (4, 6, 10) + + def test_total_derived_from_input_output(self): + # No total_tokens key -> total is inp + out + assert extract_usage_from_dict( + {"prompt_tokens": 2, "completion_tokens": 3} + ) == (2, 3, 5) + + def test_zero_total_returns_none(self): + assert extract_usage_from_dict({"prompt_tokens": 0, "completion_tokens": 0}) is None + + def test_empty_dict_returns_none(self): + assert extract_usage_from_dict({}) is None + + def test_non_dict_returns_none(self): + assert extract_usage_from_dict(None) is None + assert extract_usage_from_dict("not a dict") is None + + def test_non_int_values_are_coerced(self): + # int(MagicMock) yields 1 by default; non-coercible values fall back to None + class Bad: + def __int__(self): + raise ValueError("bad") + + assert extract_usage_from_dict( + {"prompt_tokens": Bad(), "completion_tokens": Bad(), "total_tokens": Bad()} + ) is None + + +class TestExtractUsageFromResponse: + def test_object_with_usage_attributes(self): + usage = SimpleNamespace(input_tokens=11, output_tokens=4, total_tokens=15) + response = SimpleNamespace(usage=usage) + assert extract_usage_from_response(response) == (11, 4, 15) + + def test_object_with_dict_usage(self): + response = SimpleNamespace( + usage={"prompt_tokens": 8, "completion_tokens": 2, "total_tokens": 10} + ) + assert extract_usage_from_response(response) == (8, 2, 10) + + def test_none_response_returns_none(self): + assert extract_usage_from_response(None) is None + + def test_response_without_usage_returns_none(self): + assert extract_usage_from_response(SimpleNamespace(usage=None)) is None + + def test_non_coercible_usage_returns_none(self): + class Bad: + def __int__(self): + raise TypeError("bad") + + response = SimpleNamespace( + usage=SimpleNamespace(input_tokens=Bad(), output_tokens=Bad(), total_tokens=Bad()) + ) + assert extract_usage_from_response(response) is None + + +class TestExtractUsageFromUpdate: + def test_contents_usage_details(self): + item = SimpleNamespace( + usage_details={"input_token_count": 5, "output_token_count": 2, "total_token_count": 7} + ) + update = SimpleNamespace(contents=[item]) + assert extract_usage_from_update(update) == (5, 2, 7) + + def test_raw_representation_usage_object(self): + raw = SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=9, completion_tokens=1, total_tokens=10) + ) + update = SimpleNamespace(contents=[], raw_representation=raw) + assert extract_usage_from_update(update) == (9, 1, 10) + + def test_raw_representation_usage_dict(self): + raw = SimpleNamespace(usage={"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}) + update = SimpleNamespace(contents=[], raw_representation=raw) + assert extract_usage_from_update(update) == (3, 4, 7) + + def test_no_usage_returns_none(self): + update = SimpleNamespace(contents=[], raw_representation=None) + assert extract_usage_from_update(update) is None + + +class TestTrackEventIfConfigured: + def test_calls_track_event_when_configured(self, monkeypatch): + monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "InstrumentationKey=x") + with patch("token_usage.track_event") as mock_track: + track_event_if_configured("E", {"k": "v"}) + mock_track.assert_called_once_with("E", {"k": "v"}) + + def test_skips_when_not_configured(self, monkeypatch): + monkeypatch.delenv("APPLICATIONINSIGHTS_CONNECTION_STRING", raising=False) + with patch("token_usage.track_event") as mock_track: + track_event_if_configured("E", {"k": "v"}) + mock_track.assert_not_called() + + def test_warns_only_once_when_not_configured(self, monkeypatch): + """Repeated calls without App Insights should warn at most once + (subsequent calls log at DEBUG to avoid flooding logs).""" + import token_usage + + monkeypatch.delenv("APPLICATIONINSIGHTS_CONNECTION_STRING", raising=False) + monkeypatch.setattr(token_usage, "_app_insights_warning_emitted", False) + + with patch("token_usage.logger.warning") as mock_warning, \ + patch("token_usage.logger.debug") as mock_debug: + track_event_if_configured("E1", {}) + track_event_if_configured("E2", {}) + track_event_if_configured("E3", {}) + + assert mock_warning.call_count == 1 + assert mock_debug.call_count == 2 + + +class TestTrackTokenUsage: + def test_no_events_when_total_zero(self): + with patch("token_usage.track_event_if_configured") as mock_emit: + track_token_usage("agent", "model", 0, 0, 0) + mock_emit.assert_not_called() + + def test_emits_three_events(self, monkeypatch): + monkeypatch.setenv("APPLICATIONINSIGHTS_CONNECTION_STRING", "InstrumentationKey=x") + with patch("token_usage.track_event_if_configured") as mock_emit: + track_token_usage("agent", "model", 10, 5, 15, "u", "c") + assert mock_emit.call_count == 3 + event_names = [call.args[0] for call in mock_emit.call_args_list] + assert event_names == [ + "LLM_Token_Usage_Summary", + "LLM_Agent_Token_Usage", + "LLM_Model_Token_Usage", + ] + + +class TestUsageAccumulator: + def test_initial_state(self): + acc = UsageAccumulator() + assert (acc.input, acc.output, acc.total) == (0, 0, 0) + + def test_add_tuple(self): + acc = UsageAccumulator() + acc.add((3, 4, 7)) + acc.add((1, 2, 3)) + assert (acc.input, acc.output, acc.total) == (4, 6, 10) + + def test_add_none_is_noop(self): + acc = UsageAccumulator() + acc.add(None) + assert (acc.input, acc.output, acc.total) == (0, 0, 0) + + def test_add_from_response(self): + acc = UsageAccumulator() + response = SimpleNamespace( + usage=SimpleNamespace(input_tokens=2, output_tokens=3, total_tokens=5) + ) + acc.add_from_response(response) + assert acc.total == 5 + + def test_add_from_update(self): + acc = UsageAccumulator() + item = SimpleNamespace( + usage_details={"input_token_count": 1, "output_token_count": 1, "total_token_count": 2} + ) + update = SimpleNamespace(contents=[item]) + acc.add_from_update(update) + assert acc.total == 2 + + def test_emit_delegates_to_track_token_usage(self): + acc = UsageAccumulator() + acc.add((10, 5, 15)) + with patch("token_usage.track_token_usage") as mock_track: + acc.emit("agent", "model", user_id="u", conversation_id="c") + mock_track.assert_called_once_with( + agent_name="agent", + model_deployment_name="model", + input_tokens=10, + output_tokens=5, + total_tokens=15, + user_id="u", + conversation_id="c", + ) + + def test_emit_with_zero_does_not_track(self): + acc = UsageAccumulator() + with patch("token_usage.track_event_if_configured") as mock_emit: + acc.emit("agent", "model") + mock_emit.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])