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

mistralai/extra/run/tools.py

@@ -1,13 +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
 
+import opentelemetry.semconv._incubating.attributes.gen_ai_attributes as gen_ai_attributes
 from griffe import (
     Docstring,
     DocstringSectionKind,
@@ -15,9 +13,13 @@ from griffe import (
     DocstringParameter,
     DocstringSection,
 )
+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
+from mistralai.extra.observability.otel import GenAISpanEnum, MistralAIAttributes, set_available_attributes
 from mistralai.extra.run.result import RunOutputEntries
 from mistralai.models import (
     FunctionResultEntry,
@@ -51,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:
@@ -191,22 +193,31 @@ async def create_function_result(
         if isinstance(function_call.arguments, str)
         else function_call.arguments
     )
-    try:
-        if isinstance(run_tool, RunFunction):
-            res = run_tool.callable(**arguments)
-        elif isinstance(run_tool, RunCoroutine):
-            res = await run_tool.awaitable(**arguments)
-        elif isinstance(run_tool, RunMCPTool):
-            res = await run_tool.mcp_client.execute_tool(function_call.name, arguments)
-    except Exception as e:
-        if continue_on_fn_error is True:
-            return FunctionResultEntry(
-                tool_call_id=function_call.tool_call_id,
-                result=f"Error while executing {function_call.name}: {str(e)}",
-            )
-        raise RunException(
-            f"Failed to execute tool {function_call.name} with arguments '{function_call.arguments}'"
-        ) from e
+    tracer = trace.get_tracer(__name__)
+    with tracer.start_as_current_span(GenAISpanEnum.function_call(function_call.name)) as span:
+        try:
+            if isinstance(run_tool, RunFunction):
+                res = run_tool.callable(**arguments)
+            elif isinstance(run_tool, RunCoroutine):
+                res = await run_tool.awaitable(**arguments)
+            elif isinstance(run_tool, RunMCPTool):
+                res = await run_tool.mcp_client.execute_tool(function_call.name, arguments)
+            function_call_attributes = {
+                    gen_ai_attributes.GEN_AI_OPERATION_NAME: gen_ai_attributes.GenAiOperationNameValues.EXECUTE_TOOL.value,
+                    gen_ai_attributes.GEN_AI_TOOL_CALL_ID: function_call.id,
+                    MistralAIAttributes.MISTRAL_AI_TOOL_CALL_ARGUMENTS: str(function_call.arguments),
+                    gen_ai_attributes.GEN_AI_TOOL_NAME: function_call.name
+                }
+            set_available_attributes(span, function_call_attributes)
+        except Exception as e:
+            if continue_on_fn_error is True:
+                return FunctionResultEntry(
+                    tool_call_id=function_call.tool_call_id,
+                    result=f"Error while executing {function_call.name}: {str(e)}",
+                )
+            raise RunException(
+                f"Failed to execute tool {function_call.name} with arguments '{function_call.arguments}'"
+            ) from e
 
     return FunctionResultEntry(
         tool_call_id=function_call.tool_call_id,

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)

mistralai/files.py

@@ -212,6 +212,7 @@ class Files(BaseSDK):
         *,
         page: Optional[int] = 0,
         page_size: Optional[int] = 100,
+        include_total: Optional[bool] = True,
         sample_type: OptionalNullable[List[models.SampleType]] = UNSET,
         source: OptionalNullable[List[models.Source]] = UNSET,
         search: OptionalNullable[str] = UNSET,
@@ -227,6 +228,7 @@ class Files(BaseSDK):
 
         :param page:
         :param page_size:
+        :param include_total:
         :param sample_type:
         :param source:
         :param search:
@@ -249,6 +251,7 @@ class Files(BaseSDK):
         request = models.FilesAPIRoutesListFilesRequest(
             page=page,
             page_size=page_size,
+            include_total=include_total,
             sample_type=sample_type,
             source=source,
             search=search,
@@ -310,6 +313,7 @@ class Files(BaseSDK):
         *,
         page: Optional[int] = 0,
         page_size: Optional[int] = 100,
+        include_total: Optional[bool] = True,
         sample_type: OptionalNullable[List[models.SampleType]] = UNSET,
         source: OptionalNullable[List[models.Source]] = UNSET,
         search: OptionalNullable[str] = UNSET,
@@ -325,6 +329,7 @@ class Files(BaseSDK):
 
         :param page:
         :param page_size:
+        :param include_total:
         :param sample_type:
         :param source:
         :param search:
@@ -347,6 +352,7 @@ class Files(BaseSDK):
         request = models.FilesAPIRoutesListFilesRequest(
             page=page,
             page_size=page_size,
+            include_total=include_total,
             sample_type=sample_type,
             source=source,
             search=search,

mistralai/fim.py

@@ -6,7 +6,7 @@ from mistralai._hooks import HookContext
 from mistralai.types import OptionalNullable, UNSET
 from mistralai.utils import eventstreaming, get_security_from_env
 from mistralai.utils.unmarshal_json_response import unmarshal_json_response
-from typing import Any, Mapping, Optional, Union
+from typing import Any, Dict, Mapping, Optional, Union
 
 
 class Fim(BaseSDK):
@@ -28,6 +28,7 @@ class Fim(BaseSDK):
             ]
         ] = None,
         random_seed: OptionalNullable[int] = UNSET,
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         suffix: OptionalNullable[str] = UNSET,
         min_tokens: OptionalNullable[int] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -39,7 +40,7 @@ class Fim(BaseSDK):
 
         FIM completion.
 
-        :param model: ID of the model to use. Only compatible for now with:   - `codestral-2405`   - `codestral-latest`
+        :param model: ID of the model with FIM to use.
         :param prompt: The text/code to complete.
         :param temperature: What sampling temperature to use, we recommend between 0.0 and 0.7. Higher values like 0.7 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both. The default value varies depending on the model you are targeting. Call the `/models` endpoint to retrieve the appropriate value.
         :param top_p: Nucleus sampling, where the model considers the results of the tokens with `top_p` probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or `temperature` but not both.
@@ -47,6 +48,7 @@ class Fim(BaseSDK):
         :param stream: Whether to stream back partial progress. If set, tokens will be sent as data-only server-side events as they become available, with the stream terminated by a data: [DONE] message. Otherwise, the server will hold the request open until the timeout or until completion, with the response containing the full result as JSON.
         :param stop: Stop generation if this token is detected. Or if one of these tokens is detected when providing an array
         :param random_seed: The seed to use for random sampling. If set, different calls will generate deterministic results.
+        :param metadata:
         :param suffix: Optional text/code that adds more context for the model. When given a `prompt` and a `suffix` the model will fill what is between them. When `suffix` is not provided, the model will simply execute completion starting with `prompt`.
         :param min_tokens: The minimum number of tokens to generate in the completion.
         :param retries: Override the default retry configuration for this method
@@ -72,6 +74,7 @@ class Fim(BaseSDK):
             stream=stream,
             stop=stop,
             random_seed=random_seed,
+            metadata=metadata,
             prompt=prompt,
             suffix=suffix,
             min_tokens=min_tokens,
@@ -152,6 +155,7 @@ class Fim(BaseSDK):
             ]
         ] = None,
         random_seed: OptionalNullable[int] = UNSET,
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         suffix: OptionalNullable[str] = UNSET,
         min_tokens: OptionalNullable[int] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -163,7 +167,7 @@ class Fim(BaseSDK):
 
         FIM completion.
 
