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

vanna/core/simple_component.py

@@ -0,0 +1,27 @@
+"""Base classes for simple UI components."""
+
+from typing import Any, Dict, Optional
+from pydantic import BaseModel, Field
+from enum import Enum
+
+
+class SimpleComponentType(str, Enum):
+    TEXT = "text"
+    IMAGE = "image"
+    LINK = "link"
+
+
+class SimpleComponent(BaseModel):
+    """A simple UI component with basic attributes."""
+
+    type: SimpleComponentType = Field(..., description="Type of the component.")
+    semantic_type: Optional[str] = Field(
+        default=None, description="Semantic type for better categorization."
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata for the component."
+    )
+
+    def serialize_for_frontend(self) -> Dict[str, Any]:
+        """Serialize simple component for API consumption."""
+        return self.model_dump()

vanna/core/storage/__init__.py

@@ -0,0 +1,14 @@
+"""
+Storage domain.
+
+This module provides the core abstractions for conversation storage in the Vanna Agents framework.
+"""
+
+from .base import ConversationStore
+from .models import Conversation, Message
+
+__all__ = [
+    "ConversationStore",
+    "Conversation",
+    "Message",
+]

vanna/core/storage/base.py

@@ -0,0 +1,46 @@
+"""
+Storage domain interface.
+
+This module contains the abstract base class for conversation storage.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from .models import Conversation
+from ..user.models import User
+
+
+class ConversationStore(ABC):
+    """Abstract base class for conversation storage."""
+
+    @abstractmethod
+    async def create_conversation(
+        self, conversation_id: str, user: User, initial_message: str
+    ) -> Conversation:
+        """Create a new conversation with the specified ID."""
+        pass
+
+    @abstractmethod
+    async def get_conversation(
+        self, conversation_id: str, user: User
+    ) -> Optional[Conversation]:
+        """Get conversation by ID, scoped to user."""
+        pass
+
+    @abstractmethod
+    async def update_conversation(self, conversation: Conversation) -> None:
+        """Update conversation with new messages."""
+        pass
+
+    @abstractmethod
+    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
+        """Delete conversation."""
+        pass
+
+    @abstractmethod
+    async def list_conversations(
+        self, user: User, limit: int = 50, offset: int = 0
+    ) -> List[Conversation]:
+        """List conversations for user."""
+        pass

vanna/core/storage/models.py

@@ -0,0 +1,46 @@
+"""
+Storage domain models.
+
+This module contains data models for conversation storage.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from ..tool.models import ToolCall
+from ..user.models import User
+
+
+class Message(BaseModel):
+    """Single message in a conversation."""
+
+    role: str = Field(description="Message role (user/assistant/system/tool)")
+    content: str = Field(description="Message content")
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    tool_calls: Optional[List[ToolCall]] = Field(default=None)
+    tool_call_id: Optional[str] = Field(
+        default=None, description="ID if this is a tool response"
+    )
+
+
+class Conversation(BaseModel):
+    """Conversation containing multiple messages."""
+
+    id: str = Field(description="Unique conversation identifier")
+    user: User = Field(description="User this conversation belongs to")
+    messages: List[Message] = Field(
+        default_factory=list, description="Messages in conversation"
+    )
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    updated_at: datetime = Field(default_factory=datetime.utcnow)
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional conversation metadata"
+    )
+
+    def add_message(self, message: Message) -> None:
+        """Add a message to the conversation."""
+        self.messages.append(message)
+        self.updated_at = datetime.utcnow()

vanna/core/system_prompt/__init__.py

@@ -0,0 +1,13 @@
+"""
+System prompt domain.
+
+This module provides the core abstractions for building system prompts in the Vanna Agents framework.
+"""
+
+from .base import SystemPromptBuilder
+from .default import DefaultSystemPromptBuilder
+
+__all__ = [
+    "SystemPromptBuilder",
+    "DefaultSystemPromptBuilder",
+]

vanna/core/system_prompt/base.py

