mistralai 1.7.1__py3-none-any.whl → 1.8.1__py3-none-any.whl

mistralai/embeddings.py

@@ -18,6 +18,8 @@ class Embeddings(BaseSDK):
         inputs: Union[
             models.EmbeddingRequestInputs, models.EmbeddingRequestInputsTypedDict
         ],
+        output_dimension: OptionalNullable[int] = UNSET,
+        output_dtype: Optional[models.EmbeddingDtype] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -29,6 +31,8 @@ class Embeddings(BaseSDK):
 
         :param model: ID of the model to use.
         :param inputs: Text to embed.
+        :param output_dimension: The dimension of the output embeddings.
+        :param output_dtype:
         :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
@@ -47,6 +51,8 @@ class Embeddings(BaseSDK):
         request = models.EmbeddingRequest(
             model=model,
             inputs=inputs,
+            output_dimension=output_dimension,
+            output_dtype=output_dtype,
         )
 
         req = self._build_request(
@@ -125,6 +131,8 @@ class Embeddings(BaseSDK):
         inputs: Union[
             models.EmbeddingRequestInputs, models.EmbeddingRequestInputsTypedDict
         ],
+        output_dimension: OptionalNullable[int] = UNSET,
+        output_dtype: Optional[models.EmbeddingDtype] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -136,6 +144,8 @@ class Embeddings(BaseSDK):
 
         :param model: ID of the model to use.
         :param inputs: Text to embed.
+        :param output_dimension: The dimension of the output embeddings.
+        :param output_dtype:
         :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
@@ -154,6 +164,8 @@ class Embeddings(BaseSDK):
         request = models.EmbeddingRequest(
             model=model,
             inputs=inputs,
+            output_dimension=output_dimension,
+            output_dtype=output_dtype,
         )
 
         req = self._build_request_async(

mistralai/extra/__init__.py

@@ -1,5 +1,13 @@
-from .struct_chat import ParsedChatCompletionResponse, convert_to_parsed_chat_completion_response
+from .struct_chat import (
+    ParsedChatCompletionResponse,
+    convert_to_parsed_chat_completion_response,
+)
 from .utils import response_format_from_pydantic_model
 from .utils.response_format import CustomPydanticModel
 
-__all__ = ["convert_to_parsed_chat_completion_response", "response_format_from_pydantic_model", "CustomPydanticModel", "ParsedChatCompletionResponse"]
+__all__ = [
+    "convert_to_parsed_chat_completion_response",
+    "response_format_from_pydantic_model",
+    "CustomPydanticModel",
+    "ParsedChatCompletionResponse",
+]

mistralai/extra/exceptions.py

@@ -0,0 +1,14 @@
+class MistralClientException(Exception):
+    """Base exception for all the client errors."""
+
+
+class RunException(MistralClientException):
+    """Exception raised for errors during a conversation run."""
+
+
+class MCPException(MistralClientException):
+    """Exception raised for errors related to MCP operations."""
+
+
+class MCPAuthException(MCPException):
+    """Exception raised for authentication errors with an MCP server."""

mistralai/extra/mcp/auth.py

@@ -0,0 +1,166 @@
+from typing import Optional
+
+from authlib.oauth2.rfc8414 import AuthorizationServerMetadata
+from authlib.integrations.httpx_client import AsyncOAuth2Client as AsyncOAuth2ClientBase
+import httpx
+import logging
+
+from mistralai.types import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class Oauth2AuthorizationScheme(BaseModel):
+    """Information about the oauth flow to perform with the authorization server."""
+
+    authorization_url: str
+    token_url: str
+    scope: list[str]
+    description: Optional[str] = None
+    refresh_url: Optional[str] = None
+
+
+class OAuthParams(BaseModel):
+    """Required params for authorization."""
+
+    scheme: Oauth2AuthorizationScheme
+    client_id: str
+    client_secret: str
+
+
+class AsyncOAuth2Client(AsyncOAuth2ClientBase):
+    """Subclass of the Async httpx oauth client which provides a constructor from OAuthParams."""
+
+    @classmethod
+    def from_oauth_params(cls, oauth_params: OAuthParams) -> "AsyncOAuth2Client":
+        return cls(
+            client_id=oauth_params.client_id,
+            client_secret=oauth_params.client_secret,
+            scope=oauth_params.scheme.scope,
+        )
+
+
+async def get_well_known_authorization_server_metadata(
+    server_url: str,
+) -> Optional[AuthorizationServerMetadata]:
+    """Fetch the metadata from the well-known location.
+
+    This should be available on MCP servers as described by the specification:
+    https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization#2-3-server-metadata-discovery.
+    """
+    well_known_url = f"{server_url}/.well-known/oauth-authorization-server"
+    response = await httpx.AsyncClient().get(well_known_url)
+    if 200 <= response.status_code < 300:
+        try:
+            server_metadata = AuthorizationServerMetadata(**response.json())
+            server_metadata.validate()
+            return server_metadata
+        except ValueError:
+            logger.exception("Failed to parse oauth well-known metadata")
+            return None
+    else:
+        logger.error(f"Failed to get oauth well-known metadata from {server_url}")
+        return None
+
+
+async def get_oauth_server_metadata(server_url: str) -> AuthorizationServerMetadata:
+    """Fetch the metadata from the authorization server to perform the oauth flow."""
+    # 1) attempt to get the metadata from the resource server at /.well-known/oauth-protected-resource
+    # TODO: new self-discovery protocol, not released yet
+
+    # 2) attempt to get the metadata from the authorization server at /.well-known/oauth-authorization-server
+    metadata = await get_well_known_authorization_server_metadata(server_url=server_url)
+    if metadata is not None:
+        return metadata
+
+    # 3) fallback on default endpoints
+    # https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization#2-3-3-fallbacks-for-servers-without-metadata-discovery
+    return AuthorizationServerMetadata(
+        issuer=server_url,
+        authorization_endpoint=f"{server_url}/authorize",
+        token_endpoint=f"{server_url}/token",
+        register_endpoint=f"{server_url}/register",
+        response_types_supported=["code"],
+        response_modes_supported=["query"],
+        grant_types_supported=["authorization_code", "refresh_token"],
+        token_endpoint_auth_methods_supported=["client_secret_basic"],
+        code_challenge_methods_supported=["S256", "plain"],
+    )
+
+
+async def dynamic_client_registration(
+    register_endpoint: str,
+    redirect_url: str,
+    async_client: httpx.AsyncClient,
+) -> tuple[str, str]:
+    """Try to register the client dynamically with an MCP server.
+
+    Returns a client_id and client_secret.
+    """
+    # Construct the registration request payload
+    registration_payload = {
+        "client_name": "MistralSDKClient",
+        "grant_types": ["authorization_code", "refresh_token"],
+        "token_endpoint_auth_method": "client_secret_basic",
+        "response_types": ["code"],
+        "redirect_uris": [redirect_url],
+    }
+
+    # Make the registration request
+    response = await async_client.post(register_endpoint, json=registration_payload)
+    try:
+        response.raise_for_status()
+        registration_info = response.json()
+        client_id = registration_info["client_id"]
+        client_secret = registration_info["client_secret"]
+    except Exception as e:
+        raise ValueError(
+            f"Client registration failed: status={response.status_code}, error={response.text}"
+        ) from e
+    return client_id, client_secret
+
+
+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,
+) -> OAuthParams:
+    """Get issuer metadata and build the oauth required params."""
+    metadata = await get_oauth_server_metadata(server_url=server_url)
+    oauth_scheme = Oauth2AuthorizationScheme(
+        authorization_url=metadata.authorization_endpoint,
+        token_url=metadata.token_endpoint,
+        scope=scope or [],
+        refresh_url=metadata.token_endpoint
+        if "refresh_token" in metadata.grant_types_supported
+        else None,
+    )
+    if client_id and client_secret:
+        return OAuthParams(
+            client_id=client_id,
+            client_secret=client_secret,
+            scheme=oauth_scheme,
+        )
+
+    # Try to dynamically register the client
+    if async_client:
+        reg_client_id, reg_client_secret = await dynamic_client_registration(
+            register_endpoint=metadata.registration_endpoint,
+            redirect_url=redirect_url,
+            async_client=async_client,
+        )
+    else:
+        async with httpx.AsyncClient() as async_client:
+            reg_client_id, reg_client_secret = await dynamic_client_registration(
+                register_endpoint=metadata.registration_endpoint,
+                redirect_url=redirect_url,
+                async_client=async_client,
+            )
+    return OAuthParams(
+        client_id=reg_client_id,
+        client_secret=reg_client_secret,
+        scheme=oauth_scheme,
+    )

mistralai/extra/mcp/base.py

@@ -0,0 +1,155 @@
+from typing import Optional, Union
+import logging
+import typing
+from contextlib import AsyncExitStack
+from typing import Protocol, Any
+
+from mcp import ClientSession
+from mcp.types import ListPromptsResult, EmbeddedResource, ImageContent, TextContent
+
+from mistralai.extra.exceptions import MCPException
+from mistralai.models import (
+    FunctionTool,
+    Function,
+    SystemMessageTypedDict,
+    AssistantMessageTypedDict,
+    TextChunkTypedDict,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MCPSystemPrompt(typing.TypedDict):
+    description: Optional[str]
+    messages: list[Union[SystemMessageTypedDict, AssistantMessageTypedDict]]
+
+
+class MCPClientProtocol(Protocol):
+    """MCP client that converts MCP artifacts to Mistral format."""
+
+    _name: str
+
+    async def initialize(self, exit_stack: Optional[AsyncExitStack]) -> None:
+        ...
+
+    async def aclose(self) -> None:
+        ...
+
+    async def get_tools(self) -> list[FunctionTool]:
+        ...
+
+    async def execute_tool(
+        self, name: str, arguments: dict
+    ) -> list[TextChunkTypedDict]:
+        ...
+
+    async def get_system_prompt(
+        self, name: str, arguments: dict[str, Any]
+    ) -> MCPSystemPrompt:
+        ...
+
+    async def list_system_prompts(self) -> ListPromptsResult:
+        ...
+
+
+class MCPClientBase(MCPClientProtocol):
+    """Base class to implement functionalities from an initialized MCP session."""
+
+    _session: ClientSession
+
+    def __init__(self, name: Optional[str] = None):
+        self._name = name or self.__class__.__name__
+        self._exit_stack: Optional[AsyncExitStack] = None
+        self._is_initialized = False
+
+    def _convert_content(
+        self, mcp_content: Union[TextContent, ImageContent, EmbeddedResource]
+    ) -> 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]]
+    ) -> list[TextChunkTypedDict]:
+        content_chunks = []
+        for mcp_content in mcp_contents:
+            content_chunks.append(self._convert_content(mcp_content))
+        return content_chunks
+
+    async def get_tools(self) -> list[FunctionTool]:
+        mcp_tools = await self._session.list_tools()
+        tools = []
+        for mcp_tool in mcp_tools.tools:
+            tools.append(
+                FunctionTool(
+                    type="function",
+                    function=Function(
+                        name=mcp_tool.name,
+                        description=mcp_tool.description,
+                        parameters=mcp_tool.inputSchema,
+                        strict=True,
+                    ),
+                )
+            )
+        return tools
+
+    async def execute_tool(
+        self, name: str, arguments: dict[str, Any]
+    ) -> list[TextChunkTypedDict]:
+        contents = await self._session.call_tool(name=name, arguments=arguments)
+        return self._convert_content_list(contents.content)
+
+    async def get_system_prompt(
+        self, name: str, arguments: dict[str, Any]
+    ) -> MCPSystemPrompt:
+        prompt_result = await self._session.get_prompt(name=name, arguments=arguments)
+        return {
+            "description": prompt_result.description,
+            "messages": [
+                typing.cast(
+                    Union[SystemMessageTypedDict, AssistantMessageTypedDict],
+                    {
+                        "role": message.role,
+                        "content": self._convert_content(mcp_content=message.content),
+                    },
+                )
+                for message in prompt_result.messages
+            ],
+        }
+
+    async def list_system_prompts(self) -> ListPromptsResult:
+        return await self._session.list_prompts()
+
+    async def initialize(self, exit_stack: Optional[AsyncExitStack] = None) -> None:
+        """Initialize the MCP session."""
+        # client is already initialized so return
+        if self._is_initialized:
+            return
+        if exit_stack is None:
+            self._exit_stack = AsyncExitStack()
+            exit_stack = self._exit_stack
+        stdio_transport = await self._get_transport(exit_stack=exit_stack)
+        mcp_session = await exit_stack.enter_async_context(
+            ClientSession(
+                read_stream=stdio_transport[0],
+                write_stream=stdio_transport[1],
+            )
+        )
+        await mcp_session.initialize()
+        self._session = mcp_session
+        self._is_initialized = True
+
+    async def aclose(self):
+        """Close the MCP session."""
+        if self._exit_stack:
+            await self._exit_stack.aclose()
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} name={self._name!r} id=0x{id(self):x}>"
+
+    def __str__(self):
+        return f"{self.__class__.__name__}(name={self._name})"
+
+    async def _get_transport(self, exit_stack: AsyncExitStack):
+        raise NotImplementedError

