Add Mistral Provider Support (#11)

Author: rajkstats

README.md

@@ -47,6 +47,7 @@ The providers are imported from [providers.py](/os_computer_use/providers.py) an
 - HuggingFace Spaces:
   - OS-Atlas (grounding)
   - ShowUI (grounding)
+- Mistral AI (Pixtral for vision, Mistral Large for actions)
 
 If you add a new model or provider, please [make a PR](../../pulls) to this repository with the updated providers.py!
 
@@ -101,6 +102,8 @@ GROQ_API_KEY=...
 GEMINI_API_KEY=...
 OPENAI_API_KEY=...
 ANTHROPIC_API_KEY=...
+# Required: Provide your Hugging Face token to bypass Gradio rate limits.
+HF_TOKEN=...
 ```
 
 ### 4. Start the web interface

os_computer_use/config.py

@@ -5,12 +5,15 @@
 grounding_model = providers.OSAtlasProvider()
 # grounding_model = providers.ShowUIProvider()
 
-# vision_model = providers.FireworksProvider("llama3.2")
+#vision_model = providers.FireworksProvider("llama3.2")
 # vision_model = providers.OpenAIProvider("gpt-4o")
 # vision_model = providers.AnthropicProvider("claude-3.5-sonnet")
 vision_model = providers.GroqProvider("llama3.2")
+#vision_model = providers.MistralProvider("pixtral")  # pixtral-large-latest has vision capabilities
 
-# action_model = providers.FireworksProvider("llama3.3")
+
+#action_model = providers.FireworksProvider("llama3.3")
 # action_model = providers.OpenAIProvider("gpt-4o")
 # action_model = providers.AnthropicProvider("claude-3.5-sonnet")
 action_model = providers.GroqProvider("llama3.3")
+#action_model = providers.MistralProvider("large")  # mistral-large-latest for non-vision tasks

os_computer_use/llm_provider.py

@@ -116,7 +116,6 @@ def create_image_block(self, base64_image):
         }
 
     def call(self, messages, functions=None):
-
         # If functions are provided, only return actions
         tools = self.create_function_schema(functions) if functions else None
         completion = self.completion(messages, tools=tools)
@@ -204,3 +203,20 @@ def call(self, messages, functions=None):
         # Only return response text
         else:
             return text
+
+
+class MistralBaseProvider(OpenAIBaseProvider):
+    def create_function_def(self, name, details, properties, required):
+        # If description is wrapped in a dict, extract the inner string
+        if isinstance(details.get("description"), dict):
+            details["description"] = details["description"].get("description", "")
+        return super().create_function_def(name, details, properties, required)
+    
+    def call(self, messages, functions=None):
+        if messages and messages[-1].get("role") == "assistant":
+            prefix = messages.pop()["content"]
+            if messages and messages[-1].get("role") == "user":
+                messages[-1]["content"] = prefix + "\n" + messages[-1].get("content", "")
+            else:
+                messages.append({"role": "user", "content": prefix})
+        return super().call(messages, functions)

os_computer_use/osatlas_provider.py

@@ -1,20 +1,24 @@
 from gradio_client import Client, handle_file
 from os_computer_use.logging import logger
-
 from os_computer_use.grounding import extract_bbox_midpoint
 
+import os
+
+
 OSATLAS_HUGGINGFACE_SOURCE = "maxiw/OS-ATLAS"
 OSATLAS_HUGGINGFACE_MODEL = "OS-Copilot/OS-Atlas-Base-7B"
 OSATLAS_HUGGINGFACE_API = "/run_example"
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 
 class OSAtlasProvider:
     """
     The OS-Atlas provider is used to make calls to OS-Atlas.
     """
 
     def __init__(self):
-        self.client = Client(OSATLAS_HUGGINGFACE_SOURCE)
+        self.client = Client(OSATLAS_HUGGINGFACE_SOURCE,  hf_token=HF_TOKEN)
 
     def call(self, prompt, image_data):
         result = self.client.predict(

os_computer_use/providers.py

@@ -1,6 +1,6 @@
 import os
 from dotenv import load_dotenv
-from os_computer_use.llm_provider import OpenAIBaseProvider, AnthropicBaseProvider
+from os_computer_use.llm_provider import OpenAIBaseProvider, AnthropicBaseProvider, MistralBaseProvider
 from os_computer_use.osatlas_provider import OSAtlasProvider
 from os_computer_use.showui_provider import ShowUIProvider
 
@@ -62,3 +62,13 @@ class GroqProvider(OpenAIBaseProvider):
         "llama3.2": "llama-3.2-90b-vision-preview",
         "llama3.3": "llama-3.3-70b-versatile",
     }
+
+class MistralProvider(MistralBaseProvider):
+    base_url = "https://api.mistral.ai/v1"
+    api_key = os.getenv("MISTRAL_API_KEY")
+    aliases = {
+        "small": "mistral-small-latest",
+        "medium": "mistral-medium-latest",
+        "large": "mistral-large-latest",
+        "pixtral": "pixtral-large-latest"
+    }

os_computer_use/sandbox_agent.py

@@ -219,4 +219,4 @@ def run(self, instruction):
 
                 self.messages.append(
                     Message(logger.log(f"OBSERVATION: {result}", "yellow"))
-                )
+                )

tests/llm_provider.py

@@ -3,6 +3,7 @@
     OpenAIProvider,
     GroqProvider,
     FireworksProvider,
+    MistralProvider,
 )
 from os_computer_use.llm_provider import Message
 
@@ -15,6 +16,7 @@
 }
 
 
+
 # Function to simulate taking a screenshot
 def take_screenshot():
     with open("./tests/test_screenshot.png", "rb") as f:
@@ -62,3 +64,23 @@ def take_screenshot():
 fireworks = FireworksProvider("llama3.2")
 print(fireworks.call(toolcall_messages, tools)[1])
 print(fireworks.call(messages))
+
+
+
+# Pixtral
+mistral = MistralProvider("pixtral") 
+print("\nTesting Mistral :")
+print(mistral.call(toolcall_messages, tools)[1])
+print(mistral.call(messages))
+
+
+# Mistral Large (non-vision) using text-only messages
+mistral_large = MistralProvider("large")  # Using mistral-large-latest for non-vision tasks
+text_messages = [Message("What is the capital of France?", role="user")]
+print("\nTesting Mistral Large with text-only:")
+print(mistral_large.call(text_messages))
+
+# Test tool calls for Mistral Large using text-only messages (no image data)
+text_tool_messages = [Message("Click on the submit button", role="user")]
+print("\nTesting Mistral Large Tool Calls with text:")
+print(mistral_large.call(text_tool_messages, tools)[1])