@@ -0,0 +1,36 @@
+"""
+System prompt builder interface.
+
+This module contains the abstract base class for system prompt builders.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from ..tool.models import ToolSchema
+    from ..user.models import User
+
+
+class SystemPromptBuilder(ABC):
+    """Abstract base class for system prompt builders.
+
+    Subclasses should implement the build_system_prompt method to generate
+    system prompts based on user context and available tools.
+    """
+
+    @abstractmethod
+    async def build_system_prompt(
+        self, user: "User", tools: List["ToolSchema"]
+    ) -> Optional[str]:
+        """
+        Build a system prompt based on user context and available tools.
+
+        Args:
+            user: The user making the request
+            tools: List of tools available to the user
+
+        Returns:
+            System prompt string, or None if no system prompt should be used
+        """
+        pass

vanna/core/system_prompt/default.py

@@ -0,0 +1,157 @@
+"""
+Default system prompt builder implementation with memory workflow support.
+
+This module provides a default implementation of the SystemPromptBuilder interface
+that automatically includes memory workflow instructions when memory tools are available.
+"""
+
+from typing import TYPE_CHECKING, List, Optional
+from datetime import datetime
+
+from .base import SystemPromptBuilder
+
+if TYPE_CHECKING:
+    from ..tool.models import ToolSchema
+    from ..user.models import User
+
+
+class DefaultSystemPromptBuilder(SystemPromptBuilder):
+    """Default system prompt builder with automatic memory workflow integration.
+
+    Dynamically generates system prompts that include memory workflow
+    instructions when memory tools (search_saved_correct_tool_uses and
+    save_question_tool_args) are available.
+    """
+
+    def __init__(self, base_prompt: Optional[str] = None):
+        """Initialize with an optional base prompt.
+
+        Args:
+            base_prompt: Optional base system prompt. If not provided, uses a default.
+        """
+        self.base_prompt = base_prompt
+
+    async def build_system_prompt(
+        self, user: "User", tools: List["ToolSchema"]
+    ) -> Optional[str]:
+        """
+        Build a system prompt with memory workflow instructions.
+
+        Args:
+            user: The user making the request
+            tools: List of tools available to the user
+
+        Returns:
+            System prompt string with memory workflow instructions if applicable
+        """
+        if self.base_prompt is not None:
+            return self.base_prompt
+
+        # Check which memory tools are available
+        tool_names = [tool.name for tool in tools]
+        has_search = "search_saved_correct_tool_uses" in tool_names
+        has_save = "save_question_tool_args" in tool_names
+        has_text_memory = "save_text_memory" in tool_names
+
+        # Get today's date
+        today_date = datetime.now().strftime("%Y-%m-%d")
+
+        # Base system prompt
+        prompt_parts = [
+            f"You are Vanna, an AI data analyst assistant created to help users with data analysis tasks. Today's date is {today_date}.",
+            "",
+            "Response Guidelines:",
+            "- Any summary of what you did or observations should be the final step.",
+            "- Use the available tools to help the user accomplish their goals.",
+            "- When you execute a query, that raw result is shown to the user outside of your response so YOU DO NOT need to include it in your response. Focus on summarizing and interpreting the results.",
+        ]
+
+        if tools:
+            prompt_parts.append(
+                f"\nYou have access to the following tools: {', '.join(tool_names)}"
+            )
+
+        # Add memory workflow instructions based on available tools
+        if has_search or has_save or has_text_memory:
+            prompt_parts.append("\n" + "=" * 60)
+            prompt_parts.append("MEMORY SYSTEM:")
+            prompt_parts.append("=" * 60)
+
+        if has_search or has_save:
+            prompt_parts.append("\n1. TOOL USAGE MEMORY (Structured Workflow):")
+            prompt_parts.append("-" * 50)
+
+        if has_search:
+            prompt_parts.extend(
+                [
+                    "",
+                    "• BEFORE executing any tool (run_sql, visualize_data, or calculator), you MUST first call search_saved_correct_tool_uses with the user's question to check if there are existing successful patterns for similar questions.",
+                    "",
+                    "• Review the search results (if any) to inform your approach before proceeding with other tool calls.",
+                ]
+            )
+
+        if has_save:
+            prompt_parts.extend(
+                [
+                    "",
+                    "• AFTER successfully executing a tool that produces correct and useful results, you MUST call save_question_tool_args to save the successful pattern for future use.",
+                ]
+            )
+
+        if has_search or has_save:
+            prompt_parts.extend(
+                [
+                    "",
+                    "Example workflow:",
+                    "  • User asks a question",
+                    f'  • First: Call search_saved_correct_tool_uses(question="user\'s question")'
+                    if has_search
+                    else "",
+                    "  • Then: Execute the appropriate tool(s) based on search results and the question",
+                    f'  • Finally: If successful, call save_question_tool_args(question="user\'s question", tool_name="tool_used", args={{the args you used}})'
+                    if has_save
+                    else "",
+                    "",
+                    "Do NOT skip the search step, even if you think you know how to answer. Do NOT forget to save successful executions."
+                    if has_search
+                    else "",
+                    "",
+                    "The only exceptions to searching first are:",
+                    '  • When the user is explicitly asking about the tools themselves (like "list the tools")',
+                    "  • When the user is testing or asking you to demonstrate the save/search functionality itself",
+                ]
+            )
+
+        if has_text_memory:
+            prompt_parts.extend(
+                [
+                    "",
+                    "2. TEXT MEMORY (Domain Knowledge & Context):",
+                    "-" * 50,
+                    "",
+                    "• save_text_memory: Save important context about the database, schema, or domain",
+                    "",
+                    "Use text memory to save:",
+                    "  • Database schema details (column meanings, data types, relationships)",
+                    "  • Company-specific terminology and definitions",
+                    "  • Query patterns or best practices for this database",
+                    "  • Domain knowledge about the business or data",
+                    "  • User preferences for queries or visualizations",
+                    "",
+                    "DO NOT save:",
+                    "  • Information already captured in tool usage memory",
+                    "  • One-time query results or temporary observations",
+                    "",
+                    "Examples:",
+                    '  • save_text_memory(content="The status column uses 1 for active, 0 for inactive")',
+                    '  • save_text_memory(content="MRR means Monthly Recurring Revenue in our schema")',
+                    "  • save_text_memory(content=\"Always exclude test accounts where email contains 'test'\")",
+                ]
+            )
+
+        if has_search or has_save or has_text_memory:
+            # Remove empty strings from the list
+            prompt_parts = [part for part in prompt_parts if part != ""]
+
+        return "\n".join(prompt_parts)

vanna/core/tool/__init__.py

@@ -0,0 +1,18 @@
+"""
+Tool domain.
+
+This module provides the core abstractions for tools in the Vanna Agents framework.
+"""
+
+from .base import T, Tool
+from .models import ToolCall, ToolContext, ToolRejection, ToolResult, ToolSchema
+
+__all__ = [
+    "Tool",
+    "T",
+    "ToolCall",
+    "ToolContext",
+    "ToolRejection",
+    "ToolResult",
+    "ToolSchema",
+]

vanna/core/tool/base.py

@@ -0,0 +1,70 @@
+"""
+Tool domain interface.
+
+This module contains the abstract base class for tools.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Generic, List, Type, TypeVar
+
+from .models import ToolContext, ToolResult, ToolSchema
+
+# Type variable for tool argument types
+T = TypeVar("T")
+
+
+class Tool(ABC, Generic[T]):
+    """Abstract base class for tools."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique name for this tool."""
+        pass
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Description of what this tool does."""
+        pass
+
+    @property
+    def access_groups(self) -> List[str]:
+        """Groups permitted to access this tool."""
+        return []
+
+    @abstractmethod
+    def get_args_schema(self) -> Type[T]:
+        """Return the Pydantic model for arguments."""
+        pass
+
+    @abstractmethod
+    async def execute(self, context: ToolContext, args: T) -> ToolResult:
+        """Execute the tool with validated arguments.
+
+        Args:
+            context: Execution context containing user, conversation_id, and request_id
+            args: Validated tool arguments
+
+        Returns:
+            ToolResult with success status, result for LLM, and optional UI component
+        """
+        pass
+
+    def get_schema(self) -> ToolSchema:
+        """Generate tool schema for LLM."""
+        from typing import Any, cast
+
+        args_model = self.get_args_schema()
+        # Get the schema - args_model should be a Pydantic model class
+        schema = (
+            cast(Any, args_model).model_json_schema()
+            if hasattr(args_model, "model_json_schema")
+            else {}
+        )
+        return ToolSchema(
+            name=self.name,
+            description=self.description,
+            parameters=schema,
+            access_groups=self.access_groups,
+        )

vanna/core/tool/models.py

@@ -0,0 +1,84 @@
+"""
+Tool domain models.
+
+This module contains data models for tool execution.
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+# Import AgentMemory at runtime for Pydantic model resolution
+from vanna.capabilities.agent_memory import AgentMemory
+
+if TYPE_CHECKING:
+    from ..components import UiComponent
+    from ..user.models import User
+    from ..observability import ObservabilityProvider
+
+
+class ToolCall(BaseModel):
+    """Represents a tool call from the LLM."""
+
+    id: str = Field(description="Unique identifier for this tool call")
+    name: str = Field(description="Name of the tool to execute")
+    arguments: Dict[str, Any] = Field(description="Raw arguments from LLM")
+
+
+class ToolContext(BaseModel):
+    """Context passed to all tool executions."""
+
+    user: "User"  # Forward reference to avoid circular import
+    conversation_id: str
+    request_id: str = Field(description="Unique request identifier for tracing")
+    agent_memory: AgentMemory = Field(
+        description="Agent memory for tool usage learning"
+    )
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    observability_provider: Optional["ObservabilityProvider"] = Field(
+        default=None,
+        description="Optional observability provider for metrics and spans",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class ToolResult(BaseModel):
+    """Result from tool execution.
+
+    Changes:
+    - `result_for_llm`: string that will be sent back to the LLM.
+    - `ui_component`: optional UI payload for rendering in clients.
+    """
+
+    success: bool = Field(description="Whether execution succeeded")
+    result_for_llm: str = Field(description="String content to send back to the LLM")
+    ui_component: Optional["UiComponent"] = Field(
+        default=None, description="Optional UI component for rendering"
+    )
+    error: Optional[str] = Field(default=None, description="Error message if failed")
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolSchema(BaseModel):
+    """Schema describing a tool for LLM consumption."""
+
+    name: str = Field(description="Tool name")
+    description: str = Field(description="What this tool does")
+    parameters: Dict[str, Any] = Field(description="JSON Schema of parameters")
+    access_groups: List[str] = Field(
+        default_factory=list, description="Groups permitted to access this tool"
+    )
+
+
+class ToolRejection(BaseModel):
+    """Indicates tool execution should be rejected with a message.
+
+    Used by transform_args to reject tool execution when arguments
+    cannot be appropriately transformed for the user's context.
+    """
+
+    reason: str = Field(
+        description="Explanation of why the tool execution was rejected"
+    )

vanna/core/user/__init__.py

@@ -0,0 +1,17 @@
+"""
+User domain.
+
+This module provides the core abstractions for user management in the Vanna Agents framework.
+"""
+
+from .base import UserService
+from .models import User
+from .resolver import UserResolver
+from .request_context import RequestContext
+
+__all__ = [
+    "UserService",
+    "User",
+    "UserResolver",
+    "RequestContext",
+]

vanna/core/user/base.py

@@ -0,0 +1,29 @@
+"""
+User domain interface.
+
+This module contains the abstract base class for user services.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+from .models import User
+
+
+class UserService(ABC):
+    """Service for user management and authentication."""
+
+    @abstractmethod
+    async def get_user(self, user_id: str) -> Optional[User]:
+        """Get user by ID."""
+        pass
+
+    @abstractmethod
+    async def authenticate(self, credentials: Dict[str, Any]) -> Optional[User]:
+        """Authenticate user and return User object if successful."""
+        pass
+
+    @abstractmethod
+    async def has_permission(self, user: User, permission: str) -> bool:
+        """Check if user has specific permission."""
+        pass

vanna/core/user/models.py

@@ -0,0 +1,25 @@
+"""
+User domain models.
+
+This module contains data models for user management.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class User(BaseModel):
+    """User model for authentication and scoping."""
+
+    id: str = Field(description="Unique user identifier")
+    username: Optional[str] = Field(default=None, description="Username")
+    email: Optional[str] = Field(default=None, description="User email")
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional user metadata"
+    )
+    group_memberships: List[str] = Field(
+        default_factory=list, description="Groups the user belongs to"
+    )
+
+    model_config = ConfigDict(extra="allow")

