mistralai 0.0.7__tar.gz → 0.0.9__tar.gz

{mistralai-0.0.7 → mistralai-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mistralai
-Version: 0.0.7
+Version: 0.0.9
 Summary: 
 Author: Bam4d
 Author-email: bam4d@mistral.ai
@@ -11,17 +11,25 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: aiohttp (>=3.9.1,<4.0.0)
-Requires-Dist: backoff (>=2.2.1,<3.0.0)
+Requires-Dist: httpx (>=0.25.2,<0.26.0)
 Requires-Dist: orjson (>=3.9.10,<4.0.0)
 Requires-Dist: pydantic (>=2.5.2,<3.0.0)
-Requires-Dist: requests (>=2.31.0,<3.0.0)
 Description-Content-Type: text/markdown
 
 # Mistral Python Client
 
+This client is inspired from [cohere-python](https://github.com/cohere-ai/cohere-python)
+
 You can use the Mistral Python client to interact with the Mistral AI API.
 
+## Installing
+
+```bash
+pip install mistralai
+```
+
+### From Source
+
 This client uses `poetry` as a dependency and virtual environment manager.
 
 You can install poetry with
@@ -30,8 +38,6 @@ You can install poetry with
 pip install poetry
 ```
 
-## Installing
-
 `poetry` will set up a virtual environment and install dependencies with the following command:
 
 ```bash
@@ -42,6 +48,21 @@ poetry install
 
 You can run the examples in the `examples/` directory using `poetry run` or by entering the virtual environment using `poetry shell`.
 
+### API Key Setup
+
+Running the examples requires a Mistral AI API key.
+
+1. Get your own Mistral API Key: <https://docs.mistral.ai/#api-access>
+2. Set your Mistral API Key as an environment variable. You only need to do this once.
+
+```bash
+# set Mistral API Key (using zsh for example)
+$ echo 'export MISTRAL_API_KEY=[your_key_here]' >> ~/.zshenv
+
+# reload the environment (or just quit and open a new terminal)
+$ source ~/.zshenv
+```
+
 ### Using poetry run
 
 ```bash

{mistralai-0.0.7 → mistralai-0.0.9}/README.md

@@ -1,7 +1,17 @@
 # Mistral Python Client
 
+This client is inspired from [cohere-python](https://github.com/cohere-ai/cohere-python)
+
 You can use the Mistral Python client to interact with the Mistral AI API.
 
+## Installing
+
+```bash
+pip install mistralai
+```
+
+### From Source
+
 This client uses `poetry` as a dependency and virtual environment manager.
 
 You can install poetry with
@@ -10,8 +20,6 @@ You can install poetry with
 pip install poetry
 ```
 
-## Installing
-
 `poetry` will set up a virtual environment and install dependencies with the following command:
 
 ```bash
@@ -22,6 +30,21 @@ poetry install
 
 You can run the examples in the `examples/` directory using `poetry run` or by entering the virtual environment using `poetry shell`.
 
+### API Key Setup
+
+Running the examples requires a Mistral AI API key.
+
+1. Get your own Mistral API Key: <https://docs.mistral.ai/#api-access>
+2. Set your Mistral API Key as an environment variable. You only need to do this once.
+
+```bash
+# set Mistral API Key (using zsh for example)
+$ echo 'export MISTRAL_API_KEY=[your_key_here]' >> ~/.zshenv
+
+# reload the environment (or just quit and open a new terminal)
+$ source ~/.zshenv
+```
+
 ### Using poetry run
 
 ```bash

{mistralai-0.0.7 → mistralai-0.0.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mistralai"
-version = "0.0.7"
+version = "0.0.9"
 description = ""
 authors = ["Bam4d <bam4d@mistral.ai>"]
 readme = "README.md"
@@ -24,11 +24,9 @@ exclude = ["docs", "tests", "examples", "tools", "build"]
 
 [tool.poetry.dependencies]
 python = "^3.8"
-aiohttp = "^3.9.1"
-backoff = "^2.2.1"
 orjson = "^3.9.10"
-requests = "^2.31.0"
 pydantic = "^2.5.2"
+httpx = "^0.25.2"
 
 
 [tool.poetry.group.dev.dependencies]

mistralai-0.0.9/src/mistralai/async_client.py

@@ -0,0 +1,245 @@
+import os
+import posixpath
+import time
+from json import JSONDecodeError
+from typing import Any, AsyncGenerator, Dict, List, Optional, Union
+
+from httpx import (
+    AsyncClient,
+    AsyncHTTPTransport,
+    ConnectError,
+    Limits,
+    RequestError,
+    Response,
+)
+
+from mistralai.client_base import ClientBase
+from mistralai.constants import ENDPOINT
+from mistralai.exceptions import (
+    MistralAPIException,
+    MistralAPIStatusException,
+    MistralConnectionException,
+    MistralException,
+)
+from mistralai.models.chat_completion import (
+    ChatCompletionResponse,
+    ChatCompletionStreamResponse,
+    ChatMessage,
+)
+from mistralai.models.embeddings import EmbeddingResponse
+from mistralai.models.models import ModelList
+
+
+class MistralAsyncClient(ClientBase):
+    def __init__(
+        self,
+        api_key: Optional[str] = os.environ.get("MISTRAL_API_KEY", None),
+        endpoint: str = ENDPOINT,
+        max_retries: int = 5,
+        timeout: int = 120,
+        max_concurrent_requests: int = 64,
+    ):
+        super().__init__(endpoint, api_key, max_retries, timeout)
+
+        self._client = AsyncClient(
+            follow_redirects=True,
+            timeout=timeout,
+            limits=Limits(max_connections=max_concurrent_requests),
+            transport=AsyncHTTPTransport(retries=max_retries),
+        )
+
+    async def close(self) -> None:
+        await self._client.aclose()
+
+    async def _request(
+        self,
+        method: str,
+        json: Dict[str, Any],
+        path: str,
+        stream: bool = False,
+        attempt: int = 1,
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        headers = {
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+        }
+
+        url = posixpath.join(self._endpoint, path)
+
+        self._logger.debug(f"Sending request: {method} {url} {json}")
+
+        response: Response
+
+        try:
+            if stream:
+                async with self._client.stream(
+                    method,
+                    url,
+                    headers=headers,
+                    json=json,
+                ) as response:
+                    self._check_streaming_response(response)
+
+                    async for line in response.aiter_lines():
+                        json_streamed_response = self._process_line(line)
+                        if json_streamed_response:
+                            yield json_streamed_response
+
+            else:
+                response = await self._client.request(
+                    method,
+                    url,
+                    headers=headers,
+                    json=json,
+                )
+
+                yield self._check_response(response)
+
+        except ConnectError as e:
+            raise MistralConnectionException(str(e)) from e
+        except RequestError as e:
+            raise MistralException(
+                f"Unexpected exception ({e.__class__.__name__}): {e}"
+            ) from e
+        except JSONDecodeError as e:
+            raise MistralAPIException.from_response(
+                response,
+                message=f"Failed to decode json body: {response.text}",
+            ) from e
+        except MistralAPIStatusException as e:
+            attempt += 1
+            if attempt > self._max_retries:
+                raise MistralAPIStatusException.from_response(
+                    response, message=str(e)
+                ) from e
+            backoff = 2.0**attempt  # exponential backoff
+            time.sleep(backoff)
+
+            # Retry as a generator
+            async for r in self._request(
+                method, json, path, stream=stream, attempt=attempt
+            ):
+                yield r
+
+    async def chat(
+        self,
+        model: str,
+        messages: List[ChatMessage],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_p: Optional[float] = None,
+        random_seed: Optional[int] = None,
+        safe_mode: bool = False,
+    ) -> ChatCompletionResponse:
+        """A asynchronous chat endpoint that returns a single response.
+
+        Args:
+            model (str): model the name of the model to chat with, e.g. mistral-tiny
+            messages (List[ChatMessage]): messages an array of messages to chat with, e.g.
+                [{role: 'user', content: 'What is the best French cheese?'}]
+            temperature (Optional[float], optional): temperature the temperature to use for sampling, e.g. 0.5.
+            max_tokens (Optional[int], optional): the maximum number of tokens to generate, e.g. 100. Defaults to None.
+            top_p (Optional[float], optional): the cumulative probability of tokens to generate, e.g. 0.9.
+            Defaults to None.
+            random_seed (Optional[int], optional): the random seed to use for sampling, e.g. 42. Defaults to None.
+            safe_mode (bool, optional): whether to use safe mode, e.g. true. Defaults to False.
+
+        Returns:
+            ChatCompletionResponse: a response object containing the generated text.
+        """
+        request = self._make_chat_request(
+            model,
+            messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            random_seed=random_seed,
+            stream=False,
+            safe_mode=safe_mode,
+        )
+
+        single_response = self._request("post", request, "v1/chat/completions")
+
+        async for response in single_response:
+            return ChatCompletionResponse(**response)
+
+        raise MistralException("No response received")
+
+    async def chat_stream(
+        self,
+        model: str,
+        messages: List[ChatMessage],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_p: Optional[float] = None,
+        random_seed: Optional[int] = None,
+        safe_mode: bool = False,
+    ) -> AsyncGenerator[ChatCompletionStreamResponse, None]:
+        """An Asynchronous chat endpoint that streams responses.
+
+        Args:
+            model (str): model the name of the model to chat with, e.g. mistral-tiny
+            messages (List[ChatMessage]): messages an array of messages to chat with, e.g.
+                [{role: 'user', content: 'What is the best French cheese?'}]
+            temperature (Optional[float], optional): temperature the temperature to use for sampling, e.g. 0.5.
+            max_tokens (Optional[int], optional): the maximum number of tokens to generate, e.g. 100. Defaults to None.
+            top_p (Optional[float], optional): the cumulative probability of tokens to generate, e.g. 0.9.
+            Defaults to None.
+            random_seed (Optional[int], optional): the random seed to use for sampling, e.g. 42. Defaults to None.
+            safe_mode (bool, optional): whether to use safe mode, e.g. true. Defaults to False.
+
+        Returns:
+            AsyncGenerator[ChatCompletionStreamResponse, None]:
+                An async generator that yields ChatCompletionStreamResponse objects.
+        """
+
+        request = self._make_chat_request(
+            model,
+            messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            random_seed=random_seed,
+            stream=True,
+            safe_mode=safe_mode,
+        )
+        async_response = self._request(
+            "post", request, "v1/chat/completions", stream=True
+        )
+
+        async for json_response in async_response:
+            yield ChatCompletionStreamResponse(**json_response)
+
+    async def embeddings(
+        self, model: str, input: Union[str, List[str]]
+    ) -> EmbeddingResponse:
+        """An asynchronous embeddings endpoint that returns embeddings for a single, or batch of inputs
+
+        Args:
+            model (str): The embedding model to use, e.g. mistral-embed
+            input (Union[str, List[str]]): The input to embed,
+                 e.g. ['What is the best French cheese?']
+
+        Returns:
+            EmbeddingResponse: A response object containing the embeddings.
+        """
+        request = {"model": model, "input": input}
+        single_response = self._request("post", request, "v1/embeddings")
+
+        async for response in single_response:
+            return EmbeddingResponse(**response)
+
+        raise MistralException("No response received")
+
+    async def list_models(self) -> ModelList:
+        """Returns a list of the available models
+
+        Returns:
+            ModelList: A response object containing the list of models.
+        """
+        single_response = self._request("get", {}, "v1/models")
+
+        async for response in single_response:
+            return ModelList(**response)
+
+        raise MistralException("No response received")

{mistralai-0.0.7 → mistralai-0.0.9}/src/mistralai/client.py

@@ -1,18 +1,16 @@
 import os
 import posixpath
+import time
 from json import JSONDecodeError
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Union
 
-import orjson
-import requests
-from requests import Response
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
+from httpx import Client, ConnectError, HTTPTransport, RequestError, Response
 
 from mistralai.client_base import ClientBase
-from mistralai.constants import ENDPOINT, RETRY_STATUS_CODES
+from mistralai.constants import ENDPOINT
 from mistralai.exceptions import (
     MistralAPIException,
+    MistralAPIStatusException,
     MistralConnectionException,
     MistralException,
 )
@@ -39,14 +37,22 @@ class MistralClient(ClientBase):
     ):
         super().__init__(endpoint, api_key, max_retries, timeout)
 
+        self._client = Client(
+            follow_redirects=True,
+            timeout=self._timeout,
+            transport=HTTPTransport(retries=self._max_retries))
+
+    def __del__(self) -> None:
+        self._client.close()
+
     def _request(
         self,
         method: str,
         json: Dict[str, Any],
         path: str,
         stream: bool = False,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> Union[Response, Dict[str, Any]]:
+        attempt: int = 1,
+    ) -> Iterator[Dict[str, Any]]:
         headers = {
             "Authorization": f"Bearer {self._api_key}",
             "Content-Type": "application/json",
@@ -54,49 +60,58 @@ class MistralClient(ClientBase):
 
         url = posixpath.join(self._endpoint, path)
 
-        with requests.Session() as session:
-            retries = Retry(
-                total=self._max_retries,
-                backoff_factor=0.5,
-                allowed_methods=["POST", "GET"],
-                status_forcelist=RETRY_STATUS_CODES,
-                raise_on_status=False,
-            )
-            session.mount("https://", HTTPAdapter(max_retries=retries))
-            session.mount("http://", HTTPAdapter(max_retries=retries))
+        self._logger.debug(f"Sending request: {method} {url} {json}")
+
+        response: Response
 
+        try:
             if stream:
-                return session.request(
-                    method, url, headers=headers, json=json, stream=True
-                )
+                with self._client.stream(
+                    method,
+                    url,
+                    headers=headers,
+                    json=json,
+                ) as response:
+                    self._check_streaming_response(response)
+
+                    for line in response.iter_lines():
+                        json_streamed_response = self._process_line(line)
+                        if json_streamed_response:
+                            yield json_streamed_response
 
-            try:
-                response = session.request(
+            else:
+                response = self._client.request(
                     method,
                     url,
                     headers=headers,
                     json=json,
-                    timeout=self._timeout,
-                    params=params,
                 )
-            except requests.exceptions.ConnectionError as e:
-                raise MistralConnectionException(str(e)) from e
-            except requests.exceptions.RequestException as e:
-                raise MistralException(
-                    f"Unexpected exception ({e.__class__.__name__}): {e}"
-                ) from e
 
-            try:
-                json_response: Dict[str, Any] = response.json()
-            except JSONDecodeError:
-                raise MistralAPIException.from_response(
-                    response, message=f"Failed to decode json body: {response.text}"
-                )
+                yield self._check_response(response)
+
+        except ConnectError as e:
+            raise MistralConnectionException(str(e)) from e
+        except RequestError as e:
+            raise MistralException(
+                f"Unexpected exception ({e.__class__.__name__}): {e}"
+            ) from e
+        except JSONDecodeError as e:
+            raise MistralAPIException.from_response(
+                response,
+                message=f"Failed to decode json body: {response.text}",
+            ) from e
+        except MistralAPIStatusException as e:
+            attempt += 1
+            if attempt > self._max_retries:
+                raise MistralAPIStatusException.from_response(
+                    response, message=str(e)
+                ) from e
+            backoff = 2.0**attempt  # exponential backoff
+            time.sleep(backoff)
 
-            self._check_response(
-                json_response, dict(response.headers), response.status_code
-            )
-        return json_response
+            # Retry as a generator
+            for r in self._request(method, json, path, stream=stream, attempt=attempt):
+                yield r
 
     def chat(
         self,
@@ -106,9 +121,9 @@ class MistralClient(ClientBase):
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
-        safe_mode: bool = True,
+        safe_mode: bool = False,
     ) -> ChatCompletionResponse:
-        """ A chat endpoint that returns a single response.
+        """A chat endpoint that returns a single response.
 
         Args:
             model (str): model the name of the model to chat with, e.g. mistral-tiny
@@ -135,11 +150,12 @@ class MistralClient(ClientBase):
             safe_mode=safe_mode,
         )
 
-        response = self._request("post", request, "v1/chat/completions")
+        single_response = self._request("post", request, "v1/chat/completions")
 
-        assert isinstance(response, dict), "Bad response from _request"
+        for response in single_response:
+            return ChatCompletionResponse(**response)
 
-        return ChatCompletionResponse(**response)
+        raise MistralException("No response received")
 
     def chat_stream(
         self,
@@ -149,9 +165,9 @@ class MistralClient(ClientBase):
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
-        safe_mode: bool = True,
+        safe_mode: bool = False,
     ) -> Iterable[ChatCompletionStreamResponse]:
-        """ A chat endpoint that streams responses.
+        """A chat endpoint that streams responses.
 
         Args:
             model (str): model the name of the model to chat with, e.g. mistral-tiny
@@ -181,18 +197,8 @@ class MistralClient(ClientBase):
 
         response = self._request("post", request, "v1/chat/completions", stream=True)
 
-        assert isinstance(response, Response), "Bad response from _request"
-
-        for line in response.iter_lines():
-            self._logger.debug(f"Received line: {line}")
-            if line == b"\n":
-                continue
-
-            if line.startswith(b"data: "):
-                line = line[6:].strip()
-                if line != b"[DONE]":
-                    json_response = orjson.loads(line)
-                    yield ChatCompletionStreamResponse(**json_response)
+        for json_streamed_response in response:
+            yield ChatCompletionStreamResponse(**json_streamed_response)
 
     def embeddings(self, model: str, input: Union[str, List[str]]) -> EmbeddingResponse:
         """An embeddings endpoint that returns embeddings for a single, or batch of inputs
@@ -206,9 +212,12 @@ class MistralClient(ClientBase):
             EmbeddingResponse: A response object containing the embeddings.
         """
         request = {"model": model, "input": input}
-        response = self._request("post", request, "v1/embeddings")
-        assert isinstance(response, dict), "Bad response from _request"
-        return EmbeddingResponse(**response)
+        singleton_response = self._request("post", request, "v1/embeddings")
+
+        for response in singleton_response:
+            return EmbeddingResponse(**response)
+
+        raise MistralException("No response received")
 
     def list_models(self) -> ModelList:
         """Returns a list of the available models
@@ -216,6 +225,9 @@ class MistralClient(ClientBase):
         Returns:
             ModelList: A response object containing the list of models.
         """
-        response = self._request("get", {}, "v1/models")
-        assert isinstance(response, dict), "Bad response from _request"
-        return ModelList(**response)
+        singleton_response = self._request("get", {}, "v1/models")
+
+        for response in singleton_response:
+            return ModelList(**response)
+
+        raise MistralException("No response received")

mistralai-0.0.9/src/mistralai/client_base.py

@@ -0,0 +1,109 @@
+import logging
+import os
+from abc import ABC
+from typing import Any, Dict, List, Optional
+
+import orjson
+from httpx import Response
+
+from mistralai.constants import RETRY_STATUS_CODES
+from mistralai.exceptions import (
+    MistralAPIException,
+    MistralAPIStatusException,
+    MistralException,
+)
+from mistralai.models.chat_completion import ChatMessage
+
+logging.basicConfig(
+    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    level=os.getenv("LOG_LEVEL", "ERROR"),
+)
+
+
+class ClientBase(ABC):
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: Optional[str] = None,
+        max_retries: int = 5,
+        timeout: int = 120,
+    ):
+        self._max_retries = max_retries
+        self._timeout = timeout
+
+        self._endpoint = endpoint
+        self._api_key = api_key
+        self._logger = logging.getLogger(__name__)
+
+    def _make_chat_request(
+        self,
+        model: str,
+        messages: List[ChatMessage],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_p: Optional[float] = None,
+        random_seed: Optional[int] = None,
+        stream: Optional[bool] = None,
+        safe_mode: Optional[bool] = False,
+    ) -> Dict[str, Any]:
+        request_data: Dict[str, Any] = {
+            "model": model,
+            "messages": [msg.model_dump() for msg in messages],
+            "safe_prompt": safe_mode,
+        }
+        if temperature is not None:
+            request_data["temperature"] = temperature
+        if max_tokens is not None:
+            request_data["max_tokens"] = max_tokens
+        if top_p is not None:
+            request_data["top_p"] = top_p
+        if random_seed is not None:
+            request_data["random_seed"] = random_seed
+        if stream is not None:
+            request_data["stream"] = stream
+
+        self._logger.debug(f"Chat request: {request_data}")
+
+        return request_data
+
+    def _check_response_status_codes(self, response: Response) -> None:
+        if response.status_code in RETRY_STATUS_CODES:
+            raise MistralAPIStatusException.from_response(
+                response,
+                message=f"Cannot stream response. Status: {response.status_code}",
+            )
+        elif 400 <= response.status_code < 500:
+            raise MistralAPIException.from_response(
+                response,
+                message=f"Cannot stream response. Status: {response.status_code}",
+            )
+        elif response.status_code >= 500:
+            raise MistralException(
+                message=f"Unexpected server error (status {response.status_code})"
+            )
+
+    def _check_streaming_response(self, response: Response) -> None:
+        self._check_response_status_codes(response)
+
+    def _check_response(self, response: Response) -> Dict[str, Any]:
+        self._check_response_status_codes(response)
+
+        json_response: Dict[str, Any] = response.json()
+
+        if "object" not in json_response:
+            raise MistralException(message=f"Unexpected response: {json_response}")
+        if "error" == json_response["object"]:  # has errors
+            raise MistralAPIException.from_response(
+                response,
+                message=json_response["message"],
+            )
+
+        return json_response
+
+    def _process_line(self, line: str) -> Optional[Dict[str, Any]]:
+        if line.startswith("data: "):
+            line = line[6:].strip()
+            if line != "[DONE]":
+                json_streamed_response: Dict[str, Any] = orjson.loads(line)
+                return json_streamed_response
+        return None

{mistralai-0.0.7 → mistralai-0.0.9}/src/mistralai/exceptions.py

@@ -2,8 +2,7 @@ from __future__ import annotations
 
 from typing import Any, Dict, Optional
 
-import aiohttp
-from requests import Response
+from httpx import Response
 
 
 class MistralException(Exception):
@@ -45,19 +44,11 @@ class MistralAPIException(MistralException):
             headers=dict(response.headers),
         )
 
-    @classmethod
-    def from_aio_response(
-        cls, response: aiohttp.ClientResponse, message: Optional[str] = None
-    ) -> MistralAPIException:
-        return cls(
-            message=message,
-            http_status=response.status,
-            headers=dict(response.headers),
-        )
-
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(message={str(self)}, http_status={self.http_status})"
 
+class MistralAPIStatusException(MistralAPIException):
+    """Returned when we receive a non-200 response from the API that we should retry"""
 
 class MistralConnectionException(MistralException):
     """Returned when the SDK can not reach the API server for any reason"""

mistralai-0.0.7/src/mistralai/async_client.py

@@ -1,342 +0,0 @@
-import asyncio
-import logging
-import os
-import posixpath
-import time
-from collections import defaultdict
-from json import JSONDecodeError
-from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Union
-
-import aiohttp
-import backoff
-import orjson
-
-from mistralai.client_base import ClientBase
-from mistralai.constants import ENDPOINT, RETRY_STATUS_CODES
-from mistralai.exceptions import (
-    MistralAPIException,
-    MistralConnectionException,
-    MistralException,
-)
-from mistralai.models.chat_completion import (
-    ChatCompletionResponse,
-    ChatCompletionStreamResponse,
-    ChatMessage,
-)
-from mistralai.models.embeddings import EmbeddingResponse
-from mistralai.models.models import ModelList
-
-
-class AIOHTTPBackend:
-    """HTTP backend which handles retries, concurrency limiting and logging"""
-
-    SLEEP_AFTER_FAILURE = defaultdict(lambda: 0.25, {429: 5.0})
-
-    _requester: Callable[..., Awaitable[aiohttp.ClientResponse]]
-    _semaphore: asyncio.Semaphore
-    _session: Optional[aiohttp.ClientSession]
-
-    def __init__(
-        self,
-        max_concurrent_requests: int = 64,
-        max_retries: int = 5,
-        timeout: int = 120,
-    ):
-        self._logger = logging.getLogger(__name__)
-        self._timeout = timeout
-        self._max_retries = max_retries
-        self._session = None
-        self._max_concurrent_requests = max_concurrent_requests
-
-    def build_aio_requester(
-        self,
-    ) -> Callable:  # returns a function for retryable requests
-        @backoff.on_exception(
-            backoff.expo,
-            (aiohttp.ClientError, aiohttp.ClientResponseError),
-            max_tries=self._max_retries + 1,
-            max_time=self._timeout,
-        )
-        async def make_request_fn(
-            session: aiohttp.ClientSession, *args: Any, **kwargs: Any
-        ) -> aiohttp.ClientResponse:
-            async with self._semaphore:  # this limits total concurrency by the client
-                response = await session.request(*args, **kwargs)
-            if (
-                response.status in RETRY_STATUS_CODES
-            ):  # likely temporary, raise to retry
-                self._logger.info(f"Received status {response.status}, retrying...")
-                await asyncio.sleep(self.SLEEP_AFTER_FAILURE[response.status])
-                response.raise_for_status()
-
-            return response
-
-        return make_request_fn
-
-    async def request(
-        self,
-        url: str,
-        json: Optional[Dict[str, Any]] = None,
-        method: str = "post",
-        headers: Optional[Dict[str, Any]] = None,
-        session: Optional[aiohttp.ClientSession] = None,
-        params: Optional[Dict[str, Any]] = None,
-        **kwargs: Any,
-    ) -> aiohttp.ClientResponse:
-        session = session or await self.session()
-        self._logger.debug(f"Making request to {url} with content {json}")
-
-        request_start = time.time()
-        try:
-            response = await self._requester(
-                session,
-                method,
-                url,
-                headers=headers,
-                json=json,
-                params=params,
-                **kwargs,
-            )
-        except (
-            aiohttp.ClientConnectionError
-        ) as e:  # ensure the SDK user does not have to deal with knowing aiohttp
-            self._logger.debug(
-                f"Fatal connection error after {time.time()-request_start:.1f}s: {e}"
-            )
-            raise MistralConnectionException(str(e)) from e
-        except (
-            aiohttp.ClientResponseError
-        ) as e:  # status 500 or something remains after retries
-            self._logger.debug(
-                f"Fatal ClientResponseError error after {time.time()-request_start:.1f}s: {e}"
-            )
-            raise MistralConnectionException(str(e)) from e
-        except asyncio.TimeoutError as e:
-            self._logger.debug(
-                f"Fatal timeout error after {time.time()-request_start:.1f}s: {e}"
-            )
-            raise MistralConnectionException("The request timed out") from e
-        except Exception as e:  # Anything caught here should be added above
-            self._logger.debug(
-                f"Unexpected fatal error after {time.time()-request_start:.1f}s: {e}"
-            )
-            raise MistralException(
-                f"Unexpected exception ({e.__class__.__name__}): {e}"
-            ) from e
-
-        self._logger.debug(
-            f"Received response with status {response.status} after {time.time()-request_start:.1f}s"
-        )
-        return response
-
-    async def session(self) -> aiohttp.ClientSession:
-        if self._session is None:
-            self._session = aiohttp.ClientSession(
-                timeout=aiohttp.ClientTimeout(self._timeout),
-                connector=aiohttp.TCPConnector(limit=0),
-            )
-            self._semaphore = asyncio.Semaphore(self._max_concurrent_requests)
-            self._requester = self.build_aio_requester()
-        return self._session
-
-    async def close(self) -> None:
-        if self._session is not None:
-            await self._session.close()
-            self._session = None
-
-    def __del__(self) -> None:
-        # https://stackoverflow.com/questions/54770360/how-can-i-wait-for-an-objects-del-to-finish-before-the-async-loop-closes
-        if self._session:
-            try:
-                loop = asyncio.get_event_loop()
-                if loop.is_running():
-                    loop.create_task(self.close())
-                else:
-                    loop.run_until_complete(self.close())
-            except Exception:
-                pass
-
-
-class MistralAsyncClient(ClientBase):
-    def __init__(
-        self,
-        api_key: Optional[str] = os.environ.get("MISTRAL_API_KEY", None),
-        endpoint: str = ENDPOINT,
-        max_retries: int = 5,
-        timeout: int = 120,
-        max_concurrent_requests: int = 64,
-    ):
-        super().__init__(endpoint, api_key, max_retries, timeout)
-
-        self._backend = AIOHTTPBackend(
-            max_concurrent_requests=max_concurrent_requests,
-            max_retries=max_retries,
-            timeout=timeout,
-        )
-
-    async def close(self) -> None:
-        await self._backend.close()
-
-    async def _request(
-        self,
-        method: str,
-        json: Dict[str, Any],
-        path: str,
-        stream: bool = False,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> Union[Dict[str, Any], aiohttp.ClientResponse]:
-
-        headers = {
-            "Authorization": f"Bearer {self._api_key}",
-            "Content-Type": "application/json",
-        }
-
-        url = posixpath.join(self._endpoint, path)
-
-        response = await self._backend.request(
-            url, json, method, headers, params=params
-        )
-        if stream:
-            return response
-
-        try:
-            json_response: Dict[str, Any] = await response.json()
-        except JSONDecodeError:
-            raise MistralAPIException.from_aio_response(
-                response, message=f"Failed to decode json body: {await response.text()}"
-            )
-        except aiohttp.ClientPayloadError as e:
-            raise MistralAPIException.from_aio_response(
-                response,
-                message=f"An unexpected error occurred while receiving the response: {e}",
-            )
-
-        self._logger.debug(f"JSON response: {json_response}")
-        self._check_response(json_response, dict(response.headers), response.status)
-        return json_response
-
-    async def chat(
-        self,
-        model: str,
-        messages: List[ChatMessage],
-        temperature: Optional[float] = None,
-        max_tokens: Optional[int] = None,
-        top_p: Optional[float] = None,
-        random_seed: Optional[int] = None,
-        safe_mode: bool = True,
-    ) -> ChatCompletionResponse:
-        """ A asynchronous chat endpoint that returns a single response.
-
-        Args:
-            model (str): model the name of the model to chat with, e.g. mistral-tiny
-            messages (List[ChatMessage]): messages an array of messages to chat with, e.g.
-                [{role: 'user', content: 'What is the best French cheese?'}]
-            temperature (Optional[float], optional): temperature the temperature to use for sampling, e.g. 0.5.
-            max_tokens (Optional[int], optional): the maximum number of tokens to generate, e.g. 100. Defaults to None.
-            top_p (Optional[float], optional): the cumulative probability of tokens to generate, e.g. 0.9.
-            Defaults to None.
-            random_seed (Optional[int], optional): the random seed to use for sampling, e.g. 42. Defaults to None.
-            safe_mode (bool, optional): whether to use safe mode, e.g. true. Defaults to False.
-
-        Returns:
-            ChatCompletionResponse: a response object containing the generated text.
-        """
-        request = self._make_chat_request(
-            model,
-            messages,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p,
-            random_seed=random_seed,
-            stream=False,
-            safe_mode=safe_mode,
-        )
-
-        response = await self._request("post", request, "v1/chat/completions")
-        assert isinstance(response, dict), "Bad response from _request"
-        return ChatCompletionResponse(**response)
-
-    async def chat_stream(
-        self,
-        model: str,
-        messages: List[ChatMessage],
-        temperature: Optional[float] = None,
-        max_tokens: Optional[int] = None,
-        top_p: Optional[float] = None,
-        random_seed: Optional[int] = None,
-        safe_mode: bool = False,
-    ) -> AsyncGenerator[ChatCompletionStreamResponse, None]:
-        """ An Asynchronous chat endpoint that streams responses.
-
-        Args:
-            model (str): model the name of the model to chat with, e.g. mistral-tiny
-            messages (List[ChatMessage]): messages an array of messages to chat with, e.g.
-                [{role: 'user', content: 'What is the best French cheese?'}]
-            temperature (Optional[float], optional): temperature the temperature to use for sampling, e.g. 0.5.
-            max_tokens (Optional[int], optional): the maximum number of tokens to generate, e.g. 100. Defaults to None.
-            top_p (Optional[float], optional): the cumulative probability of tokens to generate, e.g. 0.9.
-            Defaults to None.
-            random_seed (Optional[int], optional): the random seed to use for sampling, e.g. 42. Defaults to None.
-            safe_mode (bool, optional): whether to use safe mode, e.g. true. Defaults to False.
-
-        Returns:
-            AsyncGenerator[ChatCompletionStreamResponse, None]:
-                An async generator that yields ChatCompletionStreamResponse objects.
-        """
-
-        request = self._make_chat_request(
-            model,
-            messages,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p,
-            random_seed=random_seed,
-            stream=True,
-            safe_mode=safe_mode,
-        )
-        async_response = await self._request(
-            "post", request, "v1/chat/completions", stream=True
-        )
-
-        assert isinstance(
-            async_response, aiohttp.ClientResponse
-        ), "Bad response from _request"
-
-        async with async_response as response:
-            async for line in response.content:
-                if line == b"\n":
-                    continue
-
-                if line.startswith(b"data: "):
-                    line = line[6:].strip()
-                    if line != b"[DONE]":
-                        json_response = orjson.loads(line)
-                        yield ChatCompletionStreamResponse(**json_response)
-
-    async def embeddings(
-        self, model: str, input: Union[str, List[str]]
-    ) -> EmbeddingResponse:
-        """An asynchronous embeddings endpoint that returns embeddings for a single, or batch of inputs
-
-        Args:
-            model (str): The embedding model to use, e.g. mistral-embed
-            input (Union[str, List[str]]): The input to embed,
-                 e.g. ['What is the best French cheese?']
-
-        Returns:
-            EmbeddingResponse: A response object containing the embeddings.
-        """
-        request = {"model": model, "input": input}
-        response = await self._request("post", request, "v1/embeddings")
-        assert isinstance(response, dict), "Bad response from _request"
-        return EmbeddingResponse(**response)
-
-    async def list_models(self) -> ModelList:
-        """Returns a list of the available models
-
-        Returns:
-            ModelList: A response object containing the list of models.
-        """
-        response = await self._request("get", {}, "v1/models")
-        assert isinstance(response, dict), "Bad response from _request"
-        return ModelList(**response)

mistralai-0.0.7/src/mistralai/client_base.py

@@ -1,73 +0,0 @@
-import logging
-from abc import ABC
-from typing import Any, Dict, List, Optional
-
-from mistralai.exceptions import MistralAPIException, MistralException
-from mistralai.models.chat_completion import ChatMessage
-
-
-class ClientBase(ABC):
-    def __init__(
-        self,
-        endpoint: str,
-        api_key: Optional[str] = None,
-        max_retries: int = 5,
-        timeout: int = 120,
-    ):
-        self._max_retries = max_retries
-        self._timeout = timeout
-
-        self._endpoint = endpoint
-        self._api_key = api_key
-        self._logger = logging.getLogger(__name__)
-
-    @staticmethod
-    def _make_chat_request(
-        model: str,
-        messages: List[ChatMessage],
-        temperature: Optional[float] = None,
-        max_tokens: Optional[int] = None,
-        top_p: Optional[float] = None,
-        random_seed: Optional[int] = None,
-        stream: Optional[bool] = None,
-        safe_mode: Optional[bool] = True,
-    ) -> Dict[str, Any]:
-        request_data: Dict[str, Any] = {
-            "model": model,
-            "messages": [msg.model_dump() for msg in messages],
-            "safe_prompt": safe_mode,
-        }
-        if temperature is not None:
-            request_data["temperature"] = temperature
-        if max_tokens is not None:
-            request_data["max_tokens"] = max_tokens
-        if top_p is not None:
-            request_data["top_p"] = top_p
-        if random_seed is not None:
-            request_data["random_seed"] = random_seed
-        if stream is not None:
-            request_data["stream"] = stream
-
-        return request_data
-
-    def _check_response(
-        self, json_response: Dict[str, Any], headers: Dict[str, Any], status: int
-    ) -> None:
-        if "object" not in json_response:
-            raise MistralException(message=f"Unexpected response: {json_response}")
-        if "error" == json_response["object"]:  # has errors
-            raise MistralAPIException(
-                message=json_response["message"],
-                http_status=status,
-                headers=headers,
-            )
-        if 400 <= status < 500:
-            raise MistralAPIException(
-                message=f"Unexpected client error (status {status}): {json_response}",
-                http_status=status,
-                headers=headers,
-            )
-        if status >= 500:
-            raise MistralException(
-                message=f"Unexpected server error (status {status}): {json_response}"
-            )