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

vanna/capabilities/agent_memory/base.py

@@ -0,0 +1,103 @@
+"""
+Agent memory capability interface for tool usage learning.
+
+This module contains the abstract base class for agent memory operations,
+following the same pattern as the FileSystem interface.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+if TYPE_CHECKING:
+    from vanna.core.tool import ToolContext
+    from .models import (
+        ToolMemorySearchResult,
+        TextMemory,
+        TextMemorySearchResult,
+        ToolMemory,
+    )
+
+
+class AgentMemory(ABC):
+    """Abstract base class for agent memory operations."""
+
+    @abstractmethod
+    async def save_tool_usage(
+        self,
+        question: str,
+        tool_name: str,
+        args: Dict[str, Any],
+        context: "ToolContext",
+        success: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Save a tool usage pattern for future reference."""
+        pass
+
+    @abstractmethod
+    async def save_text_memory(
+        self, content: str, context: "ToolContext"
+    ) -> "TextMemory":
+        """Save a free-form text memory."""
+        pass
+
+    @abstractmethod
+    async def search_similar_usage(
+        self,
+        question: str,
+        context: "ToolContext",
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+        tool_name_filter: Optional[str] = None,
+    ) -> List[ToolMemorySearchResult]:
+        """Search for similar tool usage patterns based on a question."""
+        pass
+
+    @abstractmethod
+    async def search_text_memories(
+        self,
+        query: str,
+        context: "ToolContext",
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+    ) -> List["TextMemorySearchResult"]:
+        """Search stored text memories based on a query."""
+        pass
+
+    @abstractmethod
+    async def get_recent_memories(
+        self, context: "ToolContext", limit: int = 10
+    ) -> List[ToolMemory]:
+        """Get recently added memories. Returns most recent memories first."""
+        pass
+
+    @abstractmethod
+    async def get_recent_text_memories(
+        self, context: "ToolContext", limit: int = 10
+    ) -> List["TextMemory"]:
+        """Fetch recently stored text memories."""
+        pass
+
+    @abstractmethod
+    async def delete_by_id(self, context: "ToolContext", memory_id: str) -> bool:
+        """Delete a memory by its ID. Returns True if deleted, False if not found."""
+        pass
+
+    @abstractmethod
+    async def delete_text_memory(self, context: "ToolContext", memory_id: str) -> bool:
+        """Delete a text memory by its ID. Returns True if deleted, False if not found."""
+        pass
+
+    @abstractmethod
+    async def clear_memories(
+        self,
+        context: "ToolContext",
+        tool_name: Optional[str] = None,
+        before_date: Optional[str] = None,
+    ) -> int:
+        """Clear stored memories (tool or text). Returns number of memories deleted."""
+        pass

vanna/capabilities/agent_memory/models.py

@@ -0,0 +1,53 @@
+"""
+Memory storage models and types.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class ToolMemory(BaseModel):
+    """Represents a stored tool usage memory."""
+
+    memory_id: Optional[str] = None
+    question: str
+    tool_name: str
+    args: Dict[str, Any]
+    timestamp: Optional[str] = None
+    success: bool = True
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class TextMemory(BaseModel):
+    """Represents a stored free-form text memory."""
+
+    memory_id: Optional[str] = None
+    content: str
+    timestamp: Optional[str] = None
+
+
+class ToolMemorySearchResult(BaseModel):
+    """Represents a search result from tool memory storage."""
+
+    memory: ToolMemory
+    similarity_score: float
+    rank: int
+
+
+class TextMemorySearchResult(BaseModel):
+    """Represents a search result from text memory storage."""
+
+    memory: TextMemory
+    similarity_score: float
+    rank: int
+
+
+class MemoryStats(BaseModel):
+    """Memory storage statistics."""
+
+    total_memories: int
+    unique_tools: int
+    unique_questions: int
+    success_rate: float
+    most_used_tools: Dict[str, int]

vanna/capabilities/file_system/__init__.py

@@ -0,0 +1,14 @@
+"""
+File system capability.
+
+This module provides abstractions for file system operations used by tools.
+"""
+
+from .base import FileSystem
+from .models import CommandResult, FileSearchMatch
+
+__all__ = [
+    "FileSystem",
+    "FileSearchMatch",
+    "CommandResult",
+]

vanna/capabilities/file_system/base.py

@@ -0,0 +1,71 @@
+"""
+File system capability interface.
+
+This module contains the abstract base class for file system operations.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional
+
+from .models import CommandResult, FileSearchMatch
+
+if TYPE_CHECKING:
+    from vanna.core.tool import ToolContext
+
+
+class FileSystem(ABC):
+    """Abstract base class for file system operations."""
+
+    @abstractmethod
+    async def list_files(self, directory: str, context: "ToolContext") -> List[str]:
+        """List files in a directory."""
+        pass
+
+    @abstractmethod
+    async def read_file(self, filename: str, context: "ToolContext") -> str:
+        """Read the contents of a file."""
+        pass
+
+    @abstractmethod
+    async def write_file(
+        self,
+        filename: str,
+        content: str,
+        context: "ToolContext",
+        overwrite: bool = False,
+    ) -> None:
+        """Write content to a file."""
+        pass
+
+    @abstractmethod
+    async def exists(self, path: str, context: "ToolContext") -> bool:
+        """Check if a file or directory exists."""
+        pass
+
+    @abstractmethod
+    async def is_directory(self, path: str, context: "ToolContext") -> bool:
+        """Check if a path is a directory."""
+        pass
+
+    @abstractmethod
+    async def search_files(
+        self,
+        query: str,
+        context: "ToolContext",
+        *,
+        max_results: int = 20,
+        include_content: bool = False,
+    ) -> List[FileSearchMatch]:
+        """Search for files matching a query within the accessible namespace."""
+        pass
+
+    @abstractmethod
+    async def run_bash(
+        self,
+        command: str,
+        context: "ToolContext",
+        *,
+        timeout: Optional[float] = None,
+    ) -> CommandResult:
+        """Execute a bash command within the accessible namespace."""
+        pass

vanna/capabilities/file_system/models.py

@@ -0,0 +1,25 @@
+"""
+File system capability models.
+
+This module contains data models for file system operations.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class FileSearchMatch:
+    """Represents a single search result within a file system."""
+
+    path: str
+    snippet: Optional[str] = None
+
+
+@dataclass
+class CommandResult:
+    """Represents the result of executing a shell command."""
+
+    stdout: str
+    stderr: str
+    returncode: int

vanna/capabilities/sql_runner/__init__.py

@@ -0,0 +1,13 @@
+"""
+SQL runner capability.
+
+This module provides abstractions for SQL execution used by tools.
+"""
+
+from .base import SqlRunner
+from .models import RunSqlToolArgs
+
+__all__ = [
+    "SqlRunner",
+    "RunSqlToolArgs",
+]

vanna/capabilities/sql_runner/base.py

@@ -0,0 +1,37 @@
+"""
+SQL runner capability interface.
+
+This module contains the abstract base class for SQL execution.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import pandas as pd
+
+from .models import RunSqlToolArgs
+
+if TYPE_CHECKING:
+    from vanna.core.tool import ToolContext
+
+
+class SqlRunner(ABC):
+    """Interface for SQL execution with different implementations."""
+
+    @abstractmethod
+    async def run_sql(
+        self, args: RunSqlToolArgs, context: "ToolContext"
+    ) -> pd.DataFrame:
+        """Execute SQL query and return results as a DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            Exception: If query execution fails
+        """
+        pass

vanna/capabilities/sql_runner/models.py

@@ -0,0 +1,13 @@
+"""
+SQL runner capability models.
+
+This module contains data models for SQL execution.
+"""
+
+from pydantic import BaseModel, Field
+
+
+class RunSqlToolArgs(BaseModel):
+    """Arguments for run_sql tool."""
+
+    sql: str = Field(description="SQL query to execute")

vanna/components/__init__.py

@@ -0,0 +1,92 @@
+"""UI Component system for Vanna Agents."""
+
+# Base component
+from .base import UiComponent
+
+# Simple components
+from .simple import (
+    SimpleComponent,
+    SimpleComponentType,
+    SimpleTextComponent,
+    SimpleImageComponent,
+    SimpleLinkComponent,
+)
+
+# Rich components - re-export all
+from .rich import (
+    # Base
+    RichComponent,
+    ComponentType,
+    ComponentLifecycle,
+    # Text
+    RichTextComponent,
+    # Data
+    DataFrameComponent,
+    ChartComponent,
+    # Feedback
+    NotificationComponent,
+    StatusCardComponent,
+    ProgressBarComponent,
+    ProgressDisplayComponent,
+    StatusIndicatorComponent,
+    LogViewerComponent,
+    LogEntry,
+    BadgeComponent,
+    IconTextComponent,
+    # Interactive
+    TaskListComponent,
+    Task,
+    StatusBarUpdateComponent,
+    TaskTrackerUpdateComponent,
+    ChatInputUpdateComponent,
+    TaskOperation,
+    ButtonComponent,
+    ButtonGroupComponent,
+    # Containers
+    CardComponent,
+    # Specialized
+    ArtifactComponent,
+)
+
+__all__ = [
+    # Base
+    "UiComponent",
+    # Simple components
+    "SimpleComponent",
+    "SimpleComponentType",
+    "SimpleTextComponent",
+    "SimpleImageComponent",
+    "SimpleLinkComponent",
+    # Rich components - Base
+    "RichComponent",
+    "ComponentType",
+    "ComponentLifecycle",
+    # Rich components - Text
+    "RichTextComponent",
+    # Rich components - Data
+    "DataFrameComponent",
+    "ChartComponent",
+    # Rich components - Feedback
+    "NotificationComponent",
+    "StatusCardComponent",
+    "ProgressBarComponent",
+    "ProgressDisplayComponent",
+    "StatusIndicatorComponent",
+    "LogViewerComponent",
+    "LogEntry",
+    "BadgeComponent",
+    "IconTextComponent",
+    # Rich components - Interactive
+    "TaskListComponent",
+    "Task",
+    "StatusBarUpdateComponent",
+    "TaskTrackerUpdateComponent",
+    "ChatInputUpdateComponent",
+    "TaskOperation",
+    "ButtonComponent",
+    "ButtonGroupComponent",
+    # Rich components - Containers
+    "CardComponent",
+    # Rich components - Specialized
+    "ArtifactComponent",
+]

vanna/components/base.py

@@ -0,0 +1,11 @@
+"""
+UI components base - re-exports UiComponent from core.
+
+UiComponent lives in core/ because it's a fundamental return type for tools.
+This module provides backward compatibility by re-exporting it here.
+"""
+
+# Re-export UiComponent from core for backward compatibility
+from ..core.components import UiComponent
+
+__all__ = ["UiComponent"]

vanna/components/rich/__init__.py

@@ -0,0 +1,83 @@
+"""Rich UI components for the Vanna Agents framework."""
+
+# Base classes and enums - import from core
+from ...core.rich_component import RichComponent, ComponentType, ComponentLifecycle
+
+# Text component
+from .text import RichTextComponent
+
+# Data components
+from .data import (
+    DataFrameComponent,
+    ChartComponent,
+)
+
+# Feedback components
+from .feedback import (
+    NotificationComponent,
+    StatusCardComponent,
+    ProgressBarComponent,
+    ProgressDisplayComponent,
+    StatusIndicatorComponent,
+    LogViewerComponent,
+    LogEntry,
+    BadgeComponent,
+    IconTextComponent,
+)
+
+# Interactive components
+from .interactive import (
+    TaskListComponent,
+    Task,
+    StatusBarUpdateComponent,
+    TaskTrackerUpdateComponent,
+    ChatInputUpdateComponent,
+    TaskOperation,
+    ButtonComponent,
+    ButtonGroupComponent,
+)
+
+# Container components
+from .containers import (
+    CardComponent,
+)
+
+# Specialized components
+from .specialized import (
+    ArtifactComponent,
+)
+
+__all__ = [
+    # Base
+    "RichComponent",
+    "ComponentType",
+    "ComponentLifecycle",
+    # Text
+    "RichTextComponent",
+    # Data
+    "DataFrameComponent",
+    "ChartComponent",
+    # Feedback
+    "NotificationComponent",
+    "StatusCardComponent",
+    "ProgressBarComponent",
+    "ProgressDisplayComponent",
+    "StatusIndicatorComponent",
+    "LogViewerComponent",
+    "LogEntry",
+    "BadgeComponent",
+    "IconTextComponent",
+    # Interactive
+    "TaskListComponent",
+    "Task",
+    "StatusBarUpdateComponent",
+    "TaskTrackerUpdateComponent",
+    "ChatInputUpdateComponent",
+    "TaskOperation",
+    "ButtonComponent",
+    "ButtonGroupComponent",
+    # Containers
+    "CardComponent",
+    # Specialized
+    "ArtifactComponent",
+]

vanna/components/rich/containers/__init__.py

@@ -0,0 +1,7 @@
+"""Container components for layout."""
+
+from .card import CardComponent
+
+__all__ = [
+    "CardComponent",
+]

vanna/components/rich/containers/card.py

@@ -0,0 +1,20 @@
+"""Card component for displaying structured information."""
+
+from typing import Any, Dict, List, Optional
+from pydantic import Field
+from ....core.rich_component import RichComponent, ComponentType
+
+
+class CardComponent(RichComponent):
+    """Card component for displaying structured information."""
+
+    type: ComponentType = ComponentType.CARD
+    title: str
+    content: str
+    subtitle: Optional[str] = None
+    icon: Optional[str] = None
+    status: Optional[str] = None  # "success", "warning", "error", "info"
+    actions: List[Dict[str, Any]] = Field(default_factory=list)
+    collapsible: bool = False
+    collapsed: bool = False
+    markdown: bool = False  # Whether content should be rendered as markdown

vanna/components/rich/data/__init__.py

@@ -0,0 +1,9 @@
+"""Data display components."""
+
+from .dataframe import DataFrameComponent
+from .chart import ChartComponent
+
+__all__ = [
+    "DataFrameComponent",
+    "ChartComponent",
+]

vanna/components/rich/data/chart.py

@@ -0,0 +1,17 @@
+"""Chart component for data visualization."""
+
+from typing import Any, Dict, Optional, Union
+from pydantic import Field
+from ....core.rich_component import RichComponent, ComponentType
+
+
+class ChartComponent(RichComponent):
+    """Chart component for data visualization."""
+
+    type: ComponentType = ComponentType.CHART
+    chart_type: str  # "line", "bar", "pie", "scatter", etc.
+    data: Dict[str, Any]  # Chart data in format expected by frontend
+    title: Optional[str] = None
+    width: Optional[Union[str, int]] = None
+    height: Optional[Union[str, int]] = None
+    config: Dict[str, Any] = Field(default_factory=dict)  # Chart-specific config

vanna/components/rich/data/dataframe.py

@@ -0,0 +1,93 @@
+"""DataFrame component for displaying tabular data."""
+
+from typing import Any, Dict, List, Optional
+from pydantic import Field
+from ....core.rich_component import RichComponent, ComponentType
+
+
+class DataFrameComponent(RichComponent):
+    """DataFrame component specifically for displaying tabular data from SQL queries and similar sources."""
+
+    type: ComponentType = ComponentType.DATAFRAME
+    rows: List[Dict[str, Any]] = Field(default_factory=list)  # List of row dictionaries
+    columns: List[str] = Field(default_factory=list)  # Column names in display order
+    title: Optional[str] = None
+    description: Optional[str] = None
+    row_count: int = 0
+    column_count: int = 0
+
+    # Display options
+    max_rows_displayed: int = 100  # Limit rows shown in UI
+    searchable: bool = True
+    sortable: bool = True
+    filterable: bool = True
+    exportable: bool = True  # Allow export to CSV/Excel
+
+    # Styling options
+    striped: bool = True
+    bordered: bool = True
+    compact: bool = False
+
+    # Pagination
+    paginated: bool = True
+    page_size: int = 25
+
+    # Data types for better formatting (optional)
+    column_types: Dict[str, str] = Field(
+        default_factory=dict
+    )  # column_name -> "string"|"number"|"date"|"boolean"
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Set defaults before calling super().__init__
+        if "rows" not in kwargs:
+            kwargs["rows"] = []
+        if "columns" not in kwargs:
+            kwargs["columns"] = []
+        if "column_types" not in kwargs:
+            kwargs["column_types"] = {}
+
+        super().__init__(**kwargs)
+
+        # Auto-calculate counts if not provided
+        if self.rows and len(self.rows) > 0:
+            if "row_count" not in kwargs:
+                self.row_count = len(self.rows)
+            if not self.columns and self.rows:
+                self.columns = list(self.rows[0].keys())
+            if "column_count" not in kwargs:
+                self.column_count = len(self.columns)
+        else:
+            if "row_count" not in kwargs:
+                self.row_count = 0
+            if "column_count" not in kwargs:
+                self.column_count = len(self.columns) if self.columns else 0
+
+    @classmethod
+    def from_records(
+        cls,
+        records: List[Dict[str, Any]],
+        title: Optional[str] = None,
+        description: Optional[str] = None,
+        **kwargs: Any,
+    ) -> "DataFrameComponent":
+        """Create a DataFrame component from a list of record dictionaries."""
+        columns = list(records[0].keys()) if records else []
+
+        # Ensure we pass the required arguments correctly
+        component_data = {
+            "rows": records,
+            "columns": columns,
+            "row_count": len(records),
+            "column_count": len(columns),
+            "column_types": {},  # Initialize empty dict
+        }
+
+        if title is not None:
+            component_data["title"] = title
+        if description is not None:
+            component_data["description"] = description
+
+        # Merge with any additional kwargs
+        component_data.update(kwargs)
+
+        return cls(**component_data)