mistralai 1.10.0__py3-none-any.whl → 1.10.1__py3-none-any.whl

mistralai/extra/observability/otel.py

@@ -5,34 +5,23 @@ import os
 import traceback
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Optional, Tuple
 
 import httpx
 import opentelemetry.semconv._incubating.attributes.gen_ai_attributes as gen_ai_attributes
 import opentelemetry.semconv._incubating.attributes.http_attributes as http_attributes
 import opentelemetry.semconv.attributes.server_attributes as server_attributes
 from opentelemetry import propagate, trace
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.resources import SERVICE_NAME, Resource
-from opentelemetry.sdk.trace import SpanProcessor, TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExportResult
+from opentelemetry.sdk.trace import SpanProcessor
 from opentelemetry.trace import Span, Status, StatusCode, Tracer, set_span_in_context
 
 logger = logging.getLogger(__name__)
 
 
 OTEL_SERVICE_NAME: str = "mistralai_sdk"
-OTEL_EXPORTER_OTLP_ENDPOINT: str = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "")
-OTEL_EXPORTER_OTLP_TIMEOUT: int = int(os.getenv("OTEL_EXPORTER_OTLP_TIMEOUT", "2"))
-OTEL_EXPORTER_OTLP_MAX_EXPORT_BATCH_SIZE: int = int(os.getenv("OTEL_EXPORTER_OTLP_MAX_EXPORT_BATCH_SIZE", "512"))
-OTEL_EXPORTER_OTLP_SCHEDULE_DELAY_MILLIS: int = int(os.getenv("OTEL_EXPORTER_OTLP_SCHEDULE_DELAY_MILLIS", "1000"))
-OTEL_EXPORTER_OTLP_MAX_QUEUE_SIZE: int = int(os.getenv("OTEL_EXPORTER_OTLP_MAX_QUEUE_SIZE", "2048"))
-OTEL_EXPORTER_OTLP_EXPORT_TIMEOUT_MILLIS: int = int(os.getenv("OTEL_EXPORTER_OTLP_EXPORT_TIMEOUT_MILLIS", "5000"))
-
 MISTRAL_SDK_OTEL_TRACER_NAME: str = OTEL_SERVICE_NAME + "_tracer"
 
 MISTRAL_SDK_DEBUG_TRACING: bool = os.getenv("MISTRAL_SDK_DEBUG_TRACING", "false").lower() == "true"
-DEBUG_HINT: str = "To see detailed exporter logs, set MISTRAL_SDK_DEBUG_TRACING=true."
+DEBUG_HINT: str = "To see detailed tracing logs, set MISTRAL_SDK_DEBUG_TRACING=true."
 
 
 class MistralAIAttributes:
@@ -52,13 +41,11 @@ class MistralAINameValues(Enum):
     OCR = "ocr"
 
 class TracingErrors(Exception, Enum):
-    FAILED_TO_EXPORT_OTEL_SPANS = "Failed to export OpenTelemetry (OTEL) spans."
-    FAILED_TO_INITIALIZE_OPENTELEMETRY_TRACING = "Failed to initialize OpenTelemetry tracing."
     FAILED_TO_CREATE_SPAN_FOR_REQUEST = "Failed to create span for request."
     FAILED_TO_ENRICH_SPAN_WITH_RESPONSE = "Failed to enrich span with response."
     FAILED_TO_HANDLE_ERROR_IN_SPAN = "Failed to handle error in span."
     FAILED_TO_END_SPAN = "Failed to end span."
-    
+
     def __str__(self):
         return str(self.value)
 
