fix: fix example types and enable mypy in CI

- Rename azure examples from .py.py to .py
- Fix message types in azure and mistral examples
- Add type annotations where needed for mypy
- Enable mypy for examples in lint_custom_code.sh

Author: louis-sanna-dev

examples/azure/az_chat_no_streaming.py

@@ -0,0 +1,15 @@
+import os
+
+from mistralai_azure import MistralAzure
+from mistralai_azure.models import ChatCompletionRequestMessages, UserMessage
+
+client = MistralAzure(
+    azure_api_key=os.environ["AZURE_API_KEY"],
+    azure_endpoint=os.environ["AZURE_ENDPOINT"],
+)
+
+messages: list[ChatCompletionRequestMessages] = [
+    UserMessage(content="What is the capital of France?"),
+]
+res = client.chat.complete(messages=messages)
+print(res.choices[0].message.content)

examples/azure/az_chat_no_streaming.py.py

@@ -0,0 +1,15 @@
+import os
+
+from mistralai_azure import MistralAzure
+from mistralai_azure.models import ChatCompletionRequestMessages, UserMessage
+
+client = MistralAzure(
+    azure_api_key=os.environ["AZURE_API_KEY"],
+    azure_endpoint=os.environ["AZURE_ENDPOINT"],
+)
+
+messages: list[ChatCompletionRequestMessages] = [
+    UserMessage(content="What is the capital of France?"),
+]
+res = client.chat.complete(messages=messages)
+print(res.choices[0].message.content)

examples/azure/chat_no_streaming.py

@@ -7,6 +7,7 @@
 import os
 import readline
 import sys
+from typing import Any
 
 from mistralai.client import Mistral
 from mistralai.client.models import AssistantMessage, SystemMessage, UserMessage
@@ -21,7 +22,7 @@
 DEFAULT_TEMPERATURE = 0.7
 LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
 # A dictionary of all commands and their arguments, used for tab completion.
