Feature/custom ai roles

cortex/cli.py

@@ -37,7 +37,7 @@
 
 
 class CortexCLI:
-    def __init__(self, verbose: bool = False):
+    def __init__(self, verbose: bool = False, role: str = "default"):
         self.spinner_chars = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
         self.spinner_idx = 0
         self.verbose = verbose

cortex/confirmation.py

@@ -0,0 +1,11 @@
+def confirm_plan(steps):
+    print("\nProposed installation plan:\n")
+    for step in steps:
+        print(step)
+
+    print("\nProceed?")
+    print("[y] yes   [e] edit   [n] cancel")
+
+    choice = input("> ").lower()
+
+    return choice

cortex/llm/interpreter.py

@@ -28,19 +28,24 @@ def __init__(
         self,
         api_key: str,
         provider: str = "openai",
+        role: str = "default",
         model: str | None = None,
         cache: Optional["SemanticCache"] = None,
     ):
-        """Initialize the command interpreter.
-
-        Args:
-            api_key: API key for the LLM provider
-            provider: Provider name ("openai", "claude", or "ollama")
-            model: Optional model name override
-            cache: Optional SemanticCache instance for response caching
+        """Args:
+        api_key: API key for the LLM provider
+        provider: Provider name ("openai", "claude", or "ollama")
+        model: Optional model name override
+        cache: Optional SemanticCache instance for response caching
         """
+        from cortex.roles.loader import load_role
+
         self.api_key = api_key
         self.provider = APIProvider(provider.lower())
+        # Load role and system prompt
+        self.role_name = role
+        self.role = load_role(role)
+        self.system_prompt: str = self.role.get("system_prompt", "")
 
         if cache is None:
             try:
@@ -112,35 +117,36 @@ def _get_system_prompt(self, simplified: bool = False) -> str:
         Args:
             simplified: If True, return a shorter prompt optimized for local models
         """
-        if simplified:
-            return """You must respond with ONLY a JSON object. No explanations, no markdown, no code blocks.
-
-Format: {"commands": ["command1", "command2"]}
-
-Example input: install nginx
-Example output: {"commands": ["sudo apt update", "sudo apt install -y nginx"]}
-
-Rules:
-- Use apt for Ubuntu packages
-- Add sudo for system commands
-- Return ONLY the JSON object"""
 
-        return """You are a Linux system command expert. Convert natural language requests into safe, validated bash commands.
-
-Rules:
-1. Return ONLY a JSON array of commands
-2. Each command must be a safe, executable bash command
-3. Commands should be atomic and sequential
-4. Avoid destructive operations without explicit user confirmation
-5. Use package managers appropriate for Debian/Ubuntu systems (apt)
-6. Include necessary privilege escalation (sudo) when required
-7. Validate command syntax before returning
-
-Format:
-{"commands": ["command1", "command2", ...]}
+        if simplified:
+            base_prompt = (
+                "You must respond with ONLY a JSON object. No explanations, no markdown.\n\n"
+                "Format:\n"
+                '{"commands": ["command1", "command2"]}\n\n'
+                "Rules:\n"
+                "- Use apt for Ubuntu packages\n"
+                "- Add sudo for system commands\n"
+                "- Return ONLY the JSON object\n"
+            )
+        else:
+            base_prompt = (
+                "You are a Linux system command expert. Convert natural language requests "
+                "into safe, validated bash commands.\n\n"
+                "Rules:\n"
+                "1. Return ONLY a JSON array of commands\n"
+                "2. Each command must be safe and executable\n"
+                "3. Commands should be atomic and sequential\n"
+                "4. Avoid destructive operations without explicit user confirmation\n"
+                "5. Use apt for Debian/Ubuntu systems\n"
+                "6. Include sudo when required\n\n"
+                "Format:\n"
+                '{"commands": ["command1", "command2", ...]}\n'
+            )
 
-Example request: "install docker with nvidia support"
-Example response: {"commands": ["sudo apt update", "sudo apt install -y docker.io", "sudo apt install -y nvidia-docker2", "sudo systemctl restart docker"]}"""
+        system_prompt = getattr(self, "system_prompt", "")
+        if system_prompt:
+            return f"{self.system_prompt}\n\n{base_prompt}"
+        return base_prompt
 
     def _call_openai(self, user_input: str) -> list[str]:
         try:
