CodeToolsMixin

hub/agents/python/code/gaia_agent_code/tools/code_tools.py

from typing import Any, Dict, List, Optional

class CodeToolsMixin:
    """
    Consolidated mixin providing code generation, analysis, and helper methods.

    Provides tools for:
    - Generating Python functions, classes, and tests
    - Parsing and analyzing Python code structure
    - Validating Python syntax
    - Extracting code symbols
    - Running pylint analysis
    """

    def register_code_tools(self) -> None:
        """
        Register all code-related tools.

        Registers:
            - generate_function: Generate Python function
            - generate_class: Generate Python class
            - generate_test: Generate unit tests
            - parse_python_code: Parse code structure
            - validate_syntax: Validate Python syntax
            - list_symbols: List code symbols
            - analyze_with_pylint: Run pylint analysis
        """
        pass

    def _generate_code_for_file(
        self, filename: str, purpose: str, context: str = ""
    ) -> str:
        """
        Generate code for a specific file using LLM with streaming preview.

        Args:
            filename: Name of the file to generate
            purpose: Description of what this file should do
            context: Additional context about the project

        Returns:
            Generated code as string

        Implementation:
            - Uses streaming for live console preview
            - Retries up to MAX_CODE_RETRIES times
            - Reduces max_tokens on retry
            - Falls back to placeholder on failure
            - Extracts code from markdown blocks
        """
        pass

    def _fix_code_with_llm(
        self,
        code: str,
        file_path: str,
        error_msg: str,
        context: str = "",
        max_attempts: int = 3,
    ) -> Optional[str]:
        """
        Fix code using LLM based on error message.

        Args:
            code: Code with errors
            file_path: Path to the file (for language detection)
            error_msg: Error message from linting/execution
            context: Additional context
            max_attempts: Maximum number of fix attempts

        Returns:
            Fixed code or None if unable to fix

        Supports:
            - Python (.py files)
            - TypeScript (.ts, .tsx files)
            - Syntax validation for Python
            - Multiple fix attempts
        """
        pass

{
    "status": "success",
    "is_valid": bool,
    "errors": List[str],
    "code": str,
    "file_path": str,  # If write_to_file=True
    "message": str
}

result = generate_function(
    name="calculate_sum",
    params="a: int, b: int",
    docstring="Calculate the sum of two integers.",
    body="return a + b",
    return_type="int",
    write_to_file=True
)

{
    "name": str,
    "params": str,
    "docstring": str,
    "body": str
}

{
    "status": "success",
    "is_valid": bool,
    "errors": List[str],
    "code": str,
    "file_path": str,
    "message": str
}

{
    "status": "success",
    "is_valid": bool,
    "errors": List[str],
    "code": str,
    "functions_tested": List[str],
    "test_class": str,
    "file_path": str,
    "message": str
}

import unittest
from unittest.mock import patch, MagicMock

class TestFunctionName(unittest.TestCase):
    def setUp(self):
        """Set up test fixtures."""
        pass

    def test_function_basic(self):
        """Test basic functionality."""
        pass

    def test_function_edge_cases(self):
        """Test edge cases."""
        pass

    def test_function_invalid_input(self):
        """Test error handling."""
        pass

if __name__ == "__main__":
    unittest.main(verbosity=2)

{
    "status": "success",
    "is_valid": bool,
    "symbols": [
        {
            "name": str,
            "type": "function" | "class" | "variable",
            "line": int,
            "signature": str,
            "docstring": str
        }
    ],
    "imports": List[str],
    "errors": List[str]
}

{
    "status": "success",
    "is_valid": bool,
    "errors": List[str]
}

{
    "status": "success",
    "total_symbols": int,
    "symbols": [
        {
            "name": str,
            "type": str,
            "line": int,
            "signature": str,
            "docstring": str
        }
    ]
}

{
    "status": "success",
    "total_issues": int,
    "errors": int,
    "warnings": int,
    "conventions": int,
    "refactors": int,
    "issues": [
        {
            "type": "error" | "warning" | "convention" | "refactor",
            "symbol": str,
            "message": str,
            "line": int,
            "column": int,
            "confidence": str
        }
    ],
    "clean": bool,
    "command": str,
    "stdout": str,
    "stderr": str,
    "return_code": int
}

from gaia_agent_code import CodeAgent

agent = CodeAgent()

result = agent.generate_function(
    name="fibonacci",
    params="n: int",
    docstring="Calculate the nth Fibonacci number.",
    body="""
if n <= 0:
    return 0
elif n == 1:
    return 1
else:
    return fibonacci(n-1) + fibonacci(n-2)
""",
    return_type="int",
    write_to_file=True
)

