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

vanna/examples/extensibility_example.py

@@ -0,0 +1,262 @@
+"""
+Comprehensive example demonstrating all extensibility interfaces.
+
+This example shows how to use:
+- LlmMiddleware for caching
+- ErrorRecoveryStrategy for retry logic
+- ToolContextEnricher for adding user preferences
+- ConversationFilter for context window management
+- ObservabilityProvider for monitoring
+"""
+
+import asyncio
+import time
+from typing import Any, Dict, List, Optional
+
+from vanna.core import (
+    Agent,
+    LlmMiddleware,
+    ErrorRecoveryStrategy,
+    ToolContextEnricher,
+    ConversationFilter,
+    ObservabilityProvider,
+    User,
+    ToolContext,
+    Conversation,
+    Message,
+    LlmRequest,
+    LlmResponse,
+    Span,
+    Metric,
+)
+from vanna.core.recovery import RecoveryAction, RecoveryActionType
+from vanna.core.registry import ToolRegistry
+
+
+# 1. LlmMiddleware Example: Simple Caching
+class CachingMiddleware(LlmMiddleware):
+    """Cache LLM responses to reduce costs and latency."""
+
+    def __init__(self) -> None:
+        self.cache: Dict[str, LlmResponse] = {}
+        self.hits = 0
+        self.misses = 0
+
+    def _compute_cache_key(self, request: LlmRequest) -> str:
+        """Create cache key from request."""
+        messages_str = str([(m.role, m.content) for m in request.messages])
+        return f"{messages_str}:{request.temperature}"
+
+    async def before_llm_request(self, request: LlmRequest) -> LlmRequest:
+        """Check cache before sending request."""
+        cache_key = self._compute_cache_key(request)
+        if cache_key in self.cache:
+            self.hits += 1
+            print(f"[CACHE HIT] Cache stats: {self.hits} hits, {self.misses} misses")
+        return request
+
+    async def after_llm_response(
+        self, request: LlmRequest, response: LlmResponse
+    ) -> LlmResponse:
+        """Cache the response."""
+        cache_key = self._compute_cache_key(request)
+        if cache_key not in self.cache:
+            self.cache[cache_key] = response
+            self.misses += 1
+            print(f"[CACHE MISS] Caching response")
+        return response
+
+
+# 2. ErrorRecoveryStrategy Example: Exponential Backoff
+class ExponentialBackoffStrategy(ErrorRecoveryStrategy):
+    """Retry failed operations with exponential backoff."""
+
+    def __init__(self, max_retries: int = 3) -> None:
+        self.max_retries = max_retries
+
+    async def handle_tool_error(
+        self, error: Exception, context: ToolContext, attempt: int = 1
+    ) -> RecoveryAction:
+        """Retry tool errors with exponential backoff."""
+        if attempt < self.max_retries:
+            delay_ms = (2 ** (attempt - 1)) * 1000
+            print(
+                f"[RETRY] Tool failed, retrying in {delay_ms}ms (attempt {attempt}/{self.max_retries})"
+            )
+            return RecoveryAction(
+                action=RecoveryActionType.RETRY,
+                retry_delay_ms=delay_ms,
+                message=f"Retrying after {delay_ms}ms",
+            )
+
+        print(f"[FAIL] Max retries exceeded for tool error: {error}")
+        return RecoveryAction(
+            action=RecoveryActionType.FAIL,
+            message=f"Tool error after {self.max_retries} attempts: {str(error)}",
+        )
+
+    async def handle_llm_error(
+        self, error: Exception, request: LlmRequest, attempt: int = 1
+    ) -> RecoveryAction:
+        """Retry LLM errors with backoff."""
+        if attempt < self.max_retries:
+            delay_ms = (2 ** (attempt - 1)) * 1000
+            print(
+                f"[RETRY] LLM failed, retrying in {delay_ms}ms (attempt {attempt}/{self.max_retries})"
+            )
+            return RecoveryAction(
+                action=RecoveryActionType.RETRY,
+                retry_delay_ms=delay_ms,
+                message=f"Retrying LLM after {delay_ms}ms",
+            )
+
+        print(f"[FAIL] Max retries exceeded for LLM error: {error}")
+        return RecoveryAction(
+            action=RecoveryActionType.FAIL,
+            message=f"LLM error after {self.max_retries} attempts: {str(error)}",
+        )
+
+
+# 3. ToolContextEnricher Example: Add User Preferences
+class UserPreferencesEnricher(ToolContextEnricher):
+    """Enrich context with user preferences."""
+
+    def __init__(self) -> None:
+        # Mock user preferences database
+        self.preferences: Dict[str, Dict[str, Any]] = {
+            "user123": {
+                "timezone": "America/New_York",
+                "language": "en",
+                "theme": "dark",
+            }
+        }
+
+    async def enrich_context(self, context: ToolContext) -> ToolContext:
+        """Add user preferences to context."""
+        prefs = self.preferences.get(context.user.id, {})
+        context.metadata["user_preferences"] = prefs
+        context.metadata["timezone"] = prefs.get("timezone", "UTC")
+        print(f"[ENRICH] Added preferences for user {context.user.id}: {prefs}")
+        return context
+
+
+# 4. ConversationFilter Example: Context Window Management
+class ContextWindowFilter(ConversationFilter):
+    """Limit conversation to fit within context window."""
+
+    def __init__(self, max_messages: int = 20) -> None:
+        self.max_messages = max_messages
+
+    async def filter_messages(self, messages: List[Message]) -> List[Message]:
+        """Keep only recent messages within limit."""
+        if len(messages) <= self.max_messages:
+            return messages
+
+        # Keep system messages and recent messages
+        system_messages = [m for m in messages if m.role == "system"]
+        other_messages = [m for m in messages if m.role != "system"]
+
+        # Take the most recent messages
+        recent_messages = other_messages[-self.max_messages :]
+        filtered = system_messages + recent_messages
+
+        print(f"[FILTER] Reduced {len(messages)} messages to {len(filtered)}")
+        return filtered
+
+
+# 5. ObservabilityProvider Example: Simple Logging
+class LoggingObservabilityProvider(ObservabilityProvider):
+    """Log metrics and spans for monitoring."""
+
+    def __init__(self) -> None:
+        self.metrics: List[Metric] = []
+        self.spans: List[Span] = []
+
+    async def record_metric(
+        self,
+        name: str,
+        value: float,
+        unit: str = "",
+        tags: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Record and log a metric."""
+        metric = Metric(name=name, value=value, unit=unit, tags=tags or {})
+        self.metrics.append(metric)
+        tags_str = ", ".join(f"{k}={v}" for k, v in (tags or {}).items())
+        print(f"[METRIC] {name}: {value}{unit} {tags_str}")
+
+    async def create_span(
+        self, name: str, attributes: Optional[Dict[str, Any]] = None
+    ) -> Span:
+        """Create a span for tracing."""
+        span = Span(name=name, attributes=attributes or {})
+        print(f"[SPAN START] {name}")
+        return span
+
+    async def end_span(self, span: Span) -> None:
+        """End and record a span."""
+        span.end()
+        self.spans.append(span)
+        duration = span.duration_ms() or 0
+        print(f"[SPAN END] {span.name}: {duration:.2f}ms")
+
+
+async def run_example() -> None:
+    """
+    Example showing all extensibility interfaces working together.
+    """
+    from vanna.integrations.anthropic import AnthropicLlmService
+
+    # Create all extensibility components
+    caching_middleware = CachingMiddleware()
+    retry_strategy = ExponentialBackoffStrategy(max_retries=3)
+    preferences_enricher = UserPreferencesEnricher()
+    context_filter = ContextWindowFilter(max_messages=20)
+    observability = LoggingObservabilityProvider()
+
+    # Mock conversation store
+    class MockStore:
+        async def get_conversation(self, cid: str, uid: str) -> Optional[Conversation]:
+            return None
+
+        async def create_conversation(
+            self, cid: str, uid: str, title: str
+        ) -> Conversation:
+            return Conversation(
+                id=cid, user_id=uid, messages=[Message(role="user", content=title)]
+            )
+
+        async def update_conversation(self, conv: Conversation) -> None:
+            pass
+
+        async def delete_conversation(self, cid: str, uid: str) -> bool:
+            return False
+
+        async def list_conversations(
+            self, uid: str, limit: int = 50, offset: int = 0
+        ) -> List[Conversation]:
+            return []
+
+    # Create agent with all extensibility components
+    agent = Agent(
+        llm_service=AnthropicLlmService(api_key="test-key"),
+        tool_registry=ToolRegistry(),
+        conversation_store=MockStore(),  # type: ignore
+        llm_middlewares=[caching_middleware],
+        error_recovery_strategy=retry_strategy,
+        context_enrichers=[preferences_enricher],
+        conversation_filters=[context_filter],
+        observability_provider=observability,
+    )
+
+    print("✓ Agent created with all extensibility components:")
+    print(f"  - LLM Middleware: {len(agent.llm_middlewares)} middlewares")
+    print(f"  - Error Recovery: {type(agent.error_recovery_strategy).__name__}")
+    print(f"  - Context Enrichers: {len(agent.context_enrichers)} enrichers")
+    print(f"  - Conversation Filters: {len(agent.conversation_filters)} filters")
+    print(f"  - Observability: {type(agent.observability_provider).__name__}")
+    print("\n🎉 All extensibility interfaces integrated successfully!")
+
+
+if __name__ == "__main__":
+    asyncio.run(run_example())

vanna/examples/minimal_example.py

@@ -0,0 +1,67 @@
+"""Minimal Claude + SQLite example ready for FastAPI."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from vanna import AgentConfig, Agent
+from vanna.core.registry import ToolRegistry
+from vanna.integrations.anthropic import AnthropicLlmService
+from vanna.integrations.sqlite import SqliteRunner
+from vanna.integrations.local import LocalFileSystem
+from vanna.tools import (
+    RunSqlTool,
+    # Visualization
+    VisualizeDataTool,
+    # Python execution
+    RunPythonFileTool,
+    PipInstallTool,
+    # File system (for coding agents)
+    SearchFilesTool,
+    ListFilesTool,
+    ReadFileTool,
+    WriteFileTool,
+)
+
+_DB = Path(__file__).resolve().parents[2] / "Chinook.sqlite"
+
+
+def create_demo_agent() -> Agent:
+    # Load environment variables from .env file
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
+    llm = AnthropicLlmService(model=os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5"))
+
+    # Shared file system for all tools
+    file_system = LocalFileSystem("./claude_data")
+
+    tools = ToolRegistry()
+
+    # 1. Basic SQL agent - query databases
+    tools.register(
+        RunSqlTool(
+            sql_runner=SqliteRunner(database_path=str(_DB)),
+            file_system=file_system,
+        )
+    )
+
+    # 2. Add visualization - create charts from data
+    tools.register(VisualizeDataTool(file_system=file_system))
+
+    # 3. Add Python execution - build dashboards with artifacts
+    # tools.register(RunPythonFileTool(file_system=file_system))
+    # tools.register(PipInstallTool(file_system=file_system))
+
+    # 4. Full coding agent - read, write, search files
+    # tools.register(SearchFilesTool(file_system=file_system))
+    # tools.register(ListFilesTool(file_system=file_system))
+    # tools.register(ReadFileTool(file_system=file_system))
+    # tools.register(WriteFileTool(file_system=file_system))
+
+    return Agent(
+        llm_service=llm,
+        tool_registry=tools,
+    )

vanna/examples/mock_auth_example.py

@@ -0,0 +1,227 @@
+"""
+Mock authentication example to verify user resolution is working.
+
+This example demonstrates the new UserResolver architecture where:
+1. UserResolver is a required parameter of Agent
+2. Agent.send_message() accepts RequestContext (not User directly)
+3. The Agent resolves the user internally using the UserResolver
+
+The agent uses an LLM middleware to inject user info into the response,
+so we can verify the authentication is working correctly.
+
+Usage:
+    python -m vanna.examples.mock_auth_example
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+from vanna import AgentConfig, Agent
+from vanna.core.registry import ToolRegistry
+from vanna.core.llm import LlmRequest, LlmResponse
+from vanna.core.middleware import LlmMiddleware
+from vanna.integrations.mock import MockLlmService
+from vanna.core.user import CookieEmailUserResolver, RequestContext
+
+
+class UserEchoMiddleware(LlmMiddleware):
+    """Middleware that injects user email into LLM responses."""
+
+    async def after_llm_response(
+        self, request: LlmRequest, response: LlmResponse
+    ) -> LlmResponse:
+        """Inject user email into response."""
+        # Extract user email from request user_id (which is set to user.id in the agent)
+        user_id = request.user_id
+
+        # Create a new response with user info
+        new_content = f"Hello! You are authenticated as: {user_id}"
+
+        return LlmResponse(
+            content=new_content,
+            finish_reason=response.finish_reason,
+            usage=response.usage,
+        )
+
+
+def create_demo_agent() -> Agent:
+    """Create a demo agent for server usage.
+
+    Returns:
+        Configured Agent instance with cookie-based authentication
+    """
+    # Create a mock LLM
+    llm_service = MockLlmService(response_content="Mock response")
+
+    # Empty tool registry
+    tool_registry = ToolRegistry()
+
+    # Cookie-based user resolver
+    user_resolver = CookieEmailUserResolver(cookie_name="vanna_email")
+
+    # User echo middleware
+    middleware = UserEchoMiddleware()
+
+    # Create agent with user resolver and middleware
+    agent = Agent(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        user_resolver=user_resolver,
+        llm_middlewares=[middleware],
+        config=AgentConfig(
+            stream_responses=True,
+            include_thinking_indicators=False,
+        ),
+    )
+
+    return agent
+
+
+async def demo_authentication():
+    """Demonstrate authentication with different request contexts."""
+    agent = create_demo_agent()
+
+    print("=== Mock Authentication Demo ===")
+    print("This example verifies that user resolution is working correctly.\n")
+
+    # Test 1: Request with email cookie
+    print("🔹 Test 1: Authenticated user (alice@example.com)")
+    request_context = RequestContext(
+        cookies={"vanna_email": "alice@example.com"},
+        headers={},
+        remote_addr="127.0.0.1",
+    )
+
+    print(
+        "Request context:",
+        {
+            "cookies": request_context.cookies,
+            "headers": request_context.headers,
+            "remote_addr": request_context.remote_addr,
+        },
+    )
+
+    # Send message - Agent will resolve user internally
+    agent_response = ""
+    async for component in agent.send_message(
+        request_context=request_context,
+        message="Who am I?",
+        conversation_id="test_conv_1",
+    ):
+        # Extract and display user info from the resolved user
+        if hasattr(component, "rich_component"):
+            rich = component.rich_component
+            # Check if it's a text component
+            if rich.type.value == "text":
+                # Access content directly from the component (before serialization)
+                if hasattr(rich, "content"):
+                    agent_response = rich.content
+
+    print(f"Agent response: {agent_response}")
+
+    # Verify user was resolved by checking the conversation store
+    user_resolver = agent.user_resolver
+    resolved_user = await user_resolver.resolve_user(request_context)
+    print(
+        f"✅ Resolved user: {resolved_user.email} (username: {resolved_user.username}, id: {resolved_user.id})"
+    )
+    print(f"   Permissions: {resolved_user.permissions}")
+    print(f"   Metadata: {resolved_user.metadata}")
+
+    print("\n" + "=" * 60 + "\n")
+
+    # Test 2: Request without email cookie (anonymous)
+    print("🔹 Test 2: Anonymous user (no cookie)")
+    anonymous_context = RequestContext(cookies={}, headers={}, remote_addr="127.0.0.1")
+
+    print(
+        "Request context:",
+        {
+            "cookies": anonymous_context.cookies,
+            "headers": anonymous_context.headers,
+            "remote_addr": anonymous_context.remote_addr,
+        },
+    )
+
+    agent_response = ""
+    async for component in agent.send_message(
+        request_context=anonymous_context,
+        message="Who am I?",
+        conversation_id="test_conv_2",
+    ):
+        if hasattr(component, "rich_component"):
+            rich = component.rich_component
+            if rich.type.value == "text" and hasattr(rich, "content"):
+                agent_response = rich.content
+
+    print(f"Agent response: {agent_response}")
+
+    resolved_user = await user_resolver.resolve_user(anonymous_context)
+    print(
+        f"✅ Resolved user: {resolved_user.email or 'None'} (username: {resolved_user.username}, id: {resolved_user.id})"
+    )
+    print(f"   Permissions: {resolved_user.permissions}")
+    print(f"   Metadata: {resolved_user.metadata}")
+
+    print("\n" + "=" * 60 + "\n")
+
+    # Test 3: Different user
+    print("🔹 Test 3: Different authenticated user (bob@company.com)")
+    bob_context = RequestContext(
+        cookies={"vanna_email": "bob@company.com"},
+        headers={"User-Agent": "Mozilla/5.0"},
+        remote_addr="192.168.1.100",
+    )
+
+    print(
+        "Request context:",
+        {
+            "cookies": bob_context.cookies,
+            "headers": bob_context.headers,
+            "remote_addr": bob_context.remote_addr,
+        },
+    )
+
+    agent_response = ""
+    async for component in agent.send_message(
+        request_context=bob_context, message="Who am I?", conversation_id="test_conv_3"
+    ):
+        if hasattr(component, "rich_component"):
+            rich = component.rich_component
+            if rich.type.value == "text" and hasattr(rich, "content"):
+                agent_response = rich.content
+
+    print(f"Agent response: {agent_response}")
+
+    resolved_user = await user_resolver.resolve_user(bob_context)
+    print(
+        f"✅ Resolved user: {resolved_user.email} (username: {resolved_user.username}, id: {resolved_user.id})"
+    )
+    print(f"   Permissions: {resolved_user.permissions}")
+    print(f"   Metadata: {resolved_user.metadata}")
+
+    print("\n" + "=" * 60)
+    print("\n✅ Authentication demo complete!")
+    print("\nKey Features Verified:")
+    print("• UserResolver is part of Agent")
+    print("• Agent.send_message() accepts RequestContext")
+    print("• User resolution happens internally in Agent")
+    print("• CookieEmailUserResolver extracts email from vanna_email cookie")
+    print("• Anonymous users are created when no cookie is present")
+    print("• Different users can be resolved from different request contexts")
+
+
+async def main():
+    """Run the authentication example."""
+    await demo_authentication()
+
+
+def run_interactive():
+    """Entry point for interactive usage."""
+    print("Starting mock authentication example...")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()