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

mistralai/documents.py

@@ -6,7 +6,7 @@ from mistralai._hooks import HookContext
 from mistralai.types import OptionalNullable, UNSET
 from mistralai.utils import 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 Documents(BaseSDK):
@@ -19,6 +19,7 @@ class Documents(BaseSDK):
         search: OptionalNullable[str] = UNSET,
         page_size: Optional[int] = 100,
         page: Optional[int] = 0,
+        filters_attributes: OptionalNullable[str] = UNSET,
         sort_by: Optional[str] = "created_at",
         sort_order: Optional[str] = "desc",
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -26,7 +27,7 @@ class Documents(BaseSDK):
         timeout_ms: Optional[int] = None,
         http_headers: Optional[Mapping[str, str]] = None,
     ) -> models.ListDocumentOut:
-        r"""List document in a given library.
+        r"""List documents in a given library.
 
         Given a library, lists the document that have been uploaded to that library.
 
@@ -34,6 +35,7 @@ class Documents(BaseSDK):
         :param search:
         :param page_size:
         :param page:
+        :param filters_attributes:
         :param sort_by:
         :param sort_order:
         :param retries: Override the default retry configuration for this method
@@ -56,6 +58,7 @@ class Documents(BaseSDK):
             search=search,
             page_size=page_size,
             page=page,
+            filters_attributes=filters_attributes,
             sort_by=sort_by,
             sort_order=sort_order,
         )
@@ -123,6 +126,7 @@ class Documents(BaseSDK):
         search: OptionalNullable[str] = UNSET,
         page_size: Optional[int] = 100,
         page: Optional[int] = 0,
+        filters_attributes: OptionalNullable[str] = UNSET,
         sort_by: Optional[str] = "created_at",
         sort_order: Optional[str] = "desc",
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
@@ -130,7 +134,7 @@ class Documents(BaseSDK):
         timeout_ms: Optional[int] = None,
         http_headers: Optional[Mapping[str, str]] = None,
     ) -> models.ListDocumentOut:
-        r"""List document in a given library.
+        r"""List documents in a given library.
 
         Given a library, lists the document that have been uploaded to that library.
 
@@ -138,6 +142,7 @@ class Documents(BaseSDK):
         :param search:
         :param page_size:
         :param page:
+        :param filters_attributes:
         :param sort_by:
         :param sort_order:
         :param retries: Override the default retry configuration for this method
@@ -160,6 +165,7 @@ class Documents(BaseSDK):
             search=search,
             page_size=page_size,
             page=page,
+            filters_attributes=filters_attributes,
             sort_by=sort_by,
             sort_order=sort_order,
         )
@@ -612,6 +618,9 @@ class Documents(BaseSDK):
         library_id: str,
         document_id: str,
         name: OptionalNullable[str] = UNSET,
+        attributes: OptionalNullable[
+            Union[Dict[str, models.Attributes], Dict[str, models.AttributesTypedDict]]
+        ] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -624,6 +633,7 @@ class Documents(BaseSDK):
         :param library_id:
         :param document_id:
         :param name:
+        :param attributes:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -644,6 +654,7 @@ class Documents(BaseSDK):
             document_id=document_id,
             document_update_in=models.DocumentUpdateIn(
                 name=name,
+                attributes=attributes,
             ),
         )
 
@@ -716,6 +727,9 @@ class Documents(BaseSDK):
         library_id: str,
         document_id: str,
         name: OptionalNullable[str] = UNSET,
+        attributes: OptionalNullable[
+            Union[Dict[str, models.Attributes], Dict[str, models.AttributesTypedDict]]
+        ] = UNSET,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -728,6 +742,7 @@ class Documents(BaseSDK):
         :param library_id:
         :param document_id:
         :param name:
+        :param attributes:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -748,6 +763,7 @@ class Documents(BaseSDK):
             document_id=document_id,
             document_update_in=models.DocumentUpdateIn(
                 name=name,
+                attributes=attributes,
             ),
         )
 

mistralai/embeddings.py

@@ -6,7 +6,7 @@ from mistralai._hooks import HookContext
 from mistralai.types import OptionalNullable, UNSET
 from mistralai.utils import 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 Embeddings(BaseSDK):
