mistral.py•3.07 kB
from __future__ import annotations as _annotations
import os
from typing import overload
import httpx
from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import cached_async_http_client
from pydantic_ai.profiles.mistral import mistral_model_profile
from pydantic_ai.providers import Provider
try:
from mistralai import Mistral
except ImportError as e: # pragma: no cover
raise ImportError(
'Please install the `mistral` package to use the Mistral provider, '
'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
) from e
class MistralProvider(Provider[Mistral]):
"""Provider for Mistral API."""
@property
def name(self) -> str:
return 'mistral'
@property
def base_url(self) -> str:
return self.client.sdk_configuration.get_server_details()[0]
@property
def client(self) -> Mistral:
return self._client
def model_profile(self, model_name: str) -> ModelProfile | None:
return mistral_model_profile(model_name)
@overload
def __init__(self, *, mistral_client: Mistral | None = None) -> None: ...
@overload
def __init__(self, *, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None: ...
def __init__(
self,
*,
api_key: str | None = None,
mistral_client: Mistral | None = None,
base_url: str | None = None,
http_client: httpx.AsyncClient | None = None,
) -> None:
"""Create a new Mistral provider.
Args:
api_key: The API key to use for authentication, if not provided, the `MISTRAL_API_KEY` environment variable
will be used if available.
mistral_client: An existing `Mistral` client to use, if provided, `api_key` and `http_client` must be `None`.
base_url: The base url for the Mistral requests.
http_client: An existing async client to use for making HTTP requests.
"""
if mistral_client is not None:
assert http_client is None, 'Cannot provide both `mistral_client` and `http_client`'
assert api_key is None, 'Cannot provide both `mistral_client` and `api_key`'
assert base_url is None, 'Cannot provide both `mistral_client` and `base_url`'
self._client = mistral_client
else:
api_key = api_key or os.getenv('MISTRAL_API_KEY')
if not api_key:
raise UserError(
'Set the `MISTRAL_API_KEY` environment variable or pass it via `MistralProvider(api_key=...)`'
'to use the Mistral provider.'
)
elif http_client is not None:
self._client = Mistral(api_key=api_key, async_client=http_client, server_url=base_url)
else:
http_client = cached_async_http_client(provider='mistral')
self._client = Mistral(api_key=api_key, async_client=http_client, server_url=base_url)