tools.py

import os
import subprocess
import tempfile
import tarfile
import json
import hashlib
from pathlib import Path
from typing import Dict, Any, List
from io import BytesIO

# Tower API imports
try:
    import tower
    from tower._context import TowerContext
    from tower._client import _env_client
    from tower.tower_api_client.api.default import list_apps as list_apps_api
    from tower.tower_api_client.api.default import deploy_app as deploy_app_api
    from tower.tower_api_client.types import File
    TOWER_AVAILABLE = True
except ImportError:
    TOWER_AVAILABLE = False

# Try to import toml for reading Towerfiles
try:
    import tomli as toml
except ImportError:
    try:
        import tomllib as toml
    except ImportError:
        try:
            import toml
        except ImportError:
            toml = None

# Tool definitions for Claude
TOOLS = [
    {
        "name": "write_python_file",
        "description": "Write Python code to a file on the local filesystem. Use this to create Python scripts, modules, or configuration files. The file will be created in a workspace directory.",
        "input_schema": {
            "type": "object",
            "properties": {
                "filename": {
                    "type": "string",
                    "description": "The name of the Python file (e.g., 'script.py', 'pipeline.py')"
                },
                "content": {
                    "type": "string",
                    "description": "The Python code to write to the file"
                },
                "directory": {
                    "type": "string",
                    "description": "Optional subdirectory within workspace (e.g., 'pipelines', 'utils')"
                }
            },
            "required": ["filename", "content"]
        }
    },
    {
        "name": "execute_python",
        "description": "Execute Python code or a Python file. Use this to run scripts, test code, or execute data processing tasks. Returns stdout, stderr, and exit code.",
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Python code to execute directly"
                },
                "filename": {
                    "type": "string",
                    "description": "Path to a Python file to execute (alternative to code)"
                },
                "args": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Command-line arguments to pass to the script"
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default: 30)"
                }
            }
        }
    },
    {
        "name": "read_file",
        "description": "Read the contents of a file from the workspace. Use this to inspect files, check outputs, or read data.",
        "input_schema": {
            "type": "object",
            "properties": {
                "filename": {
                    "type": "string",
                    "description": "The name or path of the file to read"
                }
            },
            "required": ["filename"]
        }
    },
    {
        "name": "list_files",
        "description": "List files in the workspace directory. Use this to see what files exist.",
        "input_schema": {
            "type": "object",
            "properties": {
                "directory": {
                    "type": "string",
                    "description": "Optional subdirectory to list (defaults to root workspace)"
                }
            }
        }
    },
    {
        "name": "tower_deploy",
        "description": "Deploy a Tower app using the Tower Python API. Creates a TAR package from the app directory and uploads it. The app must have a Towerfile. Authentication uses TOWER_API_KEY environment variable.",
        "input_schema": {
            "type": "object",
            "properties": {
                "app_directory": {
                    "type": "string",
                    "description": "Directory containing the Tower app to deploy (must contain a Towerfile)"
                }
            },
            "required": ["app_directory"]
        }
    },
    {
        "name": "tower_run",
        "description": "Trigger a Tower app run on the Tower platform. Use this to execute a deployed Tower app remotely. The run will be tracked and you'll receive the run ID.",
        "input_schema": {
            "type": "object",
            "properties": {
                "app_name": {
                    "type": "string",
                    "description": "Name of the Tower app to run"
                },
                "environment": {
                    "type": "string",
                    "description": "Optional environment to run in (defaults to 'default')"
                },
                "parameters": {
                    "type": "object",
                    "description": "Parameters to pass to the app as key-value pairs"
                }
            },
            "required": ["app_name"]
        }
    },
    {
        "name": "tower_list_apps",
        "description": "List all Tower apps in the current account. Use this to see what apps are deployed.",
        "input_schema": {
            "type": "object",
            "properties": {}
        }
    },
    {
        "name": "install_package",
        "description": "Install a Python package using pip. Use this when you need additional libraries for your code.",
        "input_schema": {
            "type": "object",
            "properties": {
                "package": {
                    "type": "string",
                    "description": "Package name (e.g., 'pandas', 'requests==2.28.0')"
                }
            },
            "required": ["package"]
        }
    }
]

# Workspace directory for file operations
WORKSPACE_DIR = Path(tempfile.gettempdir()) / "data-engineering-agent-workspace"
WORKSPACE_DIR.mkdir(exist_ok=True)

