vanna 0.7.9__py3-none-any.whl → 2.0.0__py3-none-any.whl

vanna/tools/python.py

@@ -0,0 +1,222 @@
+"""Python-specific tooling built on top of the file system service."""
+
+from __future__ import annotations
+
+import shlex
+import sys
+from typing import Any, List, Optional, Sequence, Type
+
+from pydantic import BaseModel, Field
+
+from vanna.components import (
+    UiComponent,
+    CardComponent,
+    ComponentType,
+    NotificationComponent,
+    SimpleTextComponent,
+)
+from vanna.core.tool import Tool, ToolContext, ToolResult
+
+from .file_system import CommandResult, FileSystem, LocalFileSystem
+
+MAX_OUTPUT_LENGTH = 4000
+
+
+class RunPythonFileArgs(BaseModel):
+    """Arguments required to execute a Python file."""
+
+    filename: str = Field(
+        description="Python file to execute (relative to the workspace root)"
+    )
+    arguments: Sequence[str] = Field(
+        default_factory=list,
+        description="Optional arguments to pass to the Python script",
+    )
+    timeout_seconds: Optional[float] = Field(
+        default=None,
+        ge=0,
+        description="Optional timeout for the command in seconds",
+    )
+
+
+class RunPythonFileTool(Tool[RunPythonFileArgs]):
+    """Execute a Python file using the provided file system service."""
+
+    def __init__(self, file_system: Optional[FileSystem] = None):
+        self.file_system = file_system or LocalFileSystem()
+
+    @property
+    def name(self) -> str:
+        return "run_python_file"
+
+    @property
+    def description(self) -> str:
+        return "Execute a Python file using the workspace interpreter"
+
+    def get_args_schema(self) -> Type[RunPythonFileArgs]:
+        return RunPythonFileArgs
+
+    async def execute(
+        self, context: ToolContext, args: RunPythonFileArgs
+    ) -> ToolResult:
+        exists = await self.file_system.exists(args.filename, context)
+        if not exists:
+            message = f"Cannot execute '{args.filename}' because it does not exist."
+            return _error_result(message)
+
+        command_parts = [sys.executable, args.filename]
+        command_parts.extend(args.arguments)
+        command = _quote_command(command_parts)
+
+        try:
+            result = await self.file_system.run_bash(
+                command,
+                context,
+                timeout=args.timeout_seconds,
+            )
+        except TimeoutError as exc:
+            message = str(exc)
+            return _error_result(message)
+
+        summary = f"Executed python {args.filename} (exit code {result.returncode})."
+        success = result.returncode == 0
+        return _result_from_command(summary, command, result, success=success)
+
+
+class PipInstallArgs(BaseModel):
+    """Arguments required to run pip install."""
+
+    packages: List[str] = Field(
+        description="Packages (with optional specifiers) to install", min_length=1
+    )
+    upgrade: bool = Field(
+        default=False,
+        description="Whether to include --upgrade in the pip invocation",
+    )
+    extra_args: Sequence[str] = Field(
+        default_factory=list,
+        description="Additional arguments to pass to pip install",
+    )
+    timeout_seconds: Optional[float] = Field(
+        default=None,
+        ge=0,
+        description="Optional timeout for the command in seconds",
+    )
+
+
+class PipInstallTool(Tool[PipInstallArgs]):
+    """Install Python packages using pip inside the workspace environment."""
+
+    def __init__(self, file_system: Optional[FileSystem] = None):
+        self.file_system = file_system or LocalFileSystem()
+
+    @property
+    def name(self) -> str:
+        return "pip_install"
+
+    @property
+    def description(self) -> str:
+        return "Install Python packages using pip"
+
+    def get_args_schema(self) -> Type[PipInstallArgs]:
+        return PipInstallArgs
+
+    async def execute(self, context: ToolContext, args: PipInstallArgs) -> ToolResult:
+        command_parts = [sys.executable, "-m", "pip", "install"]
+        if args.upgrade:
+            command_parts.append("--upgrade")
+        command_parts.extend(args.packages)
+        command_parts.extend(args.extra_args)
+        command = _quote_command(command_parts)
+
+        try:
+            result = await self.file_system.run_bash(
+                command,
+                context,
+                timeout=args.timeout_seconds,
+            )
+        except TimeoutError as exc:
+            return _error_result(str(exc))
+
+        success = result.returncode == 0
+        summary = (
+            "pip install completed successfully"
+            if success
+            else f"pip install failed (exit code {result.returncode})."
+        )
+
+        return _result_from_command(summary, command, result, success=success)
+
+
+def create_python_tools(file_system: Optional[FileSystem] = None) -> List[Tool[Any]]:
+    """Create Python-specific tools backed by a shared file system service."""
+
+    fs = file_system or LocalFileSystem()
+    return [
+        RunPythonFileTool(fs),
+        PipInstallTool(fs),
+    ]
+
+
+def _quote_command(parts: Sequence[str]) -> str:
+    return " ".join(shlex.quote(part) for part in parts)
+
+
+def _truncate(text: str, limit: int = MAX_OUTPUT_LENGTH) -> str:
+    if len(text) <= limit:
+        return text
+    return f"{text[: limit - 1]}…"
+
+
+def _result_from_command(
+    summary: str,
+    command: str,
+    result: CommandResult,
+    *,
+    success: bool = True,
+) -> ToolResult:
+    stdout = result.stdout.strip()
+    stderr = result.stderr.strip()
+
+    blocks: List[str] = [f"$ {command}"]
+    if stdout:
+        blocks.append("STDOUT:\n" + _truncate(stdout))
+    if stderr:
+        blocks.append("STDERR:\n" + _truncate(stderr))
+    if not stdout and not stderr:
+        blocks.append("(no output)")
+
+    content = "\n\n".join(blocks)
+    card_status = "success" if success else "error"
+    component = CardComponent(
+        type=ComponentType.CARD,
+        title="Command Result",
+        content=content,
+        status=card_status,
+    )
+
+    return ToolResult(
+        success=success,
+        result_for_llm=f"{summary}\n\n{content}",
+        ui_component=UiComponent(
+            rich_component=component,
+            simple_component=SimpleTextComponent(text=summary),
+        ),
+        error=None if success else content,
+    )
+
+
+def _error_result(message: str) -> ToolResult:
+    return ToolResult(
+        success=False,
+        result_for_llm=message,
+        ui_component=UiComponent(
+            rich_component=NotificationComponent(
+                type=ComponentType.NOTIFICATION,
+                level="error",
+                message=message,
+            ),
+            simple_component=SimpleTextComponent(text=message),
+        ),
+        error=message,
+    )

