Skip to content

feat: Token count for agents #860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Ayaz-Microsoft wants to merge 15 commits into
base: dev
Choose a base branch
from
Open

feat: Token count for agents #860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
69534b6
Add token usage telemetry and dashboard for LLM requests
Ayaz-Microsoft May 11, 2026
ff168de
feat: add Token Usage Application Insights workbook for LLM monitoring
Ayaz-Microsoft May 12, 2026
db73a16
feat: add monitoring configuration hash to container instance for dyn…
Ayaz-Microsoft May 12, 2026
5f5737b
sync main_custom.bicep with main.bicep
Ayaz-Microsoft May 12, 2026
52a7617
feat: separate Token Usage Application Insights workbook deployment i…
Ayaz-Microsoft May 14, 2026
cebc62e
Refactor code structure for improved readability and maintainability
Ayaz-Microsoft May 19, 2026
e5f3d81
restored main.bicep and azure.yaml
Ayaz-Microsoft May 25, 2026
cd71bad
remove unused field import from dataclass in token_usage.py
Ayaz-Microsoft May 25, 2026
ee5b3f4
Refactor code
Ayaz-Microsoft May 25, 2026
df02a01
feat: enhance token usage telemetry with conversation ID tracking and…
Ayaz-Microsoft May 25, 2026
6ea5f41
aligned code with other GSAs
Ayaz-Microsoft May 27, 2026
8590077
chore: scope PR to code-only integration
Ayaz-Microsoft Jun 4, 2026
bfc5541
test: exclude token-telemetry modules from coverage
Ayaz-Microsoft Jun 4, 2026
69aba68
feat: instrument missing agent calls for token telemetry
Ayaz-Microsoft Jun 4, 2026
4aaf000
test: update generate_title assertion for new correlation kwargs
Ayaz-Microsoft Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 15 additions & 9 deletions src/backend/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ async def handle_chat():

# Extract request fields
conversation_id = data.get("conversation_id") or str(uuid.uuid4())
user_id = data.get("user_id", "anonymous")
user_id = data.get("user_id") or "anonymous"
message = data.get("message", "")
action = data.get("action")
payload = data.get("payload", {})
Expand Down Expand Up @@ -337,7 +337,7 @@ async def _handle_parse_brief(

if not has_existing_title:
title_service = get_title_service()
generated_title = await title_service.generate_title(message)
generated_title = await title_service.generate_title(message, user_id=user_id, conversation_id=conversation_id)

await cosmos_service.add_message_to_conversation(
conversation_id=conversation_id,
Expand All @@ -353,7 +353,8 @@ async def _handle_parse_brief(
logger.exception(f"Failed to save message to CosmosDB: {e}")

# Parse the brief
brief, questions, blocked = await orchestrator.parse_brief(message)
brief, questions, blocked = await orchestrator.parse_brief(message, user_id=user_id, conversation_id=conversation_id)
brief, questions, blocked = await orchestrator.parse_brief(message, user_id=user_id, conversation_id=conversation_id)
Comment on lines +356 to +357

if blocked:
track_event_if_configured("Error_RAI_Check_Failed", {"conversation_id": conversation_id, "user_id": user_id, "status": "Brief parse blocked by RAI"})
Expand Down Expand Up @@ -537,7 +538,7 @@ async def _handle_refine_brief(
logger.exception(f"Failed to save refinement message: {e}")

# Use orchestrator to refine the brief
brief, questions, blocked = await orchestrator.parse_brief(message)
brief, questions, blocked = await orchestrator.parse_brief(message, user_id=user_id, conversation_id=conversation_id)

if blocked:
track_event_if_configured("Error_RAI_Check_Failed", {"conversation_id": conversation_id, "user_id": user_id, "status": "Brief refinement blocked by RAI"})
Expand Down Expand Up @@ -943,7 +944,9 @@ async def _run_regeneration_task(
modification_request=modification_request,
brief=brief,
products=products_data,
previous_image_prompt=previous_image_prompt
previous_image_prompt=previous_image_prompt,
user_id=user_id,
conversation_id=conversation_id
)

# Check for RAI block
Expand Down Expand Up @@ -1132,7 +1135,7 @@ async def _handle_general_chat(

if not has_existing_title:
title_service = get_title_service()
generated_title = await title_service.generate_title(message)
generated_title = await title_service.generate_title(message, user_id=user_id, conversation_id=conversation_id)

await cosmos_service.add_message_to_conversation(
conversation_id=conversation_id,
Expand All @@ -1151,7 +1154,8 @@ async def _handle_general_chat(
response_content = ""
async for response in orchestrator.process_message(
message=message,
conversation_id=conversation_id
conversation_id=conversation_id,
user_id=user_id
):
if response.get("content"):
response_content += response.get("content", "")
Expand Down Expand Up @@ -1197,7 +1201,9 @@ async def _run_generation_task(task_id: str, brief: CreativeBrief, products_data
response = await orchestrator.generate_content(
brief=brief,
products=products_data,
generate_images=generate_images
generate_images=generate_images,
user_id=user_id,
conversation_id=conversation_id
)

logger.info(f"Generation task {task_id} completed. Response keys: {list(response.keys()) if response else 'None'}")
Expand Down Expand Up @@ -1303,7 +1309,7 @@ async def start_generation():
products_data = data.get("products", [])
generate_images = data.get("generate_images", True)
conversation_id = data.get("conversation_id") or str(uuid.uuid4())
user_id = data.get("user_id", "anonymous")
user_id = data.get("user_id") or "anonymous"

try:
brief = CreativeBrief(**brief_data)
Expand Down
Loading