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

vanna/examples/mock_rich_components_demo.py

@@ -0,0 +1,396 @@
+"""
+Mock rich components demonstration example.
+
+This example shows how to create an agent that emits rich, stateful components
+including cards, task lists, and tool execution displays using a mock LLM service.
+
+Usage:
+  PYTHONPATH=. python vanna/examples/mock_rich_components_demo.py
+"""
+
+import asyncio
+import time
+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,
+    Task,
+)
+
+
+class RichComponentsAgent(Agent):
+    """Agent that demonstrates rich component capabilities."""
+
+    async def send_message(
+        self,
+        user: User,
+        message: str,
+        *,
+        conversation_id: Optional[str] = None,
+    ) -> AsyncGenerator[UiComponent, None]:
+        """Send message and yield UiComponent(rich_component=rich) components."""
+
+        # Welcome message using IconText
+        yield UiComponent(
+            rich_component=IconTextComponent(
+                id="welcome-message",
+                icon="👋",
+                text=f"Hello {user.username}! I'll demonstrate primitive components.",
+                variant="primary",
+                size="large",
+            )
+        )
+
+        # Status card showing we're processing
+        status_card = StatusCardComponent(
+            id="processing-status",
+            title="Processing Request",
+            status="running",
+            description="Processing your request...",
+            icon="⚙️",
+        )
+        yield UiComponent(rich_component=status_card)
+
+        # Simulate some processing time
+        await asyncio.sleep(1)
+
+        # Update status to success
+        yield UiComponent(
+            rich_component=status_card.set_status(
+                "success", "Request processed successfully!"
+            )
+        )
+
+        # Create a status card for overall demo progress
+        demo_card = StatusCardComponent(
+            id="demo-progress",
+            title="Demo Progress",
+            status="running",
+            description="Starting primitive components demonstration...",
+            icon="🎯",
+        )
+        yield UiComponent(rich_component=demo_card)
+
+        # Create badges for different stages
+        stages = [
+            ("Initialize", "success", "✅"),
+            ("Components", "running", "⚙️"),
+            ("Progress", "pending", "⏳"),
+            ("Logs", "pending", "📋"),
+            ("Complete", "pending", "🎉"),
+        ]
+
+        for stage_name, stage_status, stage_icon in stages:
+            yield UiComponent(
+                rich_component=BadgeComponent(
+                    id=f"stage-{stage_name.lower()}",
+                    text=stage_name,
+                    variant=stage_status if stage_status != "pending" else "default",
+                    icon=stage_icon,
+                    size="md",
+                )
+            )
+
+        # Progress display
+        progress_display = ProgressDisplayComponent(
+            id="demo-progress-bar",
+            label="Overall Progress",
+            value=0.2,
+            description="Initializing demonstration...",
+            status="info",
+            animated=True,
+        )
+        yield UiComponent(rich_component=progress_display)
+
+        # Create log viewer for detailed progress
+        log_viewer = LogViewerComponent(id="demo-logs", title="Demo Activity Log")
+        yield UiComponent(rich_component=log_viewer)
+
+        # Simulate work with updates
+        for i in range(3):
+            await asyncio.sleep(1)
+
+            # Update progress
+            progress_value = 0.2 + (i + 1) * 0.2
+            step_name = ["Creating components", "Updating progress", "Finalizing demo"][
+                i
+            ]
+
+            yield UiComponent(
+                rich_component=progress_display.update_progress(
+                    progress_value, f"Step {i + 2} of 5: {step_name}..."
+                )
+            )
+
+            # Update demo card
+            yield UiComponent(
+                rich_component=demo_card.set_status(
+                    "running",
+                    f"Step {i + 2} of 5 completed. Progress: {int(progress_value * 100)}%",
+                )
+            )
+
+            # Add log entry
+            yield UiComponent(
+                rich_component=log_viewer.add_entry(
+                    f"Completed step: {step_name}", "info"
+                )
+            )
+
+            # Update stage badges
+            if i == 0:
+                yield UiComponent(
+                    rich_component=BadgeComponent(
+                        id="stage-components",
+                        text="Components",
+                        variant="success",
+                        icon="✅",
+                        size="md",
+                    )
+                )
+            elif i == 1:
+                yield UiComponent(
+                    rich_component=BadgeComponent(
+                        id="stage-progress",
+                        text="Progress",
+                        variant="success",
+                        icon="✅",
+                        size="md",
+                    )
+                )
+                yield UiComponent(
+                    rich_component=BadgeComponent(
+                        id="stage-logs",
+                        text="Logs",
+                        variant="running",
+                        icon="📋",
+                        size="md",
+                    )
+                )
+
+        # Tool execution using primitive components
+        tool_status = StatusCardComponent(
+            id="demo-tool",
+            title="Analyze Data Tool",
+            status="running",
+            description="Running regression analysis on user_data.csv",
+            icon="🔬",
+        )
+        yield UiComponent(rich_component=tool_status)
+
+        # Tool progress
+        tool_progress = ProgressDisplayComponent(
+            id="tool-progress",
+            label="Tool Execution",
+            value=0.0,
+            description="Initializing tool...",
+            animated=True,
+        )
+        yield UiComponent(rich_component=tool_progress)
+
+        # Tool logs
+        tool_logs = LogViewerComponent(id="tool-logs", title="Tool Execution Log")
+        yield UiComponent(rich_component=tool_logs)
+
+        # Simulate tool execution steps
+        tool_steps = [
+            (0.2, "Loading dataset...", "info"),
+            (0.4, "Dataset loaded: 1000 rows, 5 columns", "info"),
+            (0.6, "Preprocessing data...", "info"),
+            (0.8, "Running regression analysis...", "info"),
+            (1.0, "Analysis complete!", "info"),
+        ]
+
+        for progress_val, log_message, log_level in tool_steps:
+            await asyncio.sleep(0.5)
+
+            yield UiComponent(
+                rich_component=tool_progress.update_progress(
+                    progress_val, f"Progress: {int(progress_val * 100)}%"
+                )
+            )
+            yield UiComponent(
+                rich_component=tool_logs.add_entry(log_message, log_level)
+            )
+
+        # Complete tool execution
+        yield UiComponent(
+            rich_component=tool_status.set_status(
+                "success",
+                "Tool completed successfully. R² = 0.85, strong correlation found.",
+            )
+        )
+
+        # Show results using IconText
+        yield UiComponent(
+            rich_component=IconTextComponent(
+                id="tool-results",
+                icon="📊",
+                text="Analysis Results: R² = 0.85 (Strong correlation)",
+                variant="success",
+                size="medium",
+            )
+        )
+
+        # Update final stage badge
+        yield UiComponent(
+            rich_component=BadgeComponent(
+                id="stage-logs", text="Logs", variant="success", icon="✅", size="md"
+            )
+        )
+        yield UiComponent(
+            rich_component=BadgeComponent(
+                id="stage-complete",
+                text="Complete",
+                variant="success",
+                icon="🎉",
+                size="md",
+            )
+        )
+
+        # Final updates
+        yield UiComponent(
+            rich_component=progress_display.update_progress(
+                1.0, "Demo completed successfully!"
+            )
+        )
+
+        yield UiComponent(
+            rich_component=demo_card.set_status(
+                "success", "Primitive components demonstration finished successfully!"
+            )
+        )
+
+        # Add final log entry
+        yield UiComponent(
+            rich_component=tool_logs.add_entry("Demo completed successfully!", "info")
+        )
+
+        # Add final text response
+        yield UiComponent(
+            rich_component=RichTextComponent(
+                content=f"""## Primitive Components Demo Complete!
+
+I've demonstrated the new primitive component system:
+
+- **Status Cards**: Domain-agnostic status displays that work for any process
+- **Progress Displays**: Reusable progress indicators with animations
+- **Log Viewers**: Structured log display for any activity
+- **Badges**: Flexible status and category indicators
+- **Icon Text**: Composable icon+text combinations
+
+### Key Benefits of Primitive Components:
+
+- **Separation of Concerns**: UI components are purely presentational
+- **Reusability**: Components work across different domains and tools
+- **Composability**: Tools build exactly the UI they need from primitives
+- **Maintainability**: Business logic changes don't affect UI components
+- **Extensibility**: New tools don't require new component types
+
+**Primitive Components**: Compose UI from domain-agnostic building blocks
+**After**: Tools compose UI from primitive `StatusCard` + `ProgressDisplay` + `LogViewer`
+
+Your message was: "{message}"
+""",
+                markdown=True,
+            )
+        )
+
+
+# CLI compatibility alias
+create_demo_agent = lambda: create_rich_demo_agent()
+
+
+def create_rich_demo_agent() -> RichComponentsAgent:
+    """Create a primitive components demo agent.
+
+    Returns:
+        Configured RichComponentsAgent instance
+    """
+    llm_service = MockLlmService(response_content="Primitive components demo response")
+
+    return RichComponentsAgent(
+        llm_service=llm_service,
+        config=AgentConfig(
+            stream_responses=True,
+            include_thinking_indicators=False,  # We'll use custom status cards
+        ),
+    )
+
+
+async def main() -> None:
+    """Run the primitive components demo."""
+
+    # Create agent
+    agent = create_rich_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 the primitive components demo!"
+
+    print(f"User: {user_message}")
+    print("Agent response (primitive components):")
+    print("=" * 50)
+
+    # 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
+        rich_comp = component.rich_component
+        component_type = getattr(rich_comp, "type", rich_comp.__class__.__name__)
+        component_id = getattr(rich_comp, "id", "N/A")
+        lifecycle = getattr(rich_comp, "lifecycle", "N/A")
+
+        print(
+            f"[{component_count:2d}] {component_type} (id: {component_id[:8]}, lifecycle: {lifecycle})"
+        )
+
+        # Show some component details
+        if hasattr(rich_comp, "title"):
+            print(f"     Title: {rich_comp.title}")
+        if hasattr(rich_comp, "content") and len(str(rich_comp.content)) < 100:
+            print(f"     Content: {rich_comp.content}")
+        if hasattr(rich_comp, "status"):
+            print(f"     Status: {rich_comp.status}")
+        if (
+            hasattr(rich_comp, "value")
+            and hasattr(rich_comp.type, "value")
+            and rich_comp.type.value == "progress_bar"
+        ):
+            print(f"     Progress: {rich_comp.value:.1%}")
+
+        print()
+
+    print("=" * 50)
+    print(f"Total components emitted: {component_count}")
+
+
+def run_interactive() -> None:
+    """Entry point for interactive usage."""
+    print("Starting Primitive Components Demo...")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()

