From 632b6d9a90a265020a361f1d9b268c16acbdde17 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Thu, 27 Nov 2025 18:51:47 +0100 Subject: [PATCH 1/3] refactor: remove explicit None assignments for usage token fields Remove unnecessary explicit None assignments for optional usage fields (input_tokens, output_tokens, billed_tokens, cached_tokens, reasoning_tokens) since they already default to None in their respective Usage classes. - Anthropic: Extract cached_tokens from cache_read_input_tokens - Google: Remove billed_tokens and cached_tokens explicit None - OpenAI: Remove billed_tokens explicit None - XAI: Remove billed_tokens explicit None - ByteDance: Remove input_tokens and output_tokens explicit None --- .../celeste_image_generation/providers/bytedance/streaming.py | 2 -- .../src/celeste_text_generation/providers/anthropic/client.py | 4 +--- .../celeste_text_generation/providers/anthropic/streaming.py | 4 +--- .../src/celeste_text_generation/providers/google/client.py | 2 -- .../src/celeste_text_generation/providers/openai/client.py | 1 - .../src/celeste_text_generation/providers/openai/streaming.py | 1 - .../src/celeste_text_generation/providers/xai/client.py | 1 - .../src/celeste_text_generation/providers/xai/streaming.py | 1 - 8 files changed, 2 insertions(+), 14 deletions(-) diff --git a/packages/image-generation/src/celeste_image_generation/providers/bytedance/streaming.py b/packages/image-generation/src/celeste_image_generation/providers/bytedance/streaming.py index 4a3fafec..28caf487 100644 --- a/packages/image-generation/src/celeste_image_generation/providers/bytedance/streaming.py +++ b/packages/image-generation/src/celeste_image_generation/providers/bytedance/streaming.py @@ -38,8 +38,6 @@ def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | Non if usage_data: self._completed_usage = ImageGenerationUsage( total_tokens=usage_data.get("total_tokens"), - input_tokens=None, - output_tokens=None, ) return None diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py index 45d957db..64367508 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py @@ -49,9 +49,7 @@ def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, - billed_tokens=None, - cached_tokens=None, - reasoning_tokens=None, + cached_tokens=usage_data.get("cache_read_input_tokens"), ) def _parse_content( diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py index 7e12bfc8..90c22af2 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py @@ -105,9 +105,7 @@ def _parse_usage_from_event( input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, - billed_tokens=None, - cached_tokens=None, - reasoning_tokens=None, + cached_tokens=usage_data.get("cache_read_input_tokens"), ) def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/client.py b/packages/text-generation/src/celeste_text_generation/providers/google/client.py index 044dc073..d547464b 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/google/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/google/client.py @@ -48,8 +48,6 @@ def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: output_tokens=usage_metadata.get("candidatesTokenCount"), total_tokens=usage_metadata.get("totalTokenCount"), reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), - billed_tokens=None, - cached_tokens=None, ) def _parse_content( diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/client.py b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py index 81dba849..3621c3a8 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py @@ -44,7 +44,6 @@ def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: total_tokens=usage_data.get("total_tokens"), cached_tokens=input_tokens_details.get("cached_tokens"), reasoning_tokens=output_tokens_details.get("reasoning_tokens"), - billed_tokens=None, ) def _parse_content( diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py index 2dc15302..8bf708d4 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py @@ -59,7 +59,6 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: total_tokens=usage_data.get("total_tokens"), cached_tokens=input_tokens_details.get("cached_tokens"), reasoning_tokens=output_tokens_details.get("reasoning_tokens"), - billed_tokens=None, ) finish_reason: TextGenerationFinishReason | None = None diff --git a/packages/text-generation/src/celeste_text_generation/providers/xai/client.py b/packages/text-generation/src/celeste_text_generation/providers/xai/client.py index 0e5a3b3a..a4ef602f 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/xai/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/xai/client.py @@ -51,7 +51,6 @@ def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: total_tokens=usage_data.get("total_tokens"), cached_tokens=prompt_tokens_details.get("cached_tokens"), reasoning_tokens=completion_tokens_details.get("reasoning_tokens"), - billed_tokens=None, ) def _parse_content( diff --git a/packages/text-generation/src/celeste_text_generation/providers/xai/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/xai/streaming.py index 1d1a74c7..daffd7b7 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/xai/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/xai/streaming.py @@ -68,7 +68,6 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: total_tokens=usage_dict.get("total_tokens"), cached_tokens=prompt_tokens_details.get("cached_tokens"), reasoning_tokens=completion_tokens_details.get("reasoning_tokens"), - billed_tokens=None, ) # Create finish reason if present From 6285868fb1b3e06c6a39072ca1d45f02e72c43ac Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Thu, 27 Nov 2025 18:52:18 +0100 Subject: [PATCH 2/3] feat: add support for Mistral Magistral reasoning models array format Handle Magistral array format in content parsing for reasoning models. When content is an array, extract text from blocks with type 'text'. --- .../providers/mistral/client.py | 15 ++++++++++++--- .../providers/mistral/streaming.py | 10 ++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py index bd2bae54..5a116243 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py @@ -62,9 +62,18 @@ def _parse_content( first_choice = choices[0] message = first_choice.get("message", {}) - content = message.get("content") or "" - - return self._transform_output(content, **parameters) + content = message.get("content") + + # Handle Magistral array format (reasoning models) + if isinstance(content, list): + for block in content: + if block.get("type") == "text": + content = block.get("text", "") + break + else: + content = "" + + return self._transform_output(content or "", **parameters) def _parse_finish_reason( self, response_data: dict[str, Any] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py index c36fe5fa..de007ccb 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py @@ -56,6 +56,16 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: # Extract content delta content_delta = delta.get("content") + + # Handle Magistral array format (reasoning models) + if isinstance(content_delta, list): + for block in content_delta: + if block.get("type") == "text": + content_delta = block.get("text", "") + break + else: + content_delta = "" + finish_reason_str = first_choice.get("finish_reason") # Extract usage from event if present (in final event) From 7d45a5bb28e45fd1caf9f8f004428df1170c0031 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Thu, 27 Nov 2025 19:06:11 +0100 Subject: [PATCH 3/3] fix(mistral): add isinstance check for magistral content blocks Prevent AttributeError on malformed responses by checking that content blocks are dictionaries before calling .get() method. Addresses PR review feedback. --- .../providers/mistral/client.py | 11 +++++------ .../providers/mistral/streaming.py | 11 +++++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py index 5a116243..669c36c8 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py @@ -64,14 +64,13 @@ def _parse_content( message = first_choice.get("message", {}) content = message.get("content") - # Handle Magistral array format (reasoning models) + # Handle magistral thinking models that return list content if isinstance(content, list): + text_parts = [] for block in content: - if block.get("type") == "text": - content = block.get("text", "") - break - else: - content = "" + if isinstance(block, dict) and block.get("type") == "text": + text_parts.append(block.get("text", "")) + content = "".join(text_parts) return self._transform_output(content or "", **parameters) diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py index de007ccb..8f848394 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py @@ -57,14 +57,13 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: # Extract content delta content_delta = delta.get("content") - # Handle Magistral array format (reasoning models) + # Handle magistral thinking models that may return list content if isinstance(content_delta, list): + text_parts = [] for block in content_delta: - if block.get("type") == "text": - content_delta = block.get("text", "") - break - else: - content_delta = "" + if isinstance(block, dict) and block.get("type") == "text": + text_parts.append(block.get("text", "")) + content_delta = "".join(text_parts) if text_parts else None finish_reason_str = first_choice.get("finish_reason")