class ToolExecutor:
    """Executes tools requested by the AI agent"""

    @staticmethod
    def execute(tool_name: str, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a tool and return the result"""
        try:
            if tool_name == "write_python_file":
                return ToolExecutor._write_python_file(tool_input)
            elif tool_name == "execute_python":
                return ToolExecutor._execute_python(tool_input)
            elif tool_name == "read_file":
                return ToolExecutor._read_file(tool_input)
            elif tool_name == "list_files":
                return ToolExecutor._list_files(tool_input)
            elif tool_name == "tower_deploy":
                return ToolExecutor._tower_deploy(tool_input)
            elif tool_name == "tower_run":
                return ToolExecutor._tower_run(tool_input)
            elif tool_name == "tower_list_apps":
                return ToolExecutor._tower_list_apps(tool_input)
            elif tool_name == "install_package":
                return ToolExecutor._install_package(tool_input)
            else:
                return {
                    "success": False,
                    "error": f"Unknown tool: {tool_name}"
                }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    @staticmethod
    def _write_python_file(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Write Python code to a file"""
        filename = tool_input["filename"]
        content = tool_input["content"]
        directory = tool_input.get("directory", "")

        # Create target directory
        target_dir = WORKSPACE_DIR / directory if directory else WORKSPACE_DIR
        target_dir.mkdir(parents=True, exist_ok=True)

        # Write file
        filepath = target_dir / filename
        filepath.write_text(content)

        return {
            "success": True,
            "message": f"File written to {filepath}",
            "filepath": str(filepath)
        }

    @staticmethod
    def _execute_python(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Execute Python code or file"""
        code = tool_input.get("code")
        filename = tool_input.get("filename")
        args = tool_input.get("args", [])
        timeout = tool_input.get("timeout", 30)

        try:
            if code:
                # Execute code directly
                result = subprocess.run(
                    ["python", "-c", code] + args,
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                    cwd=str(WORKSPACE_DIR)
                )
            elif filename:
                # Execute file
                filepath = WORKSPACE_DIR / filename
                if not filepath.exists():
                    return {
                        "success": False,
                        "error": f"File not found: {filename}"
                    }

                result = subprocess.run(
                    ["python", str(filepath)] + args,
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                    cwd=str(WORKSPACE_DIR)
                )
            else:
                return {
                    "success": False,
                    "error": "Must provide either 'code' or 'filename'"
                }

            return {
                "success": result.returncode == 0,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "exit_code": result.returncode
            }

        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "error": f"Execution timed out after {timeout} seconds"
            }

    @staticmethod
    def _read_file(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Read file contents"""
        filename = tool_input["filename"]
        filepath = WORKSPACE_DIR / filename

        if not filepath.exists():
            return {
                "success": False,
                "error": f"File not found: {filename}"
            }

        content = filepath.read_text()
        return {
            "success": True,
            "content": content,
            "filepath": str(filepath)
        }

    @staticmethod
    def _list_files(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """List files in directory"""
        directory = tool_input.get("directory", "")
        target_dir = WORKSPACE_DIR / directory if directory else WORKSPACE_DIR

        if not target_dir.exists():
            return {
                "success": False,
                "error": f"Directory not found: {directory}"
            }

        files = []
        for item in target_dir.iterdir():
            files.append({
                "name": item.name,
                "type": "directory" if item.is_dir() else "file",
                "size": item.stat().st_size if item.is_file() else None
            })

        return {
            "success": True,
            "directory": str(target_dir),
            "files": files
        }

    @staticmethod
    def _read_towerfile(app_path: Path) -> Dict[str, Any]:
        """Read and parse a Towerfile"""
        if toml is None:
            raise ImportError("toml library not available. Install with: pip install tomli")

        towerfile_path = app_path / "Towerfile"
        with open(towerfile_path, "rb") as f:
            return toml.load(f)

    @staticmethod
    def _create_package(app_path: Path, towerfile_data: Dict[str, Any]) -> BytesIO:
        """Create a TAR package for deployment"""
        # Create manifest
        manifest = {
            "version": 3,
            "invoke": towerfile_data["app"]["script"],
            "parameters": [
                {
                    "name": p["name"],
                    "description": p.get("description", ""),
                    "default": p.get("default", "")
                }
                for p in towerfile_data.get("parameters", [])
            ],
            "schedule": towerfile_data["app"].get("schedule"),
            "import_paths": [],
            "app_dir_name": "app",
            "modules_dir_name": "",
            "checksum": ""
        }

        # Create tar in memory
        tar_buffer = BytesIO()
        with tarfile.open(fileobj=tar_buffer, mode='w:gz') as tar:
            # Add manifest
            manifest_bytes = json.dumps(manifest).encode('utf-8')
            manifest_info = tarfile.TarInfo(name='MANIFEST')
            manifest_info.size = len(manifest_bytes)
            tar.addfile(manifest_info, BytesIO(manifest_bytes))

            # Add all app files
            for file_path in app_path.rglob('*'):
                if file_path.is_file():
                    # Calculate relative path
                    rel_path = file_path.relative_to(app_path)
                    arcname = f"app/{rel_path}"
                    tar.add(file_path, arcname=arcname)

        tar_buffer.seek(0)
        return tar_buffer

    @staticmethod
    def _calculate_checksum(tar_buffer: BytesIO) -> str:
        """Calculate SHA256 checksum of tar content"""
        tar_buffer.seek(0)
        sha256 = hashlib.sha256()
        sha256.update(tar_buffer.read())
        tar_buffer.seek(0)
        return sha256.hexdigest()

    @staticmethod
    def _tower_deploy(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Deploy a Tower app using Python API"""
        if not TOWER_AVAILABLE:
            return {
                "success": False,
                "error": "Tower Python package not available. Install with: pip install tower"
            }

        app_directory = tool_input["app_directory"]
        app_path = WORKSPACE_DIR / app_directory

        if not app_path.exists():
            return {
                "success": False,
                "error": f"App directory not found: {app_directory}"
            }

        if not (app_path / "Towerfile").exists():
            return {
                "success": False,
                "error": "No Towerfile found in directory"
            }

        try:
            # Read Towerfile to get app name
            towerfile_data = ToolExecutor._read_towerfile(app_path)
            app_name = towerfile_data["app"]["name"]

            # Create package
            tar_buffer = ToolExecutor._create_package(app_path, towerfile_data)
            checksum = ToolExecutor._calculate_checksum(tar_buffer)

            # Create Tower client
            ctx = TowerContext.build()
            client = _env_client(ctx)

            # Create File object for upload
            file_obj = File(
                payload=tar_buffer,
                file_name=f"{app_name}.tar.gz",
                mime_type="application/gzip"
            )

            # Deploy using API
            response = deploy_app_api.sync(
                name=app_name,
                client=client,
                body=file_obj,
                x_tower_checksum_sha256=checksum
            )

            if response is None:
                return {
                    "success": False,
                    "error": "No response from Tower API"
                }

            # Check if error
            from tower.tower_api_client.models.error_model import ErrorModel
            if isinstance(response, ErrorModel):
                return {
                    "success": False,
                    "error": f"Deployment failed: {response.title} - {response.detail}"
                }

            return {
                "success": True,
                "message": f"Successfully deployed app '{app_name}' version {response.version.number}",
                "app_name": app_name,
                "version": response.version.number
            }

        except Exception as e:
            return {
                "success": False,
                "error": f"Deployment error: {str(e)}"
            }

    @staticmethod
    def _tower_run(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Run a Tower app using the Tower Python API"""
        if not TOWER_AVAILABLE:
            return {
                "success": False,
                "error": "Tower Python package not available. Install with: pip install tower-cli"
            }

        app_name = tool_input["app_name"]
        environment = tool_input.get("environment")
        parameters = tool_input.get("parameters", {})

        try:
            # Use Tower Python API to run the app
            run = tower.run_app(
                name=app_name,
                environment=environment,
                parameters=parameters
            )

            return {
                "success": True,
                "message": f"Started run #{run.number} for app '{run.app_name}'",
                "run_id": run.number,
                "app_name": run.app_name,
                "status": run.status
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    @staticmethod
    def _tower_list_apps(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """List Tower apps using the Tower Python API"""
        if not TOWER_AVAILABLE:
            return {
                "success": False,
                "error": "Tower Python package not available. Install with: pip install tower-cli"
            }

        try:
            # Build Tower context and client
            ctx = TowerContext.build()
            client = _env_client(ctx)

            # Call list_apps API
            response = list_apps_api.sync(client=client, page_size=100)

            if response is None:
                return {
                    "success": False,
                    "error": "No response from Tower API"
                }

            # Extract app information
            apps_info = []
            for app in response.apps:
                apps_info.append({
                    "name": app.name,
                    "owner": app.owner_name if hasattr(app, 'owner_name') else None,
                    "created_at": app.created_at.isoformat() if hasattr(app, 'created_at') and app.created_at else None,
                    "is_externally_accessible": app.is_externally_accessible if hasattr(app, 'is_externally_accessible') else False
                })

            return {
                "success": True,
                "apps": apps_info,
                "total": len(apps_info)
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    @staticmethod
    def _install_package(tool_input: Dict[str, Any]) -> Dict[str, Any]:
        """Install a Python package"""
        package = tool_input["package"]

        result = subprocess.run(
            ["pip", "install", package],
            capture_output=True,
            text=True,
            timeout=120
        )

        return {
            "success": result.returncode == 0,
            "stdout": result.stdout,
            "stderr": result.stderr
        }