kubrick-cli 0.1.2__py3-none-any.whl

kubrick_cli/providers/openai_provider.py

@@ -0,0 +1,190 @@
+"""OpenAI provider adapter."""
+
+import http.client
+import json
+import ssl
+from typing import Dict, Iterator, List
+
+from .base import ProviderAdapter, ProviderMetadata
+
+
+class OpenAIProvider(ProviderAdapter):
+    """Provider adapter for OpenAI API."""
+
+    METADATA = ProviderMetadata(
+        name="openai",
+        display_name="OpenAI",
+        description="OpenAI API (GPT-4, GPT-3.5, etc.)",
+        config_fields=[
+            {
+                "key": "openai_api_key",
+                "label": "OpenAI API key",
+                "type": "password",
+                "help_text": "Get your API key from: https://platform.openai.com/api-keys",
+            },
+            {
+                "key": "openai_model",
+                "label": "Model name",
+                "type": "text",
+                "default": "gpt-4",
+            },
+        ],
+    )
+
+    def __init__(self, openai_api_key: str, openai_model: str = "gpt-4"):
+        """
+        Initialize OpenAI provider.
+
+        Args:
+            openai_api_key: OpenAI API key
+            openai_model: Model name (default: gpt-4)
+        """
+        self.api_key = openai_api_key
+        self._model_name = openai_model
+        self.base_url = "api.openai.com"
+        self.timeout = 600
+
+    def generate_streaming(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from OpenAI.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        # Build request payload
+        payload = {
+            "model": self._model_name,
+            "messages": messages,
+            "stream": True,
+        }
+
+        # Add optional parameters
+        if stream_options:
+            if "temperature" in stream_options:
+                payload["temperature"] = stream_options["temperature"]
+            if "max_tokens" in stream_options:
+                payload["max_tokens"] = stream_options["max_tokens"]
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+        body = json.dumps(payload).encode("utf-8")
+
+        # Create HTTPS connection
+        context = ssl.create_default_context()
+        conn = http.client.HTTPSConnection(
+            self.base_url, 443, timeout=self.timeout, context=context
+        )
+
+        try:
+            conn.request("POST", "/v1/chat/completions", body=body, headers=headers)
+            response = conn.getresponse()
+
+            if response.status != 200:
+                error_body = response.read().decode("utf-8")
+                raise Exception(f"OpenAI API error {response.status}: {error_body}")
+
+            # Read streaming response
+            buffer = ""
+            while True:
+                chunk = response.read(1024)
+                if not chunk:
+                    break
+
+                if isinstance(chunk, bytes):
+                    chunk = chunk.decode("utf-8")
+
+                buffer += chunk
+
+                # Process complete lines
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    line = line.strip()
+
+                    if not line:
+                        continue
+
+                    # OpenAI SSE format: "data: {json}"
+                    if line.startswith("data: "):
+                        line = line[6:]
+
+                    # Check for end signal
+                    if line == "[DONE]":
+                        return
+
+                    try:
+                        data = json.loads(line)
+                        # Extract delta content
+                        if "choices" in data and len(data["choices"]) > 0:
+                            delta = data["choices"][0].get("delta", {})
+                            content = delta.get("content", "")
+                            if content:
+                                yield content
+                    except json.JSONDecodeError:
+                        # Skip malformed JSON
+                        continue
+
+        finally:
+            conn.close()
+
+    def generate(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> str:
+        """
+        Generate non-streaming response from OpenAI.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional parameters
+
+        Returns:
+            Complete response text
+        """
+        chunks = []
+        for chunk in self.generate_streaming(messages, stream_options):
+            chunks.append(chunk)
+        return "".join(chunks)
+
+    def is_healthy(self) -> bool:
+        """
+        Check if OpenAI API is accessible.
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        try:
+            context = ssl.create_default_context()
+            conn = http.client.HTTPSConnection(
+                self.base_url, 443, timeout=10, context=context
+            )
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+            }
+            conn.request("GET", "/v1/models", headers=headers)
+            response = conn.getresponse()
+            conn.close()
+            return response.status == 200
+        except Exception:
+            return False
+
+    @property
+    def provider_name(self) -> str:
+        """Get provider name."""
+        return "openai"
+
+    @property
+    def model_name(self) -> str:
+        """Get model name."""
+        return self._model_name
+
+    def set_model(self, model_name: str):
+        """Set model name dynamically."""
+        self._model_name = model_name

kubrick_cli/providers/triton_provider.py

@@ -0,0 +1,96 @@
+"""Triton provider adapter."""
+
+from typing import Dict, Iterator, List
+
+from ..triton_client import TritonLLMClient
+from .base import ProviderAdapter, ProviderMetadata
+
+
+class TritonProvider(ProviderAdapter):
+    """Provider adapter for Triton Inference Server."""
+
+    METADATA = ProviderMetadata(
+        name="triton",
+        display_name="Triton",
+        description="Self-hosted Triton Inference Server (default)",
+        config_fields=[
+            {
+                "key": "triton_url",
+                "label": "Triton server URL",
+                "type": "url",
+                "default": "localhost:8000",
+            },
+            {
+                "key": "triton_model",
+                "label": "Triton model name",
+                "type": "text",
+                "default": "llm_decoupled",
+            },
+        ],
+    )
+
+    def __init__(self, triton_url: str = "localhost:8000", triton_model: str = "llm_decoupled"):
+        """
+        Initialize Triton provider.
+
+        Args:
+            triton_url: Triton server URL (default: localhost:8000)
+            triton_model: Triton model name (default: llm_decoupled)
+        """
+        self.client = TritonLLMClient(url=triton_url, model_name=triton_model)
+        self._url = triton_url
+        self._model_name = triton_model
+
+    def generate_streaming(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from Triton.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        yield from self.client.generate_streaming(messages, stream_options)
+
+    def generate(
+        self, messages: List[Dict[str, str]], stream_options: Dict = None
+    ) -> str:
+        """
+        Generate non-streaming response from Triton.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional parameters
+
+        Returns:
+            Complete response text
+        """
+        return self.client.generate(messages, stream_options)
+
+    def is_healthy(self) -> bool:
+        """
+        Check if Triton server is healthy.
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        return self.client.is_healthy()
+
+    @property
+    def provider_name(self) -> str:
+        """Get provider name."""
+        return "triton"
+
+    @property
+    def model_name(self) -> str:
+        """Get model name."""
+        return self._model_name
+
+    @property
+    def url(self) -> str:
+        """Get Triton server URL."""
+        return self._url

kubrick_cli/safety.py

@@ -0,0 +1,216 @@
+"""Safety manager for dangerous command detection and validation."""
+
+import re
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Confirm
+
+console = Console()
+
+
+@dataclass
+class SafetyConfig:
+    """Safety configuration settings."""
+
+    max_iterations: int = 15
+    max_tools_per_turn: int = 5
+    total_timeout_seconds: int = 600
+    tool_timeout_seconds: int = 30
+    max_file_size_mb: int = 10
+    require_dangerous_command_confirmation: bool = True
+
+    @classmethod
+    def from_config(cls, config: Dict) -> "SafetyConfig":
+        """
+        Create SafetyConfig from configuration dict.
+
+        Args:
+            config: Configuration dictionary
+
+        Returns:
+            SafetyConfig instance
+        """
+        return cls(
+            max_iterations=config.get("max_iterations", 15),
+            max_tools_per_turn=config.get("max_tools_per_turn", 5),
+            total_timeout_seconds=config.get("total_timeout_seconds", 600),
+            tool_timeout_seconds=config.get("tool_timeout_seconds", 30),
+            max_file_size_mb=config.get("max_file_size_mb", 10),
+            require_dangerous_command_confirmation=config.get(
+                "require_dangerous_command_confirmation", True
+            ),
+        )
+
+
+class SafetyManager:
+    """
+    Manages safety checks for tool execution.
+
+    Features:
+    - Dangerous command detection
+    - User confirmation prompts
+    - File size limits
+    - Timeout enforcement
+    """
+
+    # Dangerous bash command patterns
+    DANGEROUS_PATTERNS = [
+        (r"\brm\s+-rf\s+/", "Recursive delete from root"),
+        (r"\brm\s+-rf\s+~", "Recursive delete from home"),
+        (r"\brm\s+-rf\s+\*", "Recursive delete all files"),
+        (r"\brm\s+-rf", "Recursive force delete"),
+        (r"\bsudo\b", "Elevated privileges"),
+        (r"\bchmod\s+777", "Overly permissive permissions"),
+        (r"\bchmod\s+-R\s+777", "Recursive overly permissive permissions"),
+        (r">\s*/dev/", "Writing to device files"),
+        (r"\bgit\s+push\s+--force", "Force push to git"),
+        (r"\bgit\s+push\s+-f", "Force push to git (short form)"),
+        (r"\bmkfs\b", "Format filesystem"),
+        (r"\bdd\s+.*of=/dev", "Writing to block device"),
+        (r"\b:(\()\|:\1&", "Fork bomb"),
+        (r"\bcurl\b.*\|\s*bash", "Pipe curl to bash"),
+        (r"\bwget\b.*\|\s*bash", "Pipe wget to bash"),
+        (r"\beval\b.*\$\(", "Eval with command substitution"),
+    ]
+
+    def __init__(self, config: SafetyConfig):
+        """
+        Initialize safety manager.
+
+        Args:
+            config: SafetyConfig instance
+        """
+        self.config = config
+
+    def validate_bash_command(
+        self, command: str
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Validate a bash command for dangerous patterns.
+
+        Args:
+            command: The bash command to validate
+
+        Returns:
+            Tuple of (is_safe, warning_message)
+            - is_safe: True if safe, False if dangerous
+            - warning_message: Description of the danger (if any)
+        """
+        command_lower = command.lower()
+
+        # Check each dangerous pattern
+        for pattern, description in self.DANGEROUS_PATTERNS:
+            if re.search(pattern, command_lower, re.IGNORECASE):
+                return False, f"Dangerous command detected: {description}"
+
+        return True, None
+
+    def get_user_confirmation(
+        self, warning: str, command: str
+    ) -> bool:
+        """
+        Get user confirmation for a dangerous command.
+
+        Args:
+            warning: Warning message
+            command: The dangerous command
+
+        Returns:
+            True if user confirms, False otherwise
+        """
+        if not self.config.require_dangerous_command_confirmation:
+            # Auto-allow if confirmation disabled
+            return True
+
+        # Display warning
+        console.print(
+            Panel(
+                f"[bold red]⚠️  WARNING[/bold red]\n\n"
+                f"{warning}\n\n"
+                f"[bold]Command:[/bold]\n"
+                f"[yellow]{command}[/yellow]\n\n"
+                f"This command could be dangerous.",
+                border_style="red",
+                title="Safety Check",
+            )
+        )
+
+        # Ask for confirmation
+        confirmed = Confirm.ask(
+            "[bold red]Do you want to execute this command?[/bold red]",
+            default=False,
+        )
+
+        if confirmed:
+            console.print("[yellow]⚠️  Proceeding with caution...[/yellow]")
+        else:
+            console.print("[green]✓ Command cancelled[/green]")
+
+        return confirmed
+
+    def validate_file_size(self, file_path: str, size_bytes: int) -> bool:
+        """
+        Validate file size against limits.
+
+        Args:
+            file_path: Path to the file
+            size_bytes: File size in bytes
+
+        Returns:
+            True if within limits, False otherwise
+        """
+        max_bytes = self.config.max_file_size_mb * 1024 * 1024
+
+        if size_bytes > max_bytes:
+            console.print(
+                f"[yellow]⚠️  File {file_path} exceeds size limit "
+                f"({size_bytes / 1024 / 1024:.1f}MB > {self.config.max_file_size_mb}MB)[/yellow]"
+            )
+            return False
+
+        return True
+
+    def check_iteration_limit(
+        self, current: int, max_iterations: int
+    ) -> bool:
+        """
+        Check if iteration limit has been reached.
+
+        Args:
+            current: Current iteration number
+            max_iterations: Maximum allowed iterations
+
+        Returns:
+            True if within limit, False otherwise
+        """
+        if current >= max_iterations:
+            console.print(
+                f"[yellow]⚠️  Max iteration limit reached ({max_iterations})[/yellow]"
+            )
+            return False
+
+        return True
+
+    def check_tool_limit(
+        self, current: int, max_tools: int
+    ) -> bool:
+        """
+        Check if tool call limit has been reached.
+
+        Args:
+            current: Current number of tool calls
+            max_tools: Maximum allowed tools
+
+        Returns:
+            True if within limit, False otherwise
+        """
+        if current > max_tools:
+            console.print(
+                f"[yellow]⚠️  Too many tool calls ({current} > {max_tools})[/yellow]"
+            )
+            return False
+
+        return True

kubrick_cli/scheduler.py

@@ -0,0 +1,197 @@
+"""Parallel tool execution scheduler for improved performance."""
+
+import concurrent.futures
+from typing import Dict, List, Tuple
+
+from rich.console import Console
+
+console = Console()
+
+
+# Read-only tools that can be executed in parallel
+READ_ONLY_TOOLS = {
+    "read_file",
+    "list_files",
+    "search_files",
+}
+
+# Write tools that must be executed sequentially
+WRITE_TOOLS = {
+    "write_file",
+    "edit_file",
+    "create_directory",
+    "run_bash",  # Conservative: treat bash as write operation
+}
+
+
+class ToolScheduler:
+    """
+    Schedules and executes tools with intelligent parallelization.
+
+    Strategy: Conservative parallelization
+    - Read-only tools (read_file, list_files, search_files) run in parallel
+    - Write tools (write_file, edit_file, run_bash, etc.) run sequentially
+    - Graceful error handling per tool
+    """
+
+    def __init__(
+        self, tool_executor, max_workers: int = 3, enable_parallel: bool = True
+    ):
+        """
+        Initialize tool scheduler.
+
+        Args:
+            tool_executor: ToolExecutor instance
+            max_workers: Maximum parallel workers (default: 3)
+            enable_parallel: Whether to enable parallel execution
+        """
+        self.tool_executor = tool_executor
+        self.max_workers = max_workers
+        self.enable_parallel = enable_parallel
+
+    def execute_tools(
+        self, tool_calls: List[Tuple[str, Dict]]
+    ) -> List[Dict]:
+        """
+        Execute a list of tool calls with intelligent scheduling.
+
+        Args:
+            tool_calls: List of (tool_name, parameters) tuples
+
+        Returns:
+            List of result dictionaries in the same order as input
+        """
+        if not self.enable_parallel or len(tool_calls) <= 1:
+            # Execute sequentially if parallel disabled or only one tool
+            return self._execute_sequential(tool_calls)
+
+        # Separate read-only and write tools
+        read_only_calls = []
+        write_calls = []
+        call_order = []  # Track original order
+
+        for i, (tool_name, params) in enumerate(tool_calls):
+            if tool_name in READ_ONLY_TOOLS:
+                read_only_calls.append((i, tool_name, params))
+                call_order.append(("read", len(read_only_calls) - 1))
+            else:
+                write_calls.append((i, tool_name, params))
+                call_order.append(("write", len(write_calls) - 1))
+
+        # Execute read-only tools in parallel
+        read_results = {}
+        if read_only_calls:
+            console.print(
+                f"[dim]→ Executing {len(read_only_calls)} read-only tool(s) in parallel[/dim]"
+            )
+            read_results = self._execute_parallel(read_only_calls)
+
+        # Execute write tools sequentially
+        write_results = {}
+        if write_calls:
+            write_results = self._execute_sequential_indexed(write_calls)
+
+        # Reconstruct results in original order
+        results = []
+        for call_type, index in call_order:
+            if call_type == "read":
+                results.append(read_results[index])
+            else:
+                results.append(write_results[index])
+
+        return results
+
+    def _execute_parallel(
+        self, indexed_calls: List[Tuple[int, str, Dict]]
+    ) -> Dict[int, Dict]:
+        """
+        Execute tools in parallel using ThreadPoolExecutor.
+
+        Args:
+            indexed_calls: List of (index, tool_name, parameters) tuples
+
+        Returns:
+            Dict mapping index to result
+        """
+        results = {}
+
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=self.max_workers
+        ) as executor:
+            # Submit all tasks
+            future_to_index = {}
+            for index, tool_name, params in indexed_calls:
+                future = executor.submit(
+                    self._execute_single, tool_name, params
+                )
+                future_to_index[future] = index
+
+            # Collect results as they complete
+            for future in concurrent.futures.as_completed(future_to_index):
+                index = future_to_index[future]
+                try:
+                    result = future.result()
+                    results[index] = result
+                except Exception as e:
+                    # Handle execution errors gracefully
+                    results[index] = {
+                        "success": False,
+                        "error": f"Parallel execution error: {str(e)}",
+                    }
+
+        return results
+
+    def _execute_sequential(
+        self, tool_calls: List[Tuple[str, Dict]]
+    ) -> List[Dict]:
+        """
+        Execute tools sequentially.
+
+        Args:
+            tool_calls: List of (tool_name, parameters) tuples
+
+        Returns:
+            List of result dictionaries
+        """
+        results = []
+        for tool_name, params in tool_calls:
+            result = self._execute_single(tool_name, params)
+            results.append(result)
+        return results
+
+    def _execute_sequential_indexed(
+        self, indexed_calls: List[Tuple[int, str, Dict]]
+    ) -> Dict[int, Dict]:
+        """
+        Execute tools sequentially with index mapping.
+
+        Args:
+            indexed_calls: List of (index, tool_name, parameters) tuples
+
+        Returns:
+            Dict mapping index to result
+        """
+        results = {}
+        for index, tool_name, params in indexed_calls:
+            result = self._execute_single(tool_name, params)
+            results[index] = result
+        return results
+
+    def _execute_single(self, tool_name: str, params: Dict) -> Dict:
+        """
+        Execute a single tool.
+
+        Args:
+            tool_name: Name of tool
+            params: Tool parameters
+
+        Returns:
+            Result dictionary
+        """
+        try:
+            return self.tool_executor.execute(tool_name, params)
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"Tool execution error: {str(e)}",
+            }