codex Merge pull request #21 from Krafman/j2k9s2-codex/add-support-for-mistral-api

Krafman · web-flow · commit 601ab48463dc · 2025-06-06T01:14:39.000+03:00
Add Mistral LLM provider
diff --git a/README.md b/README.md
@@ -106,11 +106,12 @@ python verify_connection.py
 A successful run will print "Successfully connected to device!" followed by a dictionary of your device's information.
 
 4. LLM Configuration
-API keys for LLM providers (OpenAI, Gemini, Together.ai) are configured in `src/llm_controller/llm_interface.py`.
+API keys for LLM providers (OpenAI, Gemini, Together.ai, Mistral) are configured in `src/llm_controller/llm_interface.py`.
 The system prioritizes loading API keys from environment variables:
 -   `OPENAI_API_KEY` for OpenAI
 -   `GEMINI_API_KEY` for Gemini
 -   `TOGETHER_API_KEY` for Together.ai
+-   `MISTRAL_API_KEY` for Mistral
 -   `ANTHROPIC_API_KEY` for Anthropic (if ever fully implemented)
 
 If environment variables are not set, it falls back to `config/llm_config.yml`.
diff --git a/config/llm_config.yml b/config/llm_config.yml
@@ -26,6 +26,25 @@ openai:
   # API request timeout in seconds.
   request_timeout: 60 # seconds
 
+# --- Mistral Specific Configuration ---
+mistral:
+  # API key for Mistral
+  # IMPORTANT: Load from an environment variable.
+  # The LLMInterface will check os.getenv("MISTRAL_API_KEY") first.
+  api_key: "YOUR_MISTRAL_API_KEY_PLACEHOLDER" # Fallback if env var not set.
+
+  # Default model, e.g., "open-mistral-7b", "open-mixtral-8x7b"
+  default_model: "open-mistral-7b"
+
+  # Default temperature for model responses.
+  temperature: 0.7
+
+  # Default maximum number of tokens to generate.
+  max_tokens: 1024
+
+  # API request timeout in seconds.
+  request_timeout: 60 # seconds
+
 # --- Together.ai Specific Configuration ---
 together:
   # API key for Together.ai
diff --git a/docs/api_references.md b/docs/api_references.md
@@ -226,6 +226,23 @@ This section will cover details for the LLM providers supported by the controlle
 
 *(Add sections for other LLM providers like Anthropic Claude if they are considered for integration.)*
 
+### 2.5. Mistral API
+
+* **Official Documentation Link**: `https://docs.mistral.ai`
+* **Authentication**: API Key passed via the `MISTRAL_API_KEY` environment variable or in `llm_config.yml`.
+* **Endpoint**: `POST https://api.mistral.ai/v1/chat/completions`
+* **Example Request Body**:
+    ```json
+    {
+      "model": "open-mistral-7b",
+      "messages": [
+        {"role": "user", "content": "Hello"}
+      ],
+      "response_format": {"type": "json_object"}
+    }
+    ```
+* **Notes**: The API closely mirrors OpenAI's chat completion format. The controller sends requests asynchronously with `httpx`.
+
 ## 3. Other Relevant APIs/Libraries
 
 * **openatx/uiautomator2**:
