From b9e9a81936a7c79289bf247296f3de3a306a27f5 Mon Sep 17 00:00:00 2001
From: kamilbenkirane <benkirane.kamil@gmail.com>
Date: Thu, 27 Nov 2025 18:21:39 +0100
Subject: [PATCH] fix: Cohere usage parsing and remove deprecated models

- Fix usage parsing to use 'usage' field instead of 'meta' in API response
- Remove deprecated 'command-r' and 'command-r-plus' models
- Update streaming stream-end event handler to use correct field name
- Fixes issue where input_tokens, output_tokens, and total_tokens were None
---
 .../providers/cohere/client.py                |  8 +++---
 .../providers/cohere/models.py                | 25 -------------------
 .../providers/cohere/streaming.py             | 10 ++++----
 3 files changed, 9 insertions(+), 34 deletions(-)

diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py
index 5c7cf88a..631d6114 100644
--- a/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py
+++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py
@@ -41,10 +41,10 @@ def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]:
 
     def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage:
         """Parse usage from response."""
-        meta = response_data.get("meta", {})
+        usage_data = response_data.get("usage", {})
 
-        billed_units = meta.get("billed_units", {})
-        tokens = meta.get("tokens", {})
+        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")
@@ -54,7 +54,7 @@ def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage:
                 input_tokens=input_tokens,
                 output_tokens=output_tokens,
                 total_tokens=tokens.get("total_tokens") if tokens else None,
-                cached_tokens=meta.get("cached_tokens"),
+                cached_tokens=usage_data.get("cached_tokens"),
             )
 
         return TextGenerationUsage()
diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py
index 38f25cbc..11d0de1a 100644
--- a/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py
+++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py
@@ -18,31 +18,6 @@
             # thinking_budget: Not confirmed for this model, omit constraint
         },
     ),
-    Model(
-        id="command-r-plus",
-        provider=Provider.COHERE,
-        display_name="Command R+",
-        streaming=True,
-        parameter_constraints={
-            Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01),
-            Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1),
-            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
-            # command-r-plus supports reasoning (optimized for complex reasoning)
-            TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=31000, step=1),
-        },
-    ),
-    Model(
-        id="command-r",
-        provider=Provider.COHERE,
-        display_name="Command R",
-        streaming=True,
-        parameter_constraints={
-            Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01),
-            Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1),
-            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
-            # thinking_budget: Support unclear, omit constraint for now
-        },
-    ),
     Model(
         id="command-r7b-12-2024",
         provider=Provider.COHERE,
diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py
index 5573575f..1567bc83 100644
--- a/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py
+++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py
@@ -94,11 +94,11 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
                 else None
             )
 
-            meta = event.get("meta", {})
+            usage_data = event.get("usage", {})
             usage = None
-            if isinstance(meta, dict):
-                billed_units = meta.get("billed_units", {})
-                tokens = meta.get("tokens", {})
+            if isinstance(usage_data, dict):
+                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")
@@ -108,7 +108,7 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
                         input_tokens=input_tokens,
                         output_tokens=output_tokens,
                         total_tokens=tokens.get("total_tokens") if tokens else None,
-                        cached_tokens=meta.get("cached_tokens"),
+                        cached_tokens=usage_data.get("cached_tokens"),
                     )
 
             return TextGenerationChunk(