kairo-code 0.1.0__py3-none-any.whl

kairo_code/tools/base.py

@@ -0,0 +1,417 @@
+"""Tool system with robust registry and execution"""
+
+import json
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+
+@dataclass
+class ToolResult:
+    """Result from a tool execution."""
+    success: bool
+    output: str
+    error: str | None = None
+
+
+@dataclass
+class ParsedToolCall:
+    """A parsed tool call from LLM output."""
+    tool_name: str
+    params: dict
+    raw_text: str
+    parse_error: str | None = None
+
+
+@dataclass
+class ToolParseResult:
+    """Result of parsing tool calls from text."""
+    calls: list[ParsedToolCall] = field(default_factory=list)
+    has_thinking: bool = False
+    thinking_content: str = ""
+    raw_text: str = ""
+    parse_errors: list[str] = field(default_factory=list)
+
+
+class Tool(ABC):
+    """Base class for all tools."""
+
+    name: str
+    description: str
+    parameters: dict  # JSON schema for parameters
+
+    @abstractmethod
+    def execute(self, **kwargs) -> ToolResult:
+        """Execute the tool with given parameters."""
+        pass
+
+    def to_schema(self) -> dict:
+        """Generate schema for LLM tool calling."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": self.parameters,
+        }
+
+
+class ToolRegistry:
+    """Registry for managing available tools."""
+
+    def __init__(self):
+        self._tools: dict[str, Tool] = {}
+
+    def register(self, tool: Tool) -> None:
+        """Register a tool."""
+        self._tools[tool.name] = tool
+
+    def get(self, name: str) -> Tool | None:
+        """Get a tool by name, with strict matching only."""
+        # Normalize the name
+        name_lower = name.lower().strip()
+
+        # Exact match first
+        if name in self._tools:
+            return self._tools[name]
+
+        # Case-insensitive exact match
+        for tool_name in self._tools:
+            if tool_name.lower() == name_lower:
+                return self._tools[tool_name]
+
+        # Handle common variations (underscore vs no underscore)
+        name_normalized = name_lower.replace("_", "").replace("-", "")
+        for tool_name in self._tools:
+            if tool_name.lower().replace("_", "").replace("-", "") == name_normalized:
+                return self._tools[tool_name]
+
+        # No fuzzy matching - return None and let caller handle the error
+        # This prevents confusing matches like "write" -> "read_file"
+        return None
+
+    def get_tool_names(self) -> list[str]:
+        """Get list of all tool names."""
+        return list(self._tools.keys())
+
+    def list_tools(self) -> list[Tool]:
+        """List all registered tools."""
+        return list(self._tools.values())
+
+    def get_schemas(self) -> list[dict]:
+        """Get schemas for all tools."""
+        return [t.to_schema() for t in self._tools.values()]
+
+    def format_for_prompt(self) -> str:
+        """Format tools for inclusion in system prompt with full details."""
+        lines = ["## Available Tools\n"]
+
+        for tool in self._tools.values():
+            lines.append(f"### {tool.name}")
+            lines.append(f"{tool.description}\n")
+
+            props = tool.parameters.get("properties", {})
+            required = tool.parameters.get("required", [])
+
+            if props:
+                lines.append("Parameters:")
+                for param_name, param_info in props.items():
+                    req = "(required)" if param_name in required else "(optional)"
+                    desc = param_info.get("description", "")
+                    lines.append(f"  - {param_name} {req}: {desc}")
+
+            lines.append("")
+
+        return "\n".join(lines)
+
+
+def parse_tool_calls(text: str) -> ToolParseResult:
+    """
+    Robustly parse tool calls from LLM output.
+
+    Supports multiple formats:
+    1. XML-style: <tool>name</tool><params>{...}</params>
+    2. Bracket-style: [TOOL: name][PARAMS: {...}]
+    3. Action-style: ACTION: name\nPARAMS: {...}
+    4. Function-style: name({...})
+
+    Also extracts thinking blocks.
+    """
+    result = ToolParseResult(raw_text=text)
+
+    # Extract thinking blocks first
+    thinking_patterns = [
+        r"<thinking>(.*?)</thinking>",
+        r"<scratchpad>(.*?)</scratchpad>",
+        r"\[THINKING\](.*?)\[/THINKING\]",
+        r"```thinking\n(.*?)```",
+    ]
+
+    for pattern in thinking_patterns:
+        match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
+        if match:
+            result.has_thinking = True
+            result.thinking_content = match.group(1).strip()
+            break
+
+    # Try multiple parsing strategies
+    calls = []
+
+    # Strategy 1: XML-style (most structured)
+    # Handle various whitespace issues
+    xml_pattern = r"<tool>\s*(\w+)\s*</tool>\s*<params>\s*(.*?)\s*</params>"
+    for match in re.finditer(xml_pattern, text, re.DOTALL | re.IGNORECASE):
+        parsed = _parse_params(match.group(2), match.group(1))
+        if parsed:
+            calls.append(parsed)
+        else:
+            result.parse_errors.append(f"Failed to parse params for {match.group(1)}")
+
+    # Strategy 2: Bracket-style
+    bracket_pattern = r"\[TOOL:\s*(\w+)\s*\]\s*\[PARAMS:\s*(.*?)\s*\]"
+    for match in re.finditer(bracket_pattern, text, re.DOTALL | re.IGNORECASE):
+        parsed = _parse_params(match.group(2), match.group(1))
+        if parsed:
+            calls.append(parsed)
+
+    # Strategy 3: Action-style (newline separated)
+    action_pattern = r"ACTION:\s*(\w+)\s*\nPARAMS:\s*(\{.*?\})"
+    for match in re.finditer(action_pattern, text, re.DOTALL):
+        parsed = _parse_params(match.group(2), match.group(1))
+        if parsed:
+            calls.append(parsed)
+
+    # Strategy 4: Markdown code block with tool call
+    code_pattern = r"```(?:tool|action)\n(\w+)\n(\{.*?\})\n```"
+    for match in re.finditer(code_pattern, text, re.DOTALL):
+        parsed = _parse_params(match.group(2), match.group(1))
+        if parsed:
+            calls.append(parsed)
+
+    # Strategy 5: Simple function call style: tool_name({"param": "value"})
+    func_pattern = r"(\w+)\s*\(\s*(\{.*?\})\s*\)"
+    for match in re.finditer(func_pattern, text, re.DOTALL):
+        # Only accept if it looks like a tool name (not random function)
+        name = match.group(1).lower()
+        if name in ['read_file', 'write_file', 'edit_file', 'list_files',
+                     'search_files', 'web_search', 'bash', 'grep', 'tree']:
+            parsed = _parse_params(match.group(2), match.group(1))
+            if parsed:
+                calls.append(parsed)
+
+    # Strategy 6: Llama3-style direct tags: <tool_name>params</tool_name>
+    # e.g., <web_search>query="test"</web_search>
+    tool_names = ['read_file', 'write_file', 'edit_file', 'list_files',
+                  'search_files', 'web_search', 'bash', 'grep', 'tree']
+    for tool_name in tool_names:
+        llama_pattern = rf"<{tool_name}>(.*?)</{tool_name}>"
+        for match in re.finditer(llama_pattern, text, re.DOTALL | re.IGNORECASE):
+            params_str = match.group(1).strip()
+            parsed = _parse_llama_params(params_str, tool_name)
+            if parsed:
+                calls.append(parsed)
+
+    # Strategy 7: write_file with separate path and content blocks
+    # Handles: <write_file path="app.py">code here</write_file>
+    # or: <write_file><path>app.py</path><content>code</content></write_file>
+    write_pattern_attr = r'<write_file\s+path=["\']([^"\']+)["\']\s*>(.*?)</write_file>'
+    for match in re.finditer(write_pattern_attr, text, re.DOTALL | re.IGNORECASE):
+        calls.append(ParsedToolCall(
+            tool_name="write_file",
+            params={"path": match.group(1), "content": match.group(2).strip()},
+            raw_text=match.group(0),
+        ))
+
+    # Strategy 8: write_file with nested tags
+    write_pattern_nested = r'<write_file>\s*<path>([^<]+)</path>\s*<content>(.*?)</content>\s*</write_file>'
+    for match in re.finditer(write_pattern_nested, text, re.DOTALL | re.IGNORECASE):
+        calls.append(ParsedToolCall(
+            tool_name="write_file",
+            params={"path": match.group(1).strip(), "content": match.group(2).strip()},
+            raw_text=match.group(0),
+        ))
+
+    # Strategy 9: edit_file with nested tags
+    # <edit_file path="file.py"><old>old code</old><new>new code</new></edit_file>
+    edit_pattern_nested = r'<edit_file\s+path=["\']([^"\']+)["\']\s*>\s*<old>(.*?)</old>\s*<new>(.*?)</new>\s*</edit_file>'
+    for match in re.finditer(edit_pattern_nested, text, re.DOTALL | re.IGNORECASE):
+        calls.append(ParsedToolCall(
+            tool_name="edit_file",
+            params={
+                "path": match.group(1).strip(),
+                "old_string": match.group(2).strip(),
+                "new_string": match.group(3).strip(),
+            },
+            raw_text=match.group(0),
+        ))
+
+    # Deduplicate calls (same tool + same params)
+    seen = set()
+    unique_calls = []
+    for call in calls:
+        key = (call.tool_name, json.dumps(call.params, sort_keys=True))
+        if key not in seen:
+            seen.add(key)
+            unique_calls.append(call)
+
+    result.calls = unique_calls
+    return result
+
+
+def _parse_llama_params(params_str: str, tool_name: str) -> ParsedToolCall | None:
+    """Parse llama3-style params like: query="test" or path="file", content="code" """
+    params_str = params_str.strip()
+    params = {}
+
+    # Pattern for key="value" or key='value'
+    kv_pattern = r'(\w+)\s*=\s*["\']([^"\']*)["\']'
+    for match in re.finditer(kv_pattern, params_str):
+        params[match.group(1)] = match.group(2)
+
+    # Pattern for key={...} (JSON objects)
+    json_pattern = r'(\w+)\s*=\s*(\{[^}]*\})'
+    for match in re.finditer(json_pattern, params_str):
+        try:
+            params[match.group(1)] = json.loads(match.group(2))
+        except Exception:
+            params[match.group(1)] = match.group(2)
+
+    # If no key=value params found, treat raw text as the primary parameter
+    # based on tool type (common llama3 behavior)
+    if not params and params_str:
+        primary_param_map = {
+            'web_search': 'query',
+            'search_files': 'pattern',
+            'grep': 'pattern',
+            'bash': 'command',
+            'read_file': 'path',
+        }
+        primary_param = primary_param_map.get(tool_name.lower())
+        if primary_param:
+            params[primary_param] = params_str
+
+    if params:
+        return ParsedToolCall(
+            tool_name=tool_name.lower().strip(),
+            params=params,
+            raw_text=params_str,
+        )
+    return None
+
+
+def _parse_params(params_str: str, tool_name: str) -> ParsedToolCall | None:
+    """Parse parameters JSON with error recovery."""
+    params_str = params_str.strip()
+
+    # Fix 0: Handle Python triple-quoted strings in JSON (common LLM mistake)
+    # Convert: {"content": """code"""} to {"content": "code"}
+    triple_quote_pattern = r'"""\s*(.*?)\s*"""'
+    if '"""' in params_str:
+        def replace_triple(match):
+            content = match.group(1)
+            # Escape for JSON
+            content = content.replace('\\', '\\\\')
+            content = content.replace('"', '\\"')
+            content = content.replace('\n', '\\n')
+            content = content.replace('\t', '\\t')
+            return f'"{content}"'
+        params_str = re.sub(triple_quote_pattern, replace_triple, params_str, flags=re.DOTALL)
+
+    # Try direct JSON parse
+    try:
+        params = json.loads(params_str)
+        return ParsedToolCall(
+            tool_name=tool_name.lower().strip(),
+            params=params,
+            raw_text=params_str,
+        )
+    except json.JSONDecodeError:
+        pass
+
+    # Try fixing common issues
+
+    # Fix 1: Single quotes to double quotes
+    fixed = params_str.replace("'", '"')
+    try:
+        params = json.loads(fixed)
+        return ParsedToolCall(
+            tool_name=tool_name.lower().strip(),
+            params=params,
+            raw_text=params_str,
+        )
+    except json.JSONDecodeError:
+        pass
+
+    # Fix 2: Add missing braces
+    if not params_str.startswith("{"):
+        params_str = "{" + params_str
+    if not params_str.endswith("}"):
+        params_str = params_str + "}"
+    try:
+        params = json.loads(params_str)
+        return ParsedToolCall(
+            tool_name=tool_name.lower().strip(),
+            params=params,
+            raw_text=params_str,
+        )
+    except json.JSONDecodeError:
+        pass
+
+    # Fix 3: Try to extract key-value pairs manually
+    try:
+        # Pattern: "key": "value" or "key": value
+        kv_pattern = r'"(\w+)":\s*"([^"]*)"'
+        matches = re.findall(kv_pattern, params_str)
+        if matches:
+            params = dict(matches)
+            return ParsedToolCall(
+                tool_name=tool_name.lower().strip(),
+                params=params,
+                raw_text=params_str,
+            )
+    except Exception:
+        pass
+
+    return None
+
+
+def format_tool_error(error: str, tool_name: str = None) -> str:
+    """Format a tool error message for feeding back to the model."""
+    if tool_name:
+        return f"""[TOOL ERROR]
+Tool: {tool_name}
+Error: {error}
+
+Please fix the issue and try again. Make sure to use the correct format:
+<tool>{tool_name}</tool>
+<params>{{"param": "value"}}</params>
+"""
+    return f"""[PARSE ERROR]
+{error}
+
+Your tool call could not be parsed. Please use this exact format:
+<tool>tool_name</tool>
+<params>{{"param": "value"}}</params>
+"""
+
+
+class FunctionTool(Tool):
+    """Tool that wraps a simple function."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        func: Callable,
+        parameters: dict,
+    ):
+        self.name = name
+        self.description = description
+        self.func = func
+        self.parameters = parameters
+
+    def execute(self, **kwargs) -> ToolResult:
+        try:
+            result = self.func(**kwargs)
+            return ToolResult(success=True, output=str(result))
+        except Exception as e:
+            return ToolResult(success=False, output="", error=str(e))

kairo_code/tools/code.py

@@ -0,0 +1,58 @@
+"""Code generation and parsing utilities"""
+
+import re
+
+
+def extract_code_blocks(text: str) -> list[dict]:
+    """
+    Extract code blocks from markdown-formatted text.
+
+    Returns:
+        List of dicts with 'language' and 'code' keys
+    """
+    pattern = r"```(\w*)\n(.*?)```"
+    matches = re.findall(pattern, text, re.DOTALL)
+
+    return [
+        {"language": lang or "text", "code": code.strip()}
+        for lang, code in matches
+    ]
+
+
+def build_code_prompt(
+    request: str,
+    context: str | None = None,
+    language: str | None = None,
+) -> str:
+    """
+    Build a prompt for code generation.
+
+    Args:
+        request: What the user wants
+        context: Additional context (file contents, search results)
+        language: Preferred programming language
+
+    Returns:
+        Formatted prompt string
+    """
+    parts = []
+
+    if context:
+        parts.append(f"Context:\n{context}\n")
+
+    if language:
+        parts.append(f"Language: {language}\n")
+
+    parts.append(f"Request: {request}")
+
+    return "\n".join(parts)
+
+
+def build_explain_prompt(code: str, question: str | None = None) -> str:
+    """Build a prompt for code explanation."""
+    prompt = f"Explain this code:\n\n```\n{code}\n```"
+
+    if question:
+        prompt += f"\n\nSpecifically: {question}"
+
+    return prompt