Skip to content

feat: Add unified token usage and API call tracking - #502

Open
prateek-dagar wants to merge 2 commits into
google:mainfrom
prateek-dagar:feat/token-usage-tracking
Open

feat: Add unified token usage and API call tracking#502
prateek-dagar wants to merge 2 commits into
google:mainfrom
prateek-dagar:feat/token-usage-tracking

Conversation

@prateek-dagar

@prateek-dagar prateek-dagar commented Jul 25, 2026

Copy link
Copy Markdown

Description

This PR adds token usage and API call tracking across Gemini, OpenAI, and Ollama. The metrics are returned inside the metadata dictionary of the AnnotatedDocument.

Changes made:

  • Added token usage and request ID tracking to LLM outputs.
  • Created track_api_call_details flag to inspect individual chunk/pass metrics.
  • Added tests in tests/token_usage_test.py and updated README.md.

Fixes #155
Fixes #425
Fixes #475

Type: Feature

How Has This Been Tested?

I ran the entire test suite of the repository to verify that all existing tests pass and no regressions were introduced.

pytest tests/token_usage_test.py
pytest

Checklist:

  • I have read and acknowledged Google's Open Source Code of conduct.
  • I have read the Contributing page, and I signed the CLA.
  • I have discussed my proposed solution with code owners in the linked issue(s) and we have agreed upon the general approach.
  • I have made any needed documentation changes, or noted in the linked issue(s) that documentation elsewhere needs updating.
  • I have added tests, or I have ensured existing tests cover the changes.
  • I have followed Google's Python Style Guide and ran pylint over the affected code.

Implement unified token usage tracking and API calls count for Gemini, OpenAI, and Ollama providers, returning metrics in AnnotatedDocument.metadata.
@github-actions github-actions Bot added the size/M Pull request with 150-600 lines changed label Jul 25, 2026
@google-cla

google-cla Bot commented Jul 25, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@github-actions

Copy link
Copy Markdown

⚠️ Branch Update Required

Your branch is 1 commits behind main. Please update your branch to ensure CI checks run with the latest code:

git fetch origin main
git merge origin/main
git push

Note: Enable "Allow edits by maintainers" to allow automatic updates.

1 similar comment
@github-actions

github-actions Bot commented Aug 2, 2026

Copy link
Copy Markdown

⚠️ Branch Update Required

Your branch is 1 commits behind main. Please update your branch to ensure CI checks run with the latest code:

git fetch origin main
git merge origin/main
git push

Note: Enable "Allow edits by maintainers" to allow automatic updates.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/M Pull request with 150-600 lines changed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Feature Request: Token Usage Tracking (or model metadata?) Support Inference cost Request: Token Count on Gemini (all) APIs

1 participant