Skip to content

feat: token usage tracking in responses api#11302

Merged
RamGopalSrikar merged 17 commits intomainfrom
openai-responses-usage
Feb 25, 2026
Merged

feat: token usage tracking in responses api#11302
RamGopalSrikar merged 17 commits intomainfrom
openai-responses-usage

Conversation

@phact
Copy link
Collaborator

@phact phact commented Jan 14, 2026

Token Usage Tracking Implementation

  1. Added Usage model to Properties schema (src/lfx/src/lfx/schema/properties.py and src/backend/base/langflow/schema/properties.py):
    - New Usage class with input_tokens, output_tokens, total_tokens fields
    - Added usage field to Properties class
  2. Added usage extraction in LCModelComponent (src/lfx/src/lfx/base/models/model.py):
    - New extract_usage() method to extract token usage from AIMessage response_metadata
    - Supports OpenAI format (token_usage.prompt_tokens/completion_tokens/total_tokens)
    - Supports Anthropic format (usage.input_tokens/output_tokens)
    - Modified _get_chat_result() to set usage on returned Message for non-streaming
  3. Added streaming usage support (src/lfx/src/lfx/custom/custom_component/component.py):
    - Modified _stream_message() to return usage data from final chunk
    - Modified send_message() to set usage on message properties after streaming
  4. Updated OpenAI Responses API endpoint (src/backend/base/langflow/api/v1/openai_responses.py):
    - Non-streaming: Extract usage from component_output.results Message properties
    - Streaming: Capture usage from add_message event with state=complete
    - Added response.completed event with usage in streaming responses
  5. Fixed integration tests (src/backend/tests/integration/test_openai_responses_integration.py):
    - Added programmatic model configuration for LanguageModelComponent
    - Set API key directly in template (workaround for global variable lookup issue in tests)
    - Added usage validation assertions

Summary by CodeRabbit

Release Notes

  • New Features
    • Added token usage metrics to API responses, displaying input tokens, output tokens, and total tokens consumed for each request
    • Token usage is now captured and available in both streaming and non-streaming response formats
    • Usage data is extracted from language model responses and included in all API response payloads

✏️ Tip: You can customize this high-level summary in your review settings.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Jan 14, 2026

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Walkthrough

This pull request introduces per-response usage metrics (input_tokens, output_tokens, total_tokens) throughout the codebase. A new Usage model is added to schema definitions. Usage data is captured from OpenAI-like responses in both streaming and non-streaming paths and propagated through Message properties and API responses.

Changes

Cohort / File(s) Summary
Schema definitions
src/backend/base/langflow/schema/properties.py, src/lfx/src/lfx/schema/properties.py
New public Usage model introduced with optional fields input_tokens, output_tokens, and total_tokens. New optional usage field added to Properties model in both backend and lfx packages.
API response handling
src/backend/base/langflow/api/v1/openai_responses.py
Usage data extraction and propagation added to both streaming and non-streaming paths. Streaming path captures usage from completed messages and emits response.completed event. Non-streaming path extracts usage from result outputs. Public OpenAIResponsesResponse signature updated to accept usage field.
LLM component integration
src/lfx/src/lfx/base/models/model.py
New public method extract_usage(message: AIMessage) added to parse token usage from OpenAI and Anthropic response metadata. Usage data populated in _get_chat_result flow and attached to Message properties.
Custom component streaming
src/lfx/src/lfx/custom/custom_component/component.py
Stream handling return types updated to tuple[str, dict | None] for _stream_message and _handle_async_iterator. Usage metadata extraction added across streaming chunks with fallback to response_metadata. Captured usage propagated to stored_message properties.
Integration tests
src/backend/tests/integration/test_openai_responses_integration.py
Programmatic configuration of LanguageModelComponent added. Non-streaming test expanded to validate usage object and token fields. Streaming test adds validation for response.completed event containing usage data.

Sequence Diagram

sequenceDiagram
    participant Client
    participant OpenAIAPI as OpenAI API
    participant Handler as OpenAI Response Handler
    participant Schema as Message Properties
    participant Response as OpenAI Response Object

    Client->>OpenAIAPI: Send prompt (streaming or non-streaming)
    OpenAIAPI->>Handler: Return response with usage metadata
    
    alt Streaming Path
        loop For each chunk
            OpenAIAPI-->>Handler: Stream chunk with usage_metadata
            Handler->>Schema: Extract usage (input, output, total tokens)
            Handler->>Handler: Accumulate usage_data
        end
        Handler->>Handler: Message reaches "complete" state
        Handler->>Response: Emit completion event with usage_data
        Response-->>Client: Send response.completed with usage
    else Non-Streaming Path
        Handler->>Schema: Extract usage from result outputs
        Handler->>Handler: Construct usage_data dict
        Handler->>Response: Populate OpenAIResponsesResponse with usage field
        Response-->>Client: Send response with usage_data
    end
    
    Response-->>Client: Final response includes usage metrics
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 5 | ❌ 2
❌ Failed checks (1 error, 1 warning)
Check name Status Explanation Resolution
Test Coverage For New Implementations ❌ Error PR adds new Usage model classes and extract_usage() method but lacks unit tests for critical implementations including Usage model validation, extract_usage() method, modified streaming methods, and usage extraction logic. Add unit tests for Usage model schema validation, extract_usage() method parsing, streaming method return types, and openai_responses.py usage extraction logic to complement existing integration tests.
Test Quality And Coverage ⚠️ Warning Test coverage for token usage tracking is incomplete: no unit tests for Usage schema, extract_usage() method, or streaming usage modifications despite integration tests being added. Add unit tests for Usage model, extract_usage() method, and streaming usage; expand integration tests to include error scenarios, edge cases, and async error handling validation.
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage ✅ Passed Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.
Test File Naming And Structure ✅ Passed Test file follows all specified patterns: proper naming convention, pytest structure with decorators, descriptive test function names, logical organization with helper functions, comprehensive coverage of positive/negative scenarios and edge cases.
Excessive Mock Usage Warning ✅ Passed Integration tests use real component configuration and API interactions without evidence of excessive mock usage, following proper test design practices.
Title check ✅ Passed The title 'feat: token usage tracking in responses api' clearly and concisely describes the main change—adding token usage tracking to the responses API. It is specific, relates to the core functionality introduced across multiple files, and accurately summarizes the primary objective of the changeset.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch openai-responses-usage

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@codecov
Copy link

