kubrick-cli 0.1.4__py3-none-any.whl

kubrick_cli/tools.py

@@ -0,0 +1,400 @@
+"""Tool definitions and execution handlers."""
+
+import subprocess
+from pathlib import Path
+from typing import Any, Dict
+
+TOOL_DEFINITIONS = [
+    {
+        "name": "read_file",
+        "description": "Read the contents of a file from the filesystem",
+        "read_only": True,
+        "estimated_duration": 1,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "file_path": {
+                    "type": "string",
+                    "description": "The absolute or relative path to the file to read",
+                },
+            },
+            "required": ["file_path"],
+        },
+    },
+    {
+        "name": "write_file",
+        "description": "Write content to a file (creates or overwrites)",
+        "read_only": False,
+        "estimated_duration": 2,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "file_path": {
+                    "type": "string",
+                    "description": "The absolute or relative path to the file to write",
+                },
+                "content": {
+                    "type": "string",
+                    "description": "The content to write to the file",
+                },
+            },
+            "required": ["file_path", "content"],
+        },
+    },
+    {
+        "name": "edit_file",
+        "description": "Edit a file by replacing a specific string with new content",
+        "read_only": False,
+        "estimated_duration": 2,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "file_path": {
+                    "type": "string",
+                    "description": "The path to the file to edit",
+                },
+                "old_string": {
+                    "type": "string",
+                    "description": "The exact string to find and replace",
+                },
+                "new_string": {
+                    "type": "string",
+                    "description": "The new string to replace with",
+                },
+            },
+            "required": ["file_path", "old_string", "new_string"],
+        },
+    },
+    {
+        "name": "list_files",
+        "description": "List files matching a glob pattern (*.py, **/*.py, etc)",
+        "read_only": True,
+        "estimated_duration": 1,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "pattern": {
+                    "type": "string",
+                    "description": "Glob pattern to match files (e.g., '*.py', 'src/**/*.ts')",
+                },
+                "directory": {
+                    "type": "string",
+                    "description": "Directory to search in (defaults to current directory)",
+                },
+            },
+            "required": ["pattern"],
+        },
+    },
+    {
+        "name": "search_files",
+        "description": "Search for text content within files using grep-like functionality",
+        "read_only": True,
+        "estimated_duration": 2,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "pattern": {
+                    "type": "string",
+                    "description": "Text pattern or regex to search for",
+                },
+                "file_pattern": {
+                    "type": "string",
+                    "description": "File glob pattern to search within (e.g., '*.py')",
+                },
+                "directory": {
+                    "type": "string",
+                    "description": "Directory to search in (defaults to current directory)",
+                },
+            },
+            "required": ["pattern"],
+        },
+    },
+    {
+        "name": "run_bash",
+        "description": "Execute a bash command and return its output",
+        "read_only": False,
+        "estimated_duration": 5,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "command": {
+                    "type": "string",
+                    "description": "The bash command to execute",
+                },
+            },
+            "required": ["command"],
+        },
+    },
+    {
+        "name": "create_directory",
+        "description": "Create a new directory",
+        "read_only": False,
+        "estimated_duration": 1,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "The path of the directory to create",
+                },
+            },
+            "required": ["path"],
+        },
+    },
+]
+
+
+class ToolExecutor:
+    """Handles execution of tools."""
+
+    def __init__(self, working_dir: str = None, safety_manager=None):
+        """
+        Initialize tool executor.
+
+        Args:
+            working_dir: Working directory for file operations (defaults to current directory)
+            safety_manager: Optional SafetyManager for validation
+        """
+        self.working_dir = Path(working_dir) if working_dir else Path.cwd()
+        self.safety_manager = safety_manager
+
+    def _resolve_path(self, path: str) -> Path:
+        """Resolve a path relative to working directory."""
+        p = Path(path)
+        if p.is_absolute():
+            return p
+        return (self.working_dir / p).resolve()
+
+    def execute(self, tool_name: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Execute a tool with given parameters.
+
+        Args:
+            tool_name: Name of the tool to execute
+            parameters: Tool parameters
+
+        Returns:
+            Dict with 'success', 'result', and optionally 'error' keys
+        """
+        try:
+            if tool_name == "read_file":
+                return self._read_file(parameters)
+            elif tool_name == "write_file":
+                return self._write_file(parameters)
+            elif tool_name == "edit_file":
+                return self._edit_file(parameters)
+            elif tool_name == "list_files":
+                return self._list_files(parameters)
+            elif tool_name == "search_files":
+                return self._search_files(parameters)
+            elif tool_name == "run_bash":
+                return self._run_bash(parameters)
+            elif tool_name == "create_directory":
+                return self._create_directory(parameters)
+            else:
+                return {
+                    "success": False,
+                    "error": f"Unknown tool: {tool_name}",
+                }
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e),
+            }
+
+    def _read_file(self, params: Dict) -> Dict:
+        """Read a file."""
+        file_path = self._resolve_path(params["file_path"])
+
+        if not file_path.exists():
+            return {"success": False, "error": f"File not found: {file_path}"}
+
+        content = file_path.read_text()
+        return {
+            "success": True,
+            "result": f"Content of {file_path}:\n\n{content}",
+        }
+
+    def _write_file(self, params: Dict) -> Dict:
+        """Write to a file."""
+        file_path = self._resolve_path(params["file_path"])
+        content = params["content"]
+
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        file_path.write_text(content)
+        return {
+            "success": True,
+            "result": f"Successfully wrote {len(content)} characters to {file_path}",
+        }
+
+    def _edit_file(self, params: Dict) -> Dict:
+        """Edit a file by replacing text."""
+        file_path = self._resolve_path(params["file_path"])
+        old_string = params["old_string"]
+        new_string = params["new_string"]
+
+        if not file_path.exists():
+            return {"success": False, "error": f"File not found: {file_path}"}
+
+        content = file_path.read_text()
+
+        if old_string not in content:
+            return {
+                "success": False,
+                "error": f"String not found in file: {old_string[:100]}...",
+            }
+
+        new_content = content.replace(old_string, new_string, 1)
+        file_path.write_text(new_content)
+
+        return {
+            "success": True,
+            "result": f"Successfully edited {file_path}",
+        }
+
+    def _list_files(self, params: Dict) -> Dict:
+        """List files matching a pattern."""
+        pattern = params["pattern"]
+        directory = params.get("directory", ".")
+        search_dir = self._resolve_path(directory)
+
+        files = []
+        directories = []
+
+        for match in search_dir.glob(pattern):
+            rel_path = match.relative_to(search_dir)
+            if match.is_file():
+                files.append(str(rel_path))
+            elif match.is_dir():
+                directories.append(str(rel_path) + "/")
+
+        files.sort()
+        directories.sort()
+
+        matches = directories + files
+
+        if not matches:
+            result = f"No files or directories found matching pattern: {pattern}"
+        else:
+            dir_count = len(directories)
+            file_count = len(files)
+            result = f"Found {dir_count} directories and {file_count} files:\n"
+
+            if directories:
+                result += "\nDirectories:\n" + "\n".join(directories)
+            if files:
+                result += "\n\nFiles:\n" + "\n".join(files)
+
+        return {"success": True, "result": result}
+
+    def _search_files(self, params: Dict) -> Dict:
+        """Search for text in files."""
+        pattern = params["pattern"]
+        file_pattern = params.get("file_pattern", "**/*")
+        directory = params.get("directory", ".")
+        search_dir = self._resolve_path(directory)
+
+        results = []
+
+        for file_path in search_dir.glob(file_pattern):
+            if not file_path.is_file():
+                continue
+
+            try:
+                content = file_path.read_text()
+                lines = content.split("\n")
+
+                for line_num, line in enumerate(lines, 1):
+                    if pattern in line:
+                        rel_path = file_path.relative_to(search_dir)
+                        results.append(f"{rel_path}:{line_num}: {line.strip()}")
+            except (UnicodeDecodeError, PermissionError):
+                continue
+
+        if not results:
+            result = f"No matches found for pattern: {pattern}"
+        else:
+            result = f"Found {len(results)} matches:\n" + "\n".join(results[:50])
+            if len(results) > 50:
+                result += f"\n... and {len(results) - 50} more"
+
+        return {"success": True, "result": result}
+
+    def _run_bash(self, params: Dict) -> Dict:
+        """Run a bash command."""
+        command = params["command"]
+
+        if self.safety_manager:
+            is_safe, warning = self.safety_manager.validate_bash_command(command)
+
+            if not is_safe:
+                confirmed = self.safety_manager.get_user_confirmation(warning, command)
+
+                if not confirmed:
+                    return {
+                        "success": False,
+                        "error": "Command cancelled by user (dangerous command)",
+                    }
+
+        try:
+            timeout = 30
+            if self.safety_manager and hasattr(self.safety_manager, "config"):
+                timeout = self.safety_manager.config.tool_timeout_seconds
+
+            # nosec B602: shell=True is required for bash command execution
+            # Commands are validated by SafetyManager before execution
+            result = subprocess.run(
+                command,
+                shell=True,  # nosec
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+                cwd=self.working_dir,
+            )
+
+            output = result.stdout
+            if result.stderr:
+                output += f"\nSTDERR:\n{result.stderr}"
+
+            return {
+                "success": True,
+                "result": f"Command: {command}\nExit code: {result.returncode}\n\n{output}",
+            }
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "error": f"Command timed out after {timeout} seconds",
+            }
+
+    def _create_directory(self, params: Dict) -> Dict:
+        """Create a directory."""
+        path = self._resolve_path(params["path"])
+        path.mkdir(parents=True, exist_ok=True)
+        return {
+            "success": True,
+            "result": f"Created directory: {path}",
+        }
+
+
+def get_tools_prompt() -> str:
+    """Get a formatted string describing available tools for the LLM."""
+    tools_desc = ""
+
+    for tool in TOOL_DEFINITIONS:
+        tools_desc += f"### {tool['name']}\n"
+        tools_desc += f"{tool['description']}\n\n"
+
+        props = tool["parameters"].get("properties", {})
+        required = tool["parameters"].get("required", [])
+
+        if props:
+            tools_desc += "**Parameters:**\n"
+            for param_name, param_info in props.items():
+                req_marker = " (required)" if param_name in required else " (optional)"
+                param_desc = param_info.get("description", "No description")
+                tools_desc += f"- `{param_name}`: {param_desc}{req_marker}\n"
+
+        tools_desc += "\n"
+
+    return tools_desc

kubrick_cli/triton_client.py

@@ -0,0 +1,177 @@
+"""Triton client for streaming LLM requests using HTTP only (no dependencies)."""
+
+import http.client
+import json
+import ssl
+from typing import Dict, Iterator, List
+from urllib.parse import urlparse
+
+
+class TritonLLMClient:
+    """Client for interacting with Triton LLM backend using HTTP (no extra dependencies)."""
+
+    def __init__(
+        self,
+        url: str = "localhost:8000",
+        model_name: str = "llm_decoupled",
+    ):
+        """
+        Initialize Triton LLM client.
+
+        Args:
+            url: Triton server URL (host:port, default: localhost:8000)
+            model_name: Name of the Triton model to use
+        """
+        self.model_name = model_name
+
+        if not url.startswith("http://") and not url.startswith("https://"):
+            url = f"http://{url}"
+
+        parsed = urlparse(url)
+        self.is_https = parsed.scheme == "https"
+        self.host = parsed.hostname or "localhost"
+        self.port = parsed.port or (443 if self.is_https else 8000)
+        self.timeout = 600
+        self.url = f"{parsed.scheme}://{self.host}:{self.port}"
+
+    def _get_connection(self) -> http.client.HTTPConnection:
+        """Create an HTTP(S) connection."""
+        if self.is_https:
+            context = ssl.create_default_context()
+            return http.client.HTTPSConnection(
+                self.host, self.port, timeout=self.timeout, context=context
+            )
+        else:
+            return http.client.HTTPConnection(
+                self.host, self.port, timeout=self.timeout
+            )
+
+    def generate_streaming(
+        self,
+        messages: List[Dict[str, str]],
+        stream_options: Dict = None,
+    ) -> Iterator[str]:
+        """
+        Generate streaming response from LLM.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            stream_options: Optional streaming parameters
+
+        Yields:
+            Text chunks as they arrive
+        """
+        if stream_options is None:
+            stream_options = {"streaming": True}
+        else:
+            stream_options = {"streaming": True, **stream_options}
+
+        payload = {
+            "text_input": json.dumps(messages),
+            "parameters": stream_options,
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        }
+
+        body = json.dumps(payload).encode("utf-8")
+        path = f"/v2/models/{self.model_name}/generate_stream"
+
+        conn = None
+        try:
+            conn = self._get_connection()
+            conn.request("POST", path, body=body, headers=headers)
+            response = conn.getresponse()
+
+            if response.status not in (200, 201):
+                error_body = response.read().decode("utf-8")
+                raise Exception(f"Server returned {response.status}: {error_body}")
+
+            byte_buffer = b""
+            while True:
+                chunk = response.read(1024)
+                if not chunk:
+                    break
+
+                byte_buffer += chunk
+
+                while b"\n" in byte_buffer:
+                    line_bytes, byte_buffer = byte_buffer.split(b"\n", 1)
+
+                    try:
+                        line = line_bytes.decode("utf-8").strip()
+                    except UnicodeDecodeError:
+                        continue
+
+                    if not line:
+                        continue
+
+                    if line.startswith("data: "):
+                        line = line[6:]
+
+                    if line == "[DONE]":
+                        return
+
+                    try:
+                        data = json.loads(line)
+
+                        output_data = None
+                        if "text_output" in data:
+                            output_data = data["text_output"]
+                        elif "outputs" in data and len(data["outputs"]) > 0:
+                            output_data = data["outputs"][0].get("data", [""])[0]
+
+                        if output_data:
+                            try:
+                                chunk_data = json.loads(output_data)
+
+                                if chunk_data.get("type") == "chunk":
+                                    yield chunk_data.get("content", "")
+                                elif chunk_data.get("type") == "complete":
+                                    return
+                                elif chunk_data.get("type") == "error":
+                                    raise Exception(
+                                        f"LLM error: {chunk_data.get('content')}"
+                                    )
+                            except (json.JSONDecodeError, TypeError):
+                                yield output_data
+
+                    except json.JSONDecodeError:
+                        continue
+
+        finally:
+            if conn:
+                conn.close()
+
+    def generate(
+        self,
+        messages: List[Dict[str, str]],
+        stream_options: Dict = None,
+    ) -> str:
+        """
+        Generate non-streaming response from LLM.
+
+        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 Triton server is healthy."""
+        try:
+            conn = self._get_connection()
+            conn.request("GET", "/v2/health/live")
+            response = conn.getresponse()
+            conn.close()
+            return response.status == 200
+        except Exception:
+            return False