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

vanna/examples/primitive_components_demo.py

@@ -0,0 +1,305 @@
+"""
+Demonstration of the new primitive component system.
+
+This example shows how tools compose UI from primitive, domain-agnostic
+components like StatusCardComponent, ProgressDisplayComponent, etc.
+
+Usage:
+  PYTHONPATH=. python vanna/examples/primitive_components_demo.py
+"""
+
+import asyncio
+import uuid
+from datetime import datetime
+from typing import AsyncGenerator, Optional
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    MemoryConversationStore,
+    MockLlmService,
+    User,
+)
+from vanna.core.components import UiComponent
+from vanna.core.rich_components import (
+    StatusCardComponent,
+    ProgressDisplayComponent,
+    LogViewerComponent,
+    BadgeComponent,
+    IconTextComponent,
+    RichTextComponent,
+)
+
+
+class PrimitiveComponentsAgent(Agent):
+    """Agent that demonstrates the new primitive component system."""
+
+    async def send_message(
+        self,
+        user: User,
+        message: str,
+        *,
+        conversation_id: Optional[str] = None,
+    ) -> AsyncGenerator[UiComponent, None]:
+        """Send message and demonstrate primitive component composition."""
+
+        session_id = str(uuid.uuid4())[:8]
+
+        # Demo 1: Tool execution using primitive components
+        yield UiComponent(
+            rich_component=RichTextComponent(
+                content="## Primitive Components Demo\n\nShowing how tools now compose UI from primitive components:",
+                markdown=True,
+            )
+        )
+
+        # Status card for overall operation
+        operation_status = StatusCardComponent(
+            id=f"operation-{session_id}",
+            title="Data Analysis Pipeline",
+            status="running",
+            description="Processing user data through multiple analysis stages",
+            icon="⚙️",
+        )
+        yield UiComponent(rich_component=operation_status)
+
+        # Progress display for overall progress
+        overall_progress = ProgressDisplayComponent(
+            id=f"progress-{session_id}",
+            label="Overall Progress",
+            value=0.0,
+            description="Starting analysis...",
+            animated=True,
+        )
+        yield UiComponent(rich_component=overall_progress)
+
+        # Log viewer for detailed output
+        log_viewer = LogViewerComponent(
+            id=f"logs-{session_id}",
+            title="Analysis Log",
+            entries=[],
+            show_timestamps=True,
+            auto_scroll=True,
+        )
+        yield UiComponent(rich_component=log_viewer)
+
+        # Simulate analysis stages
+        stages = [
+            ("Data Loading", "📊", 0.2),
+            ("Data Validation", "✅", 0.4),
+            ("Statistical Analysis", "🧮", 0.6),
+            ("Report Generation", "📄", 0.8),
+            ("Finalization", "🎯", 1.0),
+        ]
+
+        for i, (stage_name, stage_icon, progress_value) in enumerate(stages):
+            await asyncio.sleep(0.8)
+
+            # Update overall status
+            status = "success" if progress_value == 1.0 else "running"
+            yield UiComponent(
+                rich_component=operation_status.set_status(
+                    status, f"Executing: {stage_name}"
+                )
+            )
+
+            # Update progress
+            yield UiComponent(
+                rich_component=overall_progress.update_progress(
+                    progress_value, f"Executing {stage_name}..."
+                )
+            )
+
+            # Add log entry
+            yield UiComponent(
+                rich_component=log_viewer.add_entry(f"Starting {stage_name}", "info")
+            )
+
+            # Create a status card for this specific stage
+            stage_status = StatusCardComponent(
+                id=f"stage-{i}-{session_id}",
+                title=stage_name,
+                status="running" if progress_value < 1.0 else "success",
+                description=f"Processing stage {i + 1} of {len(stages)}",
+                icon=stage_icon,
+            )
+            yield UiComponent(rich_component=stage_status)
+
+            await asyncio.sleep(0.5)
+
+            # Complete the stage
+            final_stage_status = "success" if progress_value < 1.0 else "completed"
+            yield UiComponent(
+                rich_component=stage_status.set_status(
+                    final_stage_status, f"{stage_name} completed successfully"
+                )
+            )
+            yield UiComponent(
+                rich_component=log_viewer.add_entry(f"Completed {stage_name}", "info")
+            )
+
+        # Demo 2: Badge and IconText primitives
+        yield UiComponent(
+            rich_component=RichTextComponent(
+                content="\n### Primitive Component Examples\n\nShowing individual primitive components:",
+                markdown=True,
+            )
+        )
+
+        # Various badge examples
+        badges = [
+            BadgeComponent(text="Processing", variant="primary", size="small"),
+            BadgeComponent(text="Complete", variant="success", size="medium"),
+            BadgeComponent(text="Warning", variant="warning", size="large", icon="⚠️"),
+            BadgeComponent(text="Error", variant="error", size="medium", icon="❌"),
+        ]
+
+        for badge in badges:
+            yield UiComponent(rich_component=badge)
+
+        # IconText examples
+        icon_texts = [
+            IconTextComponent(
+                icon="📊",
+                text="Data Analysis Complete",
+                variant="primary",
+                size="large",
+            ),
+            IconTextComponent(
+                icon="✅", text="All tests passed", variant="default", size="medium"
+            ),
+            IconTextComponent(
+                icon="⏱️",
+                text="Processing time: 2.3s",
+                variant="secondary",
+                size="small",
+            ),
+        ]
+
+        for icon_text in icon_texts:
+            yield UiComponent(rich_component=icon_text)
+
+        # Demo 3: Comparison with old approach
+        yield UiComponent(
+            rich_component=RichTextComponent(
+                content=f"""
+## Key Benefits of Primitive Components
+
+**Primitive Component Approach:**
+```python
+# Tool composes UI from primitives
+status_card = StatusCardComponent(
+    title="Data Analysis",
+    status="running",  # Pure UI state
+    icon="📊"
+)
+progress = ProgressDisplayComponent(
+    label="Analysis Progress",
+    value=0.5
+)
+logs = LogViewerComponent(
+    title="Analysis Log",
+    entries=log_entries
+)
+```
+
+### Benefits:
+- **Separation of Concerns**: UI components are purely presentational
+- **Reusability**: Status cards work for any process, not just tools
+- **Composability**: Tools build exactly the UI they need
+- **Maintainability**: Changes to business logic don't affect UI components
+- **Extensibility**: New tools don't require new component types
+
+Your message was: "{message}"
+""",
+                markdown=True,
+            )
+        )
+
+
+def create_primitive_demo_agent() -> PrimitiveComponentsAgent:
+    """Create a primitive components demo agent.
+
+    Returns:
+        Configured PrimitiveComponentsAgent instance
+    """
+    llm_service = MockLlmService(response_content="Primitive components demo response")
+
+    return PrimitiveComponentsAgent(
+        llm_service=llm_service,
+        config=AgentConfig(
+            stream_responses=True,
+            include_thinking_indicators=False,
+        ),
+    )
+
+
+async def main() -> None:
+    """Run the primitive components demo."""
+
+    # Create agent
+    agent = create_primitive_demo_agent()
+
+    # Create a test user
+    user = User(
+        id="user123", username="demo_user", email="demo@example.com", permissions=[]
+    )
+
+    # Start a conversation
+    conversation_id = "primitive_demo_123"
+    user_message = "Show me how the new primitive component system works!"
+
+    print(f"User: {user_message}")
+    print("Agent response (primitive components):")
+    print("=" * 60)
+
+    # Send message and display components
+    component_count = 0
+    async for component in agent.send_message(
+        user=user, message=user_message, conversation_id=conversation_id
+    ):
+        component_count += 1
+        component_type = getattr(component, "type", component.__class__.__name__)
+        component_id = getattr(component, "id", "N/A")
+
+        print(
+            f"[{component_count:2d}] {component_type.value if hasattr(component_type, 'value') else component_type} (id: {component_id[:12] if len(str(component_id)) > 12 else component_id})"
+        )
+
+        rich_comp = component.rich_component
+
+        # Show component details
+        if hasattr(rich_comp, "title"):
+            print(f"     Title: {rich_comp.title}")
+        if hasattr(rich_comp, "status"):
+            print(f"     Status: {rich_comp.status}")
+        if hasattr(rich_comp, "description") and rich_comp.description:
+            desc = (
+                rich_comp.description[:60] + "..."
+                if len(rich_comp.description) > 60
+                else rich_comp.description
+            )
+            print(f"     Description: {desc}")
+        if (
+            hasattr(rich_comp, "value")
+            and hasattr(rich_comp.type, "value")
+            and rich_comp.type.value == "progress_display"
+        ):
+            print(f"     Progress: {rich_comp.value:.1%}")
+
+        print()
+
+    print("=" * 60)
+    print(f"Total components emitted: {component_count}")
+    print("\nThis demonstrates how tools can now compose rich UIs")
+    print("from primitive, reusable components without semantic coupling!")
+
+
+def run_interactive() -> None:
+    """Entry point for interactive usage."""
+    print("Starting Primitive Components Demo...")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()

vanna/examples/quota_lifecycle_example.py

@@ -0,0 +1,139 @@
+"""
+Example demonstrating lifecycle hooks for user quota management.
+
+This example shows how to use lifecycle hooks to add custom functionality
+like quota management without creating custom agent runner subclasses.
+"""
+
+from typing import Any, Dict, Optional
+from vanna.core import Agent, LifecycleHook, User
+from vanna.core.errors import AgentError
+
+
+class QuotaExceededError(AgentError):
+    """Raised when a user exceeds their message quota."""
+
+    pass
+
+
+class QuotaCheckHook(LifecycleHook):
+    """Lifecycle hook that enforces user-based message quotas."""
+
+    def __init__(self, default_quota: int = 10) -> None:
+        """Initialize quota hook.
+
+        Args:
+            default_quota: Default quota per user if not specifically set
+        """
+        self._user_quotas: Dict[str, int] = {}
+        self._user_usage: Dict[str, int] = {}
+        self._default_quota = default_quota
+
+    def set_user_quota(self, user_id: str, quota: int) -> None:
+        """Set a specific quota for a user."""
+        self._user_quotas[user_id] = quota
+
+    def get_user_quota(self, user_id: str) -> int:
+        """Get the quota for a user."""
+        return self._user_quotas.get(user_id, self._default_quota)
+
+    def get_user_usage(self, user_id: str) -> int:
+        """Get current usage count for a user."""
+        return self._user_usage.get(user_id, 0)
+
+    def get_user_remaining(self, user_id: str) -> int:
+        """Get remaining messages for a user."""
+        return self.get_user_quota(user_id) - self.get_user_usage(user_id)
+
+    def reset_user_usage(self, user_id: str) -> None:
+        """Reset usage count for a user."""
+        self._user_usage[user_id] = 0
+
+    async def before_message(self, user: User, message: str) -> Optional[str]:
+        """Check quota before processing message.
+
+        Raises:
+            QuotaExceededError: If user has exceeded their quota
+        """
+        usage = self.get_user_usage(user.id)
+        quota = self.get_user_quota(user.id)
+
+        if usage >= quota:
+            raise QuotaExceededError(
+                f"User {user.username} has exceeded their quota of {quota} messages. "
+                f"Current usage: {usage}"
+            )
+
+        # Increment usage count
+        current_usage = self._user_usage.get(user.id, 0)
+        self._user_usage[user.id] = current_usage + 1
+
+        # Don't modify the message
+        return None
+
+
+class LoggingHook(LifecycleHook):
+    """Example logging hook for demonstration."""
+
+    async def before_message(self, user: User, message: str) -> Optional[str]:
+        """Log incoming messages."""
+        print(f"[LOG] User {user.username} ({user.id}) sent message: {message[:50]}...")
+        return None
+
+    async def after_message(self, result: Any) -> None:
+        """Log message completion."""
+        print(f"[LOG] Message processing completed")
+
+
+async def run_example() -> None:
+    """
+    Example showing how to use lifecycle hooks with Agent.
+
+    Instead of creating a custom subclass, we compose
+    the behavior using lifecycle hooks.
+    """
+    from vanna.core.registry import ToolRegistry
+    from vanna.integrations.anthropic import AnthropicLlmService
+    from vanna.integrations.local import MemoryConversationStore
+
+    # Create quota hook
+    quota_hook = QuotaCheckHook(default_quota=10)
+    quota_hook.set_user_quota("user123", 5)  # Set custom quota for specific user
+
+    # Create logging hook
+    logging_hook = LoggingHook()
+
+    # Create agent with multiple hooks
+    agent = Agent(
+        llm_service=AnthropicLlmService(api_key="your-api-key"),
+        tool_registry=ToolRegistry(),
+        conversation_store=MemoryConversationStore(),
+        lifecycle_hooks=[
+            logging_hook,  # Logs will happen first
+            quota_hook,  # Then quota check
+        ],
+    )
+
+    # Create a test user
+    user = User(
+        id="user123", username="test_user", email="test@example.com", permissions=[]
+    )
+
+    # Send messages - will track quota
+    try:
+        async for component in agent.send_message(user=user, message="Hello, agent!"):
+            # Process UI components
+            pass
+
+        # Check remaining quota
+        remaining = quota_hook.get_user_remaining(user.id)
+        print(f"Remaining messages: {remaining}/{quota_hook.get_user_quota(user.id)}")
+
+    except QuotaExceededError as e:
+        print(f"Quota exceeded: {e}")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(run_example())