-COMMAND_LIST = {
+COMMAND_LIST: dict[str, Any] = {
     "/new": {},
     "/help": {},
     "/model": {model: {} for model in MODEL_LIST},  # Nested completions for models

examples/azure/chat_no_streaming.py.py

@@ -1,16 +1,20 @@
 import functools
 import json
 import os
-from typing import Dict, List
+from typing import Any
 
 from mistralai.client import Mistral
-from mistralai.client.models.assistantmessage import AssistantMessage
-from mistralai.client.models.function import Function
-from mistralai.client.models.toolmessage import ToolMessage
-from mistralai.client.models.usermessage import UserMessage
+from mistralai.client.models import (
+    AssistantMessage,
+    ChatCompletionRequestMessage,
+    Function,
+    Tool,
+    ToolMessage,
+    UserMessage,
+)
 
 # Assuming we have the following data
-data = {
+data: dict[str, list[Any]] = {
     "transaction_id": ["T1001", "T1002", "T1003", "T1004", "T1005"],
     "customer_id": ["C001", "C002", "C003", "C002", "C001"],
     "payment_amount": [125.50, 89.99, 120.00, 54.30, 210.20],
@@ -25,31 +29,28 @@
 }
 
 
-def retrieve_payment_status(data: Dict[str, List], transaction_id: str) -> str:
+def retrieve_payment_status(data: dict[str, list[Any]], transaction_id: str) -> str:
     for i, r in enumerate(data["transaction_id"]):
         if r == transaction_id:
             return json.dumps({"status": data["payment_status"][i]})
-        else:
-            return json.dumps({"status": "Error - transaction id not found"})
+    return json.dumps({"status": "Error - transaction id not found"})
 
 
-def retrieve_payment_date(data: Dict[str, List], transaction_id: str) -> str:
+def retrieve_payment_date(data: dict[str, list[Any]], transaction_id: str) -> str:
     for i, r in enumerate(data["transaction_id"]):
         if r == transaction_id:
             return json.dumps({"date": data["payment_date"][i]})
-        else:
-            return json.dumps({"status": "Error - transaction id not found"})
+    return json.dumps({"status": "Error - transaction id not found"})
 
 
 names_to_functions = {
     "retrieve_payment_status": functools.partial(retrieve_payment_status, data=data),
     "retrieve_payment_date": functools.partial(retrieve_payment_date, data=data),
 }
 
-tools = [
-    {
-        "type": "function",
-        "function": Function(
+tools: list[Tool] = [
+    Tool(
+        function=Function(
             name="retrieve_payment_status",
             description="Get payment status of a transaction id",
             parameters={
@@ -63,10 +64,9 @@ def retrieve_payment_date(data: Dict[str, List], transaction_id: str) -> str:
                 },
             },
         ),
-    },
-    {
-        "type": "function",
-        "function": Function(
+    ),
+    Tool(
+        function=Function(
             name="retrieve_payment_date",
             description="Get payment date of a transaction id",
             parameters={
@@ -80,36 +80,35 @@ def retrieve_payment_date(data: Dict[str, List], transaction_id: str) -> str:
                 },
             },
         ),
-    },
+    ),
 ]
 
 api_key = os.environ["MISTRAL_API_KEY"]
 model = "mistral-small-latest"
 
 client = Mistral(api_key=api_key)
 
-messages = [UserMessage(content="What's the status of my transaction?")]
+messages: list[ChatCompletionRequestMessage] = [
+    UserMessage(content="What's the status of my transaction?")
+]
 
-response = client.chat.complete(
-    model=model, messages=messages, tools=tools, temperature=0
-)
+response = client.chat.complete(model=model, messages=messages, tools=tools, temperature=0)
 
 print(response.choices[0].message.content)
 
 messages.append(AssistantMessage(content=response.choices[0].message.content))
 messages.append(UserMessage(content="My transaction ID is T1001."))
 
-response = client.chat.complete(
-    model=model, messages=messages, tools=tools, temperature=0
-)
+response = client.chat.complete(model=model, messages=messages, tools=tools, temperature=0)
 
-tool_call = response.choices[0].message.tool_calls[0]
+tool_calls = response.choices[0].message.tool_calls
+if not tool_calls:
+    raise RuntimeError("Expected tool calls")
+tool_call = tool_calls[0]
 function_name = tool_call.function.name
-function_params = json.loads(tool_call.function.arguments)
+function_params = json.loads(str(tool_call.function.arguments))
 
-print(
-    f"calling function_name: {function_name}, with function_params: {function_params}"
-)
+print(f"calling function_name: {function_name}, with function_params: {function_params}")
 
 function_result = names_to_functions[function_name](**function_params)
 
@@ -128,8 +127,6 @@ def retrieve_payment_date(data: Dict[str, List], transaction_id: str) -> str:
 )
 print(messages)
 
-response = client.chat.complete(
-    model=model, messages=messages, tools=tools, temperature=0
-)
+response = client.chat.complete(model=model, messages=messages, tools=tools, temperature=0)
 
 print(f"{response.choices[0].message.content}")

examples/mistral/chat/chatbot_with_streaming.py

@@ -3,6 +3,7 @@
 from pprint import pprint
 import asyncio
 from mistralai.client import Mistral, TrainingFile, ClassifierTrainingParametersIn
+from mistralai.client.models import ClassifierJobOut
 
 import os
 
@@ -26,7 +27,7 @@ async def upload_files(client: Mistral, file_names: list[str]) -> list[str]:
     return file_ids
 
 
-async def train_classifier(client: Mistral,training_file_ids: list[str]) -> str:
+async def train_classifier(client: Mistral, training_file_ids: list[str]) -> str | None:
     print("Creating job...")
     job = await client.fine_tuning.jobs.create_async(
         model="ministral-3b-latest",
@@ -40,6 +41,9 @@ async def train_classifier(client: Mistral,training_file_ids: list[str]) -> str:
         ),
         auto_start=True,
     )
+    if not isinstance(job, ClassifierJobOut):
+        print("Unexpected job type returned")
+        return None
 
     print(f"Job created ({job.id})")
 
@@ -62,6 +66,9 @@ async def train_classifier(client: Mistral,training_file_ids: list[str]) -> str:
         print("Training failed")
         raise Exception(f"Job failed {detailed_job.status}")
 
+    if not detailed_job.fine_tuned_model:
+        print("No fine-tuned model returned")
+        return None
     print(f"Training succeed: {detailed_job.fine_tuned_model}")
 
     return detailed_job.fine_tuned_model

examples/mistral/chat/function_calling.py

@@ -11,8 +11,8 @@ else
 fi
 
 echo "Running mypy..."
-# TODO: Uncomment once the examples are fixed
-# uv run mypy examples/ || ERRORS=1
+echo "-> running on examples"
+uv run mypy examples/ --ignore-missing-imports || ERRORS=1
 echo "-> running on extra"
 uv run mypy src/mistralai/extra/ || ERRORS=1
 echo "-> running on hooks"