@@ -180,6 +167,7 @@ def enrich_span_from_response(tracer: trace.Tracer, span: Span, operation_id: st
                 start_ns = parse_time_to_nanos(output["created_at"])
                 end_ns = parse_time_to_nanos(output["completed_at"])
                 child_span = tracer.start_span("Tool Execution", start_time=start_ns, context=parent_context)
+                child_span.set_attributes({"agent.trace.public": ""})
                 tool_attributes = {
                     gen_ai_attributes.GEN_AI_OPERATION_NAME: gen_ai_attributes.GenAiOperationNameValues.EXECUTE_TOOL.value,
                     gen_ai_attributes.GEN_AI_TOOL_CALL_ID: output.get("id", ""),
@@ -192,6 +180,7 @@ def enrich_span_from_response(tracer: trace.Tracer, span: Span, operation_id: st
                 start_ns = parse_time_to_nanos(output["created_at"])
                 end_ns = parse_time_to_nanos(output["completed_at"])
                 child_span = tracer.start_span("Message Output", start_time=start_ns, context=parent_context)
+                child_span.set_attributes({"agent.trace.public": ""})
                 message_attributes = {
                     gen_ai_attributes.GEN_AI_OPERATION_NAME: gen_ai_attributes.GenAiOperationNameValues.CHAT.value,
                     gen_ai_attributes.GEN_AI_PROVIDER_NAME: gen_ai_attributes.GenAiProviderNameValues.MISTRAL_AI.value,
@@ -217,70 +206,47 @@ class GenAISpanProcessor(SpanProcessor):
         span.set_attributes({"agent.trace.public": ""})
 
 
-class QuietOTLPSpanExporter(OTLPSpanExporter):
-    def export(self, spans):
-        try:
-            return super().export(spans)
-        except Exception:
-            logger.warning(f"{TracingErrors.FAILED_TO_EXPORT_OTEL_SPANS} {(traceback.format_exc() if MISTRAL_SDK_DEBUG_TRACING else DEBUG_HINT)}")
-            return SpanExportResult.FAILURE
-
-
-def get_or_create_otel_tracer() -> Tuple[bool, Tracer]:
+def get_or_create_otel_tracer() -> tuple[bool, Tracer]:
     """
-    3 possible cases:
+    Get a tracer from the current TracerProvider.
 
-    -> [SDK in a Workflow / App] If there is already a tracer provider set -> use that one
+    The SDK does not set up its own TracerProvider - it relies on the application
+    to configure OpenTelemetry. This follows OTEL best practices where:
+    - Libraries/SDKs get tracers from the global provider
+    - Applications configure the TracerProvider
 
-    -> [SDK standalone] If no tracer provider is set but the OTEL_EXPORTER_OTLP_ENDPOINT is set -> create a new tracer provider that exports to the OTEL_EXPORTER_OTLP_ENDPOINT
+    If no TracerProvider is configured, the ProxyTracerProvider (default) will
+    return a NoOp tracer, effectively disabling tracing. Once the application
+    sets up a real TracerProvider, subsequent spans will be recorded.
 
-    -> Else tracing is disabled
+    Returns:
+        Tuple[bool, Tracer]: (tracing_enabled, tracer)
+            - tracing_enabled is True if a real TracerProvider is configured
+            - tracer is always valid (may be NoOp if no provider configured)
     """
-    tracing_enabled = True
     tracer_provider = trace.get_tracer_provider()
-
-    if isinstance(tracer_provider, trace.ProxyTracerProvider):
-        if OTEL_EXPORTER_OTLP_ENDPOINT:
-            # SDK standalone: No tracer provider but OTEL_EXPORTER_OTLP_ENDPOINT is set -> create a new tracer provider that exports to the OTEL_EXPORTER_OTLP_ENDPOINT
-            try:
-                exporter = QuietOTLPSpanExporter(
-                    endpoint=OTEL_EXPORTER_OTLP_ENDPOINT,
-                    timeout=OTEL_EXPORTER_OTLP_TIMEOUT
-                )
-                resource = Resource.create(attributes={SERVICE_NAME: OTEL_SERVICE_NAME})
-                tracer_provider = TracerProvider(resource=resource)
-
-                span_processor = BatchSpanProcessor(
-                    exporter,
-                    export_timeout_millis=OTEL_EXPORTER_OTLP_EXPORT_TIMEOUT_MILLIS,
-                    max_export_batch_size=OTEL_EXPORTER_OTLP_MAX_EXPORT_BATCH_SIZE,
-                    schedule_delay_millis=OTEL_EXPORTER_OTLP_SCHEDULE_DELAY_MILLIS,
-                    max_queue_size=OTEL_EXPORTER_OTLP_MAX_QUEUE_SIZE
-                )
-
-                tracer_provider.add_span_processor(span_processor)
-                tracer_provider.add_span_processor(GenAISpanProcessor())
-                trace.set_tracer_provider(tracer_provider)
-
-            except Exception:
-                logger.warning(f"{TracingErrors.FAILED_TO_INITIALIZE_OPENTELEMETRY_TRACING} {(traceback.format_exc() if MISTRAL_SDK_DEBUG_TRACING else DEBUG_HINT)}")
-                tracing_enabled = False
-        else:
-            # No tracer provider nor OTEL_EXPORTER_OTLP_ENDPOINT set -> tracing is disabled
-            tracing_enabled = False
-
     tracer = tracer_provider.get_tracer(MISTRAL_SDK_OTEL_TRACER_NAME)
 
+    # Tracing is considered enabled if we have a real TracerProvider (not the default proxy)
+    tracing_enabled = not isinstance(tracer_provider, trace.ProxyTracerProvider)
+
     return tracing_enabled, tracer
 
-def get_traced_request_and_span(tracing_enabled: bool, tracer: Tracer, span: Optional[Span], operation_id: str, request: httpx.Request) -> Tuple[httpx.Request, Optional[Span]]:
+def get_traced_request_and_span(
+    tracing_enabled: bool,
+    tracer: Tracer,
+    span: Span | None,
+    operation_id: str,
+    request: httpx.Request,
+) -> tuple[httpx.Request, Span | None]:
         if not tracing_enabled:
             return request, span
 
         try:
             span = tracer.start_span(name=operation_id)
+            span.set_attributes({"agent.trace.public": ""})
             # Inject the span context into the request headers to be used by the backend service to continue the trace
-            propagate.inject(request.headers)
+            propagate.inject(request.headers, context=set_span_in_context(span))
             span = enrich_span_from_request(span, request)
         except Exception:
             logger.warning(
@@ -295,7 +261,13 @@ def get_traced_request_and_span(tracing_enabled: bool, tracer: Tracer, span: Opt
         return request, span
 
 
-def get_traced_response(tracing_enabled: bool, tracer: Tracer, span: Optional[Span], operation_id: str, response: httpx.Response) -> httpx.Response:
+def get_traced_response(
+    tracing_enabled: bool,
+    tracer: Tracer,
+    span: Span | None,
+    operation_id: str,
+    response: httpx.Response,
+) -> httpx.Response:
     if not tracing_enabled or not span:
         return response
     try:
@@ -315,7 +287,14 @@ def get_traced_response(tracing_enabled: bool, tracer: Tracer, span: Optional[Sp
         end_span(span=span)
     return response
 
-def get_response_and_error(tracing_enabled: bool, tracer: Tracer, span: Optional[Span], operation_id: str, response: httpx.Response, error: Optional[Exception]) -> Tuple[httpx.Response, Optional[Exception]]:
+def get_response_and_error(
+    tracing_enabled: bool,
+    tracer: Tracer,
+    span: Span | None,
+    operation_id: str,
+    response: httpx.Response,
+    error: Exception | None,
+) -> tuple[httpx.Response, Exception | None]:
         if not tracing_enabled or not span:
             return response, error
         try:
@@ -366,7 +345,7 @@ class TracedResponse(httpx.Response):
 
     This hack allows ending the span only once the stream is fully consumed.
     """
-    def __init__(self, *args, span: Optional[Span], **kwargs) -> None:
+    def __init__(self, *args, span: Span | None, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.span = span
 
@@ -381,7 +360,7 @@ class TracedResponse(httpx.Response):
         await super().aclose()
 
     @classmethod
-    def from_response(cls, resp: httpx.Response, span: Optional[Span]) -> "TracedResponse":
+    def from_response(cls, resp: httpx.Response, span: Span | None) -> "TracedResponse":
         traced_resp = cls.__new__(cls)
         traced_resp.__dict__ = copy.copy(resp.__dict__)
         traced_resp.span = span

mistralai/extra/run/context.py

@@ -1,47 +1,41 @@
 import asyncio
 import inspect
 import typing
-from contextlib import AsyncExitStack
-from functools import wraps
 from collections.abc import Callable
-
+from contextlib import AsyncExitStack
 from dataclasses import dataclass, field
-from typing import Union, Optional
+from functools import wraps
+from logging import getLogger
 
 import pydantic
 
-from mistralai.extra import (
-    response_format_from_pydantic_model,
-)
+from mistralai.extra import response_format_from_pydantic_model
 from mistralai.extra.exceptions import RunException
 from mistralai.extra.mcp.base import MCPClientProtocol
 from mistralai.extra.run.result import RunResult
-from mistralai.types.basemodel import OptionalNullable, BaseModel, UNSET
+from mistralai.extra.run.tools import (
+    RunCoroutine,
+    RunFunction,
+    RunMCPTool,
+    RunTool,
+    create_function_result,
+    create_tool_call,
+)
 from mistralai.models import (
-    ResponseFormat,
-    FunctionCallEntry,
-    Tools,
-    ToolsTypedDict,
     CompletionArgs,
     CompletionArgsTypedDict,
-    FunctionResultEntry,
     ConversationInputs,
     ConversationInputsTypedDict,
+    FunctionCallEntry,
+    FunctionResultEntry,
     FunctionTool,
-    MessageInputEntry,
     InputEntries,
+    MessageInputEntry,
+    ResponseFormat,
+    Tools,
+    ToolsTypedDict,
 )
-
-from logging import getLogger
-
-from mistralai.extra.run.tools import (
-    create_function_result,
-    RunFunction,
-    create_tool_call,
-    RunTool,
-    RunMCPTool,
-    RunCoroutine,
-)
+from mistralai.types.basemodel import BaseModel, OptionalNullable, UNSET
 
 if typing.TYPE_CHECKING:
     from mistralai import Beta, OptionalNullable
@@ -56,8 +50,8 @@ class AgentRequestKwargs(typing.TypedDict):
 class ModelRequestKwargs(typing.TypedDict):
     model: str
     instructions: OptionalNullable[str]
-    tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]]
-    completion_args: OptionalNullable[Union[CompletionArgs, CompletionArgsTypedDict]]
+    tools: OptionalNullable[list[Tools] | list[ToolsTypedDict]]
+    completion_args: OptionalNullable[CompletionArgs | CompletionArgsTypedDict]
 
 
 @dataclass
@@ -72,7 +66,7 @@ class RunContext:
           passed if the user wants to continue an existing conversation.
         model (Options[str]): The model name to be used for the conversation. Can't be used along with 'agent_id'.
         agent_id (Options[str]): The agent id to be used for the conversation. Can't be used along with 'model'.
-        output_format (Optional[type[BaseModel]]): The output format expected from the conversation. It represents
+        output_format (type[BaseModel] | None): The output format expected from the conversation. It represents
           the `response_format` which is part of the `CompletionArgs`.
         request_count (int): The number of requests made in the current `RunContext`.
         continue_on_fn_error (bool): Flag to determine if the conversation should continue when function execution
@@ -83,10 +77,10 @@ class RunContext:
     _callable_tools: dict[str, RunTool] = field(init=False, default_factory=dict)
     _mcp_clients: list[MCPClientProtocol] = field(init=False, default_factory=list)
 
-    conversation_id: Optional[str] = field(default=None)
-    model: Optional[str] = field(default=None)
-    agent_id: Optional[str] = field(default=None)
-    output_format: Optional[type[BaseModel]] = field(default=None)
+    conversation_id: str | None = field(default=None)
+    model: str | None = field(default=None)
+    agent_id: str | None = field(default=None)
+    output_format: type[BaseModel] | None = field(default=None)
     request_count: int = field(default=0)
     continue_on_fn_error: bool = field(default=False)
 
@@ -215,10 +209,8 @@ class RunContext:
 
     async def prepare_model_request(
         self,
-        tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]] = UNSET,
-        completion_args: OptionalNullable[
-            Union[CompletionArgs, CompletionArgsTypedDict]
-        ] = UNSET,
+        tools: OptionalNullable[list[Tools] | list[ToolsTypedDict]] = UNSET,
+        completion_args: OptionalNullable[CompletionArgs | CompletionArgsTypedDict] = UNSET,
         instructions: OptionalNullable[str] = None,
     ) -> ModelRequestKwargs:
         if self.model is None:
@@ -254,14 +246,12 @@ async def _validate_run(
     *,
     beta_client: "Beta",
     run_ctx: RunContext,
-    inputs: Union[ConversationInputs, ConversationInputsTypedDict],
+    inputs: ConversationInputs | ConversationInputsTypedDict,
     instructions: OptionalNullable[str] = UNSET,
-    tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]] = UNSET,
-    completion_args: OptionalNullable[
-        Union[CompletionArgs, CompletionArgsTypedDict]
-    ] = UNSET,
+    tools: OptionalNullable[list[Tools] | list[ToolsTypedDict]] = UNSET,
+    completion_args: OptionalNullable[CompletionArgs | CompletionArgsTypedDict] = UNSET,
 ) -> tuple[
-    Union[AgentRequestKwargs, ModelRequestKwargs], RunResult, list[InputEntries]
+    AgentRequestKwargs | ModelRequestKwargs, RunResult, list[InputEntries]
 ]:
     input_entries: list[InputEntries] = []
     if isinstance(inputs, str):
@@ -277,7 +267,7 @@ async def _validate_run(
         output_model=run_ctx.output_format,
         conversation_id=run_ctx.conversation_id,
     )
-    req: Union[AgentRequestKwargs, ModelRequestKwargs]
+    req: AgentRequestKwargs | ModelRequestKwargs
     if run_ctx.agent_id:
         if tools or completion_args:
             raise RunException("Can't set tools or completion_args when using an agent")

mistralai/extra/run/result.py

@@ -1,9 +1,10 @@
 import datetime
 import json
 import typing
-from typing import Union, Annotated, Optional, Literal
 from dataclasses import dataclass, field
-from pydantic import Discriminator, Tag, BaseModel
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Discriminator, Tag
 
 from mistralai.extra.utils.response_format import pydantic_model_from_json
 from mistralai.models import (
@@ -35,15 +36,15 @@ from mistralai.models import (
 )
 from mistralai.utils import get_discriminator
 
-RunOutputEntries = typing.Union[
-    MessageOutputEntry,
-    FunctionCallEntry,
-    FunctionResultEntry,
-    AgentHandoffEntry,
-    ToolExecutionEntry,
-]
+RunOutputEntries = (
+    MessageOutputEntry
+    | FunctionCallEntry
+    | FunctionResultEntry
+    | AgentHandoffEntry
+    | ToolExecutionEntry
+)
 
-RunEntries = typing.Union[RunOutputEntries, MessageInputEntry]
+RunEntries = RunOutputEntries | MessageInputEntry
 
 
 def as_text(entry: RunOutputEntries) -> str:
@@ -140,12 +141,12 @@ class RunFiles:
 @dataclass
 class RunResult:
     input_entries: list[InputEntries]
-    conversation_id: Optional[str] = field(default=None)
+    conversation_id: str | None = field(default=None)
     output_entries: list[RunOutputEntries] = field(default_factory=list)
     files: dict[str, RunFiles] = field(default_factory=dict)
-    output_model: Optional[type[BaseModel]] = field(default=None)
+    output_model: type[BaseModel] | None = field(default=None)
 
-    def get_file(self, file_id: str) -> Optional[RunFiles]:
+    def get_file(self, file_id: str) -> RunFiles | None:
         return self.files.get(file_id)
 
     @property
@@ -172,36 +173,34 @@ class RunResult:
 
 
 class FunctionResultEvent(BaseModel):
-    id: Optional[str] = None
+    id: str | None = None
 
-    type: Optional[Literal["function.result"]] = "function.result"
+    type: Literal["function.result"] | None = "function.result"
 
     result: str
 
     tool_call_id: str
 
-    created_at: Optional[datetime.datetime] = datetime.datetime.now(
+    created_at: datetime.datetime | None = datetime.datetime.now(
         tz=datetime.timezone.utc
     )
 
-    output_index: Optional[int] = 0
+    output_index: int | None = 0
 
 
-RunResultEventsType = typing.Union[SSETypes, Literal["function.result"]]
+RunResultEventsType = SSETypes | Literal["function.result"]
 
 RunResultEventsData = typing.Annotated[
-    Union[
-        Annotated[AgentHandoffDoneEvent, Tag("agent.handoff.done")],
-        Annotated[AgentHandoffStartedEvent, Tag("agent.handoff.started")],
-        Annotated[ResponseDoneEvent, Tag("conversation.response.done")],
-        Annotated[ResponseErrorEvent, Tag("conversation.response.error")],
-        Annotated[ResponseStartedEvent, Tag("conversation.response.started")],
-        Annotated[FunctionCallEvent, Tag("function.call.delta")],
-        Annotated[MessageOutputEvent, Tag("message.output.delta")],
-        Annotated[ToolExecutionDoneEvent, Tag("tool.execution.done")],
-        Annotated[ToolExecutionStartedEvent, Tag("tool.execution.started")],
-        Annotated[FunctionResultEvent, Tag("function.result")],
-    ],
+    Annotated[AgentHandoffDoneEvent, Tag("agent.handoff.done")]
+    | Annotated[AgentHandoffStartedEvent, Tag("agent.handoff.started")]
+    | Annotated[ResponseDoneEvent, Tag("conversation.response.done")]
+    | Annotated[ResponseErrorEvent, Tag("conversation.response.error")]
+    | Annotated[ResponseStartedEvent, Tag("conversation.response.started")]
+    | Annotated[FunctionCallEvent, Tag("function.call.delta")]
+    | Annotated[MessageOutputEvent, Tag("message.output.delta")]
+    | Annotated[ToolExecutionDoneEvent, Tag("tool.execution.done")]
+    | Annotated[ToolExecutionStartedEvent, Tag("tool.execution.started")]
+    | Annotated[FunctionResultEvent, Tag("function.result")],
     Discriminator(lambda m: get_discriminator(m, "type", "type")),
 ]
 

mistralai/extra/run/tools.py

@@ -1,14 +1,11 @@
+import inspect
 import itertools
+import json
 import logging
 from dataclasses import dataclass
-import inspect
-
-from pydantic import Field, create_model
-from pydantic.fields import FieldInfo
-import json
-from typing import cast, Callable, Sequence, Any, ForwardRef, get_type_hints, Union
+from typing import Any, Callable, ForwardRef, Sequence, cast, get_type_hints
 
-from opentelemetry import trace
+import opentelemetry.semconv._incubating.attributes.gen_ai_attributes as gen_ai_attributes
 from griffe import (
     Docstring,
     DocstringSectionKind,
@@ -16,7 +13,9 @@ from griffe import (
     DocstringParameter,
     DocstringSection,
 )
-import opentelemetry.semconv._incubating.attributes.gen_ai_attributes as gen_ai_attributes
+from opentelemetry import trace
+from pydantic import Field, create_model
+from pydantic.fields import FieldInfo
 
 from mistralai.extra.exceptions import RunException
 from mistralai.extra.mcp.base import MCPClientProtocol
@@ -54,7 +53,7 @@ class RunMCPTool:
     mcp_client: MCPClientProtocol
 
 
-RunTool = Union[RunFunction, RunCoroutine, RunMCPTool]
+RunTool = RunFunction | RunCoroutine | RunMCPTool
 
 
 def _get_function_description(docstring_sections: list[DocstringSection]) -> str:

mistralai/extra/struct_chat.py

@@ -1,19 +1,26 @@
-from ..models import ChatCompletionResponse, ChatCompletionChoice, AssistantMessage
-from .utils.response_format import  CustomPydanticModel, pydantic_model_from_json
-from typing import List, Optional, Type, Generic
-from pydantic import BaseModel
 import json
+from typing import Generic
+
+from ..models import AssistantMessage, ChatCompletionChoice, ChatCompletionResponse
+from .utils.response_format import CustomPydanticModel, pydantic_model_from_json
+
 
 class ParsedAssistantMessage(AssistantMessage, Generic[CustomPydanticModel]):
-    parsed: Optional[CustomPydanticModel]
+    parsed: CustomPydanticModel | None
+
 
 class ParsedChatCompletionChoice(ChatCompletionChoice, Generic[CustomPydanticModel]):
-    message: Optional[ParsedAssistantMessage[CustomPydanticModel]]  # type: ignore
+    message: ParsedAssistantMessage[CustomPydanticModel] | None  # type: ignore
+
 
 class ParsedChatCompletionResponse(ChatCompletionResponse, Generic[CustomPydanticModel]):
-    choices: Optional[List[ParsedChatCompletionChoice[CustomPydanticModel]]] # type: ignore
+    choices: list[ParsedChatCompletionChoice[CustomPydanticModel]] | None  # type: ignore
+
 
-def convert_to_parsed_chat_completion_response(response: ChatCompletionResponse, response_format: Type[BaseModel]) -> ParsedChatCompletionResponse:
+def convert_to_parsed_chat_completion_response(
+    response: ChatCompletionResponse,
+    response_format: type[CustomPydanticModel],
+) -> ParsedChatCompletionResponse[CustomPydanticModel]:
     parsed_choices = []
 
     if response.choices:

mistralai/extra/utils/response_format.py

@@ -1,5 +1,6 @@
+from typing import Any, TypeVar
+
 from pydantic import BaseModel
-from typing import TypeVar, Any, Type, Dict
 from ...models import JSONSchema, ResponseFormat
 from ._pydantic_helper import rec_strict_json_schema
 
@@ -7,7 +8,7 @@ CustomPydanticModel = TypeVar("CustomPydanticModel", bound=BaseModel)
 
 
 def response_format_from_pydantic_model(
-    model: Type[CustomPydanticModel],
+    model: type[CustomPydanticModel],
 ) -> ResponseFormat:
     """Generate a strict JSON schema from a pydantic model."""
     model_schema = rec_strict_json_schema(model.model_json_schema())
@@ -18,7 +19,8 @@ def response_format_from_pydantic_model(
 
 
 def pydantic_model_from_json(
-    json_data: Dict[str, Any], pydantic_model: Type[CustomPydanticModel]
+    json_data: dict[str, Any],
+    pydantic_model: type[CustomPydanticModel],
 ) -> CustomPydanticModel:
     """Parse a JSON schema into a pydantic model."""
     return pydantic_model.model_validate(json_data)