print(f"Generated {result['file_path']}")
print(result['code'])

result = agent.generate_class(
    name="Calculator",
    docstring="A simple calculator class.",
    methods=[
        {
            "name": "add",
            "params": "a: int, b: int",
            "docstring": "Add two numbers.",
            "body": "return a + b"
        },
        {
            "name": "subtract",
            "params": "a: int, b: int",
            "docstring": "Subtract two numbers.",
            "body": "return a - b"
        }
    ],
    write_to_file=True
)

source_code = '''
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b

def divide(a: int, b: int) -> float:
    """Divide two numbers."""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b
'''

result = agent.generate_test(
    source_code=source_code,
    module_name="calculator",
    write_to_file=True
)

print(f"Generated test file: {result['file_path']}")
print(f"Functions tested: {result['functions_tested']}")

result = agent.analyze_with_pylint(
    file_path="/path/to/module.py",
    confidence="HIGH"
)

if result["clean"]:
    print("Code passes pylint checks!")
else:
    print(f"Found {result['total_issues']} issues:")
    for issue in result["issues"]:
        print(f"  Line {issue['line']}: [{issue['symbol']}] {issue['message']}")

code = '''
import os
from typing import List

class DataProcessor:
    """Process data from files."""

    def __init__(self, path: str):
        self.path = path

    def load_data(self) -> List[str]:
        """Load data from file."""
        with open(self.path) as f:
            return f.readlines()
'''

result = agent.parse_python_code(code)

print(f"Valid: {result['is_valid']}")
print(f"Imports: {result['imports']}")
for symbol in result["symbols"]:
    print(f"  {symbol['type']} {symbol['name']} at line {symbol['line']}")

# Maximum number of retry attempts for code generation
MAX_CODE_RETRIES = 2

# Token limits for generation attempts
FIRST_ATTEMPT_MAX_TOKENS = 4096
RETRY_MAX_TOKENS = 2048

import pytest
from gaia_agent_code.tools.code_tools import CodeToolsMixin

def test_generate_function():
    """Test function generation."""
    mixin = CodeToolsMixin()
    result = mixin.generate_function(
        name="test_func",
        params="x: int",
        body="return x * 2",
        return_type="int",
        write_to_file=False
    )

    assert result["status"] == "success"
    assert result["is_valid"]
    assert "def test_func(x: int) -> int:" in result["code"]

def test_validate_syntax():
    """Test syntax validation."""
    mixin = CodeToolsMixin()

    # Valid code
    result = mixin.validate_syntax("x = 1 + 2")
    assert result["is_valid"]

    # Invalid code
    result = mixin.validate_syntax("x = 1 +")
    assert not result["is_valid"]
    assert len(result["errors"]) > 0

def test_list_symbols():
    """Test symbol extraction."""
    code = '''
def func1():
    pass

class Class1:
    def method1(self):
        pass
'''

    mixin = CodeToolsMixin()
    result = mixin.list_symbols(code)

    assert result["total_symbols"] == 3  # func1, Class1, method1
    assert any(s["name"] == "func1" and s["type"] == "function" for s in result["symbols"])
    assert any(s["name"] == "Class1" and s["type"] == "class" for s in result["symbols"])

# pyproject.toml
[project]
dependencies = [
    "pylint>=2.0",
]

import ast
import json
import logging
import os
import shlex
import subprocess
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional

prompt = f"""Generate complete, production-ready Python code for: {filename}

Purpose: {purpose}

{context}

Requirements:
1. Include proper copyright header
2. Add comprehensive docstrings
3. Use type hints
4. Follow best practices
5. Include error handling
6. Make it complete and runnable

Generate ONLY the code, no explanations."""

retry_count = 0
while retry_count <= MAX_CODE_RETRIES:
    try:
        max_tokens = FIRST_ATTEMPT_MAX_TOKENS if retry_count == 0 else RETRY_MAX_TOKENS

        for chunk in self.chat.send_stream(prompt, max_tokens=max_tokens):
            if chunk.is_complete:
                continue
            code_chunks.append(chunk.text)
            self.console.update_file_preview(chunk.text)

        break  # Success

    except Exception as e:
        retry_count += 1
        if "timeout" in str(e).lower():
            if retry_count <= MAX_CODE_RETRIES:
                self.console.print_warning(f"⚠️ Timeout, retrying ({retry_count}/{MAX_CODE_RETRIES})...")
            else:
                code_chunks = [self._get_timeout_placeholder(filename, purpose)]
                break