vanna/tools/run_sql.py

@@ -0,0 +1,165 @@
+"""Generic SQL query execution tool with dependency injection."""
+
+from typing import Any, Dict, List, Optional, Type, cast
+import uuid
+from vanna.core.tool import Tool, ToolContext, ToolResult
+from vanna.components import (
+    UiComponent,
+    DataFrameComponent,
+    NotificationComponent,
+    ComponentType,
+    SimpleTextComponent,
+)
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.capabilities.file_system import FileSystem
+from vanna.integrations.local import LocalFileSystem
+
+
+class RunSqlTool(Tool[RunSqlToolArgs]):
+    """Tool that executes SQL queries using an injected SqlRunner implementation."""
+
+    def __init__(
+        self,
+        sql_runner: SqlRunner,
+        file_system: Optional[FileSystem] = None,
+        custom_tool_name: Optional[str] = None,
+        custom_tool_description: Optional[str] = None,
+    ):
+        """Initialize the tool with a SqlRunner implementation.
+
+        Args:
+            sql_runner: SqlRunner implementation that handles actual query execution
+            file_system: FileSystem implementation for saving results (defaults to LocalFileSystem)
+            custom_tool_name: Optional custom name for the tool (overrides default "run_sql")
+            custom_tool_description: Optional custom description for the tool (overrides default description)
+        """
+        self.sql_runner = sql_runner
+        self.file_system = file_system or LocalFileSystem()
+        self._custom_name = custom_tool_name
+        self._custom_description = custom_tool_description
+
+    @property
+    def name(self) -> str:
+        return self._custom_name if self._custom_name else "run_sql"
+
+    @property
+    def description(self) -> str:
+        return (
+            self._custom_description
+            if self._custom_description
+            else "Execute SQL queries against the configured database"
+        )
+
+    def get_args_schema(self) -> Type[RunSqlToolArgs]:
+        return RunSqlToolArgs
+
+    async def execute(self, context: ToolContext, args: RunSqlToolArgs) -> ToolResult:
+        """Execute a SQL query using the injected SqlRunner."""
+        try:
+            # Use the injected SqlRunner to execute the query
+            df = await self.sql_runner.run_sql(args, context)
+
+            # Determine query type
+            query_type = args.sql.strip().upper().split()[0]
+
+            if query_type == "SELECT":
+                # Handle SELECT queries with results
+                if df.empty:
+                    result = "Query executed successfully. No rows returned."
+                    ui_component = UiComponent(
+                        rich_component=DataFrameComponent(
+                            rows=[],
+                            columns=[],
+                            title="Query Results",
+                            description="No rows returned",
+                        ),
+                        simple_component=SimpleTextComponent(text=result),
+                    )
+                    metadata = {
+                        "row_count": 0,
+                        "columns": [],
+                        "query_type": query_type,
+                        "results": [],
+                    }
+                else:
+                    # Convert DataFrame to records
+                    results_data = df.to_dict("records")
+                    columns = df.columns.tolist()
+                    row_count = len(df)
+
+                    # Write DataFrame to CSV file for downstream tools
+                    file_id = str(uuid.uuid4())[:8]
+                    filename = f"query_results_{file_id}.csv"
+                    csv_content = df.to_csv(index=False)
+                    await self.file_system.write_file(
+                        filename, csv_content, context, overwrite=True
+                    )
+
+                    # Create result text for LLM with truncated results
+                    results_preview = csv_content
+                    if len(results_preview) > 1000:
+                        results_preview = (
+                            results_preview[:1000]
+                            + "\n(Results truncated to 1000 characters. FOR LARGE RESULTS YOU DO NOT NEED TO SUMMARIZE THESE RESULTS OR PROVIDE OBSERVATIONS. THE NEXT STEP SHOULD BE A VISUALIZE_DATA CALL)"
+                        )
+
+                    result = f"{results_preview}\n\nResults saved to file: {filename}\n\n**IMPORTANT: FOR VISUALIZE_DATA USE FILENAME: {filename}**"
+
+                    # Create DataFrame component for UI
+                    dataframe_component = DataFrameComponent.from_records(
+                        records=cast(List[Dict[str, Any]], results_data),
+                        title="Query Results",
+                        description=f"SQL query returned {row_count} rows with {len(columns)} columns",
+                    )
+
+                    ui_component = UiComponent(
+                        rich_component=dataframe_component,
+                        simple_component=SimpleTextComponent(text=result),
+                    )
+
+                    metadata = {
+                        "row_count": row_count,
+                        "columns": columns,
+                        "query_type": query_type,
+                        "results": results_data,
+                        "output_file": filename,
+                    }
+            else:
+                # For non-SELECT queries (INSERT, UPDATE, DELETE, etc.)
+                # The SqlRunner should return a DataFrame with affected row count
+                rows_affected = len(df) if not df.empty else 0
+                result = (
+                    f"Query executed successfully. {rows_affected} row(s) affected."
+                )
+
+                metadata = {"rows_affected": rows_affected, "query_type": query_type}
+                ui_component = UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION, level="success", message=result
+                    ),
+                    simple_component=SimpleTextComponent(text=result),
+                )
+
+            return ToolResult(
+                success=True,
+                result_for_llm=result,
+                ui_component=ui_component,
+                metadata=metadata,
+            )
+
+        except Exception as e:
+            error_message = f"Error executing query: {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=error_message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=error_message,
+                    ),
+                    simple_component=SimpleTextComponent(text=error_message),
+                ),
+                error=str(e),
+                metadata={"error_type": "sql_error"},
+            )