codecov bot commented Jan 14, 2026

Codecov Report

❌ Patch coverage is 28.23529% with 61 lines in your changes missing coverage. Please review.
✅ Project coverage is 35.45%. Comparing base (5369b63) to head (f44d1fa).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
...c/lfx/src/lfx/custom/custom_component/component.py 22.22% 25 Missing and 3 partials ⚠️
src/lfx/src/lfx/base/models/model.py 9.09% 20 Missing ⚠️
...c/backend/base/langflow/api/v1/openai_responses.py 26.66% 11 Missing ⚠️
src/lfx/src/lfx/base/models/unified_models.py 0.00% 2 Missing ⚠️

❌ Your patch status has failed because the patch coverage (28.23%) is below the target coverage (40.00%). You can increase the patch coverage or adjust the target coverage.
❌ Your project status has failed because the head coverage (41.87%) is below the target coverage (60.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main   #11302      +/-   ##
==========================================
+ Coverage   35.38%   35.45%   +0.07%     
==========================================
  Files        1527     1527              
  Lines       73484    73563      +79     
  Branches    11041    11059      +18     
==========================================
+ Hits        25999    26080      +81     
+ Misses      46073    46068       -5     
- Partials     1412     1415       +3     
Flag Coverage Δ
backend 56.29% <45.00%> (+0.29%) ⬆️
frontend 16.98% <ø> (ø)
lfx 41.87% <23.07%> (-0.06%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
src/backend/base/langflow/schema/properties.py 91.42% <100.00%> (+1.42%) ⬆️
src/lfx/src/lfx/schema/properties.py 82.85% <100.00%> (+2.85%) ⬆️
src/lfx/src/lfx/base/models/unified_models.py 30.27% <0.00%> (-0.13%) ⬇️
...c/backend/base/langflow/api/v1/openai_responses.py 41.53% <26.66%> (+29.08%) ⬆️
src/lfx/src/lfx/base/models/model.py 20.00% <9.09%> (-1.31%) ⬇️
...c/lfx/src/lfx/custom/custom_component/component.py 57.94% <22.22%> (-1.43%) ⬇️

... and 9 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Copy link
Collaborator

@edwinjosechittilappilly edwinjosechittilappilly left a comment

Choose a reason for hiding this comment

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

Tested in LF. LGTM

@github-actions github-actions bot added the lgtm This PR has been approved by a maintainer label Jan 16, 2026
@edwinjosechittilappilly edwinjosechittilappilly changed the title Feat: token usage tracking in responses api feat: token usage tracking in responses api Jan 16, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Jan 16, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Jan 16, 2026
@github-actions
Copy link
Contributor

github-actions bot commented Jan 16, 2026

Frontend Unit Test Coverage Report

Coverage Summary

Lines Statements Branches Functions
Coverage: 19%
18.8% (6095/32404) 12.25% (3097/25275) 12.64% (879/6952)

Unit Test Results

Tests Skipped Failures Errors Time
2310 0 💤 0 ❌ 0 🔥 32.971s ⏱️

@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Jan 29, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 19, 2026
Update test_openai_error_propagation to ignore response.completed events emitted by the stream and only validate response.chunk entries. The stream can send a response.completed event containing a nested response id, but per OpenAI spec only response.chunk messages have top-level id/object/delta. Filter chunks accordingly, require at least one response.chunk, and assert the presence of id/object fields.
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 19, 2026
@github-actions github-actions bot added the enhancement New feature or request label Feb 19, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 20, 2026
@edwinjosechittilappilly edwinjosechittilappilly added this pull request to the merge queue Feb 20, 2026
@edwinjosechittilappilly edwinjosechittilappilly removed this pull request from the merge queue due to a manual request Feb 20, 2026
@edwinjosechittilappilly
Copy link
Collaborator

@HimavarshaVS @RamGopalSrikar Please coordinate with QA and Merge this PR as soon as available?

@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 24, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 24, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 24, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 25, 2026
@github-actions github-actions bot added enhancement New feature or request and removed enhancement New feature or request labels Feb 25, 2026
@RamGopalSrikar RamGopalSrikar added this pull request to the merge queue Feb 25, 2026
Merged via the queue into main with commit 1ba1f09 Feb 25, 2026
95 checks passed
@RamGopalSrikar RamGopalSrikar deleted the openai-responses-usage branch February 25, 2026 15:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request lgtm This PR has been approved by a maintainer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: [LF] token usage tracking in responses api

3 participants