@@ -19,6 +19,7 @@ class Embeddings(BaseSDK):
         inputs: Union[
             models.EmbeddingRequestInputs, models.EmbeddingRequestInputsTypedDict
         ],
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         output_dimension: OptionalNullable[int] = UNSET,
         output_dtype: Optional[models.EmbeddingDtype] = None,
         encoding_format: Optional[models.EncodingFormat] = None,
@@ -31,9 +32,10 @@ class Embeddings(BaseSDK):
 
         Embeddings
 
-        :param model: ID of the model to use.
-        :param inputs: Text to embed.
-        :param output_dimension: The dimension of the output embeddings.
+        :param model: The ID of the model to be used for embedding.
+        :param inputs: The text content to be embedded, can be a string or an array of strings for fast processing in bulk.
+        :param metadata:
+        :param output_dimension: The dimension of the output embeddings when feature available. If not provided, a default output dimension will be used.
         :param output_dtype:
         :param encoding_format:
         :param retries: Override the default retry configuration for this method
@@ -53,6 +55,7 @@ class Embeddings(BaseSDK):
 
         request = models.EmbeddingRequest(
             model=model,
+            metadata=metadata,
             inputs=inputs,
             output_dimension=output_dimension,
             output_dtype=output_dtype,
@@ -125,6 +128,7 @@ class Embeddings(BaseSDK):
         inputs: Union[
             models.EmbeddingRequestInputs, models.EmbeddingRequestInputsTypedDict
         ],
+        metadata: OptionalNullable[Dict[str, Any]] = UNSET,
         output_dimension: OptionalNullable[int] = UNSET,
         output_dtype: Optional[models.EmbeddingDtype] = None,
         encoding_format: Optional[models.EncodingFormat] = None,
@@ -137,9 +141,10 @@ class Embeddings(BaseSDK):
 
         Embeddings
 
-        :param model: ID of the model to use.
-        :param inputs: Text to embed.
-        :param output_dimension: The dimension of the output embeddings.
+        :param model: The ID of the model to be used for embedding.
+        :param inputs: The text content to be embedded, can be a string or an array of strings for fast processing in bulk.
+        :param metadata:
+        :param output_dimension: The dimension of the output embeddings when feature available. If not provided, a default output dimension will be used.
         :param output_dtype:
         :param encoding_format:
         :param retries: Override the default retry configuration for this method
@@ -159,6 +164,7 @@ class Embeddings(BaseSDK):
 
         request = models.EmbeddingRequest(
             model=model,
+            metadata=metadata,
             inputs=inputs,
             output_dimension=output_dimension,
             output_dtype=output_dtype,

mistralai/extra/README.md

@@ -34,7 +34,7 @@ class Chat(BaseSDK):
 
 3. Now build the SDK with the custom code:
 ```bash
-rm -rf dist; poetry build; python3 -m pip install ~/client-python/dist/mistralai-1.4.1-py3-none-any.whl --force-reinstall
+rm -rf dist; uv build; uv pip install --reinstall ~/client-python/dist/mistralai-1.4.1-py3-none-any.whl
 ```
 
 4. And now you should be able to call the custom method:

mistralai/extra/mcp/auth.py

@@ -1,9 +1,8 @@
-from typing import Optional
+import logging
 
-from authlib.oauth2.rfc8414 import AuthorizationServerMetadata
-from authlib.integrations.httpx_client import AsyncOAuth2Client as AsyncOAuth2ClientBase
 import httpx
-import logging
+from authlib.integrations.httpx_client import AsyncOAuth2Client as AsyncOAuth2ClientBase
+from authlib.oauth2.rfc8414 import AuthorizationServerMetadata
 
 from mistralai.types import BaseModel
 
@@ -16,8 +15,8 @@ class Oauth2AuthorizationScheme(BaseModel):
     authorization_url: str
     token_url: str
     scope: list[str]
-    description: Optional[str] = None
-    refresh_url: Optional[str] = None
+    description: str | None = None
+    refresh_url: str | None = None
 
 
 class OAuthParams(BaseModel):
@@ -42,7 +41,7 @@ class AsyncOAuth2Client(AsyncOAuth2ClientBase):
 
 async def get_well_known_authorization_server_metadata(
     server_url: str,
-) -> Optional[AuthorizationServerMetadata]:
+) -> AuthorizationServerMetadata | None:
     """Fetch the metadata from the well-known location.
 
     This should be available on MCP servers as described by the specification:
@@ -123,10 +122,10 @@ async def dynamic_client_registration(
 async def build_oauth_params(
     server_url: str,
     redirect_url: str,
-    client_id: Optional[str] = None,
-    client_secret: Optional[str] = None,
-    scope: Optional[list[str]] = None,
-    async_client: Optional[httpx.AsyncClient] = None,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    scope: list[str] | None = None,
+    async_client: httpx.AsyncClient | None = None,
 ) -> OAuthParams:
     """Get issuer metadata and build the oauth required params."""
     metadata = await get_oauth_server_metadata(server_url=server_url)

mistralai/extra/mcp/base.py

@@ -1,11 +1,14 @@
-from typing import Optional, Union
 import logging
 import typing
+from collections.abc import Sequence
 from contextlib import AsyncExitStack
-from typing import Protocol, Any
+from typing import Any, Protocol
 
-from mcp import ClientSession
-from mcp.types import ListPromptsResult, EmbeddedResource, ImageContent, TextContent
+from mcp import ClientSession  # pyright: ignore[reportMissingImports]
+from mcp.types import (  # pyright: ignore[reportMissingImports]
+    ContentBlock,
+    ListPromptsResult,
+)
 
 from mistralai.extra.exceptions import MCPException
 from mistralai.models import (
@@ -20,8 +23,8 @@ logger = logging.getLogger(__name__)
 
 
 class MCPSystemPrompt(typing.TypedDict):
-    description: Optional[str]
-    messages: list[Union[SystemMessageTypedDict, AssistantMessageTypedDict]]
+    description: str | None
+    messages: list[SystemMessageTypedDict | AssistantMessageTypedDict]
 
 
 class MCPClientProtocol(Protocol):
@@ -29,7 +32,7 @@ class MCPClientProtocol(Protocol):
 
     _name: str
 
-    async def initialize(self, exit_stack: Optional[AsyncExitStack]) -> None:
+    async def initialize(self, exit_stack: AsyncExitStack | None) -> None:
         ...
 
     async def aclose(self) -> None:
@@ -39,7 +42,7 @@ class MCPClientProtocol(Protocol):
         ...
 
     async def execute_tool(
-        self, name: str, arguments: dict
+        self, name: str, arguments: dict[str, Any]
     ) -> list[TextChunkTypedDict]:
         ...
 
@@ -57,20 +60,18 @@ class MCPClientBase(MCPClientProtocol):
 
     _session: ClientSession
 
-    def __init__(self, name: Optional[str] = None):
+    def __init__(self, name: str | None = None):
         self._name = name or self.__class__.__name__
-        self._exit_stack: Optional[AsyncExitStack] = None
+        self._exit_stack: AsyncExitStack | None = None
         self._is_initialized = False
 
-    def _convert_content(
-        self, mcp_content: Union[TextContent, ImageContent, EmbeddedResource]
-    ) -> TextChunkTypedDict:
+    def _convert_content(self, mcp_content: ContentBlock) -> TextChunkTypedDict:
         if not mcp_content.type == "text":
             raise MCPException("Only supporting text tool responses for now.")
         return {"type": "text", "text": mcp_content.text}
 
     def _convert_content_list(
-        self, mcp_contents: list[Union[TextContent, ImageContent, EmbeddedResource]]
+        self, mcp_contents: Sequence[ContentBlock]
     ) -> list[TextChunkTypedDict]:
         content_chunks = []
         for mcp_content in mcp_contents:
@@ -108,7 +109,7 @@ class MCPClientBase(MCPClientProtocol):
             "description": prompt_result.description,
             "messages": [
                 typing.cast(
-                    Union[SystemMessageTypedDict, AssistantMessageTypedDict],
+                    SystemMessageTypedDict | AssistantMessageTypedDict,
                     {
                         "role": message.role,
                         "content": self._convert_content(mcp_content=message.content),
@@ -121,7 +122,7 @@ class MCPClientBase(MCPClientProtocol):
     async def list_system_prompts(self) -> ListPromptsResult:
         return await self._session.list_prompts()
 
-    async def initialize(self, exit_stack: Optional[AsyncExitStack] = None) -> None:
+    async def initialize(self, exit_stack: AsyncExitStack | None = None) -> None:
         """Initialize the MCP session."""
         # client is already initialized so return
         if self._is_initialized:

mistralai/extra/mcp/sse.py

@@ -1,22 +1,20 @@
 import http
 import logging
-import typing
-from typing import Any, Optional
 from contextlib import AsyncExitStack
 from functools import cached_property
+from typing import Any
 
 import httpx
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from authlib.oauth2.rfc6749 import OAuth2Token
+from mcp.client.sse import sse_client  # pyright: ignore[reportMissingImports]
+from mcp.shared.message import SessionMessage  # pyright: ignore[reportMissingImports]
 
 from mistralai.extra.exceptions import MCPAuthException
 from mistralai.extra.mcp.base import (
     MCPClientBase,
 )
 from mistralai.extra.mcp.auth import OAuthParams, AsyncOAuth2Client
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-
-from mcp.client.sse import sse_client
-from mcp.shared.message import SessionMessage
-from authlib.oauth2.rfc6749 import OAuth2Token
 
 from mistralai.types import BaseModel
 
@@ -27,7 +25,7 @@ class SSEServerParams(BaseModel):
     """Parameters required for a MCPClient with SSE transport"""
 
     url: str
-    headers: Optional[dict[str, Any]] = None
+    headers: dict[str, Any] | None = None
     timeout: float = 5
     sse_read_timeout: float = 60 * 5
 
@@ -41,20 +39,20 @@ class MCPClientSSE(MCPClientBase):
     This is possibly going to change in the future since the protocol has ongoing discussions.
     """
 
-    _oauth_params: Optional[OAuthParams]
+    _oauth_params: OAuthParams | None
     _sse_params: SSEServerParams
 
     def __init__(
         self,
         sse_params: SSEServerParams,
-        name: Optional[str] = None,
-        oauth_params: Optional[OAuthParams] = None,
-        auth_token: Optional[OAuth2Token] = None,
+        name: str | None = None,
+        oauth_params: OAuthParams | None = None,
+        auth_token: OAuth2Token | None = None,
     ):
         super().__init__(name=name)
         self._sse_params = sse_params
-        self._oauth_params: Optional[OAuthParams] = oauth_params
-        self._auth_token: Optional[OAuth2Token] = auth_token
+        self._oauth_params: OAuthParams | None = oauth_params
+        self._auth_token: OAuth2Token | None = auth_token
 
     @cached_property
     def base_url(self) -> str:
@@ -142,7 +140,7 @@ class MCPClientSSE(MCPClientBase):
     async def _get_transport(
         self, exit_stack: AsyncExitStack
     ) -> tuple[
-        MemoryObjectReceiveStream[typing.Union[SessionMessage, Exception]],
+        MemoryObjectReceiveStream[SessionMessage | Exception],
         MemoryObjectSendStream[SessionMessage],
     ]:
         try:

mistralai/extra/mcp/stdio.py

@@ -1,12 +1,9 @@
-from typing import Optional
 import logging
 from contextlib import AsyncExitStack
 
-from mistralai.extra.mcp.base import (
-    MCPClientBase,
-)
+from mcp import StdioServerParameters, stdio_client  # pyright: ignore[reportMissingImports]
 
-from mcp import stdio_client, StdioServerParameters
+from mistralai.extra.mcp.base import MCPClientBase
 
 logger = logging.getLogger(__name__)
 
@@ -14,7 +11,9 @@ logger = logging.getLogger(__name__)
 class MCPClientSTDIO(MCPClientBase):
     """MCP client that uses stdio for communication."""
 
-    def __init__(self, stdio_params: StdioServerParameters, name: Optional[str] = None):
+    def __init__(
+        self, stdio_params: StdioServerParameters, name: str | None = None
+    ):
         super().__init__(name=name)
         self._stdio_params = stdio_params
 

mistralai/extra/observability/__init__.py

@@ -0,0 +1,15 @@
+from contextlib import contextmanager
+
+from opentelemetry import trace as otel_trace
+
+from .otel import MISTRAL_SDK_OTEL_TRACER_NAME
+
+
+@contextmanager
+def trace(name: str, **kwargs):
+    tracer = otel_trace.get_tracer(MISTRAL_SDK_OTEL_TRACER_NAME)
+    with tracer.start_as_current_span(name, **kwargs) as span:
+        yield span
+
+
+__all__ = ["trace"]