diff --git a/src/llm_controller/llm_interface.py b/src/llm_controller/llm_interface.py
@@ -13,7 +13,7 @@
 class LLMInterface:
     """
     Handles communication with the Large Language Model (LLM).
-    Supports Gemini, OpenAI, Together.ai, and AWS Bedrock providers.
+    Supports Gemini, OpenAI, Together.ai, Mistral, and AWS Bedrock providers.
     """
     def __init__(self, config_loader: ConfigLoader):
         """
@@ -103,6 +103,27 @@ def __init__(self, config_loader: ConfigLoader):
                 self.together_client = AsyncTogether(api_key=self.api_key)
             else:
                 self.together_client = None
+        elif self.provider == "mistral":
+            mistral_settings = llm_configs.get("mistral", {})
+            self.api_key = os.getenv("MISTRAL_API_KEY")
+            if self.api_key:
+                logger.info("Loaded Mistral API key from MISTRAL_API_KEY environment variable.")
+            else:
+                self.api_key = mistral_settings.get("api_key")
+                if self.api_key and self.api_key != "YOUR_MISTRAL_API_KEY_PLACEHOLDER":
+                    logger.info("Loaded Mistral API key from llm_config.yml.")
+                    logger.warning("For production, it's recommended to set the MISTRAL_API_KEY environment variable.")
+                elif self.api_key == "YOUR_MISTRAL_API_KEY_PLACEHOLDER" or not self.api_key:
+                    self.api_key = ""
+                    logger.warning("Using empty API key for Mistral. Ensure MISTRAL_API_KEY env var is set or llm_config.yml has a valid key.")
+            if not mistral_settings and not self.api_key:
+                raise ConfigError("Mistral configuration missing and MISTRAL_API_KEY environment variable not set.")
+
+            self.model_name = mistral_settings.get("default_model", "open-mistral-7b")
+            self.api_url = "https://api.mistral.ai/v1/chat/completions"
+            self.temperature = mistral_settings.get("temperature", 0.7)
+            self.max_tokens = mistral_settings.get("max_tokens", 1024)
+            self.mistral_settings = mistral_settings
         elif self.provider == "anthropic": # Example, not fully implemented
             anthropic_settings = llm_configs.get("anthropic", {})
             self.api_key = os.getenv("ANTHROPIC_API_KEY")
@@ -329,6 +350,55 @@ async def get_llm_action_json(self, messages: list[dict]) -> dict:
                 logger.error(f"Failed to parse Together LLM response as JSON: {e}")
                 logger.error(f"Together LLM response string was: {action_json_str}")
                 raise LLMInterfaceError(f"Together LLM response was not valid JSON: {action_json_str}")
+        elif self.provider == "mistral":
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "response_format": {"type": "json_object"}
+            }
+            if hasattr(self, 'temperature'):
+                payload["temperature"] = self.temperature
+            if hasattr(self, 'max_tokens'):
+                payload["max_tokens"] = self.max_tokens
+
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json"
+            }
+
+            logger.debug(f"Sending request to Mistral API: {self.api_url}")
+            logger.debug(f"Payload (first 200 chars of messages): {json.dumps(messages, indent=2)[:200]}")
+
+            try:
+                async with httpx.AsyncClient() as client:
+                    timeout = self.mistral_settings.get("request_timeout", 60.0)
+                    response = await client.post(self.api_url, headers=headers, json=payload, timeout=timeout)
+
+                response.raise_for_status()
+                result = response.json()
+                logger.debug(f"Raw Mistral API response: {json.dumps(result, indent=2)}")
+
+            except httpx.HTTPStatusError as e:
+                error_content = e.response.text
+                logger.error(f"Mistral API request failed with status {e.response.status_code}: {error_content}")
+                raise LLMInterfaceError(f"Mistral API request failed: {e.response.status_code} - {error_content}")
+            except Exception as e:
+                logger.error(f"Error during Mistral API call: {e}")
+                raise LLMInterfaceError(f"Error during Mistral API call: {e}")
+
+            if result.get("choices") and result["choices"][0].get("message") and result["choices"][0]["message"].get("content"):
+                action_json_str = result["choices"][0]["message"]["content"]
+                logger.info("Successfully received response from Mistral LLM.")
+                logger.debug(f"Mistral LLM response text (potential JSON): {action_json_str}")
+                try:
+                    return json.loads(action_json_str)
+                except json.JSONDecodeError as e:
+                    logger.error(f"Failed to parse Mistral LLM response as JSON: {e}")
+                    logger.error(f"Mistral LLM response string was: {action_json_str}")
+                    raise LLMInterfaceError(f"Mistral LLM response was not valid JSON: {action_json_str}")
+            else:
+                logger.error(f"Unexpected Mistral LLM response structure: {result}")
+                raise LLMInterfaceError(f"Unexpected Mistral LLM response structure. Full response: {result}")
         elif self.provider == "bedrock":
             logger.warning("Bedrock get_llm_action_json called, but using mock response for now.")
             mock_response_str = """
