refactor(agents-api): centralize usage tracking and remove deprecated metrics

[abhinav-1305]

User description

resolves: #1468

---

PR Type

Enhancement

---

Description

• Centralize usage tracking logic into dedicated module

• Remove deprecated metrics from individual router files

• Consolidate Prometheus metrics and database tracking

• Improve code organization and maintainability

---

Changes diagram

flowchart LR
  A["Individual Router Files"] --> B["Centralized Usage Tracker"]
  C["Deprecated Metrics"] --> D["Consolidated Metrics"]
  B --> E["Database Tracking"]
  B --> F["Prometheus Metrics"]

Loading

---

Changes walkthrough 📝

Relevant files

Enhancement

litellm.py Update LiteLLM client to use centralized tracking                src/agents-api/agents_api/clients/litellm.py Update import to use centralized track_completion_usage Replace deprecated track_usage call with new function Update comment to reflect centralized tracking +4/-3      usage_tracker.py Add centralized usage tracking utilities                                  src/agents-api/agents_api/common/utils/usage_tracker.py Create new centralized usage tracking module Implement track_completion_usage for non-streaming responses Implement track_streaming_usage for streaming responses Consolidate Prometheus metrics and database tracking +118/-0  create_response.py Migrate response router to centralized tracking                    src/agents-api/agents_api/routers/responses/create_response.py Remove deprecated total_tokens_per_user import Replace direct Prometheus metric call with centralized tracker Add comprehensive metadata to usage tracking +9/-3      chat.py Migrate chat router to centralized tracking                            src/agents-api/agents_api/routers/sessions/chat.py Remove deprecated imports and manual usage tracking Replace streaming usage logic with centralized function Simplify code by removing duplicate tracking logic Remove non-streaming usage tracking (handled elsewhere) +7/-28    metrics.py Remove deprecated metrics module                                                  src/agents-api/agents_api/routers/sessions/metrics.py Delete entire deprecated metrics file Remove total_tokens_per_user Counter definition +0/-7

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-merge-for-open-source]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Error Handling The new centralized tracking functions lack error handling. If database tracking or Prometheus metrics fail, exceptions could propagate and potentially break the main API flow. Consider adding try-catch blocks around tracking operations. # Track Prometheus metrics if response.usage and response.usage.total_tokens > 0: total_tokens_per_user.labels(str(developer_id)).inc( amount=response.usage.total_tokens ) # Track usage in database await track_usage( developer_id=developer_id, model=model, messages=messages, response=response, custom_api_used=custom_api_used, metadata=metadata, connection_pool=connection_pool, ) Data Validation The track_streaming_usage function constructs a ModelResponse object without validating the structure of collected_output. If the output format is unexpected, this could cause runtime errors when accessing dictionary keys. response=ModelResponse( id=response_id, choices=[ Choices( message=Message( content=choice.get("content", ""), tool_calls=choice.get("tool_calls"), ), ) for choice in collected_output ], usage=usage_data, ), custom_api_used=custom_api_used, metadata=metadata, connection_pool=connection_pool, ) Missing Import The code removes imports for Choices and Message from litellm.utils but these classes are still used in the centralized usage tracker. Verify that the usage tracker module properly imports these dependencies. from litellm.utils import ModelResponse from starlette.status import HTTP_201_CREATED

[qodo-merge-for-open-source]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	✅ Validate collected output before processing Suggestion Impact: The core validation logic from the suggestion was implemented - the commit added the exact check "if not collected_output: return" to prevent processing empty output code diff: + # Only track usage in database if we have collected output + if not collected_output: + return Add validation to ensure collected_output is not empty before creating the ModelResponse. An empty list could cause issues when constructing the response object for database tracking. src/agents-api/agents_api/common/utils/usage_tracker.py [65-76] async def track_streaming_usage( *, developer_id: UUID, model: str, messages: list[dict], usage_data: dict[str, Any] | None, collected_output: list[dict], response_id: str, custom_api_used: bool = False, metadata: dict[str, Any] = {}, connection_pool: Any = None, ) -> None: + """ + Tracks usage for streaming responses. + + Args: + developer_id: The developer ID for usage tracking + model: The model name used for the response + messages: The original messages sent to the model + usage_data: Usage data from the streaming response + collected_output: The complete collected output from streaming + response_id: The response ID + custom_api_used: Whether a custom API key was used + metadata: Additional metadata for tracking + connection_pool: Connection pool for testing purposes + """ + # Track Prometheus metrics if usage data is available + if usage_data and usage_data.get("total_tokens", 0) > 0: + total_tokens_per_user.labels(str(developer_id)).inc( + amount=usage_data.get("total_tokens", 0) + ) + # Only track usage in database if we have collected output + if not collected_output: + return + + # Track usage in database + await track_usage( + developer_id=developer_id, + model=model, + messages=messages, + response=ModelResponse( + id=response_id, + choices=[ + Choices( + message=Message( + content=choice.get("content", ""), + tool_calls=choice.get("tool_calls"), + ), + ) + for choice in collected_output + ], + usage=usage_data, + ), + custom_api_used=custom_api_used, + metadata=metadata, + connection_pool=connection_pool, + ) + [Suggestion processed] Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies that collected_output could be empty and proposes adding a defensive check, which improves the function's robustness against this edge case.	Low
Update

[ellipsis-dev]

Caution

Changes requested ❌

Reviewed everything up to dd913f1 in 2 minutes and 4 seconds. Click for details.

• Reviewed 267 lines of code in 5 files
• Skipped 0 files when reviewing.
• Skipped posting 2 draft comments. View those below.
• Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.

1. src/agents-api/agents_api/clients/litellm.py:121

• Draft comment:
  Instead of using print for logging errors in the usage tracking try/except block, consider using a logging framework.
• Reason this comment was not posted:
  Comment was not on a location in the diff, so it can't be submitted as a review comment.

2. src/agents-api/agents_api/clients/litellm.py:15

• Draft comment:
  There's an unexpected, stray closing parenthesis at the beginning of the file ()). This appears to be a typographical error and should be removed.
• Reason this comment was not posted:
  Comment was not on a location in the diff, so it can't be submitted as a review comment.

Workflow ID: wflow_q29I29lQyFgHjxEa

^{You can customize by changing your verbosity settings, reacting with 👍 or 👎, replying to comments, or adding code review rules.}

[abhinav-1305]

please review - @Ahmad-mtos