diff --git a/infra/dashboards/token-usage-queries.kql b/infra/dashboards/token-usage-queries.kql new file mode 100644 index 00000000..54253bd8 --- /dev/null +++ b/infra/dashboards/token-usage-queries.kql @@ -0,0 +1,349 @@ +// ============================================================================= +// Token Usage Dashboard Queries — BYOCC Customer Chatbot +// ============================================================================= +// Source events (Application Insights `customEvents` table): +// * LLM_Token_Usage_Summary — one event per request (aggregated totals) +// * LLM_Agent_Token_Usage — one event per agent involved in a request +// * LLM_Model_Token_Usage — one event per model deployment +// +// All token counts are sent as strings in `customDimensions`; cast with toint(). +// Emitted by: src/api/app/utils/token_usage_utils.py +// ============================================================================= + + +// ----------------------------------------------------------------------------- +// 1. Overall token usage summary (last 24 hours) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(24h) +| where name == "LLM_Token_Usage_Summary" +| extend + input_tokens = toint(customDimensions.total_input_tokens), + output_tokens = toint(customDimensions.total_output_tokens), + total_tokens = toint(customDimensions.total_tokens) +| summarize + Requests = count(), + TotalInputTokens = sum(input_tokens), + TotalOutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + AvgTokensPerReq = avg(total_tokens) + + +// ----------------------------------------------------------------------------- +// 2. Token usage per agent (chat + speech, last 7 days) +// +// Unions LLM chat agent usage (LLM_Agent_Token_Usage) and Voice Live realtime +// speech usage (Speech_Usage) so all agent-level token consumption is visible. +// NOTE: Sub-agent (tool) token breakdown is not available — the Azure AI SDK +// aggregates all tokens at the orchestrator level when using .as_tool(). +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name in ("LLM_Agent_Token_Usage", "Speech_Usage") +| extend + agent_name = tostring(customDimensions.agent_name), + input_tokens = toint(customDimensions.input_tokens), + output_tokens = toint(customDimensions.output_tokens), + total_tokens = toint(customDimensions.total_tokens), + agent_kind = iff(name == "Speech_Usage", "speech", "chat") +| summarize + Requests = count(), + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens) + by agent_name, agent_kind +| order by TotalTokens desc + + +// ----------------------------------------------------------------------------- +// 3. Token usage per model deployment (chat + speech, last 7 days) +// Unions LLM chat models (LLM_Model_Token_Usage) and Voice Live realtime +// model (Speech_Usage) so all model token consumption is visible in one view. +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name in ("LLM_Model_Token_Usage", "Speech_Usage") +| extend + model_name = tostring(customDimensions.model_deployment_name), + input_tokens = toint(customDimensions.input_tokens), + output_tokens = toint(customDimensions.output_tokens), + total_tokens = toint(customDimensions.total_tokens), + model_kind = iff(name == "Speech_Usage", "speech", "chat") +| summarize + Calls = count(), + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens) + by model_name, model_kind +| order by TotalTokens desc + + +// ----------------------------------------------------------------------------- +// 4. Top users by token consumption (last 30 days) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(30d) +| where name == "LLM_Token_Usage_Summary" +| extend + user_id = tostring(customDimensions.user_id), + total_tokens = toint(customDimensions.total_tokens) +| where isnotempty(user_id) +| summarize + Requests = count(), + TotalTokens = sum(total_tokens) + by user_id +| order by TotalTokens desc +| take 25 + + +// ----------------------------------------------------------------------------- +// 5. Hourly token usage trend (area chart, last 24h) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(24h) +| where name == "LLM_Token_Usage_Summary" +| extend + input_tokens = toint(customDimensions.total_input_tokens), + 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 + + +// ----------------------------------------------------------------------------- +// 6. Token-usage distribution per request (percentiles, last 7 days) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name == "LLM_Token_Usage_Summary" +| 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) + + +// ----------------------------------------------------------------------------- +// 7. Estimated cost per model (chat + speech) — gpt-4o-mini pricing ($0.15 / $0.60 per 1M tokens) +// ----------------------------------------------------------------------------- +// NOTE: adjust the rates below if you change the deployed model SKU. +// Includes both LLM_Model_Token_Usage (chat) and Speech_Usage (speech) events. +let InputRatePerToken = 0.00000015; // $0.15 / 1,000,000 +let OutputRatePerToken = 0.00000060; // $0.60 / 1,000,000 +customEvents +| where timestamp > ago(30d) +| where name in ("LLM_Model_Token_Usage", "Speech_Usage") +| extend + model_name = tostring(customDimensions.model_deployment_name), + input_tokens = toint(customDimensions.input_tokens), + output_tokens = toint(customDimensions.output_tokens), + model_kind = iff(name == "Speech_Usage", "speech", "chat") +| summarize + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens) + by model_name, model_kind +| extend EstimatedCostUSD = + round(InputTokens * InputRatePerToken + OutputTokens * OutputRatePerToken, 4) +| order by EstimatedCostUSD desc + + +// ----------------------------------------------------------------------------- +// 8. Daily cost trend (chat + speech, last 30 days) +// ----------------------------------------------------------------------------- +let InputRatePerToken = 0.00000015; +let OutputRatePerToken = 0.00000060; +customEvents +| where timestamp > ago(30d) +| where name in ("LLM_Model_Token_Usage", "Speech_Usage") +| extend + input_tokens = toint(customDimensions.input_tokens), + output_tokens = toint(customDimensions.output_tokens) +| summarize + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens) + by bin(timestamp, 1d) +| extend EstimatedCostUSD = + round(InputTokens * InputRatePerToken + OutputTokens * OutputRatePerToken, 4) +| order by timestamp asc +| render columnchart + + +// ----------------------------------------------------------------------------- +// 9. Agent ↔ model attribution (which model each agent invoked, last 7 days) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name == "LLM_Agent_Token_Usage" +| extend + agent_name = tostring(customDimensions.agent_name), + model_name = tostring(customDimensions.model_deployment_name), + total_tokens = toint(customDimensions.total_tokens) +| summarize Calls = count(), TotalTokens = sum(total_tokens) by agent_name, model_name +| order by TotalTokens desc + + +// ----------------------------------------------------------------------------- +// 10. Agent token-usage share (pie, last 24h) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(24h) +| where name == "LLM_Agent_Token_Usage" +| extend + agent_name = tostring(customDimensions.agent_name), + total_tokens = toint(customDimensions.total_tokens) +| summarize TotalTokens = sum(total_tokens) by agent_name +| render piechart + + +// ----------------------------------------------------------------------------- +// 11. OpenTelemetry cross-check — token usage from `dependencies` (gen_ai.*) +// ----------------------------------------------------------------------------- +// Useful for validating that our custom events agree with auto-instrumented +// OpenTelemetry traces emitted by the agent_framework / Azure AI SDKs. +dependencies +| where timestamp > ago(24h) +| where isnotempty(customDimensions["gen_ai.usage.input_tokens"]) +| extend + model = tostring(customDimensions["gen_ai.request.model"]), + input_tokens = toint(customDimensions["gen_ai.usage.input_tokens"]), + output_tokens = toint(customDimensions["gen_ai.usage.output_tokens"]) +| summarize + Calls = count(), + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens) + by model +| order by InputTokens + OutputTokens desc + + +// ============================================================================= +// SPEECH / VOICE LIVE — realtime model token usage +// ============================================================================= +// Source event: `Speech_Usage` — emitted by src/api/app/utils/speech_usage_utils.py +// from `voice_live.py` on every `response.done` event from the Voice Live +// realtime model (e.g. `gpt-realtime-mini`). The chat agents invoked from voice +// (via `call_foundry_agent`) already emit the standard LLM_* events, so these +// queries cover ONLY the realtime audio I/O layer that is NOT visible in the +// chat-completion telemetry. +// +// Fields (all in customDimensions, stringified): +// * source — "voice_chat" (WebSocket) or "tts" (HTTP /tts) +// * model_deployment_name — e.g. "gpt-realtime-mini" +// * session_id, user_id — voice session correlation id (client_id) +// * input_tokens / output_tokens / total_tokens +// * input_audio_tokens, input_text_tokens, input_cached_tokens +// * output_audio_tokens, output_text_tokens +// ============================================================================= + + +// ----------------------------------------------------------------------------- +// 12. Speech token usage by model & source (last 7 days) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name == "Speech_Usage" +| extend + model = tostring(customDimensions.model_deployment_name), + source = tostring(customDimensions.source), + input_tokens = toint(customDimensions.input_tokens), + output_tokens = toint(customDimensions.output_tokens), + total_tokens = toint(customDimensions.total_tokens), + input_audio_tokens = toint(customDimensions.input_audio_tokens), + input_text_tokens = toint(customDimensions.input_text_tokens), + input_cached_tokens = toint(customDimensions.input_cached_tokens), + output_audio_tokens = toint(customDimensions.output_audio_tokens), + output_text_tokens = toint(customDimensions.output_text_tokens) +| summarize + Responses = count(), + InputTokens = sum(input_tokens), + OutputTokens = sum(output_tokens), + TotalTokens = sum(total_tokens), + InputAudioTokens = sum(input_audio_tokens), + InputTextTokens = sum(input_text_tokens), + InputCachedTokens = sum(input_cached_tokens), + OutputAudioTokens = sum(output_audio_tokens), + OutputTextTokens = sum(output_text_tokens) + by model, source +| order by TotalTokens desc + + +// ----------------------------------------------------------------------------- +// 13. Speech token usage trend (hourly, last 24h) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(24h) +| where name == "Speech_Usage" +| extend + source = tostring(customDimensions.source), + input_audio_tokens = toint(customDimensions.input_audio_tokens), + output_audio_tokens = toint(customDimensions.output_audio_tokens), + total_tokens = toint(customDimensions.total_tokens) +| summarize + InputAudioTokens = sum(input_audio_tokens), + OutputAudioTokens = sum(output_audio_tokens), + TotalTokens = sum(total_tokens) + by bin(timestamp, 1h), source +| render timechart + + +// ----------------------------------------------------------------------------- +// 14. Speech usage per session (top 50 sessions, last 7 days) +// ----------------------------------------------------------------------------- +customEvents +| where timestamp > ago(7d) +| where name == "Speech_Usage" +| extend + session_id = tostring(customDimensions.session_id), + source = tostring(customDimensions.source), + input_audio_tokens = toint(customDimensions.input_audio_tokens), + output_audio_tokens = toint(customDimensions.output_audio_tokens), + total_tokens = toint(customDimensions.total_tokens) +| summarize + Responses = count(), + InputAudioTokens = sum(input_audio_tokens), + OutputAudioTokens = sum(output_audio_tokens), + TotalTokens = sum(total_tokens) + by session_id, source +| top 50 by TotalTokens desc + + +// ----------------------------------------------------------------------------- +// 15. Speech cost estimate (last 7 days) — UPDATE RATES BEFORE USE +// ----------------------------------------------------------------------------- +// Placeholder per-1K-token rates for the realtime model. Replace with current +// Azure pricing for `gpt-realtime-mini` (or whichever model is configured). +// As of writing, realtime audio tokens are priced separately from text tokens. +// See: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/ +let rate_input_audio_per_1k = 0.0; // TODO: set current $/1K input audio tokens +let rate_output_audio_per_1k = 0.0; // TODO: set current $/1K output audio tokens +let rate_input_text_per_1k = 0.0; // TODO: set current $/1K input text tokens +let rate_output_text_per_1k = 0.0; // TODO: set current $/1K output text tokens +let rate_cached_per_1k = 0.0; // TODO: set current $/1K cached input tokens +customEvents +| where timestamp > ago(7d) +| where name == "Speech_Usage" +| extend + model = tostring(customDimensions.model_deployment_name), + input_audio_tokens = toint(customDimensions.input_audio_tokens), + input_text_tokens = toint(customDimensions.input_text_tokens), + input_cached_tokens = toint(customDimensions.input_cached_tokens), + output_audio_tokens = toint(customDimensions.output_audio_tokens), + output_text_tokens = toint(customDimensions.output_text_tokens) +| summarize + InputAudioTokens = sum(input_audio_tokens), + InputTextTokens = sum(input_text_tokens), + InputCachedTokens = sum(input_cached_tokens), + OutputAudioTokens = sum(output_audio_tokens), + OutputTextTokens = sum(output_text_tokens) + by model +| extend EstimatedCostUSD = + (InputAudioTokens / 1000.0) * rate_input_audio_per_1k + + (OutputAudioTokens / 1000.0) * rate_output_audio_per_1k + + (InputTextTokens / 1000.0) * rate_input_text_per_1k + + (OutputTextTokens / 1000.0) * rate_output_text_per_1k + + (InputCachedTokens / 1000.0) * rate_cached_per_1k +| order by EstimatedCostUSD desc diff --git a/src/api/app/routers/chat.py b/src/api/app/routers/chat.py index dbfb274f..49e1192a 100644 --- a/src/api/app/routers/chat.py +++ b/src/api/app/routers/chat.py @@ -400,6 +400,38 @@ async def send_message_legacy( question = message.content result = await retrieved_agent.run(question) track_event_if_configured("Agent_Response_Received", {"session_id": session_id, "user_id": user_id}) + + # Emit token-usage telemetry (non-fatal) + try: + # Detect which sub-agent tools were actually invoked by inspecting + # function_call content items in the result messages. Only attribute + # token usage to sub-agents that were actually called. + invoked_tool_names: set[str] = set() + for _msg in (getattr(result, "messages", None) or []): + for _c in (getattr(_msg, "contents", None) or []): + if getattr(_c, "type", None) == "function_call": + _name = getattr(_c, "name", None) + if _name: + invoked_tool_names.add(_name) + + additional_agents: dict[str, str] = {} + if "product_agent" in invoked_tool_names: + additional_agents[product_agent_name] = settings.azure_openai_deployment_name + if "policy_agent" in invoked_tool_names: + additional_agents[policy_agent_name] = settings.azure_openai_deployment_name + + from ..utils.token_usage_utils import extract_and_track_usage + extract_and_track_usage( + result, + agent_name=chat_agent_name, + model_deployment_name=settings.azure_openai_deployment_name, + user_id=user_id, + session_id=session_id, + additional_agents=additional_agents, + ) + except Exception: + logger.debug("Token usage tracking failed (non-fatal)", exc_info=True) + break # Success, exit retry loop except Exception as e: diff --git a/src/api/app/routers/voice_live.py b/src/api/app/routers/voice_live.py index 6672c357..b2471ce6 100644 --- a/src/api/app/routers/voice_live.py +++ b/src/api/app/routers/voice_live.py @@ -20,6 +20,7 @@ try: from ..config import settings from ..utils.foundry_agent_utils import call_foundry_agent + from ..utils.token_usage_utils import extract_and_track_speech_usage from ..utils.voice_utils import ( clean_text_for_speech, is_valid_realtime_endpoint, @@ -30,6 +31,7 @@ except ImportError: from app.config import settings from app.utils.foundry_agent_utils import call_foundry_agent + from app.utils.token_usage_utils import extract_and_track_speech_usage from app.utils.voice_utils import ( clean_text_for_speech, is_valid_realtime_endpoint, @@ -86,7 +88,7 @@ ) -async def _call_foundry_agent(question: str) -> str: +async def _call_foundry_agent(question: str, *, user_id: Optional[str] = None, session_id: Optional[str] = None) -> str: """Delegate to foundry_agent_utils.""" client_id = str(settings.azure_client_id) if settings.azure_client_id else None return await call_foundry_agent( @@ -96,6 +98,9 @@ async def _call_foundry_agent(question: str) -> str: product_agent_name=settings.foundry_product_agent, policy_agent_name=settings.foundry_policy_agent, azure_client_id=client_id, + user_id=user_id, + session_id=session_id, + model_deployment_name=settings.azure_openai_deployment_name, ) @@ -367,7 +372,13 @@ async def _handle_event(self, event, connection) -> None: question = args.get("question", "") if name == "ask_customer_service": # Run Foundry agent with keep-alive pings to prevent WS timeout - agent_task = asyncio.create_task(_call_foundry_agent(question)) + agent_task = asyncio.create_task( + _call_foundry_agent( + question, + user_id=self.client_id, + session_id=self.client_id, + ) + ) while not agent_task.done(): await asyncio.sleep(2) if not agent_task.done(): @@ -480,6 +491,20 @@ async def _handle_event(self, event, connection) -> None: else: logger.info("[%s] Response done. No response object.", self.client_id) + # Emit Speech_Usage telemetry for the realtime model (audio tokens + # are NOT captured by the standard LLM_* events — those only see + # the downstream Foundry agent chat completion). + try: + extract_and_track_speech_usage( + response_obj, + model_deployment_name=settings.voicelive_model, + source="voice_chat", + user_id=self.client_id, + session_id=self.client_id, + ) + except Exception as exc: # pragma: no cover — telemetry must never break the flow + logger.debug("[%s] Speech usage tracking failed: %s", self.client_id, exc) + if self._assistant_transcript or self._assistant_text_response: # Prefer the text response (actual agent output) over audio transcript (paraphrase) display_text = self._assistant_text_response or self._assistant_transcript @@ -595,6 +620,14 @@ async def text_to_speech(request: Request): audio_chunks.append(base64.b64decode(delta)) elif event_type == ServerEventType.RESPONSE_DONE.value: + try: + extract_and_track_speech_usage( + getattr(event, "response", None), + model_deployment_name=settings.voicelive_model, + source="tts", + ) + except Exception as exc: # pragma: no cover + logger.debug("TTS speech usage tracking failed: %s", exc) break elif event_type == ServerEventType.ERROR.value: diff --git a/src/api/app/utils/foundry_agent_utils.py b/src/api/app/utils/foundry_agent_utils.py index 2980ae9b..c45be8d7 100644 --- a/src/api/app/utils/foundry_agent_utils.py +++ b/src/api/app/utils/foundry_agent_utils.py @@ -14,19 +14,30 @@ async def call_foundry_agent( product_agent_name: str, policy_agent_name: str, azure_client_id: Optional[str] = None, + user_id: Optional[str] = None, + session_id: Optional[str] = None, + model_deployment_name: Optional[str] = None, ) -> str: """ Call the Foundry multi-agent pipeline (chat → product/policy agents → Azure AI Search). Returns the grounded text response. + + When the underlying SDK reports token usage, emits LLM_Token_Usage_Summary, + LLM_Agent_Token_Usage and LLM_Model_Token_Usage events to Application + Insights for dashboarding (see infra/dashboards/token-usage-queries.kql). """ try: from agent_framework_azure_ai import AzureAIProjectAgentProvider from azure.ai.projects.aio import AIProjectClient try: + from ..config import settings as _settings from ..utils.azure_credential_utils import get_azure_credential_async + from ..utils.token_usage_utils import extract_and_track_usage except ImportError: + from app.config import settings as _settings from app.utils.azure_credential_utils import get_azure_credential_async + from app.utils.token_usage_utils import extract_and_track_usage if not foundry_endpoint: return "Foundry endpoint not configured." @@ -57,6 +68,39 @@ async def call_foundry_agent( result = await retrieved_agent.run(question) + # Emit token-usage telemetry (best-effort; never breaks the response) + try: + model_name = model_deployment_name or getattr( + _settings, "azure_openai_deployment_name", "" + ) + + # Only attribute usage to sub-agents that were actually invoked + # (inspect function_call items in the result messages). + invoked_tool_names: set[str] = set() + for _msg in (getattr(result, "messages", None) or []): + for _c in (getattr(_msg, "contents", None) or []): + if getattr(_c, "type", None) == "function_call": + _name = getattr(_c, "name", None) + if _name: + invoked_tool_names.add(_name) + + additional_agents: dict[str, str] = {} + if "product_agent" in invoked_tool_names: + additional_agents[product_agent_name] = model_name + if "policy_agent" in invoked_tool_names: + additional_agents[policy_agent_name] = model_name + + extract_and_track_usage( + result, + agent_name=chat_agent_name, + model_deployment_name=model_name, + user_id=user_id, + session_id=session_id, + additional_agents=additional_agents, + ) + except Exception: + logger.debug("Token usage tracking failed (non-fatal)", exc_info=True) + if result and hasattr(result, "text"): return result.text elif result: diff --git a/src/api/app/utils/token_usage_utils.py b/src/api/app/utils/token_usage_utils.py new file mode 100644 index 00000000..37a7c700 --- /dev/null +++ b/src/api/app/utils/token_usage_utils.py @@ -0,0 +1,371 @@ +"""Token usage extraction and Application Insights tracking helpers. + +Extracts LLM token usage from agent_framework run results and emits custom +events to Application Insights for monitoring per-agent, per-model, per-user +token consumption. + +Custom events emitted (visible in App Insights `customEvents` table): + - LLM_Token_Usage_Summary (one per request, aggregated totals) + - LLM_Agent_Token_Usage (one per agent involved in a request) + - LLM_Model_Token_Usage (one per model deployment involved) + +This is adapted for the BYOCC single-agent-with-tools chatbot from the +multi-agent orchestration reference implementation +(microsoft/Multi-Agent-Custom-Automation-Engine-Solution-Accelerator, +branch psl-token-usage). +""" +from __future__ import annotations + +import logging +from typing import Any, Dict, Optional, Tuple + +try: + from .event_utils import track_event_if_configured +except ImportError: + from app.utils.event_utils import track_event_if_configured + +logger = logging.getLogger(__name__) + + +def _coerce_int(value: Any) -> int: + """Best-effort int conversion; returns 0 on failure.""" + try: + if value is None: + return 0 + return int(value) + except (TypeError, ValueError): + return 0 + + +# Token-count field aliases used by various model providers / SDK versions. +_INPUT_KEYS = ( + "input_token_count", + "input_tokens", + "prompt_tokens", + "promptTokens", +) +_OUTPUT_KEYS = ( + "output_token_count", + "output_tokens", + "completion_tokens", + "completionTokens", +) +_TOTAL_KEYS = ( + "total_token_count", + "total_tokens", + "totalTokens", +) + + +def _read_usage_obj(usage_obj: Any) -> Optional[Tuple[int, int, int]]: + """Read input/output/total counts from a usage-bearing object or dict.""" + if usage_obj is None: + return None + + # dict-like + if isinstance(usage_obj, dict): + getter = usage_obj.get + else: + def getter(key, default=None): + return getattr(usage_obj, key, default) + + inp = 0 + out = 0 + tot = 0 + for k in _INPUT_KEYS: + v = getter(k) + if v: + inp = _coerce_int(v) + break + for k in _OUTPUT_KEYS: + v = getter(k) + if v: + out = _coerce_int(v) + break + for k in _TOTAL_KEYS: + v = getter(k) + if v: + tot = _coerce_int(v) + break + if tot == 0 and (inp or out): + tot = inp + out + if inp == 0 and out == 0 and tot == 0: + return None + return inp, out, tot + + +def extract_usage_from_agent_result(result: Any) -> Optional[Tuple[int, int, int]]: + """Extract (input_tokens, output_tokens, total_tokens) from an + agent_framework AgentRunResponse (or similar). + + Tries the following locations in order: + 1. result.usage_details / result.usage + 2. result.raw_representation.usage (OpenAI-style) + 3. Aggregated message contents (.messages[*].contents[*].usage_details) + + Returns None if no usage information is found. + """ + if result is None: + return None + + # 1. direct attribute + for attr in ("usage_details", "usage"): + usage = getattr(result, attr, None) + found = _read_usage_obj(usage) + if found: + return found + + # 2. raw_representation.usage (OpenAI ChatCompletion-style) + raw = getattr(result, "raw_representation", None) + if raw is not None: + found = _read_usage_obj(getattr(raw, "usage", None)) + if found: + return found + + # 3. aggregate over messages -> contents -> usage_details + messages = getattr(result, "messages", None) or [] + total_inp = 0 + total_out = 0 + total_tot = 0 + for msg in messages: + contents = getattr(msg, "contents", None) or [] + for content in contents: + usage = getattr(content, "usage_details", None) or getattr(content, "usage", None) + found = _read_usage_obj(usage) + if found: + total_inp += found[0] + total_out += found[1] + total_tot += found[2] + if total_inp or total_out or total_tot: + if total_tot == 0: + total_tot = total_inp + total_out + return total_inp, total_out, total_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: Optional[str] = None, + session_id: Optional[str] = None, + additional_agents: Optional[Dict[str, str]] = None, +) -> None: + """Emit summary, per-agent and per-model token usage events. + + Args: + agent_name: Primary agent that produced the response (e.g. chat agent). + model_deployment_name: Deployment name of the underlying model. + input_tokens / output_tokens / total_tokens: Request-level totals. + user_id / session_id: Optional context, included on every event. + additional_agents: Optional mapping {agent_name -> model_deployment_name} + for sub-agents/tools that participated in the request. Used for + agent_count/model_count in the summary event and for per-model events. + """ + if total_tokens <= 0 and input_tokens <= 0 and output_tokens <= 0: + return + + props_common = { + "user_id": user_id or "", + "session_id": session_id or "", + } + + # Summary + try: + agents = {agent_name: model_deployment_name} + if additional_agents: + agents.update({k: v for k, v in additional_agents.items() if k}) + models = {m for m in agents.values() if m} + track_event_if_configured( + "LLM_Token_Usage_Summary", + { + **props_common, + "total_input_tokens": str(input_tokens), + "total_output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + "agent_count": str(len(agents)), + "model_count": str(len(models)), + }, + ) + + # Per-agent event for the primary (orchestrator) agent only. + # NOTE: Sub-agent (tool) token breakdown is not available — the Azure + # AI SDK aggregates all tokens at the orchestrator level when using + # .as_tool(), so we only emit one event with the full token counts. + track_event_if_configured( + "LLM_Agent_Token_Usage", + { + **props_common, + "agent_name": agent_name, + "model_deployment_name": model_deployment_name or "", + "input_tokens": str(input_tokens), + "output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + }, + ) + + # Per-model (one event per distinct model) + for model in models: + track_event_if_configured( + "LLM_Model_Token_Usage", + { + **props_common, + "model_deployment_name": model, + "input_tokens": str(input_tokens), + "output_tokens": str(output_tokens), + "total_tokens": str(total_tokens), + }, + ) + + logger.info( + "[TOKEN USAGE] agent=%s model=%s input=%d output=%d total=%d user=%s session=%s", + agent_name, model_deployment_name, + input_tokens, output_tokens, total_tokens, + user_id or "-", session_id or "-", + ) + except Exception as exc: # never let telemetry break the request + logger.warning("track_token_usage failed: %s", exc) + + +def extract_and_track_usage( + result: Any, + *, + agent_name: str, + model_deployment_name: str, + user_id: Optional[str] = None, + session_id: Optional[str] = None, + additional_agents: Optional[Dict[str, str]] = None, +) -> Optional[Tuple[int, int, int]]: + """Convenience wrapper: extract usage from result and emit events. + + Returns the (input, output, total) tuple if found, else None. + Safe to call when telemetry is not configured — extraction still occurs + so callers may persist the result. + """ + usage = extract_usage_from_agent_result(result) + if not usage: + logger.debug( + "No token usage found on agent result for agent=%s", agent_name + ) + return None + inp, out, tot = usage + + track_token_usage( + agent_name=agent_name, + model_deployment_name=model_deployment_name, + input_tokens=inp, + output_tokens=out, + total_tokens=tot, + user_id=user_id, + session_id=session_id, + additional_agents=additional_agents, + ) + return usage + + +# --------------------------------------------------------------------------- +# Realtime / Speech (Azure AI Voice Live) token usage +# --------------------------------------------------------------------------- + +def _get(obj: Any, key: str, default: Any = None) -> Any: + """Access ``key`` on a dict or an object (SDK models expose attributes).""" + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_realtime_usage(response_obj: Any) -> Optional[Dict[str, int]]: + """Extract token counts from a Voice Live ``response.done`` payload. + + Returns a flat dict of token counts, or ``None`` if no usage data is + present (or all counts are zero). + """ + usage = _get(response_obj, "usage") + if usage is None: + return None + inp = _coerce_int(_get(usage, "input_tokens")) + out = _coerce_int(_get(usage, "output_tokens")) + tot = _coerce_int(_get(usage, "total_tokens")) + if tot == 0 and (inp or out): + tot = inp + out + in_details = _get(usage, "input_token_details") or {} + out_details = _get(usage, "output_token_details") or {} + counts = { + "input_tokens": inp, + "output_tokens": out, + "total_tokens": tot, + "input_audio_tokens": _coerce_int(_get(in_details, "audio_tokens")), + "input_text_tokens": _coerce_int(_get(in_details, "text_tokens")), + "input_cached_tokens": _coerce_int(_get(in_details, "cached_tokens")), + "output_audio_tokens": _coerce_int(_get(out_details, "audio_tokens")), + "output_text_tokens": _coerce_int(_get(out_details, "text_tokens")), + } + return counts if any(counts.values()) else None + + +def track_speech_usage( + *, + model_deployment_name: str, + source: str, + counts: Dict[str, int], + user_id: Optional[str] = None, + session_id: Optional[str] = None, +) -> None: + """Emit a ``Speech_Usage`` custom event. Never raises.""" + try: + track_event_if_configured( + "Speech_Usage", + { + "user_id": user_id or "", + "session_id": session_id or "", + "model_deployment_name": model_deployment_name or "", + "source": source or "", + "input_tokens": str(counts.get("input_tokens", 0)), + "output_tokens": str(counts.get("output_tokens", 0)), + "total_tokens": str(counts.get("total_tokens", 0)), + "input_audio_tokens": str(counts.get("input_audio_tokens", 0)), + "input_text_tokens": str(counts.get("input_text_tokens", 0)), + "input_cached_tokens": str(counts.get("input_cached_tokens", 0)), + "output_audio_tokens": str(counts.get("output_audio_tokens", 0)), + "output_text_tokens": str(counts.get("output_text_tokens", 0)), + }, + ) + logger.info( + "[SPEECH USAGE] source=%s model=%s in=%d (audio=%d) out=%d (audio=%d) total=%d", + source, model_deployment_name, + counts.get("input_tokens", 0), counts.get("input_audio_tokens", 0), + counts.get("output_tokens", 0), counts.get("output_audio_tokens", 0), + counts.get("total_tokens", 0), + ) + except Exception as exc: # never let telemetry break the voice path + logger.warning("track_speech_usage failed: %s", exc) + + +def extract_and_track_speech_usage( + response_obj: Any, + *, + model_deployment_name: str, + source: str, + user_id: Optional[str] = None, + session_id: Optional[str] = None, +) -> Optional[Dict[str, int]]: + """Extract realtime token usage from ``response_obj`` and emit a + ``Speech_Usage`` event. Returns the counts dict or ``None``.""" + counts = extract_realtime_usage(response_obj) + if not counts: + logger.debug("No realtime usage found on response (source=%s)", source) + return None + track_speech_usage( + model_deployment_name=model_deployment_name, + source=source, + counts=counts, + user_id=user_id, + session_id=session_id, + ) + return counts