vanna/examples/visualization_example.py

@@ -0,0 +1,251 @@
+"""
+Example demonstrating SQL query execution with automatic visualization.
+
+This example shows the integration of RunSqlTool and VisualizeDataTool,
+demonstrating how SQL results are saved to CSV files and can be visualized
+using the visualization tool with dependency injection.
+
+Usage:
+  PYTHONPATH=. python vanna/examples/visualization_example.py
+"""
+
+import asyncio
+import os
+import sys
+import uuid
+from typing import AsyncGenerator, List, Optional
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    ToolRegistry,
+    User,
+)
+from vanna.core import LlmService
+from vanna.core import (
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+    ToolCall,
+    ToolSchema,
+)
+from vanna.integrations.sqlite import SqliteRunner
+from vanna.tools import (
+    RunSqlTool,
+    VisualizeDataTool,
+    LocalFileSystem,
+)
+
+
+class VisualizationDemoLlmService(LlmService):
+    """Mock LLM that demonstrates SQL query and visualization workflow."""
+
+    def __init__(self) -> None:
+        self.step = 0
+        self.csv_filename: Optional[str] = None
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Handle non-streaming requests."""
+        await asyncio.sleep(0.1)
+        return self._build_response(request)
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Handle streaming requests."""
+        await asyncio.sleep(0.1)
+        response = self._build_response(request)
+
+        if response.tool_calls:
+            yield LlmStreamChunk(tool_calls=response.tool_calls)
+        if response.content:
+            yield LlmStreamChunk(
+                content=response.content, finish_reason=response.finish_reason
+            )
+        else:
+            yield LlmStreamChunk(finish_reason=response.finish_reason)
+
+    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
+        """Validate tools - no errors."""
+        return []
+
+    def _build_response(self, request: LlmRequest) -> LlmResponse:
+        """Build response based on conversation state."""
+        last_message = request.messages[-1] if request.messages else None
+
+        # If we got a tool result, process it
+        if last_message and last_message.role == "tool":
+            tool_result = last_message.content or ""
+
+            # Check if this was a SQL query result with a CSV file
+            if "Results saved to" in tool_result and ".csv" in tool_result:
+                # Extract filename from result
+                import re
+
+                match = re.search(r"'([^']*\.csv)'", tool_result)
+                if match:
+                    self.csv_filename = match.group(1)
+                    # Now visualize the data
+                    return LlmResponse(
+                        content=f"Great! I've saved the query results. Now let me create a visualization of the data.",
+                        tool_calls=[
+                            ToolCall(
+                                id=f"call_{uuid.uuid4().hex[:8]}",
+                                name="visualize_data",
+                                arguments={"filename": self.csv_filename},
+                            )
+                        ],
+                        finish_reason="tool_calls",
+                    )
+
+            # If this was a visualization result, acknowledge it
+            if "Created visualization" in tool_result:
+                return LlmResponse(
+                    content=f"Perfect! I've created a visualization of the data. {tool_result}",
+                    finish_reason="stop",
+                )
+
+            # Default acknowledgment
+            return LlmResponse(
+                content=f"I've completed the operation. {tool_result}",
+                finish_reason="stop",
+            )
+
+        # Initial request - run SQL query
+        if self.step == 0:
+            self.step += 1
+            return LlmResponse(
+                content="I'll query the database for you and then create a visualization.",
+                tool_calls=[
+                    ToolCall(
+                        id=f"call_{uuid.uuid4().hex[:8]}",
+                        name="run_sql",
+                        arguments={
+                            "sql": "SELECT Name, Milliseconds, Bytes FROM Track LIMIT 20"
+                        },
+                    )
+                ],
+                finish_reason="tool_calls",
+            )
+
+        # Default response
+        return LlmResponse(
+            content="I can help you query databases and visualize the results.",
+            finish_reason="stop",
+        )
+
+
+def create_demo_agent() -> Agent:
+    """
+    Create a demo agent with SQL and visualization tools.
+
+    This function is called by the vanna server framework.
+
+    Returns:
+        Configured Agent with SQL and visualization tools
+    """
+    # Check for Chinook database
+    database_path = os.path.join(
+        os.path.dirname(__file__), "..", "..", "Chinook.sqlite"
+    )
+    database_path = os.path.abspath(database_path)
+
+    if not os.path.exists(database_path):
+        raise FileNotFoundError(
+            f"Chinook database not found at {database_path}. "
+            "Please download it from https://vanna.ai/Chinook.sqlite"
+        )
+
+    # Create shared FileSystem for both tools
+    file_system = LocalFileSystem(working_directory="./data_storage")
+
+    # Create SQL tool with FileSystem
+    sqlite_runner = SqliteRunner(database_path=database_path)
+    sql_tool = RunSqlTool(sql_runner=sqlite_runner, file_system=file_system)
+
+    # Create visualization tool with same FileSystem
+    viz_tool = VisualizeDataTool(file_system=file_system)
+
+    # Create tool registry
+    tool_registry = ToolRegistry()
+    tool_registry.register(sql_tool)
+    tool_registry.register(viz_tool)
+
+    # Create LLM service
+    llm_service = VisualizationDemoLlmService()
+
+    # Create agent with streaming enabled for web interface
+    return Agent(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        config=AgentConfig(
+            stream_responses=True,
+            include_thinking_indicators=False,
+        ),
+    )
+
+
+async def main() -> None:
+    """Demonstrate SQL query execution with automatic visualization."""
+    print("🎨 SQL + Visualization Demo")
+    print("=" * 60)
+    print("This example demonstrates:")
+    print("1. Running SQL queries that save results to CSV files")
+    print("2. Automatically visualizing the CSV data")
+    print("3. User isolation for file storage")
+    print("=" * 60)
+    print()
+
+    # Create agent using factory function
+    agent = create_demo_agent()
+
+    # Create test user
+    user = User(id="demo-user", username="demo")
+
+    # Show available tools
+    tools = await agent.get_available_tools(user)
+    print(f"Available tools: {[tool.name for tool in tools]}")
+    print()
+
+    # Run conversation
+    conversation_id = "viz-demo"
+
+    print("User: Show me some track data and visualize it")
+    print()
+
+    async for component in agent.send_message(
+        user=user,
+        message="Show me some track data and visualize it",
+        conversation_id=conversation_id,
+    ):
+        if (
+            component.simple_component
+            and hasattr(component.simple_component, "text")
+            and component.simple_component.text
+        ):
+            print(f"Agent: {component.simple_component.text}")
+        elif component.simple_component and hasattr(component.simple_component, "text"):
+            print(f"Agent: {component.simple_component.text}")
+        elif hasattr(component.rich_component, "content"):
+            if isinstance(component.rich_component.content, dict):
+                # This is the chart
+                print(
+                    f"Agent: [Chart Generated - Plotly figure with {len(str(component.rich_component.content))} chars]"
+                )
+            else:
+                print(f"Agent: {component.rich_component.content}")
+
+    print()
+    print("=" * 60)
+    print("Demo complete!")
+    print()
+    print("Key features demonstrated:")
+    print("✅ SQL queries save results to user-isolated CSV files")
+    print("✅ Visualization tool reads CSV files using FileSystem")
+    print("✅ Automatic chart type selection based on data shape")
+    print("✅ Dependency injection allows customization")
+    print()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())