mistralai 1.3.1__py3-none-any.whl → 1.5.0__py3-none-any.whl

mistralai/embeddings.py

@@ -21,7 +21,7 @@ class Embeddings(BaseSDK):
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
         http_headers: Optional[Mapping[str, str]] = None,
-    ) -> Optional[models.EmbeddingResponse]:
+    ) -> models.EmbeddingResponse:
         r"""Embeddings
 
         Embeddings
@@ -90,13 +90,16 @@ class Embeddings(BaseSDK):
 
         data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
-            return utils.unmarshal_json(
-                http_res.text, Optional[models.EmbeddingResponse]
-            )
+            return utils.unmarshal_json(http_res.text, models.EmbeddingResponse)
         if utils.match_response(http_res, "422", "application/json"):
             data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
             raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = utils.stream_to_text(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = utils.stream_to_text(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -121,7 +124,7 @@ class Embeddings(BaseSDK):
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
         http_headers: Optional[Mapping[str, str]] = None,
-    ) -> Optional[models.EmbeddingResponse]:
+    ) -> models.EmbeddingResponse:
         r"""Embeddings
 
         Embeddings
@@ -190,13 +193,16 @@ class Embeddings(BaseSDK):
 
         data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
-            return utils.unmarshal_json(
-                http_res.text, Optional[models.EmbeddingResponse]
-            )
+            return utils.unmarshal_json(http_res.text, models.EmbeddingResponse)
         if utils.match_response(http_res, "422", "application/json"):
             data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
             raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = await utils.stream_to_text_async(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = await utils.stream_to_text_async(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res

mistralai/extra/README.md

@@ -0,0 +1,56 @@
+## Context
+
+The extra package contains the custom logic which is too complex to be generated by Speakeasy from the OpenAPI specs. It was introduced to add the Structured Outputs feature.
+
+## Development / Contributing
+
+To add custom code in the SDK, you need to use [Speakeasy custom code regions](https://www.speakeasy.com/docs/customize/code/code-regions/overview) as below.
+
+### Runbook of SDK customization
+
+1. Add the code you want to import in the `src/mistralai/extra/` package. To have it importable from the SDK, you need to add it in the `__init__.py` file:
+```python
+from .my_custom_file import my_custom_function
+
+__all__ = ["my_custom_function"]
+```
+
+2. Add a new custom code region in the SDK files, e.g in `src/mistralai/chat.py`:
+```python
+# region imports
+from typing import Type
+from mistralai.extra import my_custom_function
+# endregion imports
+
+class Chat(BaseSDK):
+    r"""Chat Completion API."""
+
+    # region sdk-class-body
+    def my_custom_method(self, param: str) -> Type[some_type]:
+        output = my_custom_function(param1)
+        return output
+    # endregion sdk-class-body
+```
+
+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
+```
+
+4. And now you should be able to call the custom method:
+```python
+import os
+from mistralai import Mistral
+
+api_key = os.environ["MISTRAL_API_KEY"]
+client = Mistral(api_key=api_key)
+
+client.chat.my_custom_method(param="test")
+```
+
+### Run the unit tests
+
+To run the unit tests for the `extra` package, you can run the following command from the root of the repository:
+```bash
+python3.12 -m unittest discover -s src/mistralai/extra/tests -t src
+```

mistralai/extra/__init__.py

@@ -0,0 +1,5 @@
+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"]

mistralai/extra/struct_chat.py

@@ -0,0 +1,41 @@
+from ..models import ChatCompletionResponse, ChatCompletionChoice, AssistantMessage
+from .utils.response_format import  CustomPydanticModel, pydantic_model_from_json
+from typing import List, Optional, Type, Generic
+from pydantic import BaseModel
+import json
+
+class ParsedAssistantMessage(AssistantMessage, Generic[CustomPydanticModel]):
+    parsed: Optional[CustomPydanticModel]
+
+class ParsedChatCompletionChoice(ChatCompletionChoice, Generic[CustomPydanticModel]):
+    message: Optional[ParsedAssistantMessage[CustomPydanticModel]]  # type: ignore
+
+class ParsedChatCompletionResponse(ChatCompletionResponse, Generic[CustomPydanticModel]):
+    choices: Optional[List[ParsedChatCompletionChoice[CustomPydanticModel]]] # type: ignore
+
+def convert_to_parsed_chat_completion_response(response: ChatCompletionResponse, response_format: Type[BaseModel]) -> ParsedChatCompletionResponse:
+    parsed_choices = []
+
+    if response.choices:
+        for choice in response.choices:
+            if choice.message:
+                parsed_message: ParsedAssistantMessage = ParsedAssistantMessage(
+                    **choice.message.model_dump(),
+                    parsed=None
+                )
+                if isinstance(parsed_message.content, str):
+                    parsed_message.parsed = pydantic_model_from_json(json.loads(parsed_message.content), response_format)
+                elif parsed_message.content is None:
+                    parsed_message.parsed = None
+                else:
+                    raise TypeError(f"Unexpected type for message.content: {type(parsed_message.content)}")
+                choice_dict = choice.model_dump()
+                choice_dict["message"] = parsed_message
+                parsed_choice: ParsedChatCompletionChoice = ParsedChatCompletionChoice(**choice_dict)
+                parsed_choices.append(parsed_choice)
+            else:
+                parsed_choice = ParsedChatCompletionChoice(**choice.model_dump())
+                parsed_choices.append(parsed_choice)
+    response_dict = response.model_dump()
+    response_dict["choices"] = parsed_choices
+    return ParsedChatCompletionResponse(**response_dict)

mistralai/extra/tests/test_struct_chat.py

@@ -0,0 +1,103 @@
+import unittest
+from ..struct_chat import (
+    convert_to_parsed_chat_completion_response,
+    ParsedChatCompletionResponse,
+    ParsedChatCompletionChoice,
+    ParsedAssistantMessage,
+)
+from ...models import (
+    ChatCompletionResponse,
+    UsageInfo,
+    ChatCompletionChoice,
+    AssistantMessage,
+)
+from pydantic import BaseModel
+
+
+class Explanation(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathDemonstration(BaseModel):
+    steps: list[Explanation]
+    final_answer: str
+
+
+mock_cc_response = ChatCompletionResponse(
+    id="c0271b2098954c6094231703875ca0bc",
+    object="chat.completion",
+    model="mistral-large-latest",
+    usage=UsageInfo(prompt_tokens=75, completion_tokens=220, total_tokens=295),
+    created=1737727558,
+    choices=[
+        ChatCompletionChoice(
+            index=0,
+            message=AssistantMessage(
+                content='{\n  "final_answer": "x = -4",\n  "steps": [\n    {\n      "explanation": "Start with the given equation.",\n      "output": "8x + 7 = -23"\n    },\n    {\n      "explanation": "Subtract 7 from both sides to isolate the term with x.",\n      "output": "8x = -23 - 7"\n    },\n    {\n      "explanation": "Simplify the right side of the equation.",\n      "output": "8x = -30"\n    },\n    {\n      "explanation": "Divide both sides by 8 to solve for x.",\n      "output": "x = -30 / 8"\n    },\n    {\n      "explanation": "Simplify the fraction to get the final answer.",\n      "output": "x = -4"\n    }\n  ]\n}',
+                tool_calls=None,
+                prefix=False,
+                role="assistant",
+            ),
+            finish_reason="stop",
+        )
+    ],
+)
+
+
+expected_response = ParsedChatCompletionResponse(
+    choices=[
+        ParsedChatCompletionChoice(
+            index=0,
+            message=ParsedAssistantMessage(
+                content='{\n  "final_answer": "x = -4",\n  "steps": [\n    {\n      "explanation": "Start with the given equation.",\n      "output": "8x + 7 = -23"\n    },\n    {\n      "explanation": "Subtract 7 from both sides to isolate the term with x.",\n      "output": "8x = -23 - 7"\n    },\n    {\n      "explanation": "Simplify the right side of the equation.",\n      "output": "8x = -30"\n    },\n    {\n      "explanation": "Divide both sides by 8 to solve for x.",\n      "output": "x = -30 / 8"\n    },\n    {\n      "explanation": "Simplify the fraction to get the final answer.",\n      "output": "x = -4"\n    }\n  ]\n}',
+                tool_calls=None,
+                prefix=False,
+                role="assistant",
+                parsed=MathDemonstration(
+                    steps=[
+                        Explanation(
+                            explanation="Start with the given equation.",
+                            output="8x + 7 = -23",
+                        ),
+                        Explanation(
+                            explanation="Subtract 7 from both sides to isolate the term with x.",
+                            output="8x = -23 - 7",
+                        ),
+                        Explanation(
+                            explanation="Simplify the right side of the equation.",
+                            output="8x = -30",
+                        ),
+                        Explanation(
+                            explanation="Divide both sides by 8 to solve for x.",
+                            output="x = -30 / 8",
+                        ),
+                        Explanation(
+                            explanation="Simplify the fraction to get the final answer.",
+                            output="x = -4",
+                        ),
+                    ],
+                    final_answer="x = -4",
+                ),
+            ),
+            finish_reason="stop",
+        )
+    ],
+    created=1737727558,
+    id="c0271b2098954c6094231703875ca0bc",
+    model="mistral-large-latest",
+    object="chat.completion",
+    usage=UsageInfo(prompt_tokens=75, completion_tokens=220, total_tokens=295),
+)
+
+
+class TestConvertToParsedChatCompletionResponse(unittest.TestCase):
+    def test_convert_to_parsed_chat_completion_response(self):
+        output = convert_to_parsed_chat_completion_response(
+            mock_cc_response, MathDemonstration
+        )
+        self.assertEqual(output, expected_response)
+
+
+if __name__ == "__main__":
+    unittest.main()

mistralai/extra/tests/test_utils.py

@@ -0,0 +1,162 @@
+from ..utils.response_format import (
+    pydantic_model_from_json,
+    response_format_from_pydantic_model,
+    rec_strict_json_schema,
+)
+from pydantic import BaseModel, ValidationError
+
+from ...models import ResponseFormat, JSONSchema
+from ...types.basemodel import Unset
+
+import unittest
+
+
+class Student(BaseModel):
+    name: str
+    age: int
+
+
+class Explanation(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathDemonstration(BaseModel):
+    steps: list[Explanation]
+    final_answer: str
+
+
+mathdemo_schema = {
+    "$defs": {
+        "Explanation": {
+            "properties": {
+                "explanation": {"title": "Explanation", "type": "string"},
+                "output": {"title": "Output", "type": "string"},
+            },
+            "required": ["explanation", "output"],
+            "title": "Explanation",
+            "type": "object",
+        }
+    },
+    "properties": {
+        "steps": {
+            "items": {"$ref": "#/$defs/Explanation"},
+            "title": "Steps",
+            "type": "array",
+        },
+        "final_answer": {"title": "Final Answer", "type": "string"},
+    },
+    "required": ["steps", "final_answer"],
+    "title": "MathDemonstration",
+    "type": "object",
+}
+
+mathdemo_strict_schema = mathdemo_schema.copy()
+mathdemo_strict_schema["$defs"]["Explanation"]["additionalProperties"] = False # type: ignore
+mathdemo_strict_schema["additionalProperties"] = False
+
+mathdemo_response_format = ResponseFormat(
+    type="json_schema",
+    json_schema=JSONSchema(
+        name="MathDemonstration",
+        schema_definition=mathdemo_strict_schema,
+        description=Unset(),
+        strict=True,
+    ),
+)
+
+
+class TestResponseFormat(unittest.TestCase):
+    def test_pydantic_model_from_json(self):
+        missing_json_data = {"name": "Jean Dupont"}
+        good_json_data = {"name": "Jean Dupont", "age": 25}
+        extra_json_data = {
+            "name": "Jean Dupont",
+            "age": 25,
+            "extra_field": "extra_value",
+        }
+        complex_json_data = {
+            "final_answer": "x = -4",
+            "steps": [
+                {
+                    "explanation": "Start with the given equation.",
+                    "output": "8x + 7 = -23",
+                },
+                {
+                    "explanation": "Subtract 7 from both sides to isolate the term with x.",
+                    "output": "8x = -23 - 7",
+                },
+                {
+                    "explanation": "Simplify the right side of the equation.",
+                    "output": "8x = -30",
+                },
+                {
+                    "explanation": "Divide both sides by 8 to solve for x.",
+                    "output": "x = -30 / 8",
+                },
+                {
+                    "explanation": "Simplify the fraction to get the final answer.",
+                    "output": "x = -4",
+                },
+            ],
+        }
+
+        self.assertEqual(
+            pydantic_model_from_json(good_json_data, Student),
+            Student(name="Jean Dupont", age=25),
+        )
+        self.assertEqual(
+            pydantic_model_from_json(extra_json_data, Student),
+            Student(name="Jean Dupont", age=25),
+        )
+        self.assertEqual(
+            pydantic_model_from_json(complex_json_data, MathDemonstration),
+            MathDemonstration(
+                steps=[
+                    Explanation(
+                        explanation="Start with the given equation.",
+                        output="8x + 7 = -23",
+                    ),
+                    Explanation(
+                        explanation="Subtract 7 from both sides to isolate the term with x.",
+                        output="8x = -23 - 7",
+                    ),
+                    Explanation(
+                        explanation="Simplify the right side of the equation.",
+                        output="8x = -30",
+                    ),
+                    Explanation(
+                        explanation="Divide both sides by 8 to solve for x.",
+                        output="x = -30 / 8",
+                    ),
+                    Explanation(
+                        explanation="Simplify the fraction to get the final answer.",
+                        output="x = -4",
+                    ),
+                ],
+                final_answer="x = -4",
+            ),
+        )
+
+        # Check it raises a validation error
+        with self.assertRaises(ValidationError):
+            pydantic_model_from_json(missing_json_data, Student)  # type: ignore
+
+    def test_response_format_from_pydantic_model(self):
+        self.assertEqual(
+            response_format_from_pydantic_model(MathDemonstration),
+            mathdemo_response_format,
+        )
+
+    def test_rec_strict_json_schema(self):
+        invalid_schema = mathdemo_schema | {"wrong_value": 1}
+        self.assertEqual(
+            rec_strict_json_schema(mathdemo_schema), mathdemo_strict_schema
+        )
+
+        with self.assertRaises(ValueError):
+            rec_strict_json_schema(invalid_schema)
+
+
+if __name__ == "__main__":
+    unittest.main()

mistralai/extra/utils/__init__.py

@@ -0,0 +1,3 @@
+from .response_format import response_format_from_pydantic_model
+
+__all__ = ["response_format_from_pydantic_model"]

mistralai/extra/utils/_pydantic_helper.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+def rec_strict_json_schema(schema_node: Any) -> Any:
+    """
+    Recursively set the additionalProperties property to False for all objects in the JSON Schema.
+    This makes the JSON Schema strict (i.e. no additional properties are allowed).
+    """
+    if isinstance(schema_node, (str, bool)):
+        return schema_node
+    if isinstance(schema_node, dict):
+        if "type" in schema_node and schema_node["type"] == "object":
+            schema_node["additionalProperties"] = False
+        for key, value in schema_node.items():
+            schema_node[key] = rec_strict_json_schema(value)
+    elif isinstance(schema_node, list):
+        for i, value in enumerate(schema_node):
+            schema_node[i] = rec_strict_json_schema(value)
+    else:
+        raise ValueError(f"Unexpected type: {schema_node}")
+    return schema_node

mistralai/extra/utils/response_format.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel
+from typing import TypeVar, Any, Type
+from ...models import JSONSchema, ResponseFormat
+from ._pydantic_helper import rec_strict_json_schema
+
+CustomPydanticModel = TypeVar("CustomPydanticModel", bound=BaseModel)
+
+
+def response_format_from_pydantic_model(
+    model: type[CustomPydanticModel],
+) -> ResponseFormat:
+    """Generate a strict JSON schema from a pydantic model."""
+    model_schema = rec_strict_json_schema(model.model_json_schema())
+    json_schema = JSONSchema.model_validate(
+        {"name": model.__name__, "schema": model_schema, "strict": True}
+    )
+    return ResponseFormat(type="json_schema", json_schema=json_schema)
+
+
+def pydantic_model_from_json(
+    json_data: dict[str, Any], pydantic_model: Type[CustomPydanticModel]
+) -> CustomPydanticModel:
+    """Parse a JSON schema into a pydantic model."""
+    return pydantic_model.model_validate(json_data)