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"] 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/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 " 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}, ) 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