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

tunacode/core/token_usage/cost_calculator.py

@@ -0,0 +1,58 @@
+"""
+Module: tunacode.pricing.cost_calculator
+Provides a utility for calculating the cost of model usage based on token counts.
+"""
+
+from tunacode.configuration.models import ModelRegistry
+from tunacode.types import CostAmount, ModelName, TokenCount
+
+
+class CostCalculator:
+    """
+    Calculates the cost of a model interaction based on prompt and completion tokens.
+    """
+
+    def __init__(self, registry: ModelRegistry):
+        """
+        Initializes the CostCalculator with a model registry.
+
+        Args:
+            registry (ModelRegistry): An instance of ModelRegistry that contains
+                                      the pricing information for various models.
+        """
+        self._registry = registry
+
+    def calculate_cost(
+        self,
+        model_name: ModelName,
+        prompt_tokens: TokenCount,
+        completion_tokens: TokenCount,
+    ) -> CostAmount:
+        """
+        Calculates the total cost for a given model and token usage.
+
+        Args:
+            model_name (ModelName): The identifier for the model (e.g., "openai:gpt-4o").
+            prompt_tokens (TokenCount): The number of tokens in the input/prompt.
+            completion_tokens (TokenCount): The number of tokens in the output/completion.
+
+        Returns:
+            CostAmount: The calculated cost as a float. Returns 0.0 if the model
+                        is not found in the registry.
+        """
+        model_config = self._registry.get_model(model_name)
+
+        if not model_config:
+            return 0.0
+
+        TOKENS_PER_MILLION = 1_000_000
+
+        pricing = model_config.pricing
+
+        input_cost = (prompt_tokens / TOKENS_PER_MILLION) * pricing.input
+
+        output_cost = (completion_tokens / TOKENS_PER_MILLION) * pricing.output
+
+        total_cost = input_cost + output_cost
+
+        return total_cost

tunacode/core/token_usage/usage_tracker.py

@@ -0,0 +1,98 @@
+from typing import Any
+
+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.types import UsageTrackerProtocol
+from tunacode.ui import console as ui  # Import the ui console directly
+
+
+class UsageTracker(UsageTrackerProtocol):
+    """
+    Handles parsing, calculating, storing, and displaying token usage and cost.
+    """
+
+    def __init__(
+        self,
+        parser: ApiResponseParser,
+        calculator: CostCalculator,
+        state_manager: StateManager,
+    ):
+        self.parser = parser
+        self.calculator = calculator
+        self.state_manager = state_manager
+
+    async def track_and_display(self, response_obj: Any):
+        """
+        Main method to process a model response for usage tracking.
+        """
+        try:
+            # 1. Parse the response to get token data
+            requested_model = self.state_manager.session.current_model
+            parsed_data = self.parser.parse(model=requested_model, response_obj=response_obj)
+
+            if not parsed_data:
+                return
+
+            # 2. Calculate the cost
+            cost = self._calculate_cost(parsed_data)
+
+            # 3. Update the session state
+            self._update_state(parsed_data, cost)
+
+            # 4. Display the summary if enabled
+            if self.state_manager.session.show_thoughts:
+                await self._display_summary()
+
+        except Exception as e:
+            if self.state_manager.session.show_thoughts:
+                await ui.error(f"Error during cost calculation: {e}")
+
+    def _calculate_cost(self, parsed_data: dict) -> float:
+        """Calculates the cost for the given parsed data."""
+        requested_model = self.state_manager.session.current_model
+        api_model_name = parsed_data.get("model_name", requested_model)
+        final_model_name = api_model_name
+
+        # Logic to preserve the provider prefix
+        if ":" in requested_model:
+            provider_prefix = requested_model.split(":", 1)[0]
+            if not api_model_name.startswith(provider_prefix + ":"):
+                final_model_name = f"{provider_prefix}:{api_model_name}"
+
+        return self.calculator.calculate_cost(
+            prompt_tokens=parsed_data.get("prompt_tokens", 0),
+            completion_tokens=parsed_data.get("completion_tokens", 0),
+            model_name=final_model_name,
+        )
+
+    def _update_state(self, parsed_data: dict, cost: float):
+        """Updates the last_call and session_total usage in the state."""
+        session = self.state_manager.session
+        prompt_tokens = parsed_data.get("prompt_tokens", 0)
+        completion_tokens = parsed_data.get("completion_tokens", 0)
+
+        # Update last call usage
+        session.last_call_usage["prompt_tokens"] = prompt_tokens
+        session.last_call_usage["completion_tokens"] = completion_tokens
+        session.last_call_usage["cost"] = cost
+
+        # Accumulate session totals
+        session.session_total_usage["prompt_tokens"] += prompt_tokens
+        session.session_total_usage["completion_tokens"] += completion_tokens
+        session.session_total_usage["cost"] += cost
+
+    async def _display_summary(self):
+        """Formats and prints the usage summary to the console."""
+        session = self.state_manager.session
+        prompt = session.last_call_usage["prompt_tokens"]
+        completion = session.last_call_usage["completion_tokens"]
+        last_cost = session.last_call_usage["cost"]
+        session_cost = session.session_total_usage["cost"]
+
+        usage_summary = (
+            f"[ Tokens: {prompt + completion:,} (P: {prompt:,}, C: {completion:,}) | "
+            f"Cost: ${last_cost:.4f} | "
+            f"Session Total: ${session_cost:.4f} ]"
+        )
+        await ui.muted(usage_summary)