-        :param model: ID of the model to use. Only compatible for now with:   - `codestral-2405`   - `codestral-latest`
+        :param model: ID of the model with FIM to use.
         :param prompt: The text/code to complete.
         :param temperature: What sampling temperature to use, we recommend between 0.0 and 0.7. Higher values like 0.7 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both. The default value varies depending on the model you are targeting. Call the `/models` endpoint to retrieve the appropriate value.
         :param top_p: Nucleus sampling, where the model considers the results of the tokens with `top_p` probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or `temperature` but not both.
@@ -171,6 +175,7 @@ class Fim(BaseSDK):
         :param stream: Whether to stream back partial progress. If set, tokens will be sent as data-only server-side events as they become available, with the stream terminated by a data: [DONE] message. Otherwise, the server will hold the request open until the timeout or until completion, with the response containing the full result as JSON.
         :param stop: Stop generation if this token is detected. Or if one of these tokens is detected when providing an array
         :param random_seed: The seed to use for random sampling. If set, different calls will generate deterministic results.
+        :param metadata:
         :param suffix: Optional text/code that adds more context for the model. When given a `prompt` and a `suffix` the model will fill what is between them. When `suffix` is not provided, the model will simply execute completion starting with `prompt`.
         :param min_tokens: The minimum number of tokens to generate in the completion.
         :param retries: Override the default retry configuration for this method