vanna/examples/mock_sqlite_example.py

@@ -0,0 +1,223 @@
+"""
+Mock example showing how to use the SQL query tool with the Chinook database.
+
+This example demonstrates using the RunSqlTool with SqliteRunner and a mock LLM service
+that automatically executes sample SQL queries against the Chinook database.
+
+Usage:
+  Template: Copy this file and modify for your custom database
+  Interactive: python -m vanna.examples.mock_sqlite_example
+  REPL: from vanna.examples.mock_sqlite_example import create_demo_agent
+  Server: python -m vanna.servers --example mock_sqlite_example
+"""
+
+import asyncio
+import os
+import random
+import uuid
+from typing import AsyncGenerator, Dict, List, Optional, Type
+
+from pydantic import BaseModel, Field
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    Tool,
+    ToolContext,
+    ToolRegistry,
+    ToolResult,
+    User,
+    UiComponent,
+)
+from vanna.core.interfaces import LlmService
+from vanna.core.models import (
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+    ToolCall,
+    ToolSchema,
+)
+from vanna.core.rich_components import (
+    CardComponent,
+    NotificationComponent,
+    ComponentType,
+)
+from vanna.core.simple_components import (
+    SimpleTextComponent,
+)
+from vanna.tools import RunSqlTool
+from vanna.integrations.sqlite import SqliteRunner
+
+
+class MockSqliteLlmService(LlmService):
+    """LLM service that exercises the SQLite query tool with sample queries."""
+
+    def __init__(self, seed: Optional[int] = None):
+        self._random = random.Random(seed)
+        self._sample_queries = [
+            "SELECT name FROM sqlite_master WHERE type='table'",
+            "SELECT COUNT(*) as total_customers FROM Customer",
+            "SELECT FirstName, LastName FROM Customer LIMIT 5",
+            "SELECT Name, Composer FROM Track WHERE Composer IS NOT NULL LIMIT 5",
+            "SELECT COUNT(*) as album_count FROM Album",
+            "SELECT Name FROM Artist LIMIT 10",
+            "SELECT AVG(Total) as avg_invoice_total FROM Invoice",
+            "SELECT GenreId, COUNT(*) as track_count FROM Track GROUP BY GenreId LIMIT 5",
+        ]
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Handle non-streaming SQLite interactions."""
+        await asyncio.sleep(0.1)
+        return self._build_response(request)
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Provide streaming compatibility by yielding a single chunk."""
+        await asyncio.sleep(0.1)
+        response = self._build_response(request)
+
+        if response.tool_calls:
+            yield LlmStreamChunk(tool_calls=response.tool_calls)
+        if response.content is not None:
+            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]:
+        """Mock validation - no errors."""
+        return []
+
+    def _build_response(self, request: LlmRequest) -> LlmResponse:
+        """Create a response that either calls the tool or explains its result."""
+        last_message = request.messages[-1] if request.messages else None
+
+        if last_message and last_message.role == "tool":
+            # Respond to tool result
+            result = last_message.content or "No result provided"
+            return LlmResponse(
+                content=f"Here's what I found in the database:\n\n{result}",
+                finish_reason="stop",
+                usage={
+                    "prompt_tokens": 40,
+                    "completion_tokens": 20,
+                    "total_tokens": 60,
+                },
+            )
+
+        # Generate a random SQL query
+        sql_query = self._random.choice(self._sample_queries)
+        tool_call = ToolCall(
+            id=f"call_{uuid.uuid4().hex[:8]}",
+            name="run_sql",
+            arguments={"sql": sql_query},
+        )
+
+        return LlmResponse(
+            content="Let me query the Chinook database for you...",
+            tool_calls=[tool_call],
+            finish_reason="tool_calls",
+            usage={"prompt_tokens": 30, "completion_tokens": 10, "total_tokens": 40},
+        )
+
+
+def create_demo_agent() -> Agent:
+    """Create a demo agent with SQLite query tool.
+
+    Returns:
+        Configured Agent with SQLite tool and mock LLM
+    """
+    # Get the path to the 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"
+        )
+
+    tool_registry = ToolRegistry()
+    sqlite_runner = SqliteRunner(database_path=database_path)
+    sql_tool = RunSqlTool(sql_runner=sqlite_runner)
+    tool_registry.register(sql_tool)
+
+    llm_service = MockSqliteLlmService()
+
+    return Agent(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        config=AgentConfig(
+            stream_responses=False,
+            include_thinking_indicators=False,
+        ),
+    )
+
+
+async def main() -> None:
+    """Run the mock SQLite example."""
+
+    # Create agent using factory function
+    agent = create_demo_agent()
+    tool_registry = agent.tool_registry
+
+    # Create a test user
+    user = User(id="user123", username="testuser", permissions=[])
+
+    # Test the tool directly
+    print("Testing SQL tool directly:")
+    tool_call = ToolCall(
+        id="test123",
+        name="run_sql",
+        arguments={"sql": "SELECT name FROM sqlite_master WHERE type='table'"},
+    )
+
+    context = ToolContext(user=user, conversation_id="test", request_id="test")
+
+    result = await tool_registry.execute(tool_call, context)
+    print(
+        f"Tables in database:\n{result.result_for_llm if result.success else result.error}"
+    )
+
+    # Show available tools
+    schemas = await tool_registry.get_schemas(user)
+    print(f"\nAvailable tools for user: {[schema.name for schema in schemas]}")
+
+    # Demonstrate the mock LLM triggering SQL queries
+    print("\n" + "=" * 50)
+    print("Agent conversation demo:")
+    print("=" * 50)
+
+    conversation_id = "sqlite-demo"
+
+    # Run multiple queries to show different results
+    for i in range(3):
+        print(f"\n--- Query {i + 1} ---")
+        async for component in agent.send_message(
+            user=user,
+            message=f"Show me some data from the database (query {i + 1})",
+            conversation_id=conversation_id,
+        ):
+            if (
+                hasattr(component.rich_component, "content")
+                and component.rich_component.content
+            ):
+                print(f"Assistant: {component.rich_component.content}")
+            elif component.simple_component and hasattr(
+                component.simple_component, "text"
+            ):
+                print(f"Assistant: {component.simple_component.text}")
+
+
+def run_interactive() -> None:
+    """Entry point for interactive usage."""
+    print("Starting mock SQLite example...")
+    print("This example uses the Chinook database to demonstrate SQL queries.")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()