tunacode/exceptions.py

@@ -114,3 +114,26 @@ class TooBroadPatternError(ToolExecutionError):
             f"Pattern '{pattern}' is too broad - no matches found within {timeout_seconds}s. "
             "Please use a more specific pattern.",
         )
+
+
+class ToolBatchingJSONError(TunaCodeError):
+    """Raised when JSON parsing fails during tool batching after all retries are exhausted."""
+
+    def __init__(
+        self,
+        json_content: str,
+        retry_count: int,
+        original_error: OriginalError = None,
+    ):
+        self.json_content = json_content
+        self.retry_count = retry_count
+        self.original_error = original_error
+
+        # Truncate JSON content for display if too long
+        display_content = json_content[:100] + "..." if len(json_content) > 100 else json_content
+
+        super().__init__(
+            f"The model is having issues with tool batching. "
+            f"JSON parsing failed after {retry_count} retries. "
+            f"Invalid JSON: {display_content}"
+        )

tunacode/prompts/dspy_task_planning.md

@@ -0,0 +1,45 @@
+# DSPy Task Planning Prompt
+
+This prompt is automatically generated and optimized by DSPy for TunaCode's complex task planning.
+
+## Format
+
+Given the fields `complex_request`, produce the fields `subtasks_with_tools`, `total_tool_calls`, `requires_todo`, `parallelization_opportunities`.
+
+---
+
+Follow the following format.
+
+Complex Request: A complex task that needs breakdown
+Reasoning: Let's think step by step in order to ${produce the answer}. We need to break this down into manageable subtasks...
+Subtasks With Tools: JSON array of {task, tools, priority} objects
+Total Tool Calls: Estimated total number of tool calls
+Requires Todo: Whether todo tool should be used
+Parallelization Opportunities: Number of parallel execution opportunities
+
+---
+
+## Learned Example (Selected by DSPy)
+
+### Example: Authentication System Implementation
+Complex Request: Implement user authentication system with JWT tokens
+Reasoning: Let's think step by step. This is a complex implementation task that requires multiple components. First, I need to analyze the current application structure to understand how to integrate authentication. Then I'll design the user model and database schema. Next, I'll create the authentication endpoints for registration and login. I'll need to implement JWT token generation and validation. Finally, I'll add tests to ensure everything works correctly. Several of these tasks can be done in parallel once the initial analysis is complete.
+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
+
+---
+
+## Key Patterns for Complex Tasks
+
+1. **Break Down First**: Start with analysis/exploration before implementation
+2. **Priority Levels**: High for core functionality, medium for tests/docs, low for nice-to-haves
+3. **Tool Grouping**: Group related tools together for each subtask
+4. **Todo Usage**: Use todo tool for tasks with 5+ subtasks
+5. **Parallelization**: Identify independent subtasks that can run concurrently
+
+---
+
+Complex Request: ${complex_request}
+Reasoning: Let's think step by step...

tunacode/prompts/dspy_tool_selection.md

