Update cost-tracking for OAI chatcompletions and response API

[amanjaiswal73892]

This pull request refines token usage tracking and cost calculation logic in the src/agentlab/llm/tracking.py file. The changes improve handling of cached tokens and introduce better differentiation between API types (chatcompletion vs. response) for more accurate cost computation.

Enhancements to token usage tracking:

• src/agentlab/llm/tracking.py: Updated the __call__ method to extract and include cached_tokens from prompt_tokens_details or input_tokens_details in the usage dictionary, ensuring more precise tracking of cached token usage.

Improvements to cost calculation logic:

• src/agentlab/llm/tracking.py: Refactored the get_effective_cost_from_openai_api method to handle two distinct API types (chatcompletion and response). Added logic to compute cached_input_tokens and non_cached_input_tokens separately for each type, improving the accuracy of effective cost calculations. Introduced warnings for missing usage information or unsupported API types.…ponses API

Description by Korbit AI

What change is being made?

Update cost-tracking logic for OpenAI chatcompletion and response APIs to account for token caching and improve cost calculation accuracy.

Why are these changes being made?

The change introduces handling for prompt_tokens_details and input_tokens_details to accurately assess cached tokens, addressing previous shortcomings in cost estimation by accounting for cached and non-cached token differences. This ensures more precise cost tracking and provides fallback handling for missing API usage information, enhancing robustness and logging unsupported API types to prevent miscalculations.

Is this description stale? Ask me to generate a new description by commenting /korbit-generate-pr-description

[korbit-ai]

Review by Korbit AI

Korbit automatically attempts to detect when you fix issues in new commits.

Category	Issue	Status
	Incomplete API type detection logic ▹ view	✅ Fix detected
	Inconsistent dictionary key names causing AttributeError ▹ view	✅ Fix detected
	Missing Response Context in Log ▹ view
	Redundant Token Extraction Logic ▹ view

Files scanned

​

File Path	Reviewed
src/agentlab/llm/tracking.py	✅

Explore our documentation to understand the languages and file types we support and the files we ignore.

Check out our docs on how you can make Korbit work best for you and your team.

Loving Korbit!? Share us on LinkedIn Reddit and X

[korbit-ai]

Redundant Token Extraction Logic

Tell me more

What is the issue?

Duplicated token extraction logic with only attribute names differing between API types.

Why this matters

The duplicated structure makes the code harder to maintain and obscures the fact that both branches follow the same pattern.

Suggested change ∙ Feature Preview

TOKEN_MAPPINGS = {
    'chatcompletion': {
        'total_tokens': 'prompt_tokens',
        'output_tokens': 'completion_tokens',
        'details_attr': 'prompt_tokens_details'
    },
    'response': {
        'total_tokens': 'input_tokens',
        'output_tokens': 'output_tokens',
        'details_attr': 'input_tokens_details'
    }
}

mapping = TOKEN_MAPPINGS.get(api_type)
if mapping:
    total_input_tokens = getattr(usage, mapping['total_tokens'])
    output_tokens = getattr(usage, mapping['output_tokens'])
    details = getattr(usage, mapping['details_attr'])
    cached_input_tokens = details.cached_tokens
    non_cached_input_tokens = total_input_tokens - cached_input_tokens

Provide feedback to improve future suggestions

💬 Looking for more details? Reply to this comment to chat with Korbit.