feat: add role-specific statistics to telemetry and UI (cont. #15234)

[yunaseoul]

This PR picks up the work from #15234.

Summary

Add role-specific statistics to /stats model
demo:

Details

• Granular LLM Call Categorization: Introduced a new LlmRole enum to classify LLM interactions (e.g., main agent, subagents, various utility functions like summarization, routing, autocomplete, etc.).
• Enhanced Telemetry for LLM Usage: Extended the telemetry system to record and track detailed metrics (requests, errors, latency, token usage including input, output, and cache reads) for each LlmRole.
• Improved UI for Model Statistics: Updated the CLI's ModelStatsDisplay to visually present these role-specific LLM usage breakdowns, offering a more insightful view of model activity.
• Widespread Integration of Role Tracking: The LlmRole parameter has been consistently integrated across numerous core components, utility functions, and client calls that interact with LLMs, ensuring comprehensive data collection.

UI changes are minimal and just intended as a starting point for UX discussions. Core work is the migration of all the code in packages/core to use the new tracking.

Related Issues

Fixes #14538

How to Validate

Exercise various models and trigger subagents, autocomplete, and other ways Gemini CLI uses models. Make sure they are attributed correctly.

When reviewing the code in package/core make sure that we are never setting a default value for LlmRole. My instructions for the plan for Gemini CLI was to require all code to pass in a valid LlmRole value so that static analysis could catch any issues where it wasn't passed in but there was one case I had to catch manually where Gemini CLI added a default value to make it easier for tests to pass that introduced a bug where tokens that weren't from the main event loop were tagged as main event loop.

Pre-Merge Checklist

• Updated relevant documentation and README (if needed)
• Added/updated tests (if needed)
• Noted breaking changes (if any)
• Validated on required platforms/methods:
  • MacOS
    • npm run
    • npx
    • Docker
    • Podman
    • Seatbelt
  • Windows
    • npm run
    • npx
    • Docker
  • Linux
    • npm run
    • npx
    • Docker

[gemini-code-assist]

Summary of Changes

Hello @yunaseoul, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the observability of LLM interactions within the system by introducing a robust mechanism to categorize and track model usage based on specific roles. This change provides a more detailed breakdown of how different parts of the application utilize LLMs, enabling better analysis of performance, cost, and overall system behavior. The updated UI components will visually present these new metrics, offering users and developers clearer insights into the underlying model operations.

Highlights

• Granular LLM Call Categorization: Introduced a new LlmRole enum to classify LLM interactions, covering main agents, subagents, and various utility functions like summarization, routing, and autocomplete.
• Enhanced Telemetry for LLM Usage: Extended the telemetry system to record and track detailed metrics (requests, errors, latency, token usage) for each LlmRole, providing deeper insights into model activity.
• Improved UI for Model Statistics: Updated the CLI's ModelStatsDisplay and StatsDisplay components to visually present these role-specific LLM usage breakdowns, offering a more insightful view of model performance.
• Widespread Integration of Role Tracking: The LlmRole parameter has been consistently integrated across numerous core components, utility functions, and client calls that interact with LLMs, ensuring comprehensive data collection.

Changelog

• packages/cli/src/ui/components/ModelStatsDisplay.test.tsx
  • Imported LlmRole for role-specific testing.
  • Updated mock data to include roles property for model metrics.
  • Added new test cases to verify correct display of role breakdown and handling of long metric names.
• packages/cli/src/ui/components/ModelStatsDisplay.tsx
  • Imported LlmRole and ModelMetrics types.
  • Defined RoleMetrics interface for tracking role-specific statistics.
  • Implemented logic to collect and sort all unique LlmRole values present in active models.
  • Modified rendering to display role-specific metrics within the model statistics table, including requests, input/output tokens, and cache reads.
• packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
  • Updated mock session metrics to include an empty roles object.
• packages/cli/src/ui/components/StatsDisplay.test.tsx
  • Imported LlmRole for role-specific testing.
  • Updated mock model metrics to include an empty roles object.
  • Added a new test case to verify correct rendering of role-specific statistics in the StatsDisplay component.
• packages/cli/src/ui/components/StatsDisplay.tsx
  • Modified buildModelRows function to include roles data in the generated model rows.
  • Adjusted ModelUsageTable layout to accommodate role-specific metrics, adding indentation for sub-rows.
  • Implemented rendering logic within ModelUsageTable to display detailed metrics for each role associated with a model.
• packages/cli/src/ui/components/triage/TriageDuplicates.tsx
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_TOOL to the role parameter in the chat.generateJson call for duplicate triage.
• packages/cli/src/ui/components/triage/TriageIssues.tsx
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_TOOL to the role parameter in the chat.generateJson call for issue triage.
• packages/cli/src/ui/contexts/SessionContext.test.tsx
  • Updated mock model metrics to include an empty roles object.
• packages/cli/src/ui/hooks/usePromptCompletion.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_AUTOCOMPLETE to the role parameter in the chat.generateContentStream call for prompt completion.
• packages/cli/src/ui/utils/computeStats.test.ts
  • Updated mock model metrics to include an empty roles object in various test cases.
• packages/cli/src/zed-integration/zedIntegration.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in the chat.sendMessageStream call.
• packages/cli/src/zed-integration/zedIntegration.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in the chat.sendMessageStream call.
• packages/core/src/agents/local-executor.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.SUBAGENT to the role parameter in chat.sendMessageStream mock calls.
• packages/core/src/agents/local-executor.ts
  • Imported LlmRole enum.
  • Dynamically assigned LlmRole.SUBAGENT_CODEBASE_INVESTIGATOR or LlmRole.SUBAGENT based on agent definition name to the role parameter in chat.sendMessageStream calls.
• packages/core/src/code_assist/server.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in server.generateContent and server.generateContentStream calls.
• packages/core/src/code_assist/server.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to generateContentStream and generateContent methods.
• packages/core/src/core/baseLlmClient.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_TOOL to the role parameter in generateJson and generateContent options and mock calls.
• packages/core/src/core/baseLlmClient.ts
  • Imported LlmRole enum.
  • Added role: LlmRole property to GenerateJsonOptions and GenerateContentOptions interfaces.
  • Passed the role parameter through _generateWithRetry and contentGenerator calls.
• packages/core/src/core/client.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in mockContentGenerator.generateContent and generateContentStream calls.
• packages/core/src/core/client.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to _generateContent and _generateContentStream methods.
  • Passed the role parameter to baseLlmClient.generateContentStream.
• packages/core/src/core/contentGenerator.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to the generateContent and generateContentStream interface methods.
• packages/core/src/core/fakeContentGenerator.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in generateContent and generateContentStream calls.
• packages/core/src/core/fakeContentGenerator.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to generateContent and generateContentStream methods.
• packages/core/src/core/geminiChat.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in sendMessageStream calls across various test scenarios.
• packages/core/src/core/geminiChat.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to sendMessageStream and _generateContentStream methods.
  • Passed the role parameter to contentGenerator.generateContentStream.
• packages/core/src/core/geminiChat_network_retry.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in sendMessageStream calls for network retry tests.
• packages/core/src/core/loggingContentGenerator.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN or LlmRole.SUBAGENT to the role parameter in generateContent and generateContentStream calls and expectations.
• packages/core/src/core/loggingContentGenerator.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to _logApiRequest, _logApiResponse, _logApiError, generateContent, generateContentStream, and _streamResponses methods.
  • Included the role in the telemetry events logged.
• packages/core/src/core/recordingContentGenerator.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in generateContent and generateContentStream calls.
• packages/core/src/core/recordingContentGenerator.ts
  • Imported LlmRole type.
  • Added role: LlmRole parameter to generateContent and generateContentStream methods.
• packages/core/src/core/turn.test.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.MAIN to the role parameter in the chat.sendMessageStream call.
• packages/core/src/core/turn.ts
  • Imported LlmRole enum.
  • Added role: LlmRole parameter to the _sendMessageStream method, defaulting to LlmRole.MAIN.
• packages/core/src/output/json-formatter.test.ts
  • Updated mock model metrics to include an empty roles object.
• packages/core/src/output/stream-json-formatter.test.ts
  • Updated mock model metrics to include an empty roles object.
• packages/core/src/routing/strategies/classifierStrategy.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_ROUTER to the role parameter in the llmClient.generateJson call for classification.
• packages/core/src/routing/strategies/numericalClassifierStrategy.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_ROUTER to the role parameter in the llmClient.generateJson call for numerical classification.
• packages/core/src/services/chatCompressionService.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_COMPRESSOR to the role parameter in llmClient.generateContent and llmClient.generateJson calls for chat compression.
• packages/core/src/services/loopDetectionService.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_LOOP_DETECTOR to the role parameter in the llmClient.generateJson call for loop detection.
• packages/core/src/services/sessionSummaryService.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_SUMMARIZER to the role parameter in the llmClient.generateContent call for session summarization.
• packages/core/src/telemetry/index.ts
  • Exported the new LlmRole enum from its dedicated file.
• packages/core/src/telemetry/llmRole.ts
  • Added new file defining the LlmRole enum with various roles like MAIN, SUBAGENT, UTILITY_TOOL, UTILITY_COMPRESSOR, UTILITY_SUMMARIZER, UTILITY_ROUTER, UTILITY_LOOP_DETECTOR, UTILITY_NEXT_SPEAKER, UTILITY_EDIT_CORRECTOR, and UTILITY_AUTOCOMPLETE.
• packages/core/src/telemetry/loggers.ts
  • Modified logApiRequest function to include the event.role attribute in the generated log record.
• packages/core/src/telemetry/types.ts
  • Imported and re-exported LlmRole enum.
  • Added an optional role?: LlmRole property to ApiRequestEvent, ApiErrorEvent, and ApiResponseEvent classes.
  • Updated toLogRecord methods in these event classes to include the role in their attributes if present.
• packages/core/src/telemetry/uiTelemetry.test.ts
  • Updated mock model metrics to include an empty roles object.
  • Added a new test case to verify that ApiErrorEvent with a specified role correctly updates role-specific metrics.
• packages/core/src/telemetry/uiTelemetry.ts
  • Imported LlmRole type.
  • Defined RoleMetrics interface to track requests, errors, latency, and token usage per role.
  • Added roles: Partial<Record<LlmRole, RoleMetrics>> property to the ModelMetrics interface.
  • Created createInitialRoleMetrics helper function for initializing role-specific metrics.
  • Updated processApiResponse and processApiError methods to aggregate metrics for individual LlmRole values.
• packages/core/src/tools/web-fetch.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_TOOL to the role parameter in llmClient.generateContent calls for web fetching.
• packages/core/src/tools/web-search.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_TOOL to the role parameter in the llmClient.generateContent call for web search.
• packages/core/src/utils/editCorrector.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_EDIT_CORRECTOR to the role parameter in llmClient.generateJson calls for edit correction.
• packages/core/src/utils/llm-edit-fixer.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_EDIT_CORRECTOR to the role parameter in the llmClient.generateJson call for LLM edit fixing.
• packages/core/src/utils/nextSpeakerChecker.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_NEXT_SPEAKER to the role parameter in the llmClient.generateJson call for next speaker checking.
• packages/core/src/utils/summarizer.ts
  • Imported LlmRole enum.
  • Assigned LlmRole.UTILITY_SUMMARIZER to the role parameter in the llmClient.generateContent call for summarization.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Size Change: +6.85 kB (+0.03%)

Total Size: 24.4 MB

Filename	Size	Change
./bundle/gemini.js	24.4 MB	+6.85 kB (+0.03%)

ℹ️ View Unchanged

Filename	Size
./bundle/sandbox-macos-permissive-open.sb	890 B
./bundle/sandbox-macos-permissive-proxied.sb	1.31 kB
./bundle/sandbox-macos-restrictive-open.sb	3.36 kB
./bundle/sandbox-macos-restrictive-proxied.sb	3.56 kB
./bundle/sandbox-macos-strict-open.sb	4.82 kB
./bundle/sandbox-macos-strict-proxied.sb	5.02 kB

_{compressed-size-action}

[gemini-code-assist]

Code Review

This pull request introduces the concept of LlmRole to categorize and track LLM calls based on their purpose (e.g., main agent, subagent, utility tools like summarizers or routers). The changes involve defining the LlmRole enum, integrating it into various generateContent and generateJson calls across the codebase (including GeminiChat, BaseLlmClient, ContentGenerator implementations, and specific UI components like TriageDuplicates and TriageIssues), and updating telemetry events (ApiRequestEvent, ApiErrorEvent, ApiResponseEvent) to include this role information. Crucially, the ModelStatsDisplay and StatsDisplay UI components are enhanced to visualize these role-based metrics, showing a breakdown of requests, tokens, and latency per role for each model, with new test cases added to verify this functionality and handle long role names.

[jacob314]

please address the minor comment about the incorrect sort method