vanna/tools/visualize_data.py

@@ -0,0 +1,195 @@
+"""Tool for visualizing DataFrame data from CSV files."""
+
+from typing import Optional, Type
+import logging
+import pandas as pd
+from pydantic import BaseModel, Field
+
+from vanna.core.tool import Tool, ToolContext, ToolResult
+from vanna.components import (
+    UiComponent,
+    ChartComponent,
+    NotificationComponent,
+    ComponentType,
+    SimpleTextComponent,
+)
+
+from .file_system import FileSystem, LocalFileSystem
+from vanna.integrations.plotly import PlotlyChartGenerator
+
+logger = logging.getLogger(__name__)
+
+
+class VisualizeDataArgs(BaseModel):
+    """Arguments for visualize_data tool."""
+
+    filename: str = Field(description="Name of the CSV file to visualize")
+    title: Optional[str] = Field(
+        default=None, description="Optional title for the chart"
+    )
+
+
+class VisualizeDataTool(Tool[VisualizeDataArgs]):
+    """Tool that reads CSV files and generates visualizations using dependency injection."""
+
+    def __init__(
+        self,
+        file_system: Optional[FileSystem] = None,
+        plotly_generator: Optional[PlotlyChartGenerator] = None,
+    ):
+        """Initialize the tool with FileSystem and PlotlyChartGenerator.
+
+        Args:
+            file_system: FileSystem implementation for reading CSV files (defaults to LocalFileSystem)
+            plotly_generator: PlotlyChartGenerator for creating Plotly charts (defaults to PlotlyChartGenerator())
+        """
+        self.file_system = file_system or LocalFileSystem()
+        self.plotly_generator = plotly_generator or PlotlyChartGenerator()
+
+    @property
+    def name(self) -> str:
+        return "visualize_data"
+
+    @property
+    def description(self) -> str:
+        return "Create a visualization from a CSV file. The tool automatically selects an appropriate chart type based on the data."
+
+    def get_args_schema(self) -> Type[VisualizeDataArgs]:
+        return VisualizeDataArgs
+
+    async def execute(
+        self, context: ToolContext, args: VisualizeDataArgs
+    ) -> ToolResult:
+        """Read CSV file and generate visualization."""
+        try:
+            logger.info(f"Starting visualization for file: {args.filename}")
+
+            # Read the CSV file using FileSystem
+            csv_content = await self.file_system.read_file(args.filename, context)
+            logger.info(f"Read {len(csv_content)} bytes from CSV file")
+
+            # Parse CSV into DataFrame
+            import io
+
+            df = pd.read_csv(io.StringIO(csv_content))
+            logger.info(
+                f"Parsed DataFrame with shape {df.shape}, columns: {df.columns.tolist()}, dtypes: {df.dtypes.to_dict()}"
+            )
+
+            # Generate title
+            title = args.title or f"Visualization of {args.filename}"
+
+            # Generate chart using PlotlyChartGenerator
+            logger.info("Generating chart...")
+            chart_dict = self.plotly_generator.generate_chart(df, title)
+            logger.info(
+                f"Chart generated, type: {type(chart_dict)}, keys: {list(chart_dict.keys()) if isinstance(chart_dict, dict) else 'N/A'}"
+            )
+
+            # Create result message
+            row_count = len(df)
+            col_count = len(df.columns)
+            result = f"Created visualization from '{args.filename}' ({row_count} rows, {col_count} columns)."
+
+            # Create ChartComponent
+            logger.info("Creating ChartComponent...")
+            chart_component = ChartComponent(
+                chart_type="plotly",
+                data=chart_dict,
+                title=title,
+                config={
+                    "data_shape": {"rows": row_count, "columns": col_count},
+                    "source_file": args.filename,
+                },
+            )
+            logger.info("ChartComponent created successfully")
+
+            logger.info("Creating ToolResult...")
+            tool_result = ToolResult(
+                success=True,
+                result_for_llm=result,
+                ui_component=UiComponent(
+                    rich_component=chart_component,
+                    simple_component=SimpleTextComponent(text=result),
+                ),
+                metadata={
+                    "filename": args.filename,
+                    "rows": row_count,
+                    "columns": col_count,
+                    "chart": chart_dict,
+                },
+            )
+            logger.info("ToolResult created successfully")
+            return tool_result
+
+        except FileNotFoundError as e:
+            logger.error(f"File not found: {args.filename}", exc_info=True)
+            error_message = f"File not found: {args.filename}"
+            return ToolResult(
+                success=False,
+                result_for_llm=error_message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=error_message,
+                    ),
+                    simple_component=SimpleTextComponent(text=error_message),
+                ),
+                error=str(e),
+                metadata={"error_type": "file_not_found"},
+            )
+        except pd.errors.ParserError as e:
+            logger.error(f"CSV parse error for {args.filename}", exc_info=True)
+            error_message = f"Failed to parse CSV file '{args.filename}': {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=error_message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=error_message,
+                    ),
+                    simple_component=SimpleTextComponent(text=error_message),
+                ),
+                error=str(e),
+                metadata={"error_type": "csv_parse_error"},
+            )
+        except ValueError as e:
+            logger.error(f"Visualization error for {args.filename}", exc_info=True)
+            error_message = f"Cannot visualize data: {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=error_message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=error_message,
+                    ),
+                    simple_component=SimpleTextComponent(text=error_message),
+                ),
+                error=str(e),
+                metadata={"error_type": "visualization_error"},
+            )
+        except Exception as e:
+            logger.error(
+                f"Unexpected error creating visualization for {args.filename}",
+                exc_info=True,
+            )
+            error_message = f"Error creating visualization: {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=error_message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=error_message,
+                    ),
+                    simple_component=SimpleTextComponent(text=error_message),
+                ),
+                error=str(e),
+                metadata={"error_type": "general_error"},
+            )

vanna/web_components/__init__.py

@@ -0,0 +1,44 @@
+"""
+Web components for Vanna Agents.
+
+This module provides web components built with Lit that can be embedded
+in web applications to provide rich UI for Vanna agent interactions.
+"""
+
+import os
+from pathlib import Path
+from typing import Dict
+
+
+def get_component_files() -> Dict[str, Path]:
+    """Get paths to all web component files."""
+    component_dir = Path(__file__).parent
+    return {
+        "js": component_dir / "index.js",
+        "css": component_dir / "style.css",
+    }
+
+
+def get_component_html() -> str:
+    """Get HTML template for including components."""
+    files = get_component_files()
+
+    html = """
+<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="utf-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <title>Vanna AI Chat</title>
+</head>
+<body>
+    <vanna-chat title="Vanna AI Assistant"></vanna-chat>
+    <script type="module" src="{js_file}"></script>
+</body>
+</html>
+""".format(js_file=files["js"].name)
+
+    return html
+
+
+__all__ = ["get_component_files", "get_component_html"]