feat: Token Usage Tracking implementation

infra/dashboards/token-usage-queries.kql

@@ -0,0 +1,199 @@
+// ============================================================
+// KQL Queries for LLM Token Usage Monitoring
+// Run these in Application Insights > Logs
+// ============================================================
+
+// 1. Overall token usage summary (last 7 days)
+customEvents
+| where name == 'LLM_Token_Usage_Summary'
+| where timestamp > ago(7d)
+| extend input_tokens = toint(customDimensions['total_input_tokens'])
+| extend output_tokens = toint(customDimensions['total_output_tokens'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    TotalRequests = count(),
+    TotalInputTokens = sum(input_tokens),
+    TotalOutputTokens = sum(output_tokens),
+    TotalTokens = sum(total_tokens),
+    AvgTokensPerRequest = round(avg(total_tokens), 0);
+
+// 2. Token usage by agent
+customEvents
+| where name == 'LLM_Agent_Token_Usage'
+| where timestamp > ago(7d)
+| extend agent = tostring(customDimensions['agent_name'])
+| extend input_tokens = toint(customDimensions['input_tokens'])
+| extend output_tokens = toint(customDimensions['output_tokens'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    InputTokens = sum(input_tokens),
+    OutputTokens = sum(output_tokens),
+    TotalTokens = sum(total_tokens),
+    Invocations = count()
+    by Agent = agent
+| order by TotalTokens desc;
+
+// 3. Token usage over time (hourly)
+customEvents
+| where name == 'LLM_Token_Usage_Summary'
+| where timestamp > ago(7d)
+| extend input_tokens = toint(customDimensions['total_input_tokens'])
+| extend output_tokens = toint(customDimensions['total_output_tokens'])
+| summarize InputTokens = sum(input_tokens), OutputTokens = sum(output_tokens) by bin(timestamp, 1h)
+| order by timestamp asc
+| render areachart;
+
+// 4. Estimated cost (GPT-4o pricing: $2.50/1M input, $10.00/1M output)
+let input_price_per_million = 2.50;
+let output_price_per_million = 10.00;
+customEvents
+| where name == 'LLM_Token_Usage_Summary'
+| where timestamp > ago(30d)
+| extend input_tokens = toint(customDimensions['total_input_tokens'])
+| extend output_tokens = toint(customDimensions['total_output_tokens'])
+| summarize TotalInput = sum(input_tokens), TotalOutput = sum(output_tokens) by bin(timestamp, 1d)
+| extend InputCost = round(TotalInput * input_price_per_million / 1000000.0, 4)
+| extend OutputCost = round(TotalOutput * output_price_per_million / 1000000.0, 4)
+| extend TotalCost = InputCost + OutputCost
+| project Day = timestamp, TotalInput, TotalOutput, InputCost, OutputCost, TotalCost
+| order by Day desc;
+
+// 5. Top token consumers by user
+customEvents
+| where name == 'LLM_Token_Usage_Summary'
+| where timestamp > ago(7d)
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| extend user_id = tostring(customDimensions['user_id'])
+| summarize TotalTokens = sum(total_tokens), Requests = count() by user_id
+| order by TotalTokens desc
+| take 20;
+
+// 6. Agent token distribution (pie chart)
+customEvents
+| where name == 'LLM_Agent_Token_Usage'
+| where timestamp > ago(7d)
+| extend agent = tostring(customDimensions['agent_name'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize TotalTokens = sum(total_tokens) by agent
+| render piechart;
+
+// 7. Token usage percentiles per task
+customEvents
+| where name == 'LLM_Token_Usage_Summary'
+| where timestamp > ago(7d)
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    p50 = percentile(total_tokens, 50),
+    p90 = percentile(total_tokens, 90),
+    p95 = percentile(total_tokens, 95),
+    p99 = percentile(total_tokens, 99),
+    Max = max(total_tokens);
+
+// 8. Token usage by model deployment
+customEvents
+| where name == 'LLM_Model_Token_Usage'
+| where timestamp > ago(7d)
+| extend model = tostring(customDimensions['model_deployment_name'])
+| extend input_tokens = toint(customDimensions['input_tokens'])
+| extend output_tokens = toint(customDimensions['output_tokens'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    InputTokens = sum(input_tokens),
+    OutputTokens = sum(output_tokens),
+    TotalTokens = sum(total_tokens),
+    Invocations = count()
+    by Model = model
+| order by TotalTokens desc;
+
+// 9. Token usage by model over time (hourly)
+customEvents
+| where name == 'LLM_Model_Token_Usage'
+| where timestamp > ago(7d)
+| extend model = tostring(customDimensions['model_deployment_name'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize TotalTokens = sum(total_tokens) by bin(timestamp, 1h), model
+| order by timestamp asc
+| render areachart;
+
+// 10. Model token distribution (pie chart)
+customEvents
+| where name == 'LLM_Model_Token_Usage'
+| where timestamp > ago(7d)
+| extend model = tostring(customDimensions['model_deployment_name'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize TotalTokens = sum(total_tokens) by model
+| render piechart;
+
+// 11. Estimated cost by model (adjust pricing per model)
+let gpt41_input = 2.00;
+let gpt41_output = 8.00;
+let gpt41_mini_input = 0.40;
+let gpt41_mini_output = 1.60;
+customEvents
+| where name == 'LLM_Model_Token_Usage'
+| where timestamp > ago(30d)
+| extend model = tolower(tostring(customDimensions['model_deployment_name']))
+| extend input_tokens = toint(customDimensions['input_tokens'])
+| extend output_tokens = toint(customDimensions['output_tokens'])
+| summarize TotalInput = sum(input_tokens), TotalOutput = sum(output_tokens) by model
+| extend InputPrice = case(
+    model has "4.1-mini", gpt41_mini_input,
+    model has "4.1",      gpt41_input,
+    gpt41_mini_input)
+| extend OutputPrice = case(
+    model has "4.1-mini", gpt41_mini_output,
+    model has "4.1",      gpt41_output,
+    gpt41_mini_output)
+| extend InputCost = round(TotalInput * InputPrice / 1000000.0, 4)
+| extend OutputCost = round(TotalOutput * OutputPrice / 1000000.0, 4)
+| extend TotalCost = InputCost + OutputCost
+| project Model = model, TotalInput, TotalOutput, InputCost, OutputCost, TotalCost
+| order by TotalCost desc;
+
+// 12. Agent-to-model mapping with token usage
+customEvents
+| where name == 'LLM_Agent_Token_Usage'
+| where timestamp > ago(7d)
+| extend agent = tostring(customDimensions['agent_name'])
+| extend model = tostring(customDimensions['model_deployment_name'])
+| extend input_tokens = toint(customDimensions['input_tokens'])
+| extend output_tokens = toint(customDimensions['output_tokens'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    InputTokens = sum(input_tokens),
+    OutputTokens = sum(output_tokens),
+    TotalTokens = sum(total_tokens),
+    Invocations = count()
+    by Agent = agent, Model = model
+| order by TotalTokens desc;
+
+// 13. OpenTelemetry auto-instrumented OpenAI calls (if available)
+dependencies
+| where name has "openai" or target has "openai"
+| where timestamp > ago(7d)
+| extend input_tokens = tolong(customDimensions["gen_ai.usage.input_tokens"])
+| extend output_tokens = tolong(customDimensions["gen_ai.usage.output_tokens"])
+| extend model = tostring(customDimensions["gen_ai.request.model"])
+| where isnotnull(input_tokens)
+| summarize
+    Calls = count(),
+    TotalInput = sum(input_tokens),
+    TotalOutput = sum(output_tokens)
+    by model
+| order by TotalInput desc;
+
+// 14. Total time taken per invocation (conversation)
+customEvents
+| where name == 'LLM_Model_Token_Usage'
+| where timestamp > ago(7d)
+| extend conversation_id = tostring(customDimensions['conversation_id'])
+| extend model = tostring(customDimensions['model_deployment_name'])
+| extend total_tokens = toint(customDimensions['total_tokens'])
+| summarize
+    StartTime = min(timestamp),
+    EndTime = max(timestamp),
+    TotalTokens = sum(total_tokens),
+    Calls = count()
+    by conversation_id
+| extend TotalDuration_sec = round(datetime_diff('second', EndTime, StartTime), 0)
+| order by EndTime desc;

src/api/python/chat.py

@@ -18,7 +18,6 @@
 
 # Azure SDK
 from azure.core.exceptions import HttpResponseError
-from azure.monitor.events.extension import track_event
 from azure.ai.projects.aio import AIProjectClient
 
 # Agent Framework
@@ -28,6 +27,9 @@
 from auth.auth_utils import get_authenticated_user_details
 from auth.azure_credential_utils import get_azure_credential_async
 
+# Token usage telemetry helpers (modular, reusable)
+from token_usage import track_event_if_configured, UsageAccumulator
+
 load_dotenv()
 
 # Constants
@@ -108,15 +110,6 @@ async def _delete_thread_async(self, thread_conversation_id: str):
                 await credential.close()
 
 
-def track_event_if_configured(event_name: str, event_data: dict):
-    """Track event to Application Insights if configured."""
-    instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING")
-    if instrumentation_key:
-        track_event(event_name, event_data)
-    else:
-        logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name)
-
-
 # Global thread cache
 thread_cache = None
 
@@ -141,6 +134,8 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = ""
     complete_response = ""
     credential = None
     db_connection = None
+    # Accumulator for LLM token usage across all iterations of this request
+    usage = UsageAccumulator()
 
     try:
         if not query:
@@ -184,6 +179,9 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = ""
                 extra_body={"agent_reference": {"name": os.getenv("AGENT_NAME_CHAT"), "type": "agent_reference"}}
             )
 
+            # Accumulate token usage from initial response
+            usage.add_from_response(response)
+
             # Process response - handle function calls iteratively
             max_iterations = 10
             iteration = 0
@@ -258,6 +256,9 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = ""
                     extra_body={"agent_reference": {"name": os.getenv("AGENT_NAME_CHAT"), "type": "agent_reference"}}
                 )
 
+                # Accumulate token usage from follow-up response
+                usage.add_from_response(response)
+
             if iteration >= max_iterations:
                 logger.warning("Max iterations reached for conversation %s", conversation_id)
                 yield "\n\n(Response processing reached maximum iterations)"
@@ -270,6 +271,14 @@ async def stream_openai_text(conversation_id: str, query: str, user_id: str = ""
                 "response_length": str(len(complete_response)),
             })
 
+            # Emit LLM token usage telemetry
+            usage.emit(
+                agent_name=os.getenv("AGENT_NAME_CHAT", "") or "",
+                model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "",
+                user_id=user_id,
+                conversation_id=conversation_id,
+            )
+
     except HttpResponseError as e:
         complete_response = str(e)
         if "Rate limit is exceeded" in str(e) or e.status_code == 429:
@@ -348,6 +357,8 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id:
     complete_response = ""
     credential = None
     db_connection = None
+    # Accumulator for LLM token usage across all streaming updates
+    usage = UsageAccumulator()
 
     try:
         if not query:
@@ -420,6 +431,9 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id:
                     if raw_repr:
                         _extract_mcp_from_raw(raw_repr, mcp_docs)
 
+                # Accumulate token usage from this streaming update
+                usage.add_from_update(chunk)
+
                 chunk_text = str(chunk.text) if chunk.text else ""
                 if not chunk_text:
                     continue
@@ -473,6 +487,14 @@ async def stream_openai_text_workshop(conversation_id: str, query: str, user_id:
                 "citation_count": str(len(mcp_docs)),
             })
 
+            # Emit LLM token usage telemetry
+            usage.emit(
+                agent_name=os.getenv("AGENT_NAME_CHAT", "") or "",
+                model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "",
+                user_id=user_id,
+                conversation_id=conversation_id,
+            )
+
             # Yield citations as a tool message — deduplicated by source
             citation_list = []
             if original_markers:

src/api/python/history.py

@@ -8,7 +8,6 @@
 from azure.core.exceptions import HttpResponseError
 from azure.cosmos import exceptions
 from azure.cosmos.aio import CosmosClient
-from azure.monitor.events.extension import track_event
 from fastapi import APIRouter, HTTPException, Query, Request, status
 from fastapi.responses import JSONResponse
 from opentelemetry import trace
@@ -18,6 +17,9 @@
 from auth.auth_utils import get_authenticated_user_details
 from auth.azure_credential_utils import get_azure_credential_async
 
+# Token usage telemetry helpers (modular, reusable)
+from token_usage import track_event_if_configured, UsageAccumulator
+
 router = APIRouter()
 
 logger = logging.getLogger(__name__)
@@ -39,15 +41,6 @@
 AGENT_NAME_TITLE = os.getenv("AGENT_NAME_TITLE")
 
 
-def track_event_if_configured(event_name: str, event_data: dict):
-    """Track events to Application Insights if configured."""
-    instrumentation_key = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING")
-    if instrumentation_key:
-        track_event(event_name, event_data)
-    else:
-        logging.warning("Skipping track_event for %s as Application Insights is not configured", event_name)
-
-
 class CosmosConversationClient:
     """Client for managing conversations and messages in CosmosDB."""
 
@@ -279,7 +272,7 @@ async def init_cosmosdb_client():
         raise
 
 
-async def generate_title(conversation_messages):
+async def generate_title(conversation_messages, user_id: str = "", conversation_id: str = ""):
     """Generate a title for a conversation using Azure AI Foundry agent."""
     try:
         user_messages = [msg for msg in conversation_messages if msg["role"] == "user"]
@@ -314,6 +307,15 @@ async def generate_title(conversation_messages):
                     extra_body={"agent_reference": {"name": AGENT_NAME_TITLE, "type": "agent_reference"}}
                 )
 
+                _title_usage = UsageAccumulator()
+                _title_usage.add_from_response(response)
+                _title_usage.emit(
+                    agent_name=AGENT_NAME_TITLE or "",
+                    model_deployment_name=os.getenv("AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME", "") or "",
+                    user_id=user_id or "",
+                    conversation_id=conversation_id or "",
+                )
+
                 result_text = ""
                 for item in response.output:
                     if getattr(item, 'type', None) == 'message':
@@ -363,7 +365,7 @@ async def add_conversation(user_id: str, request_json: dict):
             raise ValueError("CosmosDB is not configured or unavailable")
 
         if not conversation_id:
-            title = await generate_title(messages)
+            title = await generate_title(messages, user_id=user_id)
             conversation_dict = await cosmos_conversation_client.create_conversation(user_id, title)
             conversation_id = conversation_dict["id"]
             history_metadata["title"] = title
@@ -397,7 +399,7 @@ async def update_conversation(user_id: str, request_json: dict):
     # Retrieve or create conversation
     conversation = await cosmos_conversation_client.get_conversation(user_id, conversation_id)
     if not conversation:
-        title = await generate_title(messages)
+        title = await generate_title(messages, user_id=user_id, conversation_id=conversation_id)
         conversation = await cosmos_conversation_client.create_conversation(
             user_id=user_id, conversation_id=conversation_id, title=title
         )