mistralai 0.0.12__tar.gz → 0.1.2__tar.gz

{mistralai-0.0.12 → mistralai-0.1.2}/PKG-INFO

@@ -1,18 +1,19 @@
 Metadata-Version: 2.1
 Name: mistralai
-Version: 0.0.12
+Version: 0.1.2
 Summary: 
 Author: Bam4d
 Author-email: bam4d@mistral.ai
-Requires-Python: >=3.8,<4.0
+Requires-Python: >=3.9,<4.0
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 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: httpx (>=0.25.2,<0.26.0)
 Requires-Dist: orjson (>=3.9.10,<4.0.0)
+Requires-Dist: pandas (>=2.2.0,<3.0.0)
+Requires-Dist: pyarrow (>=15.0.0,<16.0.0)
 Requires-Dist: pydantic (>=2.5.2,<3.0.0)
 Description-Content-Type: text/markdown
 

{mistralai-0.0.12 → mistralai-0.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mistralai"
-version = "0.0.12"
+version = "0.1.2"
 description = ""
 authors = ["Bam4d <bam4d@mistral.ai>"]
 readme = "README.md"
@@ -23,10 +23,12 @@ exclude = ["docs", "tests", "examples", "tools", "build"]
 
 
 [tool.poetry.dependencies]
-python = "^3.8"
+python = "^3.9"
 orjson = "^3.9.10"
 pydantic = "^2.5.2"
 httpx = "^0.25.2"
+pandas = "^2.2.0"
+pyarrow = "^15.0.0"
 
 
 [tool.poetry.group.dev.dependencies]

{mistralai-0.0.12 → mistralai-0.1.2}/src/mistralai/async_client.py

@@ -24,7 +24,8 @@ from mistralai.exceptions import (
 from mistralai.models.chat_completion import (
     ChatCompletionResponse,
     ChatCompletionStreamResponse,
-    ChatMessage,
+    ResponseFormat,
+    ToolChoice,
 )
 from mistralai.models.embeddings import EmbeddingResponse
 from mistralai.models.models import ModelList
@@ -101,9 +102,7 @@ class MistralAsyncClient(ClientBase):
         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
+            raise MistralException(f"Unexpected exception ({e.__class__.__name__}): {e}") from e
         except JSONDecodeError as e:
             raise MistralAPIException.from_response(
                 response,
@@ -112,34 +111,33 @@ class MistralAsyncClient(ClientBase):
         except MistralAPIStatusException as e:
             attempt += 1
             if attempt > self._max_retries:
-                raise MistralAPIStatusException.from_response(
-                    response, message=str(e)
-                ) from e
+                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
-            ):
+            async for r in self._request(method, json, path, stream=stream, attempt=attempt):
                 yield r
 
     async def chat(
         self,
-        model: str,
-        messages: List[ChatMessage],
+        messages: List[Any],
+        model: Optional[str] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
         safe_mode: bool = False,
         safe_prompt: bool = False,
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        response_format: Optional[Union[Dict[str, str], ResponseFormat]] = None,
     ) -> 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.
+            messages (List[Any]): 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.
@@ -153,14 +151,17 @@ class MistralAsyncClient(ClientBase):
             ChatCompletionResponse: a response object containing the generated text.
         """
         request = self._make_chat_request(
-            model,
             messages,
+            model,
+            tools=tools,
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
             random_seed=random_seed,
             stream=False,
             safe_prompt=safe_mode or safe_prompt,
+            tool_choice=tool_choice,
+            response_format=response_format,
         )
 
         single_response = self._request("post", request, "v1/chat/completions")
@@ -172,21 +173,25 @@ class MistralAsyncClient(ClientBase):
 
     async def chat_stream(
         self,
-        model: str,
-        messages: List[ChatMessage],
+        messages: List[Any],
+        model: Optional[str] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
         safe_mode: bool = False,
         safe_prompt: bool = False,
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        response_format: Optional[Union[Dict[str, str], ResponseFormat]] = None,
     ) -> 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.
+            messages (List[Any]): messages an array of messages to chat with, e.g.
                 [{role: 'user', content: 'What is the best French cheese?'}]
+            tools (Optional[List[Function]], optional): a list of tools to use.
             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.
@@ -201,25 +206,24 @@ class MistralAsyncClient(ClientBase):
         """
 
         request = self._make_chat_request(
-            model,
             messages,
+            model,
+            tools=tools,
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
             random_seed=random_seed,
             stream=True,
             safe_prompt=safe_mode or safe_prompt,
+            tool_choice=tool_choice,
+            response_format=response_format,
         )
-        async_response = self._request(
-            "post", request, "v1/chat/completions", stream=True
-        )
+        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:
+    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:

{mistralai-0.0.12 → mistralai-0.1.2}/src/mistralai/client.py

@@ -17,7 +17,8 @@ from mistralai.exceptions import (
 from mistralai.models.chat_completion import (
     ChatCompletionResponse,
     ChatCompletionStreamResponse,
-    ChatMessage,
+    ResponseFormat,
+    ToolChoice,
 )
 from mistralai.models.embeddings import EmbeddingResponse
 from mistralai.models.models import ModelList
@@ -38,9 +39,8 @@ 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))
+            follow_redirects=True, timeout=self._timeout, transport=HTTPTransport(retries=self._max_retries)
+        )
 
     def __del__(self) -> None:
         self._client.close()
@@ -95,9 +95,7 @@ class MistralClient(ClientBase):
         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
+            raise MistralException(f"Unexpected exception ({e.__class__.__name__}): {e}") from e
         except JSONDecodeError as e:
             raise MistralAPIException.from_response(
                 response,
@@ -106,9 +104,7 @@ class MistralClient(ClientBase):
         except MistralAPIStatusException as e:
             attempt += 1
             if attempt > self._max_retries:
-                raise MistralAPIStatusException.from_response(
-                    response, message=str(e)
-                ) from e
+                raise MistralAPIStatusException.from_response(response, message=str(e)) from e
             backoff = 2.0**attempt  # exponential backoff
             time.sleep(backoff)
 
@@ -118,21 +114,25 @@ class MistralClient(ClientBase):
 
     def chat(
         self,
-        model: str,
-        messages: List[ChatMessage],
+        messages: List[Any],
+        model: Optional[str] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
         safe_mode: bool = False,
         safe_prompt: bool = False,
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        response_format: Optional[Union[Dict[str, str], ResponseFormat]] = None,
     ) -> ChatCompletionResponse:
         """A 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.
+            messages (List[Any]): messages an array of messages to chat with, e.g.
                 [{role: 'user', content: 'What is the best French cheese?'}]
+            tools (Optional[List[Function]], optional): a list of tools to use.
             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.
@@ -145,14 +145,17 @@ class MistralClient(ClientBase):
             ChatCompletionResponse: a response object containing the generated text.
         """
         request = self._make_chat_request(
-            model,
             messages,
+            model,
+            tools=tools,
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
             random_seed=random_seed,
             stream=False,
             safe_prompt=safe_mode or safe_prompt,
+            tool_choice=tool_choice,
+            response_format=response_format,
         )
 
         single_response = self._request("post", request, "v1/chat/completions")
@@ -164,21 +167,25 @@ class MistralClient(ClientBase):
 
     def chat_stream(
         self,
-        model: str,
-        messages: List[ChatMessage],
+        messages: List[Any],
+        model: Optional[str] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
         safe_mode: bool = False,
         safe_prompt: bool = False,
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        response_format: Optional[Union[Dict[str, str], ResponseFormat]] = None,
     ) -> Iterable[ChatCompletionStreamResponse]:
         """A 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.
+            messages (List[Any]): messages an array of messages to chat with, e.g.
                 [{role: 'user', content: 'What is the best French cheese?'}]
+            tools (Optional[List[Function]], optional): a list of tools to use.
             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.
@@ -192,14 +199,17 @@ class MistralClient(ClientBase):
                 A generator that yields ChatCompletionStreamResponse objects.
         """
         request = self._make_chat_request(
-            model,
             messages,
+            model,
+            tools=tools,
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
             random_seed=random_seed,
             stream=True,
             safe_prompt=safe_mode or safe_prompt,
+            tool_choice=tool_choice,
+            response_format=response_format,
         )
 
         response = self._request("post", request, "v1/chat/completions", stream=True)

{mistralai-0.0.12 → mistralai-0.1.2}/src/mistralai/client_base.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from abc import ABC
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import orjson
 from httpx import Response
@@ -12,7 +12,7 @@ from mistralai.exceptions import (
     MistralAPIStatusException,
     MistralException,
 )
-from mistralai.models.chat_completion import ChatMessage
+from mistralai.models.chat_completion import ChatMessage, Function, ResponseFormat, ToolChoice
 
 logging.basicConfig(
     format="%(asctime)s %(levelname)s %(name)s: %(message)s",
@@ -35,25 +35,76 @@ class ClientBase(ABC):
         self._api_key = api_key
         self._logger = logging.getLogger(__name__)
 
+        # For azure endpoints, we default to the mistral model
+        if "inference.azure.com" in self._endpoint:
+            self._default_model = "mistral"
+
         # This should be automatically updated by the deploy script
-        self._version = "0.0.12"
+        self._version = "0.1.2"
+
+    def _parse_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        parsed_tools: List[Dict[str, Any]] = []
+        for tool in tools:
+            if tool["type"] == "function":
+                parsed_function = {}
+                parsed_function["type"] = tool["type"]
+                if isinstance(tool["function"], Function):
+                    parsed_function["function"] = tool["function"].model_dump(exclude_none=True)
+                else:
+                    parsed_function["function"] = tool["function"]
+
+                parsed_tools.append(parsed_function)
+
+        return parsed_tools
+
+    def _parse_tool_choice(self, tool_choice: Union[str, ToolChoice]) -> str:
+        if isinstance(tool_choice, ToolChoice):
+            return tool_choice.value
+        return tool_choice
+
+    def _parse_response_format(self, response_format: Union[Dict[str, Any], ResponseFormat]) -> Dict[str, Any]:
+        if isinstance(response_format, ResponseFormat):
+            return response_format.model_dump(exclude_none=True)
+        return response_format
+
+    def _parse_messages(self, messages: List[Any]) -> List[Dict[str, Any]]:
+        parsed_messages: List[Dict[str, Any]] = []
+        for message in messages:
+            if isinstance(message, ChatMessage):
+                parsed_messages.append(message.model_dump(exclude_none=True))
+            else:
+                parsed_messages.append(message)
+
+        return parsed_messages
 
     def _make_chat_request(
         self,
-        model: str,
-        messages: List[ChatMessage],
+        messages: List[Any],
+        model: Optional[str] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
         temperature: Optional[float] = None,
         max_tokens: Optional[int] = None,
         top_p: Optional[float] = None,
         random_seed: Optional[int] = None,
         stream: Optional[bool] = None,
         safe_prompt: Optional[bool] = False,
+        tool_choice: Optional[Union[str, ToolChoice]] = None,
+        response_format: Optional[Union[Dict[str, str], ResponseFormat]] = None,
     ) -> Dict[str, Any]:
         request_data: Dict[str, Any] = {
-            "model": model,
-            "messages": [msg.model_dump() for msg in messages],
+            "messages": self._parse_messages(messages),
             "safe_prompt": safe_prompt,
         }
+
+        if model is not None:
+            request_data["model"] = model
+        else:
+            if self._default_model is None:
+                raise MistralException(message="model must be provided")
+            request_data["model"] = self._default_model
+
+        if tools is not None:
+            request_data["tools"] = self._parse_tools(tools)
         if temperature is not None:
             request_data["temperature"] = temperature
         if max_tokens is not None:
@@ -65,6 +116,11 @@ class ClientBase(ABC):
         if stream is not None:
             request_data["stream"] = stream
 
+        if tool_choice is not None:
+            request_data["tool_choice"] = self._parse_tool_choice(tool_choice)
+        if response_format is not None:
+            request_data["response_format"] = self._parse_response_format(response_format)
+
         self._logger.debug(f"Chat request: {request_data}")
 
         return request_data

{mistralai-0.0.12 → mistralai-0.1.2}/src/mistralai/models/chat_completion.py

@@ -1,22 +1,61 @@
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from pydantic import BaseModel
 
 from mistralai.models.common import UsageInfo
 
 
+class Function(BaseModel):
+    name: str
+    description: str
+    parameters: dict
+
+
+class ToolType(str, Enum):
+    function = "function"
+
+
+class FunctionCall(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel):
+    id: str = "null"
+    type: ToolType = ToolType.function
+    function: FunctionCall
+
+
+class ResponseFormats(str, Enum):
+    text: str = "text"
+    json_object: str = "json_object"
+
+
+class ToolChoice(str, Enum):
+    auto: str = "auto"
+    any: str = "any"
+    none: str = "none"
+
+
+class ResponseFormat(BaseModel):
+    type: ResponseFormats = ResponseFormats.text
+
+
 class ChatMessage(BaseModel):
     role: str
-    content: str
+    content: Union[str, List[str]]
+    name: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
 
 
 class DeltaMessage(BaseModel):
     role: Optional[str] = None
     content: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
 
 
-class FinishReason(Enum):
+class FinishReason(str, Enum):
     stop = "stop"
     length = "length"
     error = "error"