tunacode-cli 0.0.41__py3-none-any.whl → 0.0.42__py3-none-any.whl

tunacode/core/agents/dspy_tunacode.py

@@ -0,0 +1,458 @@
+#!/usr/bin/env python3
+"""
+TunaCode DSPy Production Module
+
+Optimizes tool selection and task planning for TunaCode using DSPy.
+Includes 3-4 tool batching optimization for 3x performance gains.
+"""
+
+import json
+import logging
+import os
+import re
+from typing import Dict, List
+
+import dspy
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+# Tool Categories for TunaCode
+TOOL_CATEGORIES = {
+    "read_only": ["read_file", "grep", "list_dir", "glob"],  # Parallel-executable
+    "task_management": ["todo"],  # Fast, sequential
+    "write_execute": [
+        "write_file",
+        "update_file",
+        "run_command",
+        "bash",
+    ],  # Sequential, needs confirmation
+}
+
+ALL_TOOLS = [tool for category in TOOL_CATEGORIES.values() for tool in category]
+
+
+class ToolSelectionSignature(dspy.Signature):
+    """Select optimal tools with batching awareness."""
+
+    user_request: str = dspy.InputField(desc="The user's request or task")
+    current_directory: str = dspy.InputField(desc="Current working directory context")
+    tools_json: str = dspy.OutputField(
+        desc="JSON array of tool calls with batch grouping, e.g. [[tool1, tool2, tool3], [tool4]]"
+    )
+    requires_confirmation: bool = dspy.OutputField(
+        desc="Whether any tools require user confirmation"
+    )
+    reasoning: str = dspy.OutputField(desc="Explanation of tool choice and batching strategy")
+
+
+class TaskPlanningSignature(dspy.Signature):
+    """Break down complex tasks with tool hints."""
+
+    complex_request: str = dspy.InputField(desc="A complex task that needs breakdown")
+    subtasks_with_tools: str = dspy.OutputField(
+        desc="JSON array of {task, tools, priority} objects"
+    )
+    total_tool_calls: int = dspy.OutputField(desc="Estimated total number of tool calls")
+    requires_todo: bool = dspy.OutputField(desc="Whether todo tool should be used")
+    parallelization_opportunities: int = dspy.OutputField(
+        desc="Number of parallel execution opportunities"
+    )
+
+
+class PathValidationSignature(dspy.Signature):
+    """Validate and convert paths to relative format."""
+
+    path: str = dspy.InputField(desc="Path to validate")
+    current_directory: str = dspy.InputField(desc="Current working directory")
+    is_valid: bool = dspy.OutputField(desc="Whether path is valid relative path")
+    relative_path: str = dspy.OutputField(desc="Converted relative path")
+    reason: str = dspy.OutputField(desc="Validation result explanation")
+
+
+class OptimizedToolSelector(dspy.Module):
+    """Tool selection with batching optimization"""
+
+    def __init__(self):
+        self.predict = dspy.ChainOfThought(ToolSelectionSignature)
+
+    def forward(self, user_request: str, current_directory: str = "."):
+        logger.debug(f"Tool Selection for: {user_request}")
+        result = self.predict(user_request=user_request, current_directory=current_directory)
+
+        # Parse and validate tool batches
+        try:
+            tool_batches = json.loads(result.tools_json)
+            validated_batches = self._validate_batches(tool_batches)
+            result.tool_batches = validated_batches
+        except Exception as e:
+            logger.error(f"Failed to parse tool batches: {e}")
+            result.tool_batches = []
+
+        return result
+
+    def _validate_batches(self, batches: List[List[str]]) -> List[List[Dict]]:
+        """Validate and optimize tool batches for 3-4 tool rule"""
+        validated = []
+
+        for batch in batches:
+            # Check if batch contains only read-only tools
+            all_read_only = all(
+                any(tool_name in tool for tool_name in TOOL_CATEGORIES["read_only"])
+                for tool in batch
+            )
+
+            # Optimize batch size (3-4 tools is optimal)
+            if all_read_only and len(batch) > 4:
+                # Split large batches
+                for i in range(0, len(batch), 4):
+                    sub_batch = batch[i : i + 4]
+                    validated.append(self._parse_tools(sub_batch))
+            else:
+                validated_batch = self._parse_tools(batch)
+                if validated_batch:
+                    validated.append(validated_batch)
+
+        return validated
+
+    def _parse_tools(self, tools: List[str]) -> List[Dict]:
+        """Parse tool strings into proper format"""
+        parsed = []
+        for tool_str in tools:
+            # Extract tool name and args from string like "read_file('main.py')"
+            match = re.match(r"(\w+)\((.*)\)", tool_str)
+            if match:
+                tool_name = match.group(1)
+                args_str = match.group(2)
+
+                tool_dict = {"tool": tool_name, "args": {}}
+                if args_str:
+                    # Handle simple cases like 'file.py' or "pattern", "dir"
+                    args = [arg.strip().strip("\"'") for arg in args_str.split(",")]
+                    if tool_name == "read_file":
+                        tool_dict["args"]["filepath"] = args[0]
+                    elif tool_name == "grep":
+                        tool_dict["args"]["pattern"] = args[0]
+                        if len(args) > 1:
+                            tool_dict["args"]["directory"] = args[1]
+                    elif tool_name == "list_dir":
+                        tool_dict["args"]["directory"] = args[0] if args else "."
+                    elif tool_name == "glob":
+                        tool_dict["args"]["pattern"] = args[0]
+                        if len(args) > 1:
+                            tool_dict["args"]["directory"] = args[1]
+                    elif tool_name == "todo":
+                        tool_dict["args"]["action"] = args[0] if args else "list"
+                    elif tool_name in ["write_file", "update_file"]:
+                        if args:
+                            tool_dict["args"]["filepath"] = args[0]
+                    elif tool_name in ["run_command", "bash"]:
+                        if args:
+                            tool_dict["args"]["command"] = args[0]
+
+                parsed.append(tool_dict)
+
+        return parsed
+
+
+class EnhancedTaskPlanner(dspy.Module):
+    """Task planning with tool awareness"""
+
+    def __init__(self):
+        self.predict = dspy.ChainOfThought(TaskPlanningSignature)
+
+    def forward(self, complex_request: str):
+        logger.debug(f"Task Planning for: {complex_request}")
+        result = self.predict(complex_request=complex_request)
+
+        # Parse subtasks
+        try:
+            result.subtasks = json.loads(result.subtasks_with_tools)
+        except Exception as e:
+            logger.error(f"Failed to parse subtasks: {e}")
+            result.subtasks = []
+
+        return result
+
+
+class PathValidator(dspy.Module):
+    """Ensure all paths are relative"""
+
+    def __init__(self):
+        self.predict = dspy.Predict(PathValidationSignature)
+
+    def forward(self, path: str, current_directory: str = "."):
+        return self.predict(path=path, current_directory=current_directory)
+
+
+class TunaCodeDSPy(dspy.Module):
+    """Main TunaCode DSPy agent"""
+
+    def __init__(self):
+        self.tool_selector = OptimizedToolSelector()
+        self.task_planner = EnhancedTaskPlanner()
+        self.path_validator = PathValidator()
+
+    def forward(self, user_request: str, current_directory: str = "."):
+        """Process request with optimization"""
+
+        # Detect request complexity
+        is_complex = self._is_complex_task(user_request)
+
+        result = {
+            "request": user_request,
+            "is_complex": is_complex,
+            "current_directory": current_directory,
+        }
+
+        if is_complex:
+            # Use task planner for complex tasks
+            task_plan = self.task_planner(complex_request=user_request)
+            result["subtasks"] = task_plan.subtasks
+            result["total_tool_calls"] = task_plan.total_tool_calls
+            result["requires_todo"] = task_plan.requires_todo
+            result["parallelization_opportunities"] = task_plan.parallelization_opportunities
+
+            if task_plan.requires_todo:
+                result["initial_action"] = "Use todo tool to create task list"
+        else:
+            # Direct tool selection with batching
+            tool_selection = self.tool_selector(
+                user_request=user_request, current_directory=current_directory
+            )
+            result["tool_batches"] = tool_selection.tool_batches
+            result["requires_confirmation"] = tool_selection.requires_confirmation
+            result["reasoning"] = tool_selection.reasoning
+
+        return result
+
+    def _is_complex_task(self, request: str) -> bool:
+        """Detect if task is complex based on keywords and patterns"""
+        complex_indicators = [
+            "implement",
+            "create",
+            "build",
+            "refactor",
+            "add feature",
+            "fix all",
+            "update multiple",
+            "migrate",
+            "integrate",
+            "debug",
+            "optimize performance",
+            "add authentication",
+            "setup",
+            "configure",
+            "test suite",
+        ]
+
+        request_lower = request.lower()
+
+        # Check for multiple files mentioned
+        file_pattern = r"\b\w+\.\w+\b"
+        files_mentioned = len(re.findall(file_pattern, request)) > 2
+
+        # Check for complex keywords
+        has_complex_keyword = any(indicator in request_lower for indicator in complex_indicators)
+
+        # Check for multiple operations
+        operation_words = ["and", "then", "also", "after", "before", "plus"]
+        has_multiple_ops = sum(1 for word in operation_words if word in request_lower) >= 2
+
+        return files_mentioned or has_complex_keyword or has_multiple_ops
+
+    def validate_path(self, path: str, current_directory: str = ".") -> Dict:
+        """Validate a path is relative and safe"""
+        return self.path_validator(path, current_directory)
+
+
+def get_tool_selection_examples():
+    """Training examples for tool selection"""
+    return [
+        dspy.Example(
+            user_request="Show me the authentication system implementation",
+            current_directory=".",
+            tools_json='[["grep(\\"auth\\", \\"src/\\")", "list_dir(\\"src/auth/\\")", "glob(\\"**/*auth*.py\\")"]]',
+            requires_confirmation=False,
+            reasoning="Batch 3 read-only tools for parallel search - optimal performance",
+        ).with_inputs("user_request", "current_directory"),
+        dspy.Example(
+            user_request="Read all config files and the main module",
+            current_directory=".",
+            tools_json='[["read_file(\\"config.json\\")", "read_file(\\"settings.py\\")", "read_file(\\".env\\")", "read_file(\\"main.py\\")"]]',
+            requires_confirmation=False,
+            reasoning="Batch 4 file reads together - maximum optimal batch size",
+        ).with_inputs("user_request", "current_directory"),
+        dspy.Example(
+            user_request="Find the bug in validation and fix it",
+            current_directory=".",
+            tools_json='[["grep(\\"error\\", \\"logs/\\")", "grep(\\"validation\\", \\"src/\\")", "list_dir(\\"src/validators/\\")"], ["read_file(\\"src/validators/user.py\\")"], ["update_file(\\"src/validators/user.py\\", \\"old\\", \\"new\\")"]]',
+            requires_confirmation=True,
+            reasoning="Search tools batched, then read, then sequential write operation",
+        ).with_inputs("user_request", "current_directory"),
+    ]
+
+
+def get_task_planning_examples():
+    """Training examples for task planning"""
+    return [
+        dspy.Example(
+            complex_request="Implement user authentication system with JWT tokens",
+            subtasks_with_tools='[{"task": "Analyze current app structure", "tools": ["list_dir", "grep", "read_file"], "priority": "high"}, {"task": "Design user model", "tools": ["write_file"], "priority": "high"}, {"task": "Create auth endpoints", "tools": ["write_file", "update_file"], "priority": "high"}, {"task": "Add JWT tokens", "tools": ["write_file", "grep"], "priority": "high"}, {"task": "Write tests", "tools": ["write_file", "run_command"], "priority": "medium"}]',
+            total_tool_calls=15,
+            requires_todo=True,
+            parallelization_opportunities=3,
+        ).with_inputs("complex_request"),
+    ]
+
+
+def tool_selection_metric(example, prediction):
+    """Metric for tool selection evaluation"""
+    score = 0.0
+
+    # Tool accuracy (40%)
+    if hasattr(prediction, "tool_batches") and hasattr(example, "tools_json"):
+        try:
+            expected = json.loads(example.tools_json)
+            predicted = prediction.tool_batches
+
+            # Check tool selection accuracy
+            expected_tools = set()
+            predicted_tools = set()
+
+            for batch in expected:
+                for tool in batch:
+                    tool_name = re.match(r"(\w+)\(", tool)
+                    if tool_name:
+                        expected_tools.add(tool_name.group(1))
+
+            for batch in predicted:
+                for tool in batch:
+                    if isinstance(tool, dict):
+                        predicted_tools.add(tool.get("tool", ""))
+
+            if expected_tools == predicted_tools:
+                score += 0.4
+            else:
+                overlap = len(expected_tools & predicted_tools)
+                total = len(expected_tools | predicted_tools)
+                if total > 0:
+                    score += 0.4 * (overlap / total)
+        except Exception:
+            pass
+
+    # Batching optimization (30%)
+    if hasattr(prediction, "tool_batches"):
+        batches = prediction.tool_batches
+        optimal_batching = True
+
+        for batch in batches:
+            if len(batch) > 0:
+                batch_tools = [tool.get("tool", "") for tool in batch if isinstance(tool, dict)]
+                if all(tool in TOOL_CATEGORIES["read_only"] for tool in batch_tools):
+                    if len(batch) < 3 or len(batch) > 4:
+                        optimal_batching = False
+                        break
+
+        if optimal_batching:
+            score += 0.3
+
+    # Confirmation accuracy (20%)
+    if hasattr(prediction, "requires_confirmation") and hasattr(example, "requires_confirmation"):
+        if prediction.requires_confirmation == example.requires_confirmation:
+            score += 0.2
+
+    # Reasoning quality (10%)
+    if hasattr(prediction, "reasoning") and prediction.reasoning and len(prediction.reasoning) > 20:
+        score += 0.1
+
+    return score
+
+
+def task_planning_metric(example, prediction):
+    """Metric for task planning evaluation"""
+    score = 0.0
+
+    # Subtask quality (30%)
+    if hasattr(prediction, "subtasks") and hasattr(example, "subtasks_with_tools"):
+        try:
+            expected = json.loads(example.subtasks_with_tools)
+            predicted = prediction.subtasks
+
+            if abs(len(expected) - len(predicted)) <= 1:
+                score += 0.3
+            elif abs(len(expected) - len(predicted)) <= 2:
+                score += 0.15
+        except Exception:
+            pass
+
+    # Tool estimation accuracy (30%)
+    if hasattr(prediction, "total_tool_calls") and hasattr(example, "total_tool_calls"):
+        if abs(prediction.total_tool_calls - example.total_tool_calls) <= 5:
+            score += 0.3
+        elif abs(prediction.total_tool_calls - example.total_tool_calls) <= 10:
+            score += 0.15
+
+    # Todo requirement (20%)
+    if hasattr(prediction, "requires_todo") and hasattr(example, "requires_todo"):
+        if prediction.requires_todo == example.requires_todo:
+            score += 0.2
+
+    # Parallelization awareness (20%)
+    if hasattr(prediction, "parallelization_opportunities") and hasattr(
+        example, "parallelization_opportunities"
+    ):
+        if (
+            abs(prediction.parallelization_opportunities - example.parallelization_opportunities)
+            <= 2
+        ):
+            score += 0.2
+
+    return score
+
+
+def create_optimized_agent(api_key: str = None, model: str = "openrouter/openai/gpt-4.1-mini"):
+    """Create and optimize the TunaCode DSPy agent"""
+
+    # Configure DSPy
+    if not api_key:
+        api_key = os.getenv("OPENROUTER_API_KEY")
+        if not api_key:
+            raise ValueError("Please set OPENROUTER_API_KEY environment variable")
+
+    lm = dspy.LM(
+        model,
+        api_base="https://openrouter.ai/api/v1",
+        api_key=api_key,
+        temperature=0.3,
+    )
+    dspy.configure(lm=lm)
+
+    # Create agent
+    agent = TunaCodeDSPy()
+
+    # Optimize tool selector
+    tool_examples = get_tool_selection_examples()
+    tool_optimizer = dspy.BootstrapFewShot(
+        metric=lambda ex, pred, trace: tool_selection_metric(ex, pred),
+        max_bootstrapped_demos=3,
+    )
+    agent.tool_selector = tool_optimizer.compile(agent.tool_selector, trainset=tool_examples)
+
+    # Optimize task planner
+    task_examples = get_task_planning_examples()
+    task_optimizer = dspy.BootstrapFewShot(
+        metric=lambda ex, pred, trace: task_planning_metric(ex, pred),
+        max_bootstrapped_demos=2,
+    )
+    agent.task_planner = task_optimizer.compile(agent.task_planner, trainset=task_examples)
+
+    return agent
+
+
+agent = create_optimized_agent()
+result = agent("Show me the authentication implementation", ".")