@@ -0,0 +1,58 @@
+# DSPy Tool Selection Prompt
+
+This prompt is automatically generated and optimized by DSPy for TunaCode's tool selection.
+
+## Format
+
+Given the fields `user_request`, `current_directory`, produce the fields `tools_json`, `requires_confirmation`, `reasoning`.
+
+---
+
+Follow the following format.
+
+User Request: The user's request or task
+Current Directory: Current working directory context
+Reasoning: Let's think step by step in order to ${produce the answer}. We ...
+Tools Json: JSON array of tool calls with batch grouping, e.g. [[tool1, tool2, tool3], [tool4]]
+Requires Confirmation: Whether any tools require user confirmation
+
+---
+
+## Learned Examples (Automatically Selected by DSPy)
+
+### Example 1: Searching for Implementation
+User Request: Show me the authentication system implementation
+Current Directory: .
+Reasoning: Let's think step by step. To show the authentication implementation, I need to search for auth-related files across the codebase. I'll use grep to find files containing 'auth', list the auth directory if it exists, and use glob to find all auth-related Python files. These are all read-only operations that can be executed in parallel.
+Tools Json: [["grep(\"auth\", \"src/\")", "list_dir(\"src/auth/\")", "glob(\"**/*auth*.py\")"]]
+Requires Confirmation: false
+
+### Example 2: Reading Multiple Files (Optimal Batching)
+User Request: Read all config files and the main module  
+Current Directory: .
+Reasoning: Let's think step by step. I need to read multiple specific files. All of these are read operations that can be batched together for parallel execution. I'll batch them in a group of 4 for optimal performance.
+Tools Json: [["read_file(\"config.json\")", "read_file(\"settings.py\")", "read_file(\".env\")", "read_file(\"main.py\")"]]
+Requires Confirmation: false
+
+### Example 3: Search, Read, then Modify Pattern
+User Request: Find the bug in validation and fix it
+Current Directory: .
+Reasoning: Let's think step by step. First, I need to search for validation-related code and errors. I'll use grep to search for error patterns and validation code, and list the validators directory. These search operations can be parallelized. After finding the issue, I'll need to read the specific file and then update it to fix the bug.
+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
+
+---
+
+## Key Patterns Learned by DSPy
+
+1. **3-4 Tool Batching**: Optimal batch size for parallel read-only operations
+2. **Read-Only Parallelization**: grep, list_dir, glob, read_file can run in parallel
+3. **Sequential Writes**: write_file, update_file, run_command, bash must run sequentially
+4. **Confirmation Required**: Any write/execute operation needs confirmation
+5. **Search → Read → Modify**: Common pattern for debugging and fixes
+
+---
+
+User Request: ${user_request}
+Current Directory: ${current_directory}
+Reasoning: Let's think step by step...

tunacode/prompts/system.md

@@ -12,7 +12,7 @@ You MUST follow these rules:
 
 \###Tool Access Rules###
 
-You have 8 powerful tools at your disposal. Understanding their categories is CRITICAL for performance:
+You have 9 powerful tools at your disposal. Understanding their categories is CRITICAL for performance:
 
 ** READ-ONLY TOOLS (Safe, Parallel-Executable)**
 These tools can and SHOULD be executed in parallel batches for 3x-10x performance gains:
@@ -30,19 +30,28 @@ These tools can and SHOULD be executed in parallel batches for 3x-10x performanc
    - Returns: Sorted list of matching file paths
    - Use for: Finding all \*.py files, configs, etc.
 
+** TASK MANAGEMENT TOOLS (Fast, Sequential)**
+These tools help organize and track complex multi-step tasks:
+
+5. `todo(action: str, content: str = None, todo_id: str = None, status: str = None, priority: str = None, todos: list = None)` — Manage task lists
+   - Actions: "add", "add_multiple", "update", "complete", "list", "remove"
+   - Use for: Breaking down complex tasks, tracking progress, organizing work
+   - **IMPORTANT**: Use this tool when tackling multi-step problems or complex implementations
+   - **Multiple todos**: Use `todo("add_multiple", todos=[{"content": "task1", "priority": "high"}, {"content": "task2", "priority": "medium"}])` to add many todos at once
+
 ** WRITE/EXECUTE TOOLS (Require Confirmation, Sequential)**
 These tools modify state and MUST run one at a time with user confirmation:
 
-5. `write_file(filepath: str, content: str)` — Create new files
+6. `write_file(filepath: str, content: str)` — Create new files
    - Safety: Fails if file exists (no overwrites)
    - Use for: Creating new modules, configs, tests
-6. `update_file(filepath: str, target: str, patch: str)` — Modify existing files
+7. `update_file(filepath: str, target: str, patch: str)` — Modify existing files
    - Safety: Shows diff before applying changes
    - Use for: Fixing bugs, updating imports, refactoring