vanna/core/user/request_context.py

@@ -0,0 +1,70 @@
+"""
+Request context for user resolution.
+
+This module provides the RequestContext model for passing web request
+information to UserResolver implementations.
+"""
+
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field
+
+
+class RequestContext(BaseModel):
+    """Context from a web request for user resolution.
+
+    This structured object replaces raw dictionaries for passing request
+    data to UserResolver implementations, making it easier to access
+    cookies, headers, and other request metadata.
+
+    Example:
+        context = RequestContext(
+            cookies={'vanna_email': 'alice@example.com'},
+            headers={'Authorization': 'Bearer token'},
+            remote_addr='127.0.0.1'
+        )
+        user = await resolver.resolve_user(context)
+    """
+
+    cookies: Dict[str, str] = Field(default_factory=dict, description="Request cookies")
+
+    headers: Dict[str, str] = Field(default_factory=dict, description="Request headers")
+
+    remote_addr: Optional[str] = Field(default=None, description="Remote IP address")
+
+    query_params: Dict[str, str] = Field(
+        default_factory=dict, description="Query parameters"
+    )
+
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional framework-specific metadata"
+    )
+
+    def get_cookie(self, name: str, default: Optional[str] = None) -> Optional[str]:
+        """Get cookie value by name.
+
+        Args:
+            name: Cookie name
+            default: Default value if cookie not found
+
+        Returns:
+            Cookie value or default
+        """
+        return self.cookies.get(name, default)
+
+    def get_header(self, name: str, default: Optional[str] = None) -> Optional[str]:
+        """Get header value by name (case-insensitive).
+
+        Args:
+            name: Header name
+            default: Default value if header not found
+
+        Returns:
+            Header value or default
+        """
+        # Case-insensitive header lookup
+        name_lower = name.lower()
+        for key, value in self.headers.items():
+            if key.lower() == name_lower:
+                return value
+        return default