tunacode/core/agents/main.py

@@ -30,9 +30,12 @@ except ImportError:
     STREAMING_AVAILABLE = False
 
 from tunacode.constants import READ_ONLY_TOOLS
+from tunacode.core.agents.dspy_integration import DSPyIntegration
+from tunacode.core.recursive import RecursiveTaskExecutor
 from tunacode.core.state import StateManager
 from tunacode.core.token_usage.api_response_parser import ApiResponseParser
 from tunacode.core.token_usage.cost_calculator import CostCalculator
+from tunacode.exceptions import ToolBatchingJSONError
 from tunacode.services.mcp import get_mcp_servers
 from tunacode.tools.bash import bash
 from tunacode.tools.glob import glob
@@ -471,9 +474,17 @@ async def _process_node(
         if not has_tool_calls and buffering_callback:
             for part in node.model_response.parts:
                 if hasattr(part, "content") and isinstance(part.content, str):
-                    await extract_and_execute_tool_calls(
-                        part.content, buffering_callback, state_manager
-                    )
+                    try:
+                        await extract_and_execute_tool_calls(
+                            part.content, buffering_callback, state_manager
+                        )
+                    except ToolBatchingJSONError as e:
+                        # Handle JSON parsing failure after retries
+                        logger.error(f"Tool batching JSON error: {e}")
+                        if state_manager.session.show_thoughts:
+                            await ui.error(str(e))
+                        # Continue processing other parts instead of failing completely
+                        continue
 
     # Final flush: disabled temporarily while fixing the parallel execution design
     # The buffer is not being used in the current implementation
@@ -508,6 +519,18 @@ def get_or_create_agent(model: ModelName, state_manager: StateManager) -> Pydant
                 # Use a default system prompt if neither file exists
                 system_prompt = "You are a helpful AI assistant for software development tasks."
 
+        # Enhance with DSPy optimization if enabled
+        use_dspy = state_manager.session.user_config.get("settings", {}).get(
+            "use_dspy_optimization", True
+        )
+        if use_dspy:
+            try:
+                dspy_integration = DSPyIntegration(state_manager)
+                system_prompt = dspy_integration.enhance_system_prompt(system_prompt)
+                logger.info("Enhanced system prompt with DSPy optimizations")
+            except Exception as e:
+                logger.warning(f"Failed to enhance prompt with DSPy: {e}")
+
         # Load TUNACODE.md context
         # Use sync version of get_code_style to avoid nested event loop issues
         try:
@@ -749,6 +772,110 @@ async def process_request(
     fallback_enabled = state_manager.session.user_config.get("settings", {}).get(
         "fallback_response", True
     )
+
+    # Check if DSPy optimization is enabled and if this is a complex task
+    use_dspy = state_manager.session.user_config.get("settings", {}).get(
+        "use_dspy_optimization", True
+    )
+    dspy_integration = None
+    task_breakdown = None
+
+    # Check if recursive execution is enabled
+    use_recursive = state_manager.session.user_config.get("settings", {}).get(
+        "use_recursive_execution", True
+    )
+    recursive_threshold = state_manager.session.user_config.get("settings", {}).get(
+        "recursive_complexity_threshold", 0.7
+    )
+
+    if use_dspy:
+        try:
+            dspy_integration = DSPyIntegration(state_manager)
+
+            # Check if this is a complex task that needs planning
+            if dspy_integration.should_use_task_planner(message):
+                task_breakdown = dspy_integration.get_task_breakdown(message)
+                if task_breakdown and task_breakdown.get("requires_todo"):
+                    # Auto-create todos for complex tasks
+                    from tunacode.tools.todo import TodoTool
+
+                    todo_tool = TodoTool(state_manager=state_manager)
+
+                    if state_manager.session.show_thoughts:
+                        from tunacode.ui import console as ui
+
+                        await ui.muted("DSPy: Detected complex task - creating todo list")
+
+                    # Create todos from subtasks
+                    todos = []
+                    for subtask in task_breakdown["subtasks"][:5]:  # Limit to first 5
+                        todos.append(
+                            {
+                                "content": subtask["task"],
+                                "priority": subtask.get("priority", "medium"),
+                            }
+                        )
+
+                    if todos:
+                        await todo_tool._execute(action="add_multiple", todos=todos)
+        except Exception as e:
+            logger.warning(f"DSPy task planning failed: {e}")
+
+    # Check if recursive execution should be used
+    if use_recursive and state_manager.session.current_recursion_depth == 0:
+        try:
+            # Initialize recursive executor
+            recursive_executor = RecursiveTaskExecutor(
+                state_manager=state_manager,
+                max_depth=state_manager.session.max_recursion_depth,
+                min_complexity_threshold=recursive_threshold,
+                default_iteration_budget=max_iterations,
+            )
+
+            # Analyze task complexity
+            complexity_result = await recursive_executor.decomposer.analyze_and_decompose(message)
+
+            if (
+                complexity_result.should_decompose
+                and complexity_result.total_complexity >= recursive_threshold
+            ):
+                if state_manager.session.show_thoughts:
+                    from tunacode.ui import console as ui
+
+                    await ui.muted(
+                        f"\n🔄 RECURSIVE EXECUTION: Task complexity {complexity_result.total_complexity:.2f} >= {recursive_threshold}"
+                    )
+                    await ui.muted(f"Reasoning: {complexity_result.reasoning}")
+                    await ui.muted(f"Subtasks: {len(complexity_result.subtasks)}")
+
+                # Execute recursively
+                success, result, error = await recursive_executor.execute_task(
+                    request=message, parent_task_id=None, depth=0
+                )
+
+                # Create AgentRun response
+                from datetime import datetime
+
+                if success:
+                    return AgentRun(
+                        messages=[{"role": "assistant", "content": str(result)}],
+                        timestamp=datetime.now(),
+                        model=model,
+                        iterations=1,
+                        status="success",
+                    )
+                else:
+                    return AgentRun(
+                        messages=[{"role": "assistant", "content": f"Task failed: {error}"}],
+                        timestamp=datetime.now(),
+                        model=model,
+                        iterations=1,
+                        status="error",
+                    )
+        except Exception as e:
+            logger.warning(f"Recursive execution failed, falling back to normal: {e}")
+            # Continue with normal execution
+
     from tunacode.configuration.models import ModelRegistry
     from tunacode.core.token_usage.usage_tracker import UsageTracker