Skip to content

Add token usage tracking#513

Merged
ali-sedaghatbaf merged 9 commits intomainfrom
llm-usage
Apr 17, 2026
Merged

Add token usage tracking#513
ali-sedaghatbaf merged 9 commits intomainfrom
llm-usage

Conversation

@ali-sedaghatbaf
Copy link
Copy Markdown
Contributor

@ali-sedaghatbaf ali-sedaghatbaf commented Apr 16, 2026

Description

This pull request enhances the LLM response handling across all supported providers by adding token usage tracking. The LLMResponse objects returned by each LLM now include an optional usage field that provides detailed token counts for requests and responses. This enables better monitoring and analysis of LLM costs and performance.

  • Introduced the LLMUsage type and updated imports and __all__ in src/neo4j_graphrag/llm/__init__.py to expose it across the package.
  • Updated all LLM provider modules (anthropic_llm.py, bedrock_llm.py, cohere_llm.py, mistralai_llm.py, ollama_llm.py, openai_llm.py) to extract token usage information from their respective API responses and include it in the LLMResponse via the new usage parameter.
  • Ensured that all synchronous and asynchronous invocation methods for each provider now consistently attach usage data to the response, or None if unavailable, paving the way for uniform downstream handling and analytics.

Type of Change

  • New feature
  • Bug fix
  • Breaking change
  • Documentation update
  • Project configuration change

Complexity

Complexity: low

How Has This Been Tested?

  • Unit tests
  • E2E tests
  • Manual tests

Checklist

The following requirements should have been met (depending on the changes in the branch):

  • Documentation has been updated
  • Unit tests have been updated
  • E2E tests have been updated
  • Examples have been updated
  • New files have copyright header
  • CLA (https://neo4j.com/developer/cla/) has been signed
  • CHANGELOG.md updated if appropriate

@ali-sedaghatbaf ali-sedaghatbaf marked this pull request as ready for review April 17, 2026 06:42
@ali-sedaghatbaf ali-sedaghatbaf requested a review from a team as a code owner April 17, 2026 06:42
Comment thread src/neo4j_graphrag/llm/openai_llm.py
Copy link
Copy Markdown

@serengil-neo4j serengil-neo4j left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ali-sedaghatbaf ali-sedaghatbaf merged commit 11ebd29 into main Apr 17, 2026
14 checks passed
@ali-sedaghatbaf ali-sedaghatbaf deleted the llm-usage branch April 17, 2026 09:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants