Tentative Fix for Mistral Fix

instructor/client_mistral.py

@@ -1,45 +1,48 @@
 # Future imports to ensure compatibility with Python 3.9
 from __future__ import annotations
 
-import mistralai.client
-import mistralai.async_client as mistralaiasynccli
+
+from mistralai import Mistral
 import instructor
-from typing import overload, Any
+from typing import overload, Any, Literal
 
 
 @overload
 def from_mistral(
-    client: mistralai.client.MistralClient,
+    client: Mistral,
     mode: instructor.Mode = instructor.Mode.MISTRAL_TOOLS,
+    use_async: Literal[False] = False,
     **kwargs: Any,
 ) -> instructor.Instructor: ...
 
 
 @overload
 def from_mistral(
-    client: mistralaiasynccli.MistralAsyncClient,
+    client: Mistral,
     mode: instructor.Mode = instructor.Mode.MISTRAL_TOOLS,
+    use_async: Literal[True] = True,
     **kwargs: Any,
 ) -> instructor.AsyncInstructor: ...
 
 
 def from_mistral(
-    client: mistralai.client.MistralClient | mistralaiasynccli.MistralAsyncClient,
+    client: Mistral,
     mode: instructor.Mode = instructor.Mode.MISTRAL_TOOLS,
+    use_async: bool = False,
     **kwargs: Any,
 ) -> instructor.Instructor | instructor.AsyncInstructor:
     assert mode in {
         instructor.Mode.MISTRAL_TOOLS,
     }, "Mode be one of {instructor.Mode.MISTRAL_TOOLS}"
 
     assert isinstance(
-        client, (mistralai.client.MistralClient, mistralaiasynccli.MistralAsyncClient)
-    ), "Client must be an instance of mistralai.client.MistralClient or mistralai.async_cli.MistralAsyncClient"
+        client, (Mistral)
+    ), "Client must be an instance of mistralai.Mistral"
 
-    if isinstance(client, mistralai.client.MistralClient):
+    if not use_async:
         return instructor.Instructor(
             client=client,
-            create=instructor.patch(create=client.chat, mode=mode),
+            create=instructor.patch(create=client.chat.complete, mode=mode),
             provider=instructor.Provider.MISTRAL,
             mode=mode,
             **kwargs,
@@ -48,7 +51,7 @@
     else:
         return instructor.AsyncInstructor(
             client=client,
-            create=instructor.patch(create=client.chat, mode=mode),
+            create=instructor.patch(create=client.chat.complete_async, mode=mode),
             provider=instructor.Provider.MISTRAL,
             mode=mode,
             **kwargs,