fix: Cohere usage parsing and remove deprecated models

[Kamilbenkirane]

Changes

• Fix usage parsing bug: Changed from looking for meta field to usage field in Cohere API response

  • Fixes issue where input_tokens, output_tokens, and total_tokens were None even when data existed
  • Updated both non-streaming client and streaming stream-end event handler

• Remove deprecated models: Removed command-r and command-r-plus models that were deprecated by Cohere on September 15, 2025

  • These models are no longer available in the Cohere API
  • Users should use command-a-03-2025 or command-r7b-12-2024 instead

Files Changed

• packages/text-generation/src/celeste_text_generation/providers/cohere/client.py
• packages/text-generation/src/celeste_text_generation/providers/cohere/models.py
• packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py

[claude]

Code Review: Cohere Usage Parsing Fix and Deprecated Models Removal

Summary

This PR addresses two important issues:

1. Bug Fix: Corrects usage parsing from meta to usage field in Cohere API responses
2. Maintenance: Removes deprecated command-r and command-r-plus models

✅ Positive Aspects

Code Quality

• Clean, focused changes with minimal surface area
• Consistent implementation across both streaming and non-streaming code paths
• Good variable naming (usage_data is clearer than meta for this context)
• Maintains existing code style and patterns

Bug Fix

• The fix correctly addresses the root cause - Cohere's API structure uses usage not meta
• Applied consistently in three locations:
  • client.py:44 (non-streaming)
  • streaming.py:66 (message-end event)
  • streaming.py:97 (stream-end event)

Model Removal

• Appropriate removal of deprecated models
• Clear documentation in PR description about replacement models

🔍 Observations & Suggestions

1. Test Coverage ⚠️
The most significant concern is the lack of test coverage for this bug fix. I couldn't find unit or integration tests for the Cohere provider's usage parsing logic.

Recommendations:

• Add unit tests to verify usage parsing with mock API responses containing the usage field structure
• Add regression tests to ensure input_tokens, output_tokens, and total_tokens are correctly populated
• Consider adding tests for both streaming and non-streaming scenarios
• Test edge cases: missing fields, null values, etc.

2. API Response Validation
The code assumes the Cohere API structure but doesn't validate it explicitly. Consider:

# In streaming.py:66 and 97
usage_dict = delta.get("usage", {})  # Could be message-end event
usage_data = event.get("usage", {})  # stream-end event

This inconsistency in naming (usage_dict vs usage_data) is minor but worth noting for consistency.

3. Deprecated Models Deprecation Date ⚠️
The PR description states: "deprecated by Cohere on September 15, 2025"

This is a future date (we're in November 2025 according to context, but September 2025 is still odd). Did you mean 2024? Please verify the actual deprecation date.

4. Duplication in Streaming Code
The usage parsing logic appears three times with minor variations:

• client.py:42-60 (non-streaming)
• streaming.py:66-81 (message-end)
• streaming.py:97-112 (stream-end)

Suggestion: Consider extracting this into a shared helper method to reduce duplication and improve maintainability:

def _parse_usage_from_dict(usage_data: dict[str, Any]) -> TextGenerationUsage | None:
    """Parse usage from a usage data dictionary."""
    if not isinstance(usage_data, dict):
        return None
    
    billed_units = usage_data.get("billed_units", {})
    tokens = usage_data.get("tokens", {})
    
    input_tokens = billed_units.get("input_tokens")
    output_tokens = billed_units.get("output_tokens")
    
    if input_tokens is not None or output_tokens is not None:
        return TextGenerationUsage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=tokens.get("total_tokens") if tokens else None,
            cached_tokens=usage_data.get("cached_tokens"),
        )
    
    return None

5. Migration Path
For users who may have hardcoded command-r or command-r-plus:

• Consider adding a migration note in a CHANGELOG or release notes
• The PR description mentions replacement models, which is good

🔒 Security

No security concerns identified. The changes don't introduce any new attack vectors.

⚡ Performance

No performance concerns. The changes are neutral from a performance perspective.

🐛 Potential Bugs

The code looks correct for the intended fix. However:

• Without tests, we can't verify this works with real Cohere API responses
• Backward compatibility: If there are any edge cases where Cohere still sends meta, those would now fail (unlikely, but worth considering)

📋 Action Items

Critical:

1. ✅ Fix the deprecation date in the PR description (2025 vs 2024?)
2. ⚠️ Add unit tests for usage parsing with the new usage field structure

Recommended:
3. Consider refactoring duplicate usage parsing logic into a shared helper
4. Add integration tests that verify actual Cohere API responses work correctly
5. Document the breaking change if users were relying on the old models

Verdict

The code changes are correct and well-implemented, but the PR would benefit significantly from:

1. Test coverage
2. Verification of the deprecation date
3. Minor refactoring to reduce duplication

Overall: Approve with suggestions - The bug fix is necessary and correct, but please add tests before merging.