vanna/examples/openai_quickstart.py

@@ -0,0 +1,83 @@
+"""
+OpenAI example using OpenAILlmService.
+
+Loads environment from .env (via python-dotenv), uses model 'gpt-5' by default,
+and sends a simple message through a Agent.
+
+Run:
+  PYTHONPATH=. python vanna/examples/openai_quickstart.py
+"""
+
+import asyncio
+import importlib.util
+import os
+import sys
+
+
+def ensure_env() -> None:
+    if importlib.util.find_spec("dotenv") is not None:
+        from dotenv import load_dotenv
+
+        # Load from local .env without overriding existing env
+        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"), override=False)
+    else:
+        print(
+            "[warn] python-dotenv not installed; skipping .env load. Install with: pip install python-dotenv"
+        )
+
+    if not os.getenv("OPENAI_API_KEY"):
+        print(
+            "[error] OPENAI_API_KEY is not set. Add it to your environment or .env file."
+        )
+        sys.exit(1)
+
+
+async def main() -> None:
+    ensure_env()
+
+    # Lazy import after env load to allow custom base_url/org via env
+    try:
+        from vanna.integrations.anthropic import OpenAILlmService
+    except ImportError as e:
+        print(
+            "[error] openai extra not installed. Install with: pip install -e .[openai]"
+        )
+        raise
+
+    from vanna import AgentConfig, Agent, User
+    from vanna.core.registry import ToolRegistry
+    from vanna.tools import ListFilesTool
+
+    # Default to 'gpt-5' for this demo; override via $OPENAI_MODEL if desired
+    model = os.getenv("OPENAI_MODEL", "gpt-5")
+    print(f"Using OpenAI model: {model}")
+
+    llm = OpenAILlmService(model=model)
+
+    # Create tool registry and register the list_files tool
+    tool_registry = ToolRegistry()
+    list_files_tool = ListFilesTool()
+    tool_registry.register(list_files_tool)
+
+    # Some models (e.g., reasoning/gpt-5) only support the default temperature=1.0
+    agent = Agent(
+        llm_service=llm,
+        config=AgentConfig(stream_responses=False, temperature=1.0),
+        tool_registry=tool_registry,
+    )
+
+    user = User(id="demo-user", username="demo")
+    conversation_id = "openai-demo"
+
+    print("Sending: 'List the files in the current directory'\n")
+    async for component in agent.send_message(
+        user=user,
+        message="List the files in the current directory",
+        conversation_id=conversation_id,
+    ):
+        if hasattr(component, "content") and component.content:
+            print("Assistant:", component.content)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())