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

vanna/core/user/resolver.py

@@ -0,0 +1,42 @@
+"""
+User resolver interface for web request authentication.
+
+This module provides the abstract base class for resolving web requests
+to authenticated User objects.
+"""
+
+from abc import ABC, abstractmethod
+
+from .models import User
+from .request_context import RequestContext
+
+
+class UserResolver(ABC):
+    """Resolves web requests to authenticated users.
+
+    Implementations of this interface handle the specifics of extracting
+    user identity from request context (cookies, headers, tokens, etc.)
+    and creating authenticated User objects.
+
+    Example:
+        class JwtUserResolver(UserResolver):
+            async def resolve_user(self, request_context: RequestContext) -> User:
+                token = request_context.get_header('Authorization')
+                # ... validate JWT and extract user info
+                return User(id=user_id, username=username, email=email)
+    """
+
+    @abstractmethod
+    async def resolve_user(self, request_context: RequestContext) -> User:
+        """Resolve user from request context.
+
+        Args:
+            request_context: Structured request context with cookies, headers, etc.
+
+        Returns:
+            Authenticated User object
+
+        Raises:
+            Can raise exceptions for authentication failures
+        """
+        pass

vanna/core/validation.py

@@ -0,0 +1,164 @@
+"""
+Development utilities for validating Pydantic models.
+
+This module provides utilities that can be used during development
+and testing to catch forward reference issues early.
+"""
+
+from typing import Any, Dict, List, Tuple, Type
+from pydantic import BaseModel
+import importlib
+import inspect
+
+
+def validate_pydantic_models_in_package(package_name: str) -> Dict[str, Any]:
+    """
+    Validate all Pydantic models in a package for completeness.
+
+    This function can be used in tests or development scripts to catch
+    forward reference issues before they cause runtime errors.
+
+    Args:
+        package_name: Name of the package to validate (e.g., 'vanna.core')
+
+    Returns:
+        Dictionary with validation results
+    """
+    results: Dict[str, Any] = {
+        "total_models": 0,
+        "incomplete_models": [],
+        "models": {},
+        "summary": "",
+    }
+
+    try:
+        # Import the package
+        package = importlib.import_module(package_name)
+
+        # Get all submodules
+        submodules = []
+        if hasattr(package, "__path__"):
+            import pkgutil
+
+            for _, name, _ in pkgutil.iter_modules(
+                package.__path__, package_name + "."
+            ):
+                try:
+                    submodule = importlib.import_module(name)
+                    submodules.append((name, submodule))
+                except ImportError:
+                    continue
+        else:
+            submodules = [(package_name, package)]
+
+        # Check all Pydantic models in each submodule
+        for module_name, module in submodules:
+            for name, obj in inspect.getmembers(module):
+                if (
+                    inspect.isclass(obj)
+                    and issubclass(obj, BaseModel)
+                    and obj is not BaseModel
+                ):
+                    model_key = f"{module_name}.{name}"
+                    results["total_models"] += 1
+
+                    # Check for forward references
+                    forward_refs: List[Tuple[str, str]] = []
+                    for field_name, field_info in obj.model_fields.items():
+                        annotation = field_info.annotation
+                        if annotation is not None and hasattr(
+                            annotation, "__forward_arg__"
+                        ):
+                            forward_refs.append(
+                                (field_name, annotation.__forward_arg__)
+                            )
+
+                    # Check completeness
+                    try:
+                        obj.model_json_schema()
+                        is_complete = True
+                        error = None
+                    except Exception as e:
+                        is_complete = False
+                        error = str(e)
+                        results["incomplete_models"].append(model_key)
+
+                    results["models"][model_key] = {
+                        "class": obj,
+                        "forward_references": forward_refs,
+                        "is_complete": is_complete,
+                        "error": error,
+                    }
+
+        # Generate summary
+        incomplete_models = results["incomplete_models"]
+        incomplete_count = len(incomplete_models)
+        total_models = results["total_models"]
+        if incomplete_count == 0:
+            results["summary"] = (
+                f"✓ All {total_models} Pydantic models are complete and valid!"
+            )
+        else:
+            results["summary"] = (
+                f"⚠ {incomplete_count} of {total_models} models are incomplete: "
+                f"{', '.join(incomplete_models)}"
+            )
+
+    except Exception as e:
+        results["summary"] = f"Error validating package {package_name}: {e}"
+
+    return results
+
+
+def check_models_health() -> bool:
+    """
+    Quick health check for all core Pydantic models.
+
+    Returns:
+        True if all models are healthy, False otherwise
+    """
+    core_packages = [
+        "vanna.core.tool.models",
+        "vanna.core.user.models",
+        "vanna.core.llm.models",
+        "vanna.core.storage.models",
+        "vanna.core.agent.models",
+    ]
+
+    all_healthy = True
+
+    for package in core_packages:
+        try:
+            results = validate_pydantic_models_in_package(package)
+            if results["incomplete_models"]:
+                print(f"❌ Issues in {package}: {results['incomplete_models']}")
+                all_healthy = False
+            else:
+                print(f"✅ {package}: {results['total_models']} models OK")
+        except Exception as e:
+            print(f"❌ Error checking {package}: {e}")
+            all_healthy = False
+
+    return all_healthy
+
+
+if __name__ == "__main__":
+    print("Checking Pydantic model health across core packages...")
+    print("=" * 60)
+
+    healthy = check_models_health()
+
+    print("=" * 60)
+    if healthy:
+        print("🎉 All Pydantic models are healthy!")
+    else:
+        print("⚠️  Some models need attention.")
+        print("\nTo fix forward reference issues:")
+        print("1. Ensure all referenced classes are imported")
+        print("2. Call model_rebuild() after imports")
+        print("3. Use proper TYPE_CHECKING imports for circular deps")
+
+    print("\nNote: You can also catch these issues at development time using:")
+    print("  - mypy static type checking")
+    print("  - This validation script in your test suite")
+    print("  - Pre-commit hooks")

