feat(providers): add Mistral Chat API package and migrate capability to mixins

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"]

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"]

packages/capabilities/text-generation/src/celeste_text_generation/providers/mistral/parameters.py

@@ -1,62 +1,36 @@
-"""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
 from celeste.parameters import ParameterMapper
 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(),