@@ -181,6 +187,7 @@ def _call_ollama(self, user_input: str) -> list[str]:
             enhanced_input = f"""{user_input}
 
 Respond with ONLY this JSON format (no explanations):
+
 {{\"commands\": [\"command1\", \"command2\"]}}"""
 
             response = self.client.chat.completions.create(
@@ -233,55 +240,39 @@ def _repair_json(self, content: str) -> str:
 
     def _parse_commands(self, content: str) -> list[str]:
         try:
-            # Strip markdown code blocks
-            if "```json" in content:
-                content = content.split("```json")[1].split("```")[0].strip()
-            elif "```" in content:
+            # Remove markdown code blocks if present
+            if "```" in content:
                 parts = content.split("```")
-                if len(parts) >= 3:
-                    content = parts[1].strip()
+                if len(parts) >= 2:
+                    content = parts[1]
 
-            # Try to find JSON object in the content
             import re
 
-            # Look for {"commands": [...]} pattern
-            json_match = re.search(
-                r'\{\s*["\']commands["\']\s*:\s*\[.*?\]\s*\}', content, re.DOTALL
+            match = re.search(
+                r'\{\s*"commands"\s*:\s*\[.*?\]\s*\}',
+                content,
+                re.DOTALL,
             )
-            if json_match:
-                content = json_match.group(0)
 
-            # Try to repair common JSON issues
-            content = self._repair_json(content)
+            if not match:
+                raise ValueError("No valid JSON command block found")
 
-            data = json.loads(content)
-            commands = data.get("commands", [])
+            json_text = match.group(0)
+            data = json.loads(json_text)
 
+            commands = data.get("commands")
             if not isinstance(commands, list):
-                raise ValueError("Commands must be a list")
+                raise ValueError("commands must be a list")
 
-            # Handle both formats:
-            # 1. ["cmd1", "cmd2"] - direct string array
-            # 2. [{"command": "cmd1"}, {"command": "cmd2"}] - object array
-            result = []
+            cleaned: list[str] = []
             for cmd in commands:
-                if isinstance(cmd, str):
-                    # Direct string
-                    if cmd:
-                        result.append(cmd)
-                elif isinstance(cmd, dict):
-                    # Object with "command" key
-                    cmd_str = cmd.get("command", "")
-                    if cmd_str:
-                        result.append(cmd_str)
-
-            return result
-        except (json.JSONDecodeError, ValueError) as e:
-            # Log the problematic content for debugging
-            import sys
-
-            print(f"\nDebug: Failed to parse JSON. Raw content:\n{content[:500]}", file=sys.stderr)
-            raise ValueError(f"Failed to parse LLM response: {str(e)}")
+                if isinstance(cmd, str) and cmd.strip():
+                    cleaned.append(cmd.strip())
+
+            return cleaned
+
+        except Exception as exc:
+            raise ValueError(f"Failed to parse LLM response: {exc}") from exc
 
     def _validate_commands(self, commands: list[str]) -> list[str]:
         dangerous_patterns = [
@@ -303,19 +294,6 @@ def _validate_commands(self, commands: list[str]) -> list[str]:
         return validated
 
     def parse(self, user_input: str, validate: bool = True) -> list[str]:
-        """Parse natural language input into shell commands.
-
-        Args:
-            user_input: Natural language description of desired action
-            validate: If True, validate commands for dangerous patterns
-
-        Returns:
-            List of shell commands to execute
-
-        Raises:
-            ValueError: If input is empty
-            RuntimeError: If offline mode is enabled and no cached response exists
-        """
         if not user_input or not user_input.strip():
             raise ValueError("User input cannot be empty")
 
@@ -347,23 +325,13 @@ def parse(self, user_input: str, validate: bool = True) -> list[str]:
         if validate:
             commands = self._validate_commands(commands)
 
-        if self.cache is not None and commands:
-            try:
-                self.cache.put_commands(
-                    prompt=user_input,
-                    provider=self.provider.value,
-                    model=self.model,
-                    system_prompt=cache_system_prompt,
-                    commands=commands,
-                )
-            except (OSError, sqlite3.Error):
-                # Silently fail cache writes - not critical for operation
-                pass
-
         return commands
 
     def parse_with_context(
-        self, user_input: str, system_info: dict[str, Any] | None = None, validate: bool = True
+        self,
+        user_input: str,
+        system_info: dict[str, Any] | None = None,
+        validate: bool = True,
     ) -> list[str]:
         context = ""
         if system_info:

cortex/planner.py

@@ -0,0 +1,39 @@
+from typing import Any, dict, list
+
+from cortex.llm.interpreter import CommandInterpreter
+
+
+def generate_plan(intent: str, slots: dict[str, Any]) -> list[str]:
+    """
+    Generate a human-readable installation plan using LLM (Ollama).
+    """
+
+    prompt = f"""
+You are a DevOps assistant.
+
+User intent:
+{intent}
+
+Extracted details:
+{slots}
+
+Generate a step-by-step installation plan.
+Rules:
+- High-level steps only
+- No shell commands
+- One sentence per step
+- Return as a JSON list of strings
+
+Example:
+["Install Python", "Install ML libraries", "Set up Jupyter"]
+"""
+
+    interpreter = CommandInterpreter(
+        api_key="ollama",  # dummy value, Ollama ignores it
+        provider="ollama",
+    )
+
+    # Reuse interpreter to get structured output
+    steps = interpreter.parse(prompt, validate=False)
+
+    return steps

cortex/roles/default.yaml

@@ -0,0 +1,6 @@
+name: default
+description: General-purpose Linux and software assistant
+system_prompt: |
+  You are Cortex, an AI-powered Linux command interpreter.
+  Provide clear, safe, and correct Linux commands.
+  Explain steps when helpful, but keep responses concise.

cortex/roles/devops.yaml

@@ -0,0 +1,6 @@
+name: devops
+description: DevOps and infrastructure automation expert
+system_prompt: |
+  You are a DevOps-focused AI assistant.
+  Optimize for reliability, automation, and scalability.
+  Prefer idempotent commands and infrastructure-as-code approaches.

cortex/roles/loader.py

@@ -0,0 +1,32 @@
+import os
+
+import yaml
+
+
+class RoleNotFoundError(Exception):
+    pass
+
+
+def get_roles_dir():
+    """
+    Returns the directory where built-in roles are stored.
+    """
+    return os.path.dirname(__file__)
+
+
+def load_role(role_name: str) -> dict:
+    """
+    Load a role YAML by name.
+    Falls back to default if role not found.
+    """
+    roles_dir = get_roles_dir()
+    role_file = os.path.join(roles_dir, f"{role_name}.yaml")
+
+    if not os.path.exists(role_file):
+        if role_name != "default":
+            # Fallback to default role
+            return load_role("default")
+        raise RoleNotFoundError("Default role not found")
+
+    with open(role_file, encoding="utf-8") as f:
+        return yaml.safe_load(f)

cortex/roles/security.yaml

@@ -0,0 +1,7 @@
+name: security
+description: Security auditing and hardening expert
+system_prompt: |
+  You are a security-focused AI assistant.
+  Always prioritize safety, least privilege, and risk mitigation.
+  Warn before destructive or irreversible actions.
+  Prefer secure defaults and compliance best practices.