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/mistral/client.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py index bd2bae54..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 @@ -62,9 +62,17 @@ def _parse_content( first_choice = choices[0] message = first_choice.get("message", {}) - content = message.get("content") or "" + content = message.get("content") - return self._transform_output(content, **parameters) + # Handle magistral thinking models that return list content + if isinstance(content, list): + text_parts = [] + for block in 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) 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..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 @@ -56,6 +56,15 @@ def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: # Extract content delta content_delta = delta.get("content") + + # Handle magistral thinking models that may return list content + if isinstance(content_delta, list): + text_parts = [] + for block in 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") # Extract usage from event if present (in final event) 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