From fb32a7ed787065d5810a402caae73e51859503ea Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 11:50:40 +0530 Subject: [PATCH 01/10] Added token usage taking detils code --- infra/dashboards/token-usage-queries.kql | 181 +++++++++++++++++++++++ src/api/python/chat.py | 165 +++++++++++++++++++++ src/api/python/history.py | 54 +++++++ src/api/python/history_sql.py | 54 +++++++ 4 files changed, 454 insertions(+) create mode 100644 infra/dashboards/token-usage-queries.kql diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql new file mode 100644 index 00000000..bb436518 --- /dev/null +++ b/infra/dashboards/token-usage-queries.kql @@ -0,0 +1,181 @@ +// ============================================================ +// 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 gpt4o_input = 2.50; +let gpt4o_output = 10.00; +let gpt4o_mini_input = 0.15; +let gpt4o_mini_output = 0.60; +customEvents +| where name == 'LLM_Model_Token_Usage' +| where timestamp > ago(30d) +| extend model = 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 "mini", gpt4o_mini_input, + gpt4o_input) +| extend OutputPrice = case( + model has "mini", gpt4o_mini_output, + gpt4o_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 diff --git a/src/api/python/chat.py b/src/api/python/chat.py index c47d9a03..4b7fabb2 100644 --- a/src/api/python/chat.py +++ b/src/api/python/chat.py @@ -117,6 +117,120 @@ def track_event_if_configured(event_name: str, event_data: dict): logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name) +def _extract_usage_from_dict(d: dict) -> "tuple[int, int, int] | None": + """Extract (input, output, total) token counts from a usage-like dict.""" + if not isinstance(d, dict) or not d: + return None + inp = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 + out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 + tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + return None + + +def _extract_usage_from_update(update) -> "tuple[int, int, int] | None": + """Extract (input, output, total) token counts from an agent_framework streaming update. + + Checks, in order: + 1. update.contents[*].usage_details (dict) + 2. update.raw_representation.usage (dict or object with token attributes) + """ + 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 + + raw = getattr(update, "raw_representation", None) + if raw is not None: + usage_obj = getattr(raw, "usage", None) + if usage_obj is not None: + if isinstance(usage_obj, dict): + result = _extract_usage_from_dict(usage_obj) + if result: + return result + else: + inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 + out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + return None + + +def _extract_usage_from_response(response) -> "tuple[int, int, int] | None": + """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) + inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 + out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + 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.info( + "[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) + + # Global thread cache thread_cache = None @@ -141,6 +255,10 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" complete_response = "" credential = None db_connection = None + # Accumulators for LLM token usage across all iterations of this request + total_input_tokens = 0 + total_output_tokens = 0 + total_tokens_sum = 0 try: if not query: @@ -184,6 +302,13 @@ 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 = _extract_usage_from_response(response) + if _usage: + total_input_tokens += _usage[0] + total_output_tokens += _usage[1] + total_tokens_sum += _usage[2] + # Process response - handle function calls iteratively max_iterations = 10 iteration = 0 @@ -258,6 +383,13 @@ 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 = _extract_usage_from_response(response) + if _usage: + total_input_tokens += _usage[0] + total_output_tokens += _usage[1] + total_tokens_sum += _usage[2] + if iteration >= max_iterations: logger.warning("Max iterations reached for conversation %s", conversation_id) yield "\n\n(Response processing reached maximum iterations)" @@ -270,6 +402,17 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" "response_length": str(len(complete_response)), }) + # Emit LLM token usage telemetry + _track_token_usage( + agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + input_tokens=total_input_tokens, + output_tokens=total_output_tokens, + total_tokens=total_tokens_sum, + 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 +491,10 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: complete_response = "" credential = None db_connection = None + # Accumulators for LLM token usage across all streaming updates + total_input_tokens = 0 + total_output_tokens = 0 + total_tokens_sum = 0 try: if not query: @@ -420,6 +567,13 @@ 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 = _extract_usage_from_update(chunk) + if _usage: + total_input_tokens += _usage[0] + total_output_tokens += _usage[1] + total_tokens_sum += _usage[2] + chunk_text = str(chunk.text) if chunk.text else "" if not chunk_text: continue @@ -473,6 +627,17 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: "citation_count": str(len(mcp_docs)), }) + # Emit LLM token usage telemetry + _track_token_usage( + agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", + model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", + input_tokens=total_input_tokens, + output_tokens=total_output_tokens, + total_tokens=total_tokens_sum, + 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..a7bcd10a 100644 --- a/src/api/python/history.py +++ b/src/api/python/history.py @@ -48,6 +48,58 @@ def track_event_if_configured(event_name: str, event_data: dict): logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name) +def _track_title_token_usage(response, agent_name: str, user_id: str = "", conversation_id: str = "") -> None: + """Extract usage from a Responses API result and emit LLM_* telemetry events.""" + try: + usage_obj = getattr(response, "usage", None) + if usage_obj is None: + return + if isinstance(usage_obj, dict): + inp = usage_obj.get("input_tokens", 0) or usage_obj.get("prompt_tokens", 0) or 0 + out = usage_obj.get("output_tokens", 0) or usage_obj.get("completion_tokens", 0) or 0 + tot = usage_obj.get("total_tokens", 0) or (inp + out) + else: + inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 + out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if not tot or tot <= 0: + return + inp, out, tot = int(inp), int(out), int(tot) + model_deployment_name = os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "" + track_event_if_configured("LLM_Token_Usage_Summary", { + "total_input_tokens": str(inp), + "total_output_tokens": str(out), + "total_tokens": str(tot), + "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(inp), + "output_tokens": str(out), + "total_tokens": str(tot), + "model_deployment_name": model_deployment_name, + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + track_event_if_configured("LLM_Model_Token_Usage", { + "model_deployment_name": model_deployment_name, + "input_tokens": str(inp), + "output_tokens": str(out), + "total_tokens": str(tot), + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + logger.info( + "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d", + agent_name, model_deployment_name, inp, out, tot, + ) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to emit token usage telemetry: %s", e) + + class CosmosConversationClient: """Client for managing conversations and messages in CosmosDB.""" @@ -314,6 +366,8 @@ async def generate_title(conversation_messages): extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} ) + _track_title_token_usage(response, agent_name=AGENT_NAME_TITLE or "") + result_text = "" for item in response.output: if getattr(item, 'type', None) == 'message': diff --git a/src/api/python/history_sql.py b/src/api/python/history_sql.py index 1f775c6e..3a41891a 100644 --- a/src/api/python/history_sql.py +++ b/src/api/python/history_sql.py @@ -32,6 +32,58 @@ # Database configuration +def _track_title_token_usage(response, agent_name: str, user_id: str = "", conversation_id: str = "") -> None: + """Extract usage from a Responses API result and emit LLM_* telemetry events.""" + try: + usage_obj = getattr(response, "usage", None) + if usage_obj is None: + return + if isinstance(usage_obj, dict): + inp = usage_obj.get("input_tokens", 0) or usage_obj.get("prompt_tokens", 0) or 0 + out = usage_obj.get("output_tokens", 0) or usage_obj.get("completion_tokens", 0) or 0 + tot = usage_obj.get("total_tokens", 0) or (inp + out) + else: + inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 + out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if not tot or tot <= 0: + return + inp, out, tot = int(inp), int(out), int(tot) + model_deployment_name = os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "" + track_event_if_configured("LLM_Token_Usage_Summary", { + "total_input_tokens": str(inp), + "total_output_tokens": str(out), + "total_tokens": str(tot), + "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(inp), + "output_tokens": str(out), + "total_tokens": str(tot), + "model_deployment_name": model_deployment_name, + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + track_event_if_configured("LLM_Model_Token_Usage", { + "model_deployment_name": model_deployment_name, + "input_tokens": str(inp), + "output_tokens": str(out), + "total_tokens": str(tot), + "user_id": user_id or "", + "conversation_id": conversation_id or "", + }) + logger.info( + "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d", + agent_name, model_deployment_name, inp, out, tot, + ) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to emit token usage telemetry: %s", e) + + def track_event_if_configured(event_name: str, event_data: dict): """ Track an event with Application Insights if configured. @@ -562,6 +614,8 @@ async def generate_title(conversation_messages): extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} ) + _track_title_token_usage(response, agent_name=AGENT_NAME_TITLE or "") + # Extract text from response output result_text = "" for item in response.output: From fdbb2afc195ee6284112b2db86428afa02ffb0d4 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 15:18:56 +0530 Subject: [PATCH 02/10] updated kql query --- infra/dashboards/token-usage-queries.kql | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql index bb436518..2d6be67b 100644 --- a/infra/dashboards/token-usage-queries.kql +++ b/infra/dashboards/token-usage-queries.kql @@ -125,23 +125,25 @@ customEvents | render piechart // 11. Estimated cost by model (adjust pricing per model) -let gpt4o_input = 2.50; -let gpt4o_output = 10.00; -let gpt4o_mini_input = 0.15; -let gpt4o_mini_output = 0.60; +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 = tostring(customDimensions['model_deployment_name']) +| 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 "mini", gpt4o_mini_input, - gpt4o_input) + model has "4.1-mini", gpt41_mini_input, + model has "4.1", gpt41_input, + gpt41_mini_input) | extend OutputPrice = case( - model has "mini", gpt4o_mini_output, - gpt4o_output) + 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 From b736781cee08fb7edebe28ae324508b3e1b04254 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 16:34:30 +0530 Subject: [PATCH 03/10] updated token usage --- src/api/python/chat.py | 132 ++--------------------------- src/api/python/token_usage.py | 153 ++++++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 124 deletions(-) create mode 100644 src/api/python/token_usage.py diff --git a/src/api/python/chat.py b/src/api/python/chat.py index 4b7fabb2..acffb704 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,14 @@ 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, + extract_usage_from_update as _extract_usage_from_update, + extract_usage_from_response as _extract_usage_from_response, + track_token_usage as _track_token_usage, +) + load_dotenv() # Constants @@ -108,129 +115,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) - - -def _extract_usage_from_dict(d: dict) -> "tuple[int, int, int] | None": - """Extract (input, output, total) token counts from a usage-like dict.""" - if not isinstance(d, dict) or not d: - return None - inp = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 - out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 - tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) - return None - - -def _extract_usage_from_update(update) -> "tuple[int, int, int] | None": - """Extract (input, output, total) token counts from an agent_framework streaming update. - - Checks, in order: - 1. update.contents[*].usage_details (dict) - 2. update.raw_representation.usage (dict or object with token attributes) - """ - 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 - - raw = getattr(update, "raw_representation", None) - if raw is not None: - usage_obj = getattr(raw, "usage", None) - if usage_obj is not None: - if isinstance(usage_obj, dict): - result = _extract_usage_from_dict(usage_obj) - if result: - return result - else: - inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 - out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) - return None - - -def _extract_usage_from_response(response) -> "tuple[int, int, int] | None": - """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) - inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 - out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) - 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.info( - "[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) - - # Global thread cache thread_cache = None diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py new file mode 100644 index 00000000..c22af786 --- /dev/null +++ b/src/api/python/token_usage.py @@ -0,0 +1,153 @@ +""" +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) + + +def track_event_if_configured(event_name: str, event_data: dict) -> None: + """Track event to Application Insights if a connection string is configured.""" + instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING") + if instrumentation_key: + track_event(event_name, event_data) + else: + logger.warning( + "Skipping track_event for %s as Application Insights is not configured", + event_name, + ) + + +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 + inp = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 + out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 + tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + 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.raw_representation.usage (dict or object with token attributes) + """ + 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 + + raw = getattr(update, "raw_representation", None) + if raw is not None: + usage_obj = getattr(raw, "usage", None) + if usage_obj is not None: + if isinstance(usage_obj, dict): + result = extract_usage_from_dict(usage_obj) + if result: + return result + else: + inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 + out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + 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) + inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 + out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + if tot > 0: + return (int(inp), int(out), int(tot)) + 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.info( + "[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) From abf09382e5cd391b9f8d8db7bf22cc67033f1c05 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 16:59:38 +0530 Subject: [PATCH 04/10] optimized token usage --- src/api/python/chat.py | 47 +++++---------------- src/api/python/history.py | 72 ++++---------------------------- src/api/python/history_sql.py | 77 ++++------------------------------- src/api/python/token_usage.py | 53 ++++++++++++++++++++++++ 4 files changed, 81 insertions(+), 168 deletions(-) diff --git a/src/api/python/chat.py b/src/api/python/chat.py index acffb704..e3539c55 100644 --- a/src/api/python/chat.py +++ b/src/api/python/chat.py @@ -28,12 +28,7 @@ from auth.azure_credential_utils import get_azure_credential_async # Token usage telemetry helpers (modular, reusable) -from token_usage import ( - track_event_if_configured, - extract_usage_from_update as _extract_usage_from_update, - extract_usage_from_response as _extract_usage_from_response, - track_token_usage as _track_token_usage, -) +from token_usage import track_event_if_configured, UsageAccumulator load_dotenv() @@ -139,10 +134,8 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" complete_response = "" credential = None db_connection = None - # Accumulators for LLM token usage across all iterations of this request - total_input_tokens = 0 - total_output_tokens = 0 - total_tokens_sum = 0 + # Accumulator for LLM token usage across all iterations of this request + usage = UsageAccumulator() try: if not query: @@ -187,11 +180,7 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" ) # Accumulate token usage from initial response - _usage = _extract_usage_from_response(response) - if _usage: - total_input_tokens += _usage[0] - total_output_tokens += _usage[1] - total_tokens_sum += _usage[2] + usage.add_from_response(response) # Process response - handle function calls iteratively max_iterations = 10 @@ -268,11 +257,7 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" ) # Accumulate token usage from follow-up response - _usage = _extract_usage_from_response(response) - if _usage: - total_input_tokens += _usage[0] - total_output_tokens += _usage[1] - total_tokens_sum += _usage[2] + usage.add_from_response(response) if iteration >= max_iterations: logger.warning("Max iterations reached for conversation %s", conversation_id) @@ -287,12 +272,9 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = "" }) # Emit LLM token usage telemetry - _track_token_usage( + usage.emit( agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", - input_tokens=total_input_tokens, - output_tokens=total_output_tokens, - total_tokens=total_tokens_sum, user_id=user_id, conversation_id=conversation_id, ) @@ -375,10 +357,8 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: complete_response = "" credential = None db_connection = None - # Accumulators for LLM token usage across all streaming updates - total_input_tokens = 0 - total_output_tokens = 0 - total_tokens_sum = 0 + # Accumulator for LLM token usage across all streaming updates + usage = UsageAccumulator() try: if not query: @@ -452,11 +432,7 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: _extract_mcp_from_raw(raw_repr, mcp_docs) # Accumulate token usage from this streaming update - _usage = _extract_usage_from_update(chunk) - if _usage: - total_input_tokens += _usage[0] - total_output_tokens += _usage[1] - total_tokens_sum += _usage[2] + usage.add_from_update(chunk) chunk_text = str(chunk.text) if chunk.text else "" if not chunk_text: @@ -512,12 +488,9 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id: }) # Emit LLM token usage telemetry - _track_token_usage( + usage.emit( agent_name=os.getenv("AGENT_NAME_CHAT", "") or "", model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "", - input_tokens=total_input_tokens, - output_tokens=total_output_tokens, - total_tokens=total_tokens_sum, user_id=user_id, conversation_id=conversation_id, ) diff --git a/src/api/python/history.py b/src/api/python/history.py index a7bcd10a..b118141b 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,67 +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) - - -def _track_title_token_usage(response, agent_name: str, user_id: str = "", conversation_id: str = "") -> None: - """Extract usage from a Responses API result and emit LLM_* telemetry events.""" - try: - usage_obj = getattr(response, "usage", None) - if usage_obj is None: - return - if isinstance(usage_obj, dict): - inp = usage_obj.get("input_tokens", 0) or usage_obj.get("prompt_tokens", 0) or 0 - out = usage_obj.get("output_tokens", 0) or usage_obj.get("completion_tokens", 0) or 0 - tot = usage_obj.get("total_tokens", 0) or (inp + out) - else: - inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 - out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if not tot or tot <= 0: - return - inp, out, tot = int(inp), int(out), int(tot) - model_deployment_name = os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "" - track_event_if_configured("LLM_Token_Usage_Summary", { - "total_input_tokens": str(inp), - "total_output_tokens": str(out), - "total_tokens": str(tot), - "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(inp), - "output_tokens": str(out), - "total_tokens": str(tot), - "model_deployment_name": model_deployment_name, - "user_id": user_id or "", - "conversation_id": conversation_id or "", - }) - track_event_if_configured("LLM_Model_Token_Usage", { - "model_deployment_name": model_deployment_name, - "input_tokens": str(inp), - "output_tokens": str(out), - "total_tokens": str(tot), - "user_id": user_id or "", - "conversation_id": conversation_id or "", - }) - logger.info( - "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d", - agent_name, model_deployment_name, inp, out, tot, - ) - except Exception as e: # noqa: BLE001 - logger.warning("Failed to emit token usage telemetry: %s", e) - - class CosmosConversationClient: """Client for managing conversations and messages in CosmosDB.""" @@ -366,7 +307,12 @@ async def generate_title(conversation_messages): extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} ) - _track_title_token_usage(response, agent_name=AGENT_NAME_TITLE or "") + _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 "", + ) result_text = "" for item in response.output: diff --git a/src/api/python/history_sql.py b/src/api/python/history_sql.py index 3a41891a..53ab51a0 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,72 +34,6 @@ # Database configuration -def _track_title_token_usage(response, agent_name: str, user_id: str = "", conversation_id: str = "") -> None: - """Extract usage from a Responses API result and emit LLM_* telemetry events.""" - try: - usage_obj = getattr(response, "usage", None) - if usage_obj is None: - return - if isinstance(usage_obj, dict): - inp = usage_obj.get("input_tokens", 0) or usage_obj.get("prompt_tokens", 0) or 0 - out = usage_obj.get("output_tokens", 0) or usage_obj.get("completion_tokens", 0) or 0 - tot = usage_obj.get("total_tokens", 0) or (inp + out) - else: - inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 - out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if not tot or tot <= 0: - return - inp, out, tot = int(inp), int(out), int(tot) - model_deployment_name = os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "" - track_event_if_configured("LLM_Token_Usage_Summary", { - "total_input_tokens": str(inp), - "total_output_tokens": str(out), - "total_tokens": str(tot), - "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(inp), - "output_tokens": str(out), - "total_tokens": str(tot), - "model_deployment_name": model_deployment_name, - "user_id": user_id or "", - "conversation_id": conversation_id or "", - }) - track_event_if_configured("LLM_Model_Token_Usage", { - "model_deployment_name": model_deployment_name, - "input_tokens": str(inp), - "output_tokens": str(out), - "total_tokens": str(tot), - "user_id": user_id or "", - "conversation_id": conversation_id or "", - }) - logger.info( - "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d", - agent_name, model_deployment_name, inp, out, tot, - ) - except Exception as e: # noqa: BLE001 - logger.warning("Failed to emit token usage telemetry: %s", e) - - -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(): """ @@ -614,7 +550,12 @@ async def generate_title(conversation_messages): extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}} ) - _track_title_token_usage(response, agent_name=AGENT_NAME_TITLE or "") + _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 "", + ) # Extract text from response output result_text = "" diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py index c22af786..bba83c74 100644 --- a/src/api/python/token_usage.py +++ b/src/api/python/token_usage.py @@ -151,3 +151,56 @@ def track_token_usage( ) 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, + ) From 0447fc93ca278dcb2471994dd0c3a7ba97e280d8 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 17:18:24 +0530 Subject: [PATCH 05/10] fixed lint and test issue --- src/api/python/history_sql.py | 1 - src/api/python/token_usage.py | 42 ++++++++++++++++--------- src/test/api/python/test_chat.py | 6 ++-- src/test/api/python/test_history.py | 2 +- src/test/api/python/test_history_sql.py | 4 +-- 5 files changed, 33 insertions(+), 22 deletions(-) diff --git a/src/api/python/history_sql.py b/src/api/python/history_sql.py index 53ab51a0..150db79c 100644 --- a/src/api/python/history_sql.py +++ b/src/api/python/history_sql.py @@ -34,7 +34,6 @@ # Database configuration - async def get_azure_sql_connection(): """ Get a connection to Azure SQL Server using DefaultAzureCredential. diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py index bba83c74..d01266f4 100644 --- a/src/api/python/token_usage.py +++ b/src/api/python/token_usage.py @@ -43,11 +43,15 @@ 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 - inp = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 - out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 - tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) + try: + inp = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 + out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 + tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) + inp_i, out_i, tot_i = int(inp), int(out), int(tot) + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) return None @@ -75,11 +79,15 @@ def extract_usage_from_update(update) -> Optional[UsageTuple]: if result: return result else: - inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 - out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) + try: + inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 + out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + inp_i, out_i, tot_i = int(inp), int(out), int(tot) + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) return None @@ -92,11 +100,15 @@ def extract_usage_from_response(response) -> Optional[UsageTuple]: return None if isinstance(usage_obj, dict): return extract_usage_from_dict(usage_obj) - inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 - out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - if tot > 0: - return (int(inp), int(out), int(tot)) + try: + inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 + out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 + tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) + inp_i, out_i, tot_i = int(inp), int(out), int(tot) + except (TypeError, ValueError): + return None + if tot_i > 0: + return (inp_i, out_i, tot_i) return None diff --git a/src/test/api/python/test_chat.py b/src/test/api/python/test_chat.py index a46c774c..272f5e32 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() @@ -1071,7 +1071,7 @@ def test_track_event_if_configured_without_instrumentation_key(self, monkeypatch # Ensure no instrumentation key monkeypatch.delenv("APPLICATIONINSIGHTS_CONNECTION_STRING", raising=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 +1127,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..ffc851fe 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): diff --git a/src/test/api/python/test_history_sql.py b/src/test/api/python/test_history_sql.py index 8e7d12a4..cc040057 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() From 6dbbb0b809d0cdc0aacd7b1fc3fd2313ec078781 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 17:33:21 +0530 Subject: [PATCH 06/10] fixed copilot comments --- src/api/python/history.py | 8 +- src/api/python/history_sql.py | 8 +- src/api/python/token_usage.py | 120 +++++++++++--- src/test/api/python/test_token_usage.py | 204 ++++++++++++++++++++++++ 4 files changed, 310 insertions(+), 30 deletions(-) create mode 100644 src/test/api/python/test_token_usage.py diff --git a/src/api/python/history.py b/src/api/python/history.py index b118141b..2daf1523 100644 --- a/src/api/python/history.py +++ b/src/api/python/history.py @@ -272,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"] @@ -312,6 +312,8 @@ async def generate_title(conversation_messages): _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 = "" @@ -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 150db79c..43239fc9 100644 --- a/src/api/python/history_sql.py +++ b/src/api/python/history_sql.py @@ -510,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. @@ -554,6 +556,8 @@ async def generate_title(conversation_messages): _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 @@ -750,7 +754,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 index d01266f4..831b8ad1 100644 --- a/src/api/python/token_usage.py +++ b/src/api/python/token_usage.py @@ -39,15 +39,69 @@ def track_event_if_configured(event_name: str, event_data: dict) -> None: ) +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 = d.get("input_token_count", 0) or d.get("prompt_tokens", 0) or d.get("input_tokens", 0) or 0 - out = d.get("output_token_count", 0) or d.get("completion_tokens", 0) or d.get("output_tokens", 0) or 0 - tot = d.get("total_token_count", 0) or d.get("total_tokens", 0) or (inp + out) - inp_i, out_i, tot_i = int(inp), int(out), int(tot) + 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: @@ -60,7 +114,12 @@ def extract_usage_from_update(update) -> Optional[UsageTuple]: Checks, in order: 1. update.contents[*].usage_details (dict) - 2. update.raw_representation.usage (dict or object with token attributes) + 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: @@ -69,25 +128,27 @@ def extract_usage_from_update(update) -> Optional[UsageTuple]: 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: - usage_obj = getattr(raw, "usage", None) - if usage_obj is not None: - if isinstance(usage_obj, dict): - result = extract_usage_from_dict(usage_obj) - if result: - return result - else: - try: - inp = getattr(usage_obj, "prompt_tokens", 0) or getattr(usage_obj, "input_tokens", 0) or 0 - out = getattr(usage_obj, "completion_tokens", 0) or getattr(usage_obj, "output_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - inp_i, out_i, tot_i = int(inp), int(out), int(tot) - except (TypeError, ValueError): - return None - if tot_i > 0: - return (inp_i, out_i, tot_i) + 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 @@ -101,10 +162,19 @@ def extract_usage_from_response(response) -> Optional[UsageTuple]: if isinstance(usage_obj, dict): return extract_usage_from_dict(usage_obj) try: - inp = getattr(usage_obj, "input_tokens", 0) or getattr(usage_obj, "prompt_tokens", 0) or 0 - out = getattr(usage_obj, "output_tokens", 0) or getattr(usage_obj, "completion_tokens", 0) or 0 - tot = getattr(usage_obj, "total_tokens", 0) or (inp + out) - inp_i, out_i, tot_i = int(inp), int(out), int(tot) + 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: 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..cb5b2517 --- /dev/null +++ b/src/test/api/python/test_token_usage.py @@ -0,0 +1,204 @@ +"""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() + + +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"]) From de6802d1dfd4ba726d8385ad380079e11580e611 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 17:44:12 +0530 Subject: [PATCH 07/10] comments updated --- infra/dashboards/token-usage-queries.kql | 26 ++++----- src/test/api/python/test_history.py | 72 +++++++++++++++++++++++- src/test/api/python/test_history_sql.py | 67 ++++++++++++++++++++++ 3 files changed, 151 insertions(+), 14 deletions(-) diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql index 2d6be67b..dfef47e9 100644 --- a/infra/dashboards/token-usage-queries.kql +++ b/infra/dashboards/token-usage-queries.kql @@ -15,7 +15,7 @@ customEvents TotalInputTokens = sum(input_tokens), TotalOutputTokens = sum(output_tokens), TotalTokens = sum(total_tokens), - AvgTokensPerRequest = round(avg(total_tokens), 0) + AvgTokensPerRequest = round(avg(total_tokens), 0); // 2. Token usage by agent customEvents @@ -31,7 +31,7 @@ customEvents TotalTokens = sum(total_tokens), Invocations = count() by Agent = agent -| order by TotalTokens desc +| order by TotalTokens desc; // 3. Token usage over time (hourly) customEvents @@ -41,7 +41,7 @@ customEvents | 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 +| render areachart; // 4. Estimated cost (GPT-4o pricing: $2.50/1M input, $10.00/1M output) let input_price_per_million = 2.50; @@ -56,7 +56,7 @@ customEvents | 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 +| order by Day desc; // 5. Top token consumers by user customEvents @@ -66,7 +66,7 @@ customEvents | extend user_id = tostring(customDimensions['user_id']) | summarize TotalTokens = sum(total_tokens), Requests = count() by user_id | order by TotalTokens desc -| take 20 +| take 20; // 6. Agent token distribution (pie chart) customEvents @@ -75,7 +75,7 @@ customEvents | extend agent = tostring(customDimensions['agent_name']) | extend total_tokens = toint(customDimensions['total_tokens']) | summarize TotalTokens = sum(total_tokens) by agent -| render piechart +| render piechart; // 7. Token usage percentiles per task customEvents @@ -87,7 +87,7 @@ customEvents p90 = percentile(total_tokens, 90), p95 = percentile(total_tokens, 95), p99 = percentile(total_tokens, 99), - Max = max(total_tokens) + Max = max(total_tokens); // 8. Token usage by model deployment customEvents @@ -103,7 +103,7 @@ customEvents TotalTokens = sum(total_tokens), Invocations = count() by Model = model -| order by TotalTokens desc +| order by TotalTokens desc; // 9. Token usage by model over time (hourly) customEvents @@ -113,7 +113,7 @@ customEvents | extend total_tokens = toint(customDimensions['total_tokens']) | summarize TotalTokens = sum(total_tokens) by bin(timestamp, 1h), model | order by timestamp asc -| render areachart +| render areachart; // 10. Model token distribution (pie chart) customEvents @@ -122,7 +122,7 @@ customEvents | extend model = tostring(customDimensions['model_deployment_name']) | extend total_tokens = toint(customDimensions['total_tokens']) | summarize TotalTokens = sum(total_tokens) by model -| render piechart +| render piechart; // 11. Estimated cost by model (adjust pricing per model) let gpt41_input = 2.00; @@ -148,7 +148,7 @@ customEvents | extend OutputCost = round(TotalOutput * OutputPrice / 1000000.0, 4) | extend TotalCost = InputCost + OutputCost | project Model = model, TotalInput, TotalOutput, InputCost, OutputCost, TotalCost -| order by TotalCost desc +| order by TotalCost desc; // 12. Agent-to-model mapping with token usage customEvents @@ -165,7 +165,7 @@ customEvents TotalTokens = sum(total_tokens), Invocations = count() by Agent = agent, Model = model -| order by TotalTokens desc +| order by TotalTokens desc; // 13. OpenTelemetry auto-instrumented OpenAI calls (if available) dependencies @@ -180,4 +180,4 @@ dependencies TotalInput = sum(input_tokens), TotalOutput = sum(output_tokens) by model -| order by TotalInput desc +| order by TotalInput desc; diff --git a/src/test/api/python/test_history.py b/src/test/api/python/test_history.py index ffc851fe..2ad45f05 100644 --- a/src/test/api/python/test_history.py +++ b/src/test/api/python/test_history.py @@ -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 cc040057..8e2dbec3 100644 --- a/src/test/api/python/test_history_sql.py +++ b/src/test/api/python/test_history_sql.py @@ -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.""" From 83a848c78e06ae155dc430396166ff321e32572a Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Wed, 13 May 2026 17:53:20 +0530 Subject: [PATCH 08/10] fixed as per comment --- src/api/python/history_sql.py | 62 +++++++++++++------------ src/api/python/token_usage.py | 22 +++++++-- src/test/api/python/test_chat.py | 3 ++ src/test/api/python/test_token_usage.py | 17 +++++++ 4 files changed, 71 insertions(+), 33 deletions(-) diff --git a/src/api/python/history_sql.py b/src/api/python/history_sql.py index 43239fc9..bb22ddff 100644 --- a/src/api/python/history_sql.py +++ b/src/api/python/history_sql.py @@ -538,38 +538,42 @@ async def generate_title(conversation_messages, user_id: str = "", conversation_ 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"}} + ) - _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 "", - ) + _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 + # 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) diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py index 831b8ad1..4989db57 100644 --- a/src/api/python/token_usage.py +++ b/src/api/python/token_usage.py @@ -26,17 +26,31 @@ 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: - logger.warning( - "Skipping track_event for %s as Application Insights is not configured", - event_name, - ) + 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): diff --git a/src/test/api/python/test_chat.py b/src/test/api/python/test_chat.py index 272f5e32..3c04508a 100644 --- a/src/test/api/python/test_chat.py +++ b/src/test/api/python/test_chat.py @@ -1067,9 +1067,12 @@ 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('token_usage.logger.warning') as mock_warning: track_event_if_configured("test_event", {"data": "value"}) diff --git a/src/test/api/python/test_token_usage.py b/src/test/api/python/test_token_usage.py index cb5b2517..a75a2b09 100644 --- a/src/test/api/python/test_token_usage.py +++ b/src/test/api/python/test_token_usage.py @@ -125,6 +125,23 @@ def test_skips_when_not_configured(self, monkeypatch): 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): From 694446da911a6ba1e8c0c3ab1ef7472de4c74661 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Thu, 14 May 2026 11:30:24 +0530 Subject: [PATCH 09/10] updated token usage --- src/api/python/token_usage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/api/python/token_usage.py b/src/api/python/token_usage.py index 4989db57..6e1de992 100644 --- a/src/api/python/token_usage.py +++ b/src/api/python/token_usage.py @@ -241,7 +241,7 @@ def track_token_usage( "user_id": user_id or "", "conversation_id": conversation_id or "", }) - logger.info( + 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, ) From 6a5bfcec50d94481b68bd4e00be5924a09e81100 Mon Sep 17 00:00:00 2001 From: Ajit Padhi Date: Tue, 19 May 2026 18:00:42 +0530 Subject: [PATCH 10/10] udpated doc --- infra/dashboards/token-usage-queries.kql | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql index dfef47e9..3011c33b 100644 --- a/infra/dashboards/token-usage-queries.kql +++ b/infra/dashboards/token-usage-queries.kql @@ -181,3 +181,19 @@ dependencies 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;