mistralai/extra/mcp/sse.py

@@ -0,0 +1,165 @@
+import http
+import logging
+import typing
+from typing import Any, Optional
+from contextlib import AsyncExitStack
+from functools import cached_property
+
+import httpx
+
+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
+
+logger = logging.getLogger(__name__)
+
+
+class SSEServerParams(BaseModel):
+    """Parameters required for a MCPClient with SSE transport"""
+
+    url: str
+    headers: Optional[dict[str, Any]] = None
+    timeout: float = 5
+    sse_read_timeout: float = 60 * 5
+
+
+class MCPClientSSE(MCPClientBase):
+    """MCP client that uses sse for communication.
+
+    The client provides authentication for OAuth2 protocol following the current MCP authorization spec:
+    https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization.
+
+    This is possibly going to change in the future since the protocol has ongoing discussions.
+    """
+
+    _oauth_params: Optional[OAuthParams]
+    _sse_params: SSEServerParams
+
+    def __init__(
+        self,
+        sse_params: SSEServerParams,
+        name: Optional[str] = None,
+        oauth_params: Optional[OAuthParams] = None,
+        auth_token: Optional[OAuth2Token] = None,
+    ):
+        super().__init__(name=name)
+        self._sse_params = sse_params
+        self._oauth_params: Optional[OAuthParams] = oauth_params
+        self._auth_token: Optional[OAuth2Token] = auth_token
+
+    @cached_property
+    def base_url(self) -> str:
+        return self._sse_params.url.rstrip("/sse")
+
+    def set_oauth_params(self, oauth_params: OAuthParams):
+        """Update the oauth params and client accordingly."""
+        if self._oauth_params is not None:
+            logger.warning(f"Overriding current oauth params for {self._name}")
+        self._oauth_params = oauth_params
+
+    async def get_auth_url_and_state(self, redirect_url: str) -> tuple[str, str]:
+        """Create the authorization url for client to start oauth flow."""
+        if self._oauth_params is None:
+            raise MCPAuthException(
+                "Can't generate an authorization url without oauth_params being set, "
+                "make sure the oauth params have been set."
+            )
+        oauth_client = AsyncOAuth2Client.from_oauth_params(self._oauth_params)
+        auth_url, state = oauth_client.create_authorization_url(
+            self._oauth_params.scheme.authorization_url, redirect_uri=redirect_url
+        )
+        return auth_url, state
+
+    async def get_token_from_auth_response(
+        self,
+        authorization_response: str,
+        redirect_url: str,
+        state: str,
+    ) -> OAuth2Token:
+        """Fetch the authentication token from the server."""
+        if self._oauth_params is None:
+            raise MCPAuthException(
+                "Can't fetch a token without oauth_params, make sure they have been set."
+            )
+        oauth_client = AsyncOAuth2Client.from_oauth_params(self._oauth_params)
+        oauth_token = await oauth_client.fetch_token(
+            url=self._oauth_params.scheme.token_url,
+            authorization_response=authorization_response,
+            redirect_uri=redirect_url,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            state=state,
+        )
+        return oauth_token
+
+    async def refresh_auth_token(self):
+        """Refresh an expired token."""
+        if self._oauth_params is None or self._oauth_params.scheme.refresh_url is None:
+            raise MCPAuthException(
+                "Can't refresh a token without a refresh url make sure the oauth params have been set."
+            )
+        if self._auth_token is None:
+            raise MCPAuthException(
+                "Can't refresh a token without a refresh token, use the `set_auth_token` to add a OAuth2Token."
+            )
+        oauth_client = AsyncOAuth2Client.from_oauth_params(self._oauth_params)
+        oauth_token = await oauth_client.refresh_token(
+            url=self._oauth_params.scheme.refresh_url,
+            refresh_token=self._auth_token["refresh_token"],
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+        )
+        self.set_auth_token(oauth_token)
+
+    def set_auth_token(self, token: OAuth2Token) -> None:
+        """Register the authentication token with this client."""
+        self._auth_token = token
+
+    def _format_headers(self) -> dict[str, str]:
+        headers: dict[str, str] = {}
+        if self._sse_params.headers:
+            headers |= self._sse_params.headers
+        if self._auth_token:
+            headers["Authorization"] = f"Bearer {self._auth_token['access_token']}"
+        return headers
+
+    async def requires_auth(self) -> bool:
+        """Check if the client requires authentication to communicate with the server."""
+        response = httpx.get(
+            self._sse_params.url,
+            headers=self._format_headers(),
+            timeout=self._sse_params.timeout,
+        )
+        return response.status_code == http.HTTPStatus.UNAUTHORIZED
+
+    async def _get_transport(
+        self, exit_stack: AsyncExitStack
+    ) -> tuple[
+        MemoryObjectReceiveStream[typing.Union[SessionMessage, Exception]],
+        MemoryObjectSendStream[SessionMessage],
+    ]:
+        try:
+            return await exit_stack.enter_async_context(
+                sse_client(
+                    url=self._sse_params.url,
+                    headers=self._format_headers(),
+                    timeout=self._sse_params.timeout,
+                    sse_read_timeout=self._sse_params.sse_read_timeout,
+                )
+            )
+        except Exception as e:
+            if isinstance(e, httpx.HTTPStatusError):
+                if e.response.status_code == http.HTTPStatus.UNAUTHORIZED:
+                    if self._oauth_params is None:
+                        raise MCPAuthException(
+                            "Authentication required but no auth params provided."
+                        ) from e
+                    raise MCPAuthException("Authentication required.") from e
+            raise

mistralai/extra/mcp/stdio.py

@@ -0,0 +1,22 @@
+from typing import Optional
+import logging
+from contextlib import AsyncExitStack
+
+from mistralai.extra.mcp.base import (
+    MCPClientBase,
+)
+
+from mcp import stdio_client, StdioServerParameters
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClientSTDIO(MCPClientBase):
+    """MCP client that uses stdio for communication."""
+
+    def __init__(self, stdio_params: StdioServerParameters, name: Optional[str] = None):
+        super().__init__(name=name)
+        self._stdio_params = stdio_params
+
+    async def _get_transport(self, exit_stack: AsyncExitStack):
+        return await exit_stack.enter_async_context(stdio_client(self._stdio_params))