Added Kimi K2 provider tests Fixes #40

LLM/interpreter.py

@@ -1,16 +1,24 @@
+"""Natural language to shell command interpreter backed by multiple LLMs."""
+
 import os
 import json
 from typing import List, Optional, Dict, Any
 from enum import Enum
 
 
 class APIProvider(Enum):
+    """Supported large-language-model providers for command generation."""
+
     CLAUDE = "claude"
     OPENAI = "openai"
+    KIMI = "kimi"
+    FAKE = "fake"
     OLLAMA = "ollama"
 
 
 class CommandInterpreter:
+    """Translate natural language intents into shell commands via LLMs."""
+
     def __init__(
         self,
         api_key: str,
@@ -33,6 +41,7 @@ def __init__(
         self._initialize_client()
 
     def _initialize_client(self):
+        """Instantiate the SDK client for the selected provider."""
         if self.provider == APIProvider.OPENAI:
             try:
                 from openai import OpenAI
@@ -45,12 +54,24 @@ def _initialize_client(self):
                 self.client = Anthropic(api_key=self.api_key)
             except ImportError:
                 raise ImportError("Anthropic package not installed. Run: pip install anthropic")
+        elif self.provider == APIProvider.KIMI:
+            try:
+                import requests  # type: ignore
+            except ImportError as exc:
+                raise ImportError("Requests package not installed. Run: pip install requests") from exc
+
+            self.client = requests
+            self._kimi_base_url = os.environ.get("KIMI_API_BASE_URL", "https://api.moonshot.ai")
+        elif self.provider == APIProvider.FAKE:
+            # Fake provider is used for deterministic offline or integration tests.
+            self.client = None
         elif self.provider == APIProvider.OLLAMA:
             # Ollama uses local HTTP API, no special client needed
             self.ollama_url = os.environ.get('OLLAMA_HOST', 'http://localhost:11434')
             self.client = None  # Will use requests
 
     def _get_system_prompt(self) -> str:
+        """Return the base instructions shared across all provider calls."""
         return """You are a Linux system command expert. Convert natural language requests into safe, validated bash commands.
 
 Rules:
@@ -69,6 +90,7 @@ def _get_system_prompt(self) -> str:
 Example response: {"commands": ["sudo apt update", "sudo apt install -y docker.io", "sudo apt install -y nvidia-docker2", "sudo systemctl restart docker"]}"""
 
     def _call_openai(self, user_input: str) -> List[str]:
+        """Call the OpenAI Chat Completions API and parse the response."""
         try:
             response = self.client.chat.completions.create(
                 model=self.model,
@@ -86,6 +108,7 @@ def _call_openai(self, user_input: str) -> List[str]:
             raise RuntimeError(f"OpenAI API call failed: {str(e)}")
 
     def _call_claude(self, user_input: str) -> List[str]:
+        """Call the Anthropic Messages API and parse the response."""
         try:
             response = self.client.messages.create(
                 model=self.model,
@@ -102,6 +125,75 @@ def _call_claude(self, user_input: str) -> List[str]:
         except Exception as e:
             raise RuntimeError(f"Claude API call failed: {str(e)}")
 
+    def _call_kimi(self, user_input: str) -> List[str]:
+        """Call the Kimi K2 HTTP API and parse the response body."""
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": self._get_system_prompt()},
+                {"role": "user", "content": user_input},
+            ],
+            "temperature": 0.3,
+            "max_tokens": 1000,
+        }
+
+        try:
+            import requests
+            response = requests.post(
+                f"{self._kimi_base_url.rstrip('/')}/v1/chat/completions",
+                headers=headers,
+                json=payload,
+                timeout=60,
+            )
+            response.raise_for_status()
+            data = response.json()
+            choices = data.get("choices", [])
+            if not choices:
+                raise RuntimeError("Kimi API returned no choices")
+            content = choices[0].get("message", {}).get("content", "").strip()
+            if not content:
+                raise RuntimeError("Kimi API returned empty content")
+            return self._parse_commands(content)
+        except Exception as exc:
+            raise RuntimeError(f"Kimi API call failed: {str(exc)}") from exc
+
+    def _call_fake(self, user_input: str) -> List[str]:
+        """Return predetermined commands without hitting a real provider."""
+
+        payload = os.environ.get("CORTEX_FAKE_COMMANDS")
+        if payload:
+            try:
+                data = json.loads(payload)
+            except json.JSONDecodeError as exc:
+                raise ValueError("CORTEX_FAKE_COMMANDS must contain valid JSON") from exc
+            if not isinstance(data["commands"], list):  
+                raise ValueError("'commands' must be a list in CORTEX_FAKE_COMMANDS")  
+            return data["commands"]
+
+        safe_defaults = {
+            "docker": [
+                "echo Updating package cache",
+                "echo Installing docker packages",
+                "echo Enabling docker service",
+            ],
+            "python": [
+                "echo Installing Python",
+                "echo Setting up virtual environment",
+                "echo Installing pip packages",
+            ],
+        }
+
+        for key, commands in safe_defaults.items():
+            if key in user_input.lower():
+                return commands
+
+        return ["echo Preparing environment", "echo Completed simulation"]
+
     def _call_ollama(self, user_input: str) -> List[str]:
         """Call local Ollama instance for offline/local inference"""
         import urllib.request
@@ -137,6 +229,7 @@ def _call_ollama(self, user_input: str) -> List[str]:
             raise RuntimeError(f"Ollama API call failed: {str(e)}")
 
     def _parse_commands(self, content: str) -> List[str]:
+        """Parse the JSON payload returned by an LLM into command strings."""
         try:
             if content.startswith("```json"):
                 content = content.split("```json")[1].split("```")[0].strip()
@@ -154,6 +247,7 @@ def _parse_commands(self, content: str) -> List[str]:
             raise ValueError(f"Failed to parse LLM response: {str(e)}")
 
     def _validate_commands(self, commands: List[str]) -> List[str]:
+        """Filter the provided commands to remove obviously dangerous patterns."""
         dangerous_patterns = [
             "rm -rf /",
             "dd if=",
@@ -173,13 +267,18 @@ def _validate_commands(self, commands: List[str]) -> List[str]:
         return validated
 
     def parse(self, user_input: str, validate: bool = True) -> List[str]:
+        """Parse the user's request into a list of shell commands."""
         if not user_input or not user_input.strip():
             raise ValueError("User input cannot be empty")
 
         if self.provider == APIProvider.OPENAI:
             commands = self._call_openai(user_input)
         elif self.provider == APIProvider.CLAUDE:
             commands = self._call_claude(user_input)
+        elif self.provider == APIProvider.KIMI:
+            commands = self._call_kimi(user_input)
+        elif self.provider == APIProvider.FAKE:
+            commands = self._call_fake(user_input)
         elif self.provider == APIProvider.OLLAMA:
             commands = self._call_ollama(user_input)
         else:
@@ -196,9 +295,21 @@ def parse_with_context(
         system_info: Optional[Dict[str, Any]] = None,
         validate: bool = True
     ) -> List[str]:
+        """Parse a request while appending structured system context."""
         context = ""
         if system_info:
             context = f"\n\nSystem context: {json.dumps(system_info)}"
 
         enriched_input = user_input + context
         return self.parse(enriched_input, validate=validate)
+
+    def _default_model(self) -> str:
+        """Return the default model identifier for the active provider."""
+
+        if self.provider == APIProvider.OPENAI:
+            return "gpt-4"
+        if self.provider == APIProvider.CLAUDE:
+            return "claude-3-5-sonnet-20241022"
+        if self.provider == APIProvider.KIMI:
+            return os.environ.get("KIMI_DEFAULT_MODEL", "kimi-k2-turbo-preview")
+        return "fake-local-model"

LLM/requirements.txt

@@ -1,3 +1,4 @@
 openai>=1.0.0
 anthropic>=0.18.0
 PyYAML>=6.0
+requests>=2.32.4

LLM/test_interpreter.py

@@ -1,8 +1,9 @@
-import unittest
-from unittest.mock import Mock, patch, MagicMock
 import json
-import sys
 import os
+import sys
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
@@ -13,12 +14,20 @@ class TestCommandInterpreter(unittest.TestCase):
 
     def setUp(self):
         self.api_key = "test-api-key"
+        openai_stub = SimpleNamespace(OpenAI=Mock())
+        anthropic_stub = SimpleNamespace(Anthropic=Mock())
+        self.sys_modules_patcher = patch.dict(sys.modules, {
+            'openai': openai_stub,
+            'anthropic': anthropic_stub,
+        })
+        self.sys_modules_patcher.start()
+        self.addCleanup(self.sys_modules_patcher.stop)
 
     @patch('openai.OpenAI')
     def test_initialization_openai(self, mock_openai):
         interpreter = CommandInterpreter(api_key=self.api_key, provider="openai")
         self.assertEqual(interpreter.provider, APIProvider.OPENAI)
-        self.assertEqual(interpreter.model, "gpt-4")
+        self.assertEqual(interpreter.model, "gpt-4o")
         mock_openai.assert_called_once_with(api_key=self.api_key)
 
     @patch('anthropic.Anthropic')
@@ -37,6 +46,43 @@ def test_initialization_custom_model(self, mock_openai):
         )
         self.assertEqual(interpreter.model, "gpt-4-turbo")
 
+    @patch.dict(os.environ, {}, clear=True)
+    @patch.dict(sys.modules, {'requests': Mock()})
+    def test_initialization_kimi(self):
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+        self.assertEqual(interpreter.provider, APIProvider.KIMI)
+        self.assertEqual(interpreter.model, "kimi-k2-turbo-preview")
+
+    @patch('requests.post')
+    def test_call_kimi_success(self, mock_post):
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "choices": [{"message": {"content": '{"commands": ["apt update", "apt install curl"]}'}}]
+        }
+        mock_post.return_value = mock_response
+
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+        result = interpreter._call_kimi("install curl")
+
+        self.assertEqual(result, ["apt update", "apt install curl"])
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertIn("Authorization", call_args[1]["headers"])
+        self.assertEqual(call_args[1]["headers"]["Authorization"], f"Bearer {self.api_key}")
+
+    @patch('requests.post')
+    def test_call_kimi_failure(self, mock_post):
+        mock_response = Mock()
+        mock_response.status_code = 401
+        mock_response.raise_for_status.side_effect = Exception("401 Unauthorized")
+        mock_post.return_value = mock_response
+
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+
+        with self.assertRaises(RuntimeError):
+            interpreter._call_kimi("install docker")
+
     def test_parse_commands_valid_json(self):
         interpreter = CommandInterpreter.__new__(CommandInterpreter)
 
@@ -109,7 +155,7 @@ def test_call_openai_failure(self, mock_openai):
 
         with self.assertRaises(RuntimeError):
             interpreter._call_openai("install docker")
-    
+
     @patch('anthropic.Anthropic')
     def test_call_claude_success(self, mock_anthropic):
         mock_client = Mock()
@@ -225,5 +271,74 @@ def test_parse_docker_installation(self, mock_openai):
         self.assertTrue(any("docker" in cmd.lower() for cmd in result))
 
 
+@unittest.skipUnless(
+    os.environ.get('RUN_KIMI_INTEGRATION_TESTS') == '1',
+    "Skipping Kimi K2 integration tests. Set RUN_KIMI_INTEGRATION_TESTS=1 to run them."
+)
+class TestKimiK2Integration(unittest.TestCase):
+    """Integration tests for Kimi K2 API with real API calls
+
+    To run these tests:
+    - Set environment variable: RUN_KIMI_INTEGRATION_TESTS=1
+    - Set environment variable: KIMI_API_KEY=your-api-key
+    - Run: python -m unittest LLM.test_interpreter.TestKimiK2Integration -v
+    """
+
+    def setUp(self):
+        # Use the actual API key from environment
+        self.api_key = os.environ.get('KIMI_API_KEY')
+
+        if not self.api_key:
+            self.skipTest("KIMI_API_KEY not set for integration tests")
+
+    def test_kimi_real_api_basic_request(self):
+        """Test Kimi K2 with real API - basic installation request"""
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+
+        result = interpreter.parse("Install curl on Ubuntu")
+
+        self.assertIsInstance(result, list)
+        self.assertGreater(len(result), 0)
+        self.assertTrue(any("curl" in cmd.lower() for cmd in result))
+        print(f"\n✅ Kimi K2 API Test - Generated {len(result)} commands: {result}")
+
+    def test_kimi_real_api_complex_request(self):
+        """Test Kimi K2 with real API - complex installation request"""
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+
+        result = interpreter.parse("Install nginx web server and configure it to start on boot")
+
+        self.assertIsInstance(result, list)
+        self.assertGreater(len(result), 2)
+        self.assertTrue(any("nginx" in cmd.lower() for cmd in result))
+        print(f"\n✅ Kimi K2 API Complex Test - Generated {len(result)} commands: {result}")
+
+    @patch.dict(os.environ, {'KIMI_DEFAULT_MODEL': 'kimi-k2-0905-preview'})
+    def test_kimi_real_api_with_custom_model(self):
+        """Test Kimi K2 with different model"""
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+
+        result = interpreter.parse("Install git")
+
+        self.assertIsInstance(result, list)
+        self.assertGreater(len(result), 0)
+        self.assertTrue(any("git" in cmd.lower() for cmd in result))
+        print(f"\n✅ Kimi K2 Custom Model Test - Generated {len(result)} commands: {result}")
+
+    def test_kimi_real_api_validation(self):
+        """Test Kimi K2 with command validation"""
+        interpreter = CommandInterpreter(api_key=self.api_key, provider="kimi")
+
+        result = interpreter.parse("Install docker", validate=True)
+
+        self.assertIsInstance(result, list)
+        self.assertGreater(len(result), 0)
+        # Ensure no dangerous commands passed validation
+        for cmd in result:
+            self.assertNotIn("rm -rf", cmd.lower())
+            self.assertNotIn("dd if=", cmd.lower())
+        print(f"\n✅ Kimi K2 Validation Test - Validated commands: {result}")
+
+
 if __name__ == "__main__":
     unittest.main()