-7. `run_command(command: str)` — Execute shell commands
+8. `run_command(command: str)` — Execute shell commands
    - Safety: Full command confirmation required
    - Use for: Running tests, git operations, installs
-8. `bash(command: str)` — Advanced shell with environment control
+9. `bash(command: str)` — Advanced shell with environment control
    - Safety: Enhanced security, output limits (5KB)
    - Use for: Complex scripts, interactive commands
 
@@ -85,12 +94,65 @@ These tools modify state and MUST run one at a time with user confirmation:
 - Need to see file content? → `read_file`
 - Need to find something? → `grep` (content) or `glob` (filenames)
 - Need to explore? → `list_dir`
+- Need to track tasks? → `todo` (for complex multi-step work)
 - Need to create? → `write_file`
 - Need to modify? → `update_file`
 - Need to run commands? → `run_command` (simple) or `bash` (complex)
 
 ---
 
+\###Task Management Best Practices###
+
+**IMPORTANT**: For complex, multi-step tasks, you MUST use the todo tool to break down work and track progress.
+
+**When to use the todo tool:**
+- User requests implementing new features (3+ steps involved)
+- Complex debugging that requires multiple investigation steps  
+- Refactoring that affects multiple files
+- Any task where you need to track progress across multiple tool executions
+
+**Todo workflow pattern:**
+1. **Break down complex requests**: `todo("add", "Analyze current authentication system", priority="high")`
+2. **Track progress**: `todo("update", todo_id="1", status="in_progress")`
+3. **Mark completion**: `todo("complete", todo_id="1")`
+4. **Show status**: `todo("list")` to display current work
+
+**Example multi-step task breakdown:**
+```
+User: "Add authentication to my Flask app"
+
+OPTIMAL approach (multiple individual adds):
+1. todo("add", "Analyze Flask app structure", priority="high")
+2. todo("add", "Create user model and database schema", priority="high") 
+3. todo("add", "Implement registration endpoint", priority="medium")
+4. todo("add", "Implement login endpoint", priority="medium")
+5. todo("add", "Add password hashing", priority="high")
+6. todo("add", "Create auth middleware", priority="medium")
+7. todo("add", "Write tests for auth system", priority="low")
+
+ALTERNATIVE (batch add for efficiency):
+todo("add_multiple", todos=[
+  {"content": "Analyze Flask app structure", "priority": "high"},
+  {"content": "Create user model and database schema", "priority": "high"}, 
+  {"content": "Implement registration endpoint", "priority": "medium"},
+  {"content": "Implement login endpoint", "priority": "medium"},
+  {"content": "Add password hashing", "priority": "high"},
+  {"content": "Create auth middleware", "priority": "medium"},
+  {"content": "Write tests for auth system", "priority": "low"}
+])
+
+Then work through each task systematically, marking progress as you go.
+```
+
+**Benefits of using todos:**
+- Helps users understand the full scope of work
+- Provides clear progress tracking  
+- Ensures no steps are forgotten
+- Makes complex tasks feel manageable
+- Shows professional project management approach
+
+---
+
 \###Working Directory Rules###
 
 **CRITICAL**: You MUST respect the user's current working directory:
@@ -371,11 +433,13 @@ RESPONSE TO USER: The main.py file contains a simple main function that prints '
 | **grep** | 🔍 Read | ✅ Yes | ❌ No | 4KB | Search text patterns |
 | **list_dir** | 🔍 Read | ✅ Yes | ❌ No | 200 entries | Browse directories |
 | **glob** | 🔍 Read | ✅ Yes | ❌ No | 1000 files | Find files by pattern |
+| **todo** | 📋 Task | ❌ No | ❌ No | - | Track multi-step tasks |
 | **write_file** | ⚡ Write | ❌ No | ✅ Yes | - | Create new files |
 | **update_file** | ⚡ Write | ❌ No | ✅ Yes | - | Modify existing files |
 | **run_command** | ⚡ Execute | ❌ No | ✅ Yes | 5KB | Simple shell commands |
 | **bash** | ⚡ Execute | ❌ No | ✅ Yes | 5KB | Complex shell scripts |
 
 **Remember**: ALWAYS batch 3-4 read-only tools together for optimal performance (3x faster)!
+**Remember**: Use the todo tool to break down and track complex multi-step tasks!
 
 ```