From 34748d706e0190e6094982afaba571080051f4aa Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 18:13:07 +0100 Subject: [PATCH 1/4] feat(providers): add Mistral Chat API package Add standalone provider package for Mistral Chat API with mixin pattern for capability-agnostic reuse. ## Client (MistralChatClient mixin) - HTTP POST/streaming to /v1/chat/completions endpoint - Usage parsing: prompt_tokens, completion_tokens, total_tokens - Content extraction from choices[0].message.content - Finish reason mapping (stop, length, etc.) - Structured output support via response_format ## Parameters - TemperatureMapper: temperature float [0.0-2.0] - MaxTokensMapper: max_tokens integer - OutputSchemaMapper: JSON schema via response_format - Supports single BaseModel and list[BaseModel] - Uses StrictRefResolvingJsonSchemaGenerator for schema generation ## Streaming (MistralChatStream mixin) - SSE event parsing for streaming chat completions - Text delta extraction from choices[0].delta.content - Finish reason and usage tracking in final events ## Config - API base URL: https://api.mistral.ai - Endpoint: /v1/chat/completions All clients follow the mixin pattern for reuse across capabilities. --- packages/providers/mistral/pyproject.toml | 18 +++ .../mistral/src/celeste_mistral/__init__.py | 1 + .../src/celeste_mistral/chat/__init__.py | 1 + .../src/celeste_mistral/chat/client.py | 114 ++++++++++++++++++ .../src/celeste_mistral/chat/config.py | 52 ++++++++ .../src/celeste_mistral/chat/parameters.py | 114 ++++++++++++++++++ .../src/celeste_mistral/chat/streaming.py | 67 ++++++++++ .../mistral/src/celeste_mistral/py.typed | 0 8 files changed, 367 insertions(+) create mode 100644 packages/providers/mistral/pyproject.toml create mode 100644 packages/providers/mistral/src/celeste_mistral/__init__.py create mode 100644 packages/providers/mistral/src/celeste_mistral/chat/__init__.py create mode 100644 packages/providers/mistral/src/celeste_mistral/chat/client.py create mode 100644 packages/providers/mistral/src/celeste_mistral/chat/config.py create mode 100644 packages/providers/mistral/src/celeste_mistral/chat/parameters.py create mode 100644 packages/providers/mistral/src/celeste_mistral/chat/streaming.py create mode 100644 packages/providers/mistral/src/celeste_mistral/py.typed diff --git a/packages/providers/mistral/pyproject.toml b/packages/providers/mistral/pyproject.toml new file mode 100644 index 00000000..e83438ee --- /dev/null +++ b/packages/providers/mistral/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "celeste-mistral" +version = "0.3.0" +description = "Mistral provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" +dependencies = ["celeste-ai", "httpx"] + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_mistral"] diff --git a/packages/providers/mistral/src/celeste_mistral/__init__.py b/packages/providers/mistral/src/celeste_mistral/__init__.py new file mode 100644 index 00000000..6b0ceaea --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/__init__.py @@ -0,0 +1 @@ +"""Mistral provider package for Celeste AI.""" diff --git a/packages/providers/mistral/src/celeste_mistral/chat/__init__.py b/packages/providers/mistral/src/celeste_mistral/chat/__init__.py new file mode 100644 index 00000000..fd4fa182 --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/chat/__init__.py @@ -0,0 +1 @@ +"""Mistral Chat API provider package.""" diff --git a/packages/providers/mistral/src/celeste_mistral/chat/client.py b/packages/providers/mistral/src/celeste_mistral/chat/client.py new file mode 100644 index 00000000..3ab98329 --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/chat/client.py @@ -0,0 +1,114 @@ +"""Mistral Chat API client with shared implementation.""" + +from collections.abc import AsyncIterator +from typing import Any + +import httpx + +from celeste.core import UsageField +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class MistralChatClient: + """Mixin for Mistral Chat API capabilities. + + Provides shared implementation for chat-based capabilities: + - _make_request() - HTTP POST to /v1/chat/completions + - _make_stream_request() - SSE streaming to /v1/chat/completions + + Capability clients extend parsing methods via super() to wrap/transform results. + + Usage: + class MistralTextGenerationClient(MistralChatClient, TextGenerationClient): + def _parse_content(self, response_data, **parameters): + choices = response_data.get("choices", []) + text = choices[0]["message"]["content"] + return self._transform_output(text, **parameters) + """ + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request to Mistral Chat API.""" + request_body["model"] = self.model.id # type: ignore[attr-defined] + + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + + return await self.http_client.post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{config.MistralChatEndpoint.CREATE_CHAT_COMPLETION}", + headers=headers, + json_body=request_body, + ) + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request to Mistral Chat API.""" + request_body["model"] = self.model.id # type: ignore[attr-defined] + request_body["stream"] = True + + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + + return self.http_client.stream_post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{config.MistralChatEndpoint.CREATE_CHAT_COMPLETION}", + headers=headers, + json_body=request_body, + ) + + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map Mistral usage fields to unified names. + + Shared by client and streaming across all capabilities. + """ + return { + UsageField.INPUT_TOKENS: usage_data.get("prompt_tokens"), + UsageField.OUTPUT_TOKENS: usage_data.get("completion_tokens"), + UsageField.TOTAL_TOKENS: usage_data.get("total_tokens"), + } + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from response.""" + usage_data = response_data.get("usage", {}) + return self.map_usage_fields(usage_data) + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Return choices from response.""" + choices = response_data.get("choices", []) + if not choices: + msg = "No choices in response" + raise ValueError(msg) + return choices + + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: + """Extract finish reason from choices.""" + choices = response_data.get("choices", []) + if not choices: + reason = None + else: + reason = choices[0].get("finish_reason") + return FinishReason(reason=reason) + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata dictionary, filtering out content fields.""" + content_fields = {"choices"} + filtered_data = { + k: v for k, v in response_data.items() if k not in content_fields + } + return super()._build_metadata(filtered_data) # type: ignore[misc,no-any-return] + + +__all__ = ["MistralChatClient"] diff --git a/packages/providers/mistral/src/celeste_mistral/chat/config.py b/packages/providers/mistral/src/celeste_mistral/chat/config.py new file mode 100644 index 00000000..27f3af68 --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/chat/config.py @@ -0,0 +1,52 @@ +"""Configuration for Mistral Chat API.""" + +from enum import StrEnum + + +class MistralChatEndpoint(StrEnum): + """Endpoints for Chat API.""" + + CREATE_CHAT_COMPLETION = "/v1/chat/completions" + CREATE_FIM_COMPLETION = "/v1/fim/completions" + CREATE_AGENT_COMPLETION = "/v1/agents/completions" + CREATE_EMBEDDING = "/v1/embeddings" + LIST_MODELS = "/v1/models" + GET_MODEL = "/v1/models/{model_id}" + DELETE_MODEL = "/v1/models/{model_id}" + UPLOAD_FILE = "/v1/files" + LIST_FILES = "/v1/files" + GET_FILE = "/v1/files/{file_id}" + DELETE_FILE = "/v1/files/{file_id}" + GET_FILE_CONTENT = "/v1/files/{file_id}/content" + CREATE_FINE_TUNING_JOB = "/v1/fine_tuning/jobs" + LIST_FINE_TUNING_JOBS = "/v1/fine_tuning/jobs" + GET_FINE_TUNING_JOB = "/v1/fine_tuning/jobs/{job_id}" + CANCEL_FINE_TUNING_JOB = "/v1/fine_tuning/jobs/{job_id}/cancel" + CREATE_BATCH_JOB = "/v1/batch/jobs" + LIST_BATCH_JOBS = "/v1/batch/jobs" + GET_BATCH_JOB = "/v1/batch/jobs/{job_id}" + CANCEL_BATCH_JOB = "/v1/batch/jobs/{job_id}/cancel" + CREATE_AGENT = "/v1/agents" + LIST_AGENTS = "/v1/agents" + GET_AGENT = "/v1/agents/{agent_id}" + DELETE_AGENT = "/v1/agents/{agent_id}" + UPDATE_AGENT = "/v1/agents/{agent_id}" + CREATE_OCR = "/v1/ocr" + CREATE_MODERATION = "/v1/moderations" + CREATE_CHAT_MODERATION = "/v1/chat/moderations" + + +BASE_URL = "https://api.mistral.ai" + +# Alternative +CODESTRAL_HOST = "https://codestral.mistral.ai" + +# Standard +DEFAULT_CONTENT_TYPE = "application/json" +ACCEPT_HEADER = "application/json" + +# Required +FILE_UPLOAD_CONTENT_TYPE = "multipart/form-data" + +# Server-Sent +STREAMING_DELIMITER = "data: [DONE]" diff --git a/packages/providers/mistral/src/celeste_mistral/chat/parameters.py b/packages/providers/mistral/src/celeste_mistral/chat/parameters.py new file mode 100644 index 00000000..ead366f1 --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/chat/parameters.py @@ -0,0 +1,114 @@ +"""Mistral Chat API parameter mappers.""" + +import json +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste.structured_outputs import StrictRefResolvingJsonSchemaGenerator +from celeste.types import StructuredOutput + + +class TemperatureMapper(ParameterMapper): + """Map temperature to Mistral temperature field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform temperature into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["temperature"] = validated_value + return request + + +class MaxTokensMapper(ParameterMapper): + """Map max_tokens to Mistral max_tokens field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform max_tokens into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["max_tokens"] = validated_value + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema to Mistral structured outputs format. + + Handles both single BaseModel and list[BaseModel] types. + Lists are returned as arrays directly (no wrapping needed). + """ + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + origin = get_origin(validated_value) + if origin is list: + inner_type = get_args(validated_value)[0] + inner_schema = TypeAdapter(inner_type).json_schema( + schema_generator=StrictRefResolvingJsonSchemaGenerator, + mode="serialization", + ) + schema = {"type": "array", "items": inner_schema} + name = f"{inner_type.__name__.lower()}_list" + else: + schema = TypeAdapter(validated_value).json_schema( + schema_generator=StrictRefResolvingJsonSchemaGenerator, + mode="serialization", + ) + name = validated_value.__name__.lower() + + request["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": name, + "schema": schema, + "strict": True, + }, + } + return request + + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: + """Parse JSON to BaseModel using Pydantic's TypeAdapter.""" + if value is None: + return content + + # If content is already a BaseModel, return it unchanged + if isinstance(content, BaseModel): + return content + if isinstance(content, list) and content and isinstance(content[0], BaseModel): + return content + + if isinstance(content, str): + parsed = json.loads(content) + else: + parsed = content + return TypeAdapter(value).validate_python(parsed) + + +__all__ = ["MaxTokensMapper", "OutputSchemaMapper", "TemperatureMapper"] diff --git a/packages/providers/mistral/src/celeste_mistral/chat/streaming.py b/packages/providers/mistral/src/celeste_mistral/chat/streaming.py new file mode 100644 index 00000000..c128fa04 --- /dev/null +++ b/packages/providers/mistral/src/celeste_mistral/chat/streaming.py @@ -0,0 +1,67 @@ +"""Mistral Chat SSE parsing for streaming.""" + +from typing import Any + +from .client import MistralChatClient + + +class MistralChatStream: + """Mixin for Chat API SSE parsing. + + Provides shared implementation for all capabilities using Mistral Chat API streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class MistralTextGenerationStream(MistralChatStream, TextGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return TextGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data.""" + choices = event.get("choices", []) + if not choices: + return None + + first_choice = choices[0] + if not isinstance(first_choice, dict): + return None + + delta = first_choice.get("delta", {}) + if not isinstance(delta, dict): + return None + + 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 = first_choice.get("finish_reason") + + usage = None + usage_data = event.get("usage") + if isinstance(usage_data, dict): + usage = MistralChatClient.map_usage_fields(usage_data) + + if not content_delta and not finish_reason: + return None + + return { + "content": content_delta or "", + "finish_reason": finish_reason, + "usage": usage, + "raw_event": event, + } + + +__all__ = ["MistralChatStream"] diff --git a/packages/providers/mistral/src/celeste_mistral/py.typed b/packages/providers/mistral/src/celeste_mistral/py.typed new file mode 100644 index 00000000..e69de29b From 686441987ee894b4842e84469fcde2a646a39a5b Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 18:13:19 +0100 Subject: [PATCH 2/4] refactor(text-generation): migrate Mistral to provider mixins Migrate Mistral text generation capability client to use provider package mixins, eliminating code duplication and centralizing API-specific logic. ## Changes - Mistral client now inherits from MistralChatClient mixin - Parameter mappers inherit from provider package mappers (TemperatureMapper, MaxTokensMapper, OutputSchemaMapper) - Streaming class inherits from MistralChatStream mixin - Remove duplicated HTTP request logic (_make_request, _make_stream_request) - Remove duplicated parameter mapping logic (~125 lines) - Remove duplicated streaming parsing logic (~60 lines) - Simplify _parse_usage to use mixin's implementation - Simplify _parse_content to use mixin's output parsing - Simplify _parse_finish_reason to use mixin's implementation ## Code Reduction - ~261 lines removed across client, parameters, and streaming files - Significant deduplication of HTTP request, parameter mapping, and streaming logic --- .../providers/mistral/client.py | 79 ++------ .../providers/mistral/parameters.py | 183 ++---------------- .../providers/mistral/streaming.py | 104 +++------- 3 files changed, 57 insertions(+), 309 deletions(-) diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/client.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/client.py index bcb2ed6e..cf625b28 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/client.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/client.py @@ -1,13 +1,11 @@ """Mistral client implementation for text generation.""" -from collections.abc import AsyncIterator from typing import Any, Unpack -import httpx -from pydantic import BaseModel +from celeste_mistral.chat.client import MistralChatClient -from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper +from celeste.types import StructuredOutput from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( TextGenerationFinishReason, @@ -16,12 +14,11 @@ ) from celeste_text_generation.parameters import TextGenerationParameters -from . import config from .parameters import MISTRAL_PARAMETER_MAPPERS from .streaming import MistralTextGenerationStream -class MistralTextGenerationClient(TextGenerationClient): +class MistralTextGenerationClient(MistralChatClient, TextGenerationClient): """Mistral client for text generation.""" @classmethod @@ -41,19 +38,14 @@ def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: """Parse usage from response.""" - usage_dict = response_data.get("usage", {}) - - return TextGenerationUsage( - input_tokens=usage_dict.get("prompt_tokens"), - output_tokens=usage_dict.get("completion_tokens"), - total_tokens=usage_dict.get("total_tokens"), - ) + usage = super()._parse_usage(response_data) + return TextGenerationUsage(**usage) def _parse_content( self, response_data: dict[str, Any], **parameters: Unpack[TextGenerationParameters], - ) -> str | BaseModel: + ) -> StructuredOutput: """Parse content from response.""" choices = response_data.get("choices", []) if not choices: @@ -62,7 +54,7 @@ def _parse_content( first_choice = choices[0] message = first_choice.get("message", {}) - content = message.get("content") + content = message.get("content") or "" # Handle magistral thinking models that return list content if isinstance(content, list): @@ -72,23 +64,19 @@ def _parse_content( text_parts.append(block.get("text", "")) content = "".join(text_parts) - return self._transform_output(content or "", **parameters) + return self._transform_output(content, **parameters) def _parse_finish_reason( self, response_data: dict[str, Any] - ) -> TextGenerationFinishReason | None: + ) -> TextGenerationFinishReason: """Parse finish reason from response.""" choices = response_data.get("choices", []) if not choices: - return None - - first_choice = choices[0] - finish_reason_str = first_choice.get("finish_reason") - return ( - TextGenerationFinishReason(reason=finish_reason_str) - if finish_reason_str - else None - ) + finish_reason_str = None + else: + first_choice = choices[0] + finish_reason_str = first_choice.get("finish_reason") + return TextGenerationFinishReason(reason=finish_reason_str) def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: """Build metadata dictionary from response data.""" @@ -99,48 +87,9 @@ def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: } return super()._build_metadata(filtered_data) - async def _make_request( - self, - request_body: dict[str, Any], - **parameters: Unpack[TextGenerationParameters], - ) -> httpx.Response: - """Make HTTP request(s) and return response object.""" - request_body["model"] = self.model.id - - headers = { - **self.auth.get_headers(), - "Content-Type": ApplicationMimeType.JSON, - } - - return await self.http_client.post( - f"{config.BASE_URL}{config.ENDPOINT}", - headers=headers, - json_body=request_body, - ) - def _stream_class(self) -> type[MistralTextGenerationStream]: """Return the Stream class for this client.""" return MistralTextGenerationStream - def _make_stream_request( - self, - request_body: dict[str, Any], - **parameters: Unpack[TextGenerationParameters], - ) -> AsyncIterator[dict[str, Any]]: - """Make HTTP streaming request and return async iterator of events.""" - request_body["model"] = self.model.id - request_body["stream"] = True - - headers = { - **self.auth.get_headers(), - "Content-Type": ApplicationMimeType.JSON, - } - - return self.http_client.stream_post( - f"{config.BASE_URL}{config.STREAM_ENDPOINT}", - headers=headers, - json_body=request_body, - ) - __all__ = ["MistralTextGenerationClient"] diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/parameters.py index 918f024b..79776f3c 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/parameters.py @@ -1,8 +1,16 @@ -"""Mistral parameter mappers for text generation.""" +"""Mistral Chat parameter mappers for text generation.""" -from typing import Any, get_args, get_origin +from typing import Any -from pydantic import BaseModel, TypeAdapter +from celeste_mistral.chat.parameters import ( + MaxTokensMapper as _MaxTokensMapper, +) +from celeste_mistral.chat.parameters import ( + OutputSchemaMapper as _OutputSchemaMapper, +) +from celeste_mistral.chat.parameters import ( + TemperatureMapper as _TemperatureMapper, +) from celeste.core import Parameter from celeste.models import Model @@ -10,53 +18,19 @@ from celeste_text_generation.parameters import TextGenerationParameter -class TemperatureMapper(ParameterMapper): - """Map temperature parameter to Mistral temperature field.""" - +class TemperatureMapper(_TemperatureMapper): name = Parameter.TEMPERATURE - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform temperature into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["temperature"] = validated_value - return request - - -class MaxTokensMapper(ParameterMapper): - """Map max_tokens parameter to Mistral max_tokens field.""" +class MaxTokensMapper(_MaxTokensMapper): name = Parameter.MAX_TOKENS - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform max_tokens into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["max_tokens"] = validated_value - return request - class ThinkingBudgetMapper(ParameterMapper): - """Map thinking_budget parameter to Mistral prompt_mode. + """Map thinking_budget to Mistral prompt_mode (Pattern 3: Coercion). - Maps unified thinking_budget to Mistral's prompt_mode parameter for reasoning models: - - -1: Enable reasoning (prompt_mode: "reasoning") - - 0: Disable reasoning (prompt_mode: null) - - >0: Enable reasoning (prompt_mode: "reasoning") - Note: Mistral doesn't support budget control + Mistral uses prompt_mode instead of a thinking parameter, so this creates + a unified parameter that maps to the provider-specific format. """ name = TextGenerationParameter.THINKING_BUDGET @@ -67,17 +41,13 @@ def map( value: object, model: Model, ) -> dict[str, Any]: - """Transform thinking_budget into provider request. - - Only applies to magistral reasoning models. For other models, silently ignores. - """ - if not model.id.startswith("magistral"): - return request + """Transform thinking_budget into provider request.""" validated_value = self._validate_value(value, model) if validated_value is None: return request + # Map unified values to Mistral's prompt_mode if validated_value == -1: request["prompt_mode"] = "reasoning" elif validated_value == 0: @@ -88,124 +58,9 @@ def map( return request -class OutputSchemaMapper(ParameterMapper): - """Map output_schema parameter to Mistral response_format.""" - +class OutputSchemaMapper(_OutputSchemaMapper): name = TextGenerationParameter.OUTPUT_SCHEMA - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform output_schema into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - schema = self._convert_to_mistral_schema(validated_value) - schema_name = self._get_schema_name(validated_value) - - request["response_format"] = { - "type": "json_schema", - "json_schema": { - "name": schema_name, - "description": validated_value.__doc__ - if hasattr(validated_value, "__doc__") - else "", - "schema": schema, - "strict": True, - }, - } - - return request - - def parse_output(self, content: str, value: object | None) -> str | BaseModel: - """Parse JSON string to BaseModel instance if output_schema provided. - - Args: - content: Raw text content (JSON string when output_schema is set). - value: Original output_schema parameter value. - - Returns: - BaseModel instance if value provided, otherwise str unchanged. - """ - if value is None: - return content - - return TypeAdapter(value).validate_json(content) - - def _convert_to_mistral_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 - """Convert Pydantic BaseModel or list[BaseModel] to Mistral JSON Schema format.""" - origin = get_origin(output_schema) - if origin is list: - inner_type = get_args(output_schema)[0] - items_schema = inner_type.model_json_schema() - json_schema = {"type": "array", "items": items_schema} - else: - json_schema = output_schema.model_json_schema() - - json_schema = self._resolve_refs(json_schema) - return json_schema - - def _resolve_refs(self, schema: dict[str, Any]) -> dict[str, Any]: - """Resolve all $ref references and inline definitions for reliability.""" - defs: dict[str, Any] = {} - - def collect_defs(value: Any) -> None: # noqa: ANN401 - """Recursively collect all $defs dictionaries.""" - if isinstance(value, dict): - if "$defs" in value: - defs.update(value["$defs"]) - for v in value.values(): - collect_defs(v) - elif isinstance(value, list): - for item in value: - collect_defs(item) - - collect_defs(schema) - - def remove_defs(value: Any) -> Any: # noqa: ANN401 - """Recursively remove all $defs keys.""" - if isinstance(value, dict): - result = {k: remove_defs(v) for k, v in value.items() if k != "$defs"} - return result - elif isinstance(value, list): - return [remove_defs(item) for item in value] - return value - - schema = remove_defs(schema) - - def resolve(value: Any) -> Any: # noqa: ANN401 - """Recursively resolve $ref references in schema.""" - if isinstance(value, dict): - if "$ref" in value: - ref_path = value["$ref"] - if ref_path.startswith("#/$defs/"): - ref_name = ref_path.split("/")[-1] - if ref_name in defs: - resolved = defs[ref_name].copy() - resolved.update( - {k: v for k, v in value.items() if k != "$ref"} - ) - return resolve(resolved) - return {k: resolve(v) for k, v in value.items()} - elif isinstance(value, list): - return [resolve(item) for item in value] - return value - - return resolve(schema) - - def _get_schema_name(self, output_schema: Any) -> str: # noqa: ANN401 - """Derive schema name from model class name.""" - origin = get_origin(output_schema) - if origin is list: - inner_type = get_args(output_schema)[0] - return f"{inner_type.__name__.lower()}_list" - else: - return output_schema.__name__.lower() - MISTRAL_PARAMETER_MAPPERS: list[ParameterMapper] = [ TemperatureMapper(), diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/streaming.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/streaming.py index 8f848394..47d7d752 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/streaming.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/streaming.py @@ -1,9 +1,11 @@ """Mistral streaming for text generation.""" -import logging from collections.abc import Callable from typing import Any, Unpack +from celeste_mistral.chat.streaming import MistralChatStream + +from celeste.types import StructuredOutput from celeste_text_generation.io import ( TextGenerationChunk, TextGenerationFinishReason, @@ -13,101 +15,45 @@ from celeste_text_generation.parameters import TextGenerationParameters from celeste_text_generation.streaming import TextGenerationStream -logger = logging.getLogger(__name__) - -class MistralTextGenerationStream(TextGenerationStream): +class MistralTextGenerationStream(MistralChatStream, TextGenerationStream): """Mistral streaming for text generation.""" def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[..., object], + transform_output: Callable[..., StructuredOutput], **parameters: Unpack[TextGenerationParameters], ) -> None: - """Initialize stream with output transformation support. - - Args: - sse_iterator: Server-Sent Events iterator. - transform_output: Function to transform accumulated content (e.g., JSON → BaseModel). - **parameters: Parameters passed to stream() for output transformation. - """ + """Initialize stream with output transformation support.""" super().__init__(sse_iterator, **parameters) self._transform_output = transform_output def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: - """Parse chunk from SSE event. - - Extract from choices[0].delta.content (content delta events). - Extract finish_reason and usage from final event when finish_reason is not null. - Return None if no text delta (filter lifecycle events). - """ - choices = event.get("choices", []) - if not choices: - return None - - first_choice = choices[0] - if not isinstance(first_choice, dict): + """Parse SSE event into typed Chunk.""" + raw = super()._parse_chunk(event) + if not raw: return None - delta = first_choice.get("delta", {}) - if not isinstance(delta, dict): - return 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) - usage = None - usage_dict = event.get("usage") - if isinstance(usage_dict, dict): - usage = TextGenerationUsage( - input_tokens=usage_dict.get("prompt_tokens"), - output_tokens=usage_dict.get("completion_tokens"), - total_tokens=usage_dict.get("total_tokens"), - ) - - # Create finish reason if present + usage = TextGenerationUsage(**raw["usage"]) if raw["usage"] else None finish_reason = ( - TextGenerationFinishReason(reason=finish_reason_str) - if finish_reason_str + TextGenerationFinishReason(reason=raw["finish_reason"]) + if raw["finish_reason"] else None ) - # If no content delta and no finish reason, filter this event - if not content_delta and not finish_reason: - return None - return TextGenerationChunk( - content=content_delta or "", # Empty string if no content (final event) + content=raw["content"], finish_reason=finish_reason, usage=usage, + metadata={"raw_event": raw["raw_event"]}, ) def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: - """Parse usage from chunks. - - Mistral provides usage metadata in the final event (when finish_reason is not null). - Use the last chunk that has usage metadata. - """ - if not chunks: - return TextGenerationUsage() - - # Usage metadata is typically in the final chunk (when finish_reason is set) + """Extract usage from final chunk.""" for chunk in reversed(chunks): if chunk.usage: return chunk.usage - return TextGenerationUsage() def _parse_output( @@ -115,28 +61,26 @@ def _parse_output( chunks: list[TextGenerationChunk], **parameters: Unpack[TextGenerationParameters], ) -> TextGenerationOutput: - """Assemble chunks into final output with structured output support. - - Concatenates text chunks, then applies parameter transformations - (e.g., JSON → BaseModel if output_schema provided). - """ - # Filter out empty chunks (from final events) + """Assemble chunks into final output.""" content_chunks = [chunk for chunk in chunks if chunk.content] - - # Concatenate text chunks content = "".join(chunk.content for chunk in content_chunks) - - # Apply parameter transformations (e.g., JSON → BaseModel if output_schema provided) content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None + raw_response = None + for chunk in reversed(chunks): + raw_event = chunk.metadata.get("raw_event") + if raw_event and raw_event.get("usage"): + raw_response = raw_event + break + return TextGenerationOutput( content=content, usage=usage, finish_reason=finish_reason, - metadata={}, + metadata={"raw_response": raw_response}, ) From 9598cf2ea41bcbe728f401bed31c602219b15b5a Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 18:13:27 +0100 Subject: [PATCH 3/4] chore(text-generation): remove redundant Mistral config file Remove unused config.py file from Mistral capability provider. Configuration is now centralized in the provider package (celeste_mistral.chat.config) following the established pattern from OpenAI and other providers. --- .../providers/mistral/config.py | 10 ---------- 1 file changed, 10 deletions(-) delete mode 100644 packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/config.py diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/config.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/config.py deleted file mode 100644 index 0d9030a8..00000000 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/config.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Mistral provider configuration for text generation.""" - -# HTTP Configuration -BASE_URL = "https://api.mistral.ai" -ENDPOINT = "/v1/chat/completions" -STREAM_ENDPOINT = ENDPOINT - -# Authentication -AUTH_HEADER_NAME = "Authorization" -AUTH_HEADER_PREFIX = "Bearer " From 7164e70365a5784f6aff7a1de543339ed39025f4 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 18:13:34 +0100 Subject: [PATCH 4/4] fix(capabilities): add celeste-mistral dependency to text-generation Add celeste-mistral to [tool.uv.sources] in text-generation capability package that now imports from the Mistral provider package after the mixin migration. This ensures workspace dependencies are properly declared for the refactored capability client that uses Mistral provider mixins. --- packages/capabilities/text-generation/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/capabilities/text-generation/pyproject.toml b/packages/capabilities/text-generation/pyproject.toml index 6fa4b4ac..9976d70b 100644 --- a/packages/capabilities/text-generation/pyproject.toml +++ b/packages/capabilities/text-generation/pyproject.toml @@ -29,6 +29,7 @@ Issues = "https://github.com/withceleste/celeste-python/issues" celeste-ai = { workspace = true } celeste-anthropic = { workspace = true } celeste-google = { workspace = true } +celeste-mistral = { workspace = true } celeste-openai = { workspace = true } [project.entry-points."celeste.packages"]