@@ -364,13 +434,16 @@ async def get_llm_action_json(self, messages: list[dict]) -> dict:
     if not dummy_llm_config_file.exists():
         with open(dummy_llm_config_file, "w") as f:
             f.write("""
-llm_provider: "gemini" # or "openai" or "bedrock"
+llm_provider: "gemini" # or "openai" or "mistral" or "bedrock"
 gemini:
   api_key: "YOUR_GEMINI_API_KEY_PLACEHOLDER"
   default_model: "gemini-2.0-flash"
 openai:
   api_key: "YOUR_OPENAI_API_KEY_PLACEHOLDER"
   default_model: "gpt-4o-mini"
+mistral:
+  api_key: "YOUR_MISTRAL_API_KEY_PLACEHOLDER"
+  default_model: "open-mistral-7b"
 bedrock:
   region_name: "us-east-1"
   profile_name: ""
diff --git a/tests/unit/test_llm_interface.py b/tests/unit/test_llm_interface.py
@@ -72,6 +72,16 @@ def setUp(self):
                 "request_timeout": 30.0
             }
         }
+        self.mistral_config_data = {
+            "llm_provider": "mistral",
+            "mistral": {
+                "api_key": "test_mistral_api_key",
+                "default_model": "open-mistral-7b",
+                "temperature": 0.5,
+                "max_tokens": 512,
+                "request_timeout": 30.0
+            }
+        }
         self.bedrock_config_data = {
             "llm_provider": "bedrock",
             "bedrock": {
@@ -273,6 +283,43 @@ def test_init_together_placeholder_api_key_becomes_empty(self):
         self.assertEqual(llm_interface.api_key, "")
         self.assertTrue(any("Using empty API key for Together.ai" in msg for msg in log_watcher.output))
 
+    @patch.dict(os.environ, {}, clear=True)
+    def test_init_mistral_api_key_from_env(self):
+        """Test Mistral API key loaded from environment variable."""
+        os.environ["MISTRAL_API_KEY"] = "env_mistral_key"
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+        self.assertEqual(llm_interface.provider, "mistral")
+        self.assertEqual(llm_interface.api_key, "env_mistral_key")
+        self.assertEqual(llm_interface.model_name, self.mistral_config_data["mistral"]["default_model"])
+
+    @patch.dict(os.environ, {}, clear=True)
+    def test_init_mistral_api_key_from_yaml_if_not_in_env(self):
+        """Test Mistral API key loaded from YAML when not in environment."""
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+        self.assertEqual(llm_interface.provider, "mistral")
+        self.assertEqual(llm_interface.api_key, "test_mistral_api_key")
+        self.assertEqual(llm_interface.model_name, self.mistral_config_data["mistral"]["default_model"])
+
+    @patch.dict(os.environ, {}, clear=True)
+    def test_init_mistral_config_section_missing_and_env_var_unset_raises_error(self):
+        """Test ConfigError if Mistral provider section is missing and env var not set."""
+        config_missing = {"llm_provider": "mistral"}
+        self._create_llm_config_file(config_missing)
+        with self.assertRaisesRegex(ConfigError, "Mistral configuration missing and MISTRAL_API_KEY environment variable not set."):
+            LLMInterface(config_loader=self.mock_config_loader)
+
+    @patch.dict(os.environ, {}, clear=True)
+    def test_init_mistral_placeholder_api_key_becomes_empty(self):
+        """Test placeholder Mistral API key becomes empty when env var not set."""
+        config_placeholder = {"llm_provider": "mistral", "mistral": {"api_key": "YOUR_MISTRAL_API_KEY_PLACEHOLDER", "default_model": "m"}}
+        self._create_llm_config_file(config_placeholder)
+        with self.assertLogs('llm_controller.llm_interface', level='WARNING') as log_watcher:
+            llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+        self.assertEqual(llm_interface.api_key, "")
+        self.assertTrue(any("Using empty API key for Mistral" in msg for msg in log_watcher.output))
+
     # --- End API Key Handling Tests ---
 
     # --- Gemini API Call Tests ---
@@ -519,6 +566,83 @@ async def test_get_llm_action_json_together_malformed_json(self, MockTogether):
             await llm_interface.get_llm_action_json([{"role": "user", "content": "tap"}])
     # --- End Together API Call Tests ---
 
+    # --- Mistral API Call Tests ---
+    @patch('llm_controller.llm_interface.httpx.AsyncClient')
+    async def test_get_llm_action_json_mistral_success(self, MockAsyncClient):
+        """Test successful Mistral API call and JSON response parsing."""
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+
+        mock_api_resp_structure = {
+            "choices": [
+                {"message": {"content": '{"action": "tap"}'}}
+            ]
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = mock_api_resp_structure
+
+        mock_client_instance = MockAsyncClient.return_value.__aenter__.return_value
+        mock_client_instance.post = AsyncMock(return_value=mock_response)
+
+        result_json = await llm_interface.get_llm_action_json([{"role": "user", "content": "tap"}])
+        self.assertEqual(result_json, {"action": "tap"})
+        mock_client_instance.post.assert_called_once()
+
+    @patch('llm_controller.llm_interface.httpx.AsyncClient')
+    async def test_get_llm_action_json_mistral_http_status_error(self, MockAsyncClient):
+        """Test Mistral API HTTPStatusError handling."""
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+
+        mock_http_response = MagicMock(spec=httpx.Response)
+        mock_http_response.status_code = 429
+        mock_http_response.text = "Rate limited"
+
+        mock_client_instance = MockAsyncClient.return_value.__aenter__.return_value
+        mock_client_instance.post = AsyncMock(side_effect=httpx.HTTPStatusError(
+            "Too Many Requests", request=MagicMock(), response=mock_http_response
+        ))
+
+        with self.assertRaisesRegex(LLMInterfaceError, "Mistral API request failed: 429 - Rate limited"):
+            await llm_interface.get_llm_action_json([{"role": "user", "content": "tap"}])
+
+    @patch('llm_controller.llm_interface.httpx.AsyncClient')
+    async def test_get_llm_action_json_mistral_request_error(self, MockAsyncClient):
+        """Test Mistral API request error handling."""
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+
+        mock_client_instance = MockAsyncClient.return_value.__aenter__.return_value
+        mock_client_instance.post = AsyncMock(side_effect=httpx.ConnectError("conn fail"))
+
+        with self.assertRaisesRegex(LLMInterfaceError, "Error during Mistral API call: conn fail"):
+            await llm_interface.get_llm_action_json([{"role": "user", "content": "tap"}])
+
+    @patch('llm_controller.llm_interface.httpx.AsyncClient')
+    async def test_get_llm_action_json_mistral_malformed_json(self, MockAsyncClient):
+        """Test Mistral API returning malformed JSON."""
+        self._create_llm_config_file(self.mistral_config_data)
+        llm_interface = LLMInterface(config_loader=self.mock_config_loader)
+
+        mock_api_resp_structure = {
+            "choices": [
+                {"message": {"content": 'not json'}}
+            ]
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = mock_api_resp_structure
+
+        mock_client_instance = MockAsyncClient.return_value.__aenter__.return_value
+        mock_client_instance.post = AsyncMock(return_value=mock_response)
+
+        with self.assertRaisesRegex(LLMInterfaceError, "Mistral LLM response was not valid JSON"):
+            await llm_interface.get_llm_action_json([{"role": "user", "content": "tap"}])
+    # --- End Mistral API Call Tests ---
+
     # --- Bedrock API Call Tests (current mock) ---
     async def test_get_llm_action_json_bedrock_mock_response(self):
         """Test Bedrock provider path (which currently returns a mock response)."""