vanna/core/workflow/__init__.py

@@ -0,0 +1,12 @@
+"""
+Workflow handler system for deterministic workflow execution.
+
+This module provides the WorkflowHandler interface for intercepting user messages
+and executing deterministic workflows before they reach the LLM. This is useful
+for command handling, pattern-based routing, and state-based workflows.
+"""
+
+from .base import WorkflowHandler, WorkflowResult
+from .default import DefaultWorkflowHandler
+
+__all__ = ["WorkflowHandler", "WorkflowResult", "DefaultWorkflowHandler"]

vanna/core/workflow/base.py

@@ -0,0 +1,254 @@
+"""
+Base workflow handler interface.
+
+Workflow triggers allow you to execute deterministic workflows in response to
+user messages before they are sent to the LLM. This is useful for:
+- Command handling (e.g., /help, /reset)
+- Pattern-based routing (e.g., report generation)
+- State-based workflows (e.g., onboarding flows)
+- Quota enforcement with custom responses
+"""
+
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    Union,
+    List,
+    AsyncGenerator,
+    Callable,
+    Awaitable,
+)
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from ..user.models import User
+    from ..storage import Conversation
+    from ...components import UiComponent
+    from ..agent.agent import Agent
+
+
+@dataclass
+class WorkflowResult:
+    """Result from a workflow handler attempt.
+
+    When a workflow handles a message, it can optionally return UI components to stream
+    to the user and/or mutate the conversation state.
+
+    Attributes:
+        should_skip_llm: If True, the workflow handled the message and LLM processing is skipped.
+                         If False, the message continues to the agent/LLM.
+        components: Optional UI components to stream back to the user.
+                    Can be a list or async generator for streaming responses.
+        conversation_mutation: Optional async callback to modify conversation state
+                               (e.g., clearing messages, adding system events).
+
+    Example:
+        # Simple command response
+        WorkflowResult(
+            should_skip_llm=True,
+            components=[RichTextComponent(content="Help text here")]
+        )
+
+        # With conversation mutation
+        async def clear_history(conv):
+            conv.messages.clear()
+
+        WorkflowResult(
+            should_skip_llm=True,
+            components=[StatusCardComponent(...)],
+            conversation_mutation=clear_history
+        )
+
+        # Not handled, continue to agent
+        WorkflowResult(should_skip_llm=False)
+    """
+
+    should_skip_llm: bool
+    components: Optional[
+        Union[List["UiComponent"], AsyncGenerator["UiComponent", None]]
+    ] = None
+    conversation_mutation: Optional[Callable[["Conversation"], Awaitable[None]]] = None
+
+
+class WorkflowHandler(ABC):
+    """Base class for handling deterministic workflows before LLM processing.
+
+    Implement this interface to intercept user messages and execute deterministic
+    workflows instead of sending to the LLM. This is the first extensibility point
+    in the agent's message processing pipeline, running after user resolution and
+    conversation loading but before the message is added to conversation history
+    or sent to the LLM.
+
+    Use cases:
+    - Slash commands (/help, /reset, /report)
+    - Pattern-based routing (regex matching)
+    - State-based workflows (onboarding, surveys)
+    - Custom quota enforcement with helpful messages
+    - Deterministic report generation
+    - Starter UI (buttons, welcome messages) when conversation begins
+
+    Example:
+        class CommandWorkflow(WorkflowHandler):
+            async def try_handle(self, agent, user, conversation, message):
+                if message.startswith("/help"):
+                    return WorkflowResult(
+                        should_skip_llm=True,
+                        components=[
+                            RichTextComponent(
+                                content="Available commands:\\n- /help\\n- /reset",
+                                markdown=True
+                            )
+                        ]
+                    )
+
+                # Execute tool for reports
+                if message.startswith("/report"):
+                    tool = await agent.tool_registry.get_tool("generate_report")
+                    result = await tool.execute(ToolContext(user=user), {})
+                    return WorkflowResult(should_skip_llm=True, components=[result.ui_component])
+
+                # Not handled, continue to agent
+                return WorkflowResult(should_skip_llm=False)
+
+            async def get_starter_ui(self, agent, user, conversation):
+                return [
+                    RichTextComponent(content=f"Welcome {user.username}!"),
+                    ButtonComponent(label="Generate Report", value="/report"),
+                ]
+
+        agent = Agent(
+            llm_service=...,
+            tool_registry=...,
+            user_resolver=...,
+            workflow_handler=CommandWorkflow()
+        )
+
+    Observability:
+        The agent automatically creates an "agent.workflow_handler" span when
+        a WorkflowHandler is configured, allowing you to monitor handler
+        performance and outcomes.
+    """
+
+    @abstractmethod
+    async def try_handle(
+        self, agent: "Agent", user: "User", conversation: "Conversation", message: str
+    ) -> WorkflowResult:
+        """Attempt to handle a workflow for the given message.
+
+        This method is called for every user message before it reaches the LLM.
+        Inspect the message content, user context, and conversation state to
+        decide whether to execute a deterministic workflow or allow normal
+        agent processing.
+
+        Args:
+            agent: The agent instance, providing access to tool_registry, config,
+                   and observability_provider for tool execution and logging.
+            user: The user who sent the message, including their ID, permissions,
+                  and metadata. Use this for permission checks or personalization.
+            conversation: The current conversation context, including message history.
+                          Can be inspected for state-based workflows.
+            message: The user's raw message content.
+
+        Returns:
+            WorkflowResult with should_skip_llm=True to execute a workflow and skip LLM,
+            or should_skip_llm=False to continue normal agent processing.
+
+            When should_skip_llm=True:
+            - The message is NOT added to conversation history automatically
+            - The components are streamed to the user
+            - The conversation_mutation callback (if provided) is executed
+            - The agent returns without calling the LLM
+
+            When should_skip_llm=False:
+            - The message is added to conversation history
+            - Normal agent processing continues (LLM call, tool execution, etc.)
+
+        Example:
+            async def try_handle(self, agent, user, conversation, message):
+                # Pattern matching with tool execution
+                if message.startswith("/report"):
+                    # Execute tool from registry
+                    tool = await agent.tool_registry.get_tool("generate_sales_report")
+                    context = ToolContext(user=user, conversation=conversation)
+                    result = await tool.execute(context, {})
+
+                    return WorkflowResult(
+                        should_skip_llm=True,
+                        components=[...]
+                    )
+
+                # State-based workflow
+                if user.metadata.get("needs_onboarding"):
+                    return await self._onboarding_flow(agent, user, message)
+
+                # Permission check
+                if message.startswith("/admin") and "admin" not in user.permissions:
+                    return WorkflowResult(
+                        should_skip_llm=True,
+                        components=[RichTextComponent(content="Access denied.")]
+                    )
+
+                # Continue to agent
+                return WorkflowResult(should_skip_llm=False)
+        """
+        pass
+
+    async def get_starter_ui(
+        self, agent: "Agent", user: "User", conversation: "Conversation"
+    ) -> Optional[List["UiComponent"]]:
+        """Provide UI components when a conversation starts.
+
+        Override this method to show starter buttons, welcome messages,
+        or quick actions when a new chat is opened by the user.
+
+        This is called by the frontend/server when initializing a new
+        conversation, before any user messages are sent.
+
+        Args:
+            agent: The agent instance, providing access to tool_registry, config,
+                   and observability_provider for dynamic UI generation.
+            user: The user starting the conversation
+            conversation: The new conversation (typically empty)
+
+        Returns:
+            List of UI components to display, or None for no starter UI.
+            Components can include buttons, welcome text, quick actions, etc.
+
+        Example:
+            async def get_starter_ui(self, agent, user, conversation):
+                # Show role-based quick actions
+                if "analyst" in user.permissions:
+                    # Dynamically generate buttons based on available tools
+                    report_tools = [
+                        tool for tool in agent.tool_registry.list_tools()
+                        if tool.startswith("report_")
+                    ]
+
+                    buttons = [
+                        ButtonComponent(label=f"📊 {tool}", value=f"/{tool}")
+                        for tool in report_tools
+                    ]
+
+                    return [
+                        RichTextComponent(
+                            content=f"Welcome back, {user.username}!",
+                            markdown=True
+                        ),
+                        *buttons
+                    ]
+
+                # New user onboarding
+                if user.metadata.get("is_new_user"):
+                    return [
+                        RichTextComponent(
+                            content="# Welcome to Vanna!\\n\\nTry one of these to get started:",
+                            markdown=True
+                        ),
+                        ButtonComponent(label="Show Example Query", value="/example"),
+                        ButtonComponent(label="View Tutorial", value="/tutorial"),
+                    ]
+
+                return None
+        """
+        return None