@@ -196,6 +201,7 @@ class Fim(BaseSDK):
             stream=stream,
             stop=stop,
             random_seed=random_seed,
+            metadata=metadata,
             prompt=prompt,
             suffix=suffix,
             min_tokens=min_tokens,
@@ -276,6 +282,7 @@ class Fim(BaseSDK):
             ]
         ] = None,
         random_seed: OptionalNullable[int] = UNSET,
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         suffix: OptionalNullable[str] = UNSET,
         min_tokens: OptionalNullable[int] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -287,7 +294,7 @@ class Fim(BaseSDK):
 
         Mistral AI provides the ability to stream responses back to a client in order to allow partial results for certain requests. Tokens will be sent as data-only server-sent events as they become available, with the stream terminated by a data: [DONE] message. Otherwise, the server will hold the request open until the timeout or until completion, with the response containing the full result as JSON.
 
-        :param model: ID of the model to use. Only compatible for now with:   - `codestral-2405`   - `codestral-latest`
+        :param model: ID of the model with FIM to use.
         :param prompt: The text/code to complete.
         :param temperature: What sampling temperature to use, we recommend between 0.0 and 0.7. Higher values like 0.7 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both. The default value varies depending on the model you are targeting. Call the `/models` endpoint to retrieve the appropriate value.
         :param top_p: Nucleus sampling, where the model considers the results of the tokens with `top_p` probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or `temperature` but not both.
@@ -295,6 +302,7 @@ class Fim(BaseSDK):
         :param stream:
         :param stop: Stop generation if this token is detected. Or if one of these tokens is detected when providing an array
         :param random_seed: The seed to use for random sampling. If set, different calls will generate deterministic results.
+        :param metadata:
         :param suffix: Optional text/code that adds more context for the model. When given a `prompt` and a `suffix` the model will fill what is between them. When `suffix` is not provided, the model will simply execute completion starting with `prompt`.
         :param min_tokens: The minimum number of tokens to generate in the completion.
         :param retries: Override the default retry configuration for this method
@@ -320,6 +328,7 @@ class Fim(BaseSDK):
             stream=stream,
             stop=stop,
             random_seed=random_seed,
+            metadata=metadata,
             prompt=prompt,
             suffix=suffix,
             min_tokens=min_tokens,
@@ -408,6 +417,7 @@ class Fim(BaseSDK):
             ]
         ] = None,
         random_seed: OptionalNullable[int] = UNSET,
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         suffix: OptionalNullable[str] = UNSET,
         min_tokens: OptionalNullable[int] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -419,7 +429,7 @@ class Fim(BaseSDK):
 
         Mistral AI provides the ability to stream responses back to a client in order to allow partial results for certain requests. Tokens will be sent as data-only server-sent events as they become available, with the stream terminated by a data: [DONE] message. Otherwise, the server will hold the request open until the timeout or until completion, with the response containing the full result as JSON.
 
-        :param model: ID of the model to use. Only compatible for now with:   - `codestral-2405`   - `codestral-latest`
+        :param model: ID of the model with FIM to use.
         :param prompt: The text/code to complete.
         :param temperature: What sampling temperature to use, we recommend between 0.0 and 0.7. Higher values like 0.7 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both. The default value varies depending on the model you are targeting. Call the `/models` endpoint to retrieve the appropriate value.
         :param top_p: Nucleus sampling, where the model considers the results of the tokens with `top_p` probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or `temperature` but not both.
@@ -427,6 +437,7 @@ class Fim(BaseSDK):
         :param stream:
         :param stop: Stop generation if this token is detected. Or if one of these tokens is detected when providing an array
         :param random_seed: The seed to use for random sampling. If set, different calls will generate deterministic results.
+        :param metadata:
         :param suffix: Optional text/code that adds more context for the model. When given a `prompt` and a `suffix` the model will fill what is between them. When `suffix` is not provided, the model will simply execute completion starting with `prompt`.
         :param min_tokens: The minimum number of tokens to generate in the completion.
         :param retries: Override the default retry configuration for this method
@@ -452,6 +463,7 @@ class Fim(BaseSDK):
             stream=stream,
             stop=stop,
             random_seed=random_seed,
+            metadata=metadata,
             prompt=prompt,
             suffix=suffix,
             min_tokens=min_tokens,