fix: update code for renamed model classes

- Filter UnknownAgentTool in agent update (context.py)
- Update examples for renamed classes:
  - BatchRequest, UserMessage: import from models module
  - ClassifierTrainingParametersIn -> ClassifierTrainingParameters
  - ClassifierJobOut -> ClassifierFineTuningJob
  - Add type narrowing for ClassifierFineTuningJobDetails

Author: louis-sanna-dev

examples/mistral/audio/chat_streaming.py

@@ -2,7 +2,8 @@
 
 import os
 
-from mistralai.client import Mistral, File
+from mistralai.client import Mistral
+from mistralai.client.models import File
 from mistralai.client.models import UserMessage
 
 

examples/mistral/audio/transcription_async.py

@@ -2,7 +2,8 @@
 
 import os
 import asyncio
-from mistralai.client import Mistral, File
+from mistralai.client import Mistral
+from mistralai.client.models import File
 
 
 async def main():

examples/mistral/audio/transcription_diarize_async.py

@@ -3,7 +3,8 @@
 import os
 import asyncio
 import pathlib
-from mistralai.client import Mistral, File
+from mistralai.client import Mistral
+from mistralai.client.models import File
 
 fixture_dir = pathlib.Path(__file__).parents[2] / "fixtures"
 

examples/mistral/audio/transcription_stream_async.py

@@ -2,7 +2,8 @@
 import asyncio
 import os
 
-from mistralai.client import Mistral, File
+from mistralai.client import Mistral
+from mistralai.client.models import File
 
 
 async def main():

examples/mistral/classifier/async_classifier.py

@@ -2,8 +2,8 @@
 
 from pprint import pprint
 import asyncio
-from mistralai.client import Mistral, TrainingFile, ClassifierTrainingParametersIn
-from mistralai.client.models import ClassifierJobOut
+from mistralai.client import Mistral
+from mistralai.client.models import ClassifierFineTuningJob, ClassifierFineTuningJobDetails, ClassifierTrainingParameters, TrainingFile
 
 import os
 
@@ -36,12 +36,12 @@ async def train_classifier(client: Mistral, training_file_ids: list[str]) -> str
             TrainingFile(file_id=training_file_id)
             for training_file_id in training_file_ids
         ],
-        hyperparameters=ClassifierTrainingParametersIn(
+        hyperparameters=ClassifierTrainingParameters(
             learning_rate=0.0001,
         ),
         auto_start=True,
     )
-    if not isinstance(job, ClassifierJobOut):
+    if not isinstance(job, ClassifierFineTuningJob):
         print("Unexpected job type returned")
         return None
 
@@ -51,6 +51,8 @@ async def train_classifier(client: Mistral, training_file_ids: list[str]) -> str
     while True:
         await asyncio.sleep(10)
         detailed_job = await client.fine_tuning.jobs.get_async(job_id=job.id)
+        if not isinstance(detailed_job, ClassifierFineTuningJobDetails):
+            raise Exception(f"Unexpected job type: {type(detailed_job)}")
         if detailed_job.status not in [
             "QUEUED",
             "STARTED",

examples/mistral/jobs/async_batch_job_chat_completion_inline.py

@@ -1,4 +1,5 @@
-from mistralai.client import Mistral, BatchRequest, UserMessage
+from mistralai.client import Mistral
+from mistralai.client.models import BatchRequest, UserMessage
 import os
 import asyncio
 

src/mistralai/extra/run/context.py

@@ -22,7 +22,6 @@
     create_tool_call,
 )
 from mistralai.client.models import (
-    AgentTool,
     CompletionArgs,
     CompletionArgsTypedDict,
     ConversationInputs,
@@ -35,6 +34,8 @@
     InputEntries,
     MessageInputEntry,
     ResponseFormat,
+    UnknownAgentTool,
+    UpdateAgentRequestTool,
 )
 from mistralai.client.types.basemodel import BaseModel, OptionalNullable, UNSET
 
@@ -187,8 +188,11 @@ async def prepare_agent_request(self, beta_client: "Beta") -> AgentRequestKwargs
             )
         agent = await beta_client.agents.get_async(agent_id=self.agent_id)
         agent_tools = agent.tools or []
-        updated_tools: list[AgentTool] = []
+        updated_tools: list[UpdateAgentRequestTool] = []
         for tool in agent_tools:
+            if isinstance(tool, UnknownAgentTool):
+                # Skip unknown tools - can't include them in update request
+                continue
             if not isinstance(tool, FunctionTool):
                 updated_tools.append(tool)
             elif tool.function.name in self._callable_tools: