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

vanna/core/observability/base.py

@@ -0,0 +1,88 @@
+"""
+Base observability provider interface.
+
+Observability providers allow you to collect telemetry data about
+agent execution for monitoring and debugging.
+"""
+
+from abc import ABC
+from typing import Any, Dict, Optional
+
+from .models import Span, Metric
+
+
+class ObservabilityProvider(ABC):
+    """Provider for collecting telemetry and observability data.
+
+    Subclass this to create custom observability integrations that can:
+    - Emit metrics to monitoring systems
+    - Create distributed traces
+    - Log performance data
+    - Track costs and usage
+    - Monitor error rates
+
+    Example:
+        class PrometheusProvider(ObservabilityProvider):
+            def __init__(self, registry):
+                self.registry = registry
+                self.request_counter = Counter(
+                    'agent_requests_total',
+                    'Total agent requests',
+                    registry=registry
+                )
+
+            async def record_metric(self, name: str, value: float, tags: Dict[str, str]) -> None:
+                if name == "agent.request":
+                    self.request_counter.inc()
+
+            async def create_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Span:
+                span = Span(name=name, attributes=attributes or {})
+                return span
+
+        agent = AgentRunner(
+            llm_service=...,
+            observability_provider=PrometheusProvider(registry)
+        )
+    """
+
+    async def record_metric(
+        self,
+        name: str,
+        value: float,
+        unit: str = "",
+        tags: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Record a metric measurement.
+
+        Args:
+            name: Metric name (e.g., "agent.request.duration")
+            value: Metric value
+            unit: Unit of measurement (e.g., "ms", "tokens")
+            tags: Additional tags/labels for the metric
+        """
+        pass
+
+    async def create_span(
+        self, name: str, attributes: Optional[Dict[str, Any]] = None
+    ) -> Span:
+        """Create a new span for tracing.
+
+        Args:
+            name: Span name/operation
+            attributes: Initial span attributes
+
+        Returns:
+            Span object to track the operation
+
+        Note:
+            Call span.end() when the operation completes.
+        """
+        return Span(name=name, attributes=attributes or {})
+
+    async def end_span(self, span: Span) -> None:
+        """End a span and record it.
+
+        Args:
+            span: The span to end
+        """
+        span.end()

vanna/core/observability/models.py

@@ -0,0 +1,47 @@
+"""
+Observability models for spans and metrics.
+"""
+
+import time
+from typing import Any, Dict, Optional
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+class Span(BaseModel):
+    """Represents a unit of work for distributed tracing."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()), description="Span ID")
+    name: str = Field(description="Span name/operation")
+    start_time: float = Field(default_factory=time.time, description="Start timestamp")
+    end_time: Optional[float] = Field(default=None, description="End timestamp")
+    attributes: Dict[str, Any] = Field(
+        default_factory=dict, description="Span attributes"
+    )
+    parent_id: Optional[str] = Field(default=None, description="Parent span ID")
+
+    def end(self) -> None:
+        """Mark span as ended."""
+        if self.end_time is None:
+            self.end_time = time.time()
+
+    def duration_ms(self) -> Optional[float]:
+        """Get span duration in milliseconds."""
+        if self.end_time is None:
+            return None
+        return (self.end_time - self.start_time) * 1000
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Set a span attribute."""
+        self.attributes[key] = value
+
+
+class Metric(BaseModel):
+    """Represents a metric measurement."""
+
+    name: str = Field(description="Metric name")
+    value: float = Field(description="Metric value")
+    unit: str = Field(default="", description="Unit of measurement")
+    tags: Dict[str, str] = Field(default_factory=dict, description="Metric tags")
+    timestamp: float = Field(default_factory=time.time, description="Measurement time")

vanna/core/recovery/__init__.py

@@ -0,0 +1,11 @@
+"""
+Error recovery system for handling failures gracefully.
+
+This module provides interfaces for custom error handling, retry logic,
+and fallback strategies.
+"""
+
+from .base import ErrorRecoveryStrategy
+from .models import RecoveryAction, RecoveryActionType
+
+__all__ = ["ErrorRecoveryStrategy", "RecoveryAction", "RecoveryActionType"]

vanna/core/recovery/base.py

@@ -0,0 +1,84 @@
+"""
+Base error recovery strategy interface.
+
+Recovery strategies allow you to customize how the agent handles errors
+during tool execution and LLM communication.
+"""
+
+from abc import ABC
+from typing import TYPE_CHECKING
+
+from .models import RecoveryAction, RecoveryActionType
+
+if TYPE_CHECKING:
+    from ..tool.models import ToolContext
+    from ..llm import LlmRequest
+
+
+class ErrorRecoveryStrategy(ABC):
+    """Strategy for handling errors and implementing retry logic.
+
+    Subclass this to create custom error recovery strategies that can:
+    - Retry failed operations with backoff
+    - Fallback to alternative approaches
+    - Log errors to external systems
+    - Gracefully degrade functionality
+
+    Example:
+        class ExponentialBackoffStrategy(ErrorRecoveryStrategy):
+            async def handle_tool_error(
+                self, error: Exception, context: ToolContext, attempt: int
+            ) -> RecoveryAction:
+                if attempt < 3:
+                    delay = (2 ** attempt) * 1000  # Exponential backoff
+                    return RecoveryAction(
+                        action=RecoveryActionType.RETRY,
+                        retry_delay_ms=delay,
+                        message=f"Retrying after {delay}ms"
+                    )
+                return RecoveryAction(
+                    action=RecoveryActionType.FAIL,
+                    message="Max retries exceeded"
+                )
+
+        agent = AgentRunner(
+            llm_service=...,
+            error_recovery_strategy=ExponentialBackoffStrategy()
+        )
+    """
+
+    async def handle_tool_error(
+        self, error: Exception, context: "ToolContext", attempt: int = 1
+    ) -> RecoveryAction:
+        """Handle errors during tool execution.
+
+        Args:
+            error: The exception that occurred
+            context: Tool execution context
+            attempt: Current attempt number (1-indexed)
+
+        Returns:
+            RecoveryAction indicating how to proceed
+        """
+        # Default: fail immediately
+        return RecoveryAction(
+            action=RecoveryActionType.FAIL, message=f"Tool error: {str(error)}"
+        )
+
+    async def handle_llm_error(
+        self, error: Exception, request: "LlmRequest", attempt: int = 1
+    ) -> RecoveryAction:
+        """Handle errors during LLM communication.
+
+        Args:
+            error: The exception that occurred
+            request: The LLM request that failed
+            attempt: Current attempt number (1-indexed)
+
+        Returns:
+            RecoveryAction indicating how to proceed
+        """
+        # Default: fail immediately
+        return RecoveryAction(
+            action=RecoveryActionType.FAIL, message=f"LLM error: {str(error)}"
+        )

vanna/core/recovery/models.py

@@ -0,0 +1,32 @@
+"""
+Recovery action models for error handling.
+"""
+
+from enum import Enum
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class RecoveryActionType(str, Enum):
+    """Types of recovery actions."""
+
+    RETRY = "retry"
+    FAIL = "fail"
+    FALLBACK = "fallback"
+    SKIP = "skip"
+
+
+class RecoveryAction(BaseModel):
+    """Action to take when recovering from an error."""
+
+    action: RecoveryActionType = Field(description="Type of recovery action")
+    retry_delay_ms: Optional[int] = Field(
+        default=None, description="Delay before retry in milliseconds"
+    )
+    fallback_value: Optional[Any] = Field(
+        default=None, description="Fallback value to use"
+    )
+    message: Optional[str] = Field(
+        default=None, description="Message to include with action"
+    )

vanna/core/registry.py

@@ -0,0 +1,278 @@
+"""
+Tool registry for the Vanna Agents framework.
+
+This module provides the ToolRegistry class for managing and executing tools.
+"""
+
+import time
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, Union
+
+from .tool import Tool, ToolCall, ToolContext, ToolRejection, ToolResult, ToolSchema
+from .user import User
+
+if TYPE_CHECKING:
+    from .audit import AuditLogger
+    from .agent.config import AuditConfig
+
+T = TypeVar("T")
+
+
+class _LocalToolWrapper(Tool[T]):
+    """Wrapper for tools with configurable access groups."""
+
+    def __init__(self, wrapped_tool: Tool[T], access_groups: List[str]):
+        self._wrapped_tool = wrapped_tool
+        self._access_groups = access_groups
+
+    @property
+    def name(self) -> str:
+        return self._wrapped_tool.name
+
+    @property
+    def description(self) -> str:
+        return self._wrapped_tool.description
+
+    @property
+    def access_groups(self) -> List[str]:
+        return self._access_groups
+
+    def get_args_schema(self) -> Type[T]:
+        return self._wrapped_tool.get_args_schema()
+
+    async def execute(self, context: ToolContext, args: T) -> ToolResult:
+        return await self._wrapped_tool.execute(context, args)
+
+
+class ToolRegistry:
+    """Registry for managing tools."""
+
+    def __init__(
+        self,
+        audit_logger: Optional["AuditLogger"] = None,
+        audit_config: Optional["AuditConfig"] = None,
+    ) -> None:
+        self._tools: Dict[str, Tool[Any]] = {}
+        self.audit_logger = audit_logger
+        if audit_config is not None:
+            self.audit_config = audit_config
+        else:
+            from .agent.config import AuditConfig
+
+            self.audit_config = AuditConfig()
+
+    def register_local_tool(self, tool: Tool[Any], access_groups: List[str]) -> None:
+        """Register a local tool with optional access group restrictions.
+
+        Args:
+            tool: The tool to register
+            access_groups: List of groups that can access this tool.
+                          If None or empty, tool is accessible to all users.
+        """
+        if tool.name in self._tools:
+            raise ValueError(f"Tool '{tool.name}' already registered")
+
+        if access_groups:
+            # Wrap the tool with access groups
+            wrapped_tool = _LocalToolWrapper(tool, access_groups)
+            self._tools[tool.name] = wrapped_tool
+        else:
+            # No access restrictions, register as-is
+            self._tools[tool.name] = tool
+
+    async def get_tool(self, name: str) -> Optional[Tool[Any]]:
+        """Get a tool by name."""
+        return self._tools.get(name)
+
+    async def list_tools(self) -> List[str]:
+        """List all registered tool names."""
+        return list(self._tools.keys())
+
+    async def get_schemas(self, user: Optional[User] = None) -> List[ToolSchema]:
+        """Get schemas for all tools accessible to user."""
+        schemas = []
+        for tool in self._tools.values():
+            if user is None or await self._validate_tool_permissions(tool, user):
+                schemas.append(tool.get_schema())
+        return schemas
+
+    async def _validate_tool_permissions(self, tool: Tool[Any], user: User) -> bool:
+        """Validate if user has access to tool based on group membership.
+
+        Checks for intersection between user's group memberships and tool's access groups.
+        If tool has no access groups specified, it's accessible to all users.
+        """
+        tool_access_groups = tool.access_groups
+        if not tool_access_groups:
+            return True
+
+        user_groups = set(user.group_memberships)
+        tool_groups = set(tool_access_groups)
+        # Grant access if any group in user.group_memberships exists in tool.access_groups
+        return bool(user_groups & tool_groups)
+
+    async def transform_args(
+        self,
+        tool: Tool[T],
+        args: T,
+        user: User,
+        context: ToolContext,
+    ) -> Union[T, ToolRejection]:
+        """Transform and validate tool arguments based on user context.
+
+        This method allows per-user transformation of tool arguments, such as:
+        - Applying row-level security (RLS) to SQL queries
+        - Filtering available options based on user permissions
+        - Validating required arguments are present
+        - Redacting sensitive fields
+
+        The default implementation performs no transformation (NoOp).
+        Subclasses can override this method to implement custom transformation logic.
+
+        Args:
+            tool: The tool being executed
+            args: Already Pydantic-validated arguments
+            user: The user executing the tool
+            context: Full execution context
+
+        Returns:
+            Either:
+            - Transformed arguments (may be unchanged if no transformation needed)
+            - ToolRejection with explanation of why args were rejected
+        """
+        return args  # Default: no transformation (NoOp)
+
+    async def execute(
+        self,
+        tool_call: ToolCall,
+        context: ToolContext,
+    ) -> ToolResult:
+        """Execute a tool call with validation."""
+        tool = await self.get_tool(tool_call.name)
+        if not tool:
+            msg = f"Tool '{tool_call.name}' not found"
+            return ToolResult(
+                success=False,
+                result_for_llm=msg,
+                ui_component=None,
+                error=msg,
+            )
+
+        # Validate group access
+        if not await self._validate_tool_permissions(tool, context.user):
+            msg = f"Insufficient group access for tool '{tool_call.name}'"
+
+            # Audit access denial
+            if (
+                self.audit_logger
+                and self.audit_config
+                and self.audit_config.log_tool_access_checks
+            ):
+                await self.audit_logger.log_tool_access_check(
+                    user=context.user,
+                    tool_name=tool_call.name,
+                    access_granted=False,
+                    required_groups=tool.access_groups,
+                    context=context,
+                    reason=msg,
+                )
+
+            return ToolResult(
+                success=False,
+                result_for_llm=msg,
+                ui_component=None,
+                error=msg,
+            )
+
+        # Validate and parse arguments
+        try:
+            args_model = tool.get_args_schema()
+            validated_args = args_model.model_validate(tool_call.arguments)
+        except Exception as e:
+            msg = f"Invalid arguments: {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=msg,
+                ui_component=None,
+                error=msg,
+            )
+
+        # Transform/validate arguments based on user context
+        transform_result = await self.transform_args(
+            tool=tool,
+            args=validated_args,
+            user=context.user,
+            context=context,
+        )
+
+        if isinstance(transform_result, ToolRejection):
+            return ToolResult(
+                success=False,
+                result_for_llm=transform_result.reason,
+                ui_component=None,
+                error=transform_result.reason,
+            )
+
+        # Use transformed arguments for execution
+        final_args = transform_result
+
+        # Audit successful access check
+        if (
+            self.audit_logger
+            and self.audit_config
+            and self.audit_config.log_tool_access_checks
+        ):
+            await self.audit_logger.log_tool_access_check(
+                user=context.user,
+                tool_name=tool_call.name,
+                access_granted=True,
+                required_groups=tool.access_groups,
+                context=context,
+            )
+
+        # Audit tool invocation
+        if (
+            self.audit_logger
+            and self.audit_config
+            and self.audit_config.log_tool_invocations
+        ):
+            # Get UI features if available from context
+            ui_features = context.metadata.get("ui_features_available", [])
+            await self.audit_logger.log_tool_invocation(
+                user=context.user,
+                tool_call=tool_call,
+                ui_features=ui_features,
+                context=context,
+                sanitize_parameters=self.audit_config.sanitize_tool_parameters,
+            )
+
+        # Execute tool with context-first signature
+        try:
+            start_time = time.perf_counter()
+            result = await tool.execute(context, final_args)
+            execution_time_ms = (time.perf_counter() - start_time) * 1000
+
+            # Add execution time to metadata
+            result.metadata["execution_time_ms"] = execution_time_ms
+
+            # Audit tool result
+            if (
+                self.audit_logger
+                and self.audit_config
+                and self.audit_config.log_tool_results
+            ):
+                await self.audit_logger.log_tool_result(
+                    user=context.user,
+                    tool_call=tool_call,
+                    result=result,
+                    context=context,
+                )
+
+            return result
+        except Exception as e:
+            msg = f"Execution failed: {str(e)}"
+            return ToolResult(
+                success=False,
+                result_for_llm=msg,
+                ui_component=None,
+                error=msg,
+            )

vanna/core/rich_component.py

@@ -0,0 +1,156 @@
+"""
+Base classes for rich UI components.
+
+This module provides the base RichComponent class and supporting enums
+for the component system.
+"""
+
+import uuid
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, TypeVar
+
+from pydantic import BaseModel, Field
+
+# Type variable for self-returning methods
+T = TypeVar("T", bound="RichComponent")
+
+
+class ComponentType(str, Enum):
+    """Types of rich UI components."""
+
+    # Basic components
+    TEXT = "text"
+    CARD = "card"
+    CONTAINER = "container"
+
+    # Primitive UI components (domain-agnostic)
+    STATUS_CARD = "status_card"
+    PROGRESS_DISPLAY = "progress_display"
+    LOG_VIEWER = "log_viewer"
+    BADGE = "badge"
+    ICON_TEXT = "icon_text"
+
+    # Interactive components
+    TASK_LIST = "task_list"
+    PROGRESS_BAR = "progress_bar"
+    BUTTON = "button"
+    BUTTON_GROUP = "button_group"
+
+    # Data components
+    TABLE = "table"
+    DATAFRAME = "dataframe"
+    CHART = "chart"
+    CODE_BLOCK = "code_block"
+
+    # Status components
+    STATUS_INDICATOR = "status_indicator"
+    NOTIFICATION = "notification"
+    ALERT = "alert"
+
+    # Artifact components
+    ARTIFACT = "artifact"
+
+    # UI state components
+    STATUS_BAR_UPDATE = "status_bar_update"
+    TASK_TRACKER_UPDATE = "task_tracker_update"
+    CHAT_INPUT_UPDATE = "chat_input_update"
+
+    # Legacy (deprecated - use primitives instead)
+    TOOL_EXECUTION = "tool_execution"
+
+
+class ComponentLifecycle(str, Enum):
+    """Component lifecycle operations."""
+
+    CREATE = "create"
+    UPDATE = "update"
+    REPLACE = "replace"
+    REMOVE = "remove"
+
+
+class RichComponent(BaseModel):
+    """Base class for all rich UI components."""
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    type: ComponentType
+    lifecycle: ComponentLifecycle = ComponentLifecycle.CREATE
+    data: Dict[str, Any] = Field(default_factory=dict)
+    children: List[str] = Field(default_factory=list)  # Child component IDs
+    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
+    visible: bool = True
+    interactive: bool = False
+
+    def update(self: T, **kwargs: Any) -> T:
+        """Create an updated copy of this component."""
+        updated_data = self.model_dump()
+        updated_data.update(kwargs)
+        updated_data["lifecycle"] = ComponentLifecycle.UPDATE
+        updated_data["timestamp"] = datetime.utcnow().isoformat()
+        return self.__class__(**updated_data)
+
+    def hide(self: T) -> T:
+        """Create a hidden copy of this component."""
+        return self.update(visible=False)
+
+    def show(self: T) -> T:
+        """Create a visible copy of this component."""
+        return self.update(visible=True)
+
+    def serialize_for_frontend(self) -> Dict[str, Any]:
+        """Normalize component payload for the frontend renderer.
+
+        The frontend expects component-specific fields to live under the
+        ``data`` key while the shared metadata (``id``, ``type``, layout hints,
+        etc.) remains at the top level. Pydantic's ``model_dump`` keeps
+        component attributes at the top level, so we remap them here before
+        streaming them across the wire.
+        """
+
+        # Base fields that should remain at the top level of the payload.
+        shared_fields = {
+            "id",
+            "type",
+            "lifecycle",
+            "children",
+            "timestamp",
+            "visible",
+            "interactive",
+        }
+
+        raw = self.model_dump()
+        payload: Dict[str, Any] = {}
+
+        # Preserve any existing data payload so implementations can opt-in to
+        # advanced usage without losing information.
+        raw_data = raw.get("data")
+        if raw_data is not None and isinstance(raw_data, dict):
+            component_data: Dict[str, Any] = raw_data.copy()
+        else:
+            # Handle case where data might be a sequence or other type, or None
+            component_data = {}
+
+        for key, value in raw.items():
+            if key in shared_fields:
+                payload[key] = value
+            elif key == "data":
+                # For most components, skip the base data field
+                continue
+            elif (
+                key == "rows"
+                and hasattr(self, "type")
+                and self.type.value == "dataframe"
+            ):
+                # For DataFrame components, the 'rows' field contains the actual row data
+                # which should be included in the component_data as 'data' for the frontend
+                component_data["data"] = value
+            else:
+                component_data[key] = value
+
+        payload["data"] = component_data
+
+        # Ensure enums are serialized as primitive values for the frontend.
+        payload["type"] = self.type.value
+        payload["lifecycle"] = self.lifecycle.value
+
+        return payload