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

vanna/core/component_manager.py

@@ -0,0 +1,329 @@
+"""
+Component state management and update protocol for rich components.
+"""
+
+import uuid
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Union
+
+from pydantic import BaseModel, Field
+
+from ..components.rich import ComponentLifecycle, RichComponent
+
+
+class UpdateOperation(str, Enum):
+    """Types of component update operations."""
+
+    CREATE = "create"
+    UPDATE = "update"
+    REPLACE = "replace"
+    REMOVE = "remove"
+    REORDER = "reorder"
+    BULK_UPDATE = "bulk_update"
+
+
+class Position(BaseModel):
+    """Position specification for component placement."""
+
+    index: Optional[int] = None
+    anchor_id: Optional[str] = None
+    relation: str = "after"  # "before", "after", "inside", "replace"
+
+
+class ComponentUpdate(BaseModel):
+    """Represents a change to the component tree."""
+
+    operation: UpdateOperation
+    target_id: str  # Component being affected
+    component: Optional[RichComponent] = None  # New/updated component data
+    updates: Optional[Dict[str, Any]] = None  # Partial updates for UPDATE operation
+    position: Optional[Position] = None  # For positioning operations
+    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
+    batch_id: Optional[str] = None  # For grouping related updates
+
+    def serialize_for_frontend(self) -> Dict[str, Any]:
+        """Return update payload with nested components normalized."""
+        payload = self.model_dump()
+
+        # Normalise enum values for the frontend contract.
+        payload["operation"] = self.operation.value
+
+        if self.component:
+            payload["component"] = self.component.serialize_for_frontend()
+
+        return payload
+
+
+class ComponentNode(BaseModel):
+    """Node in the component tree."""
+
+    component: RichComponent
+    children: List["ComponentNode"] = Field(default_factory=list)
+    parent_id: Optional[str] = None
+
+    def find_child(self, component_id: str) -> Optional["ComponentNode"]:
+        """Find a child node by component ID."""
+        for child in self.children:
+            if child.component.id == component_id:
+                return child
+            found = child.find_child(component_id)
+            if found:
+                return found
+        return None
+
+    def remove_child(self, component_id: str) -> bool:
+        """Remove a child component by ID."""
+        for i, child in enumerate(self.children):
+            if child.component.id == component_id:
+                self.children.pop(i)
+                return True
+            if child.remove_child(component_id):
+                return True
+        return False
+
+    def get_all_ids(self) -> Set[str]:
+        """Get all component IDs in this subtree."""
+        ids = {self.component.id}
+        for child in self.children:
+            ids.update(child.get_all_ids())
+        return ids
+
+
+class ComponentTree(BaseModel):
+    """Hierarchical structure for managing component layout."""
+
+    root: Optional[ComponentNode] = None
+    flat_index: Dict[str, ComponentNode] = Field(default_factory=dict)
+
+    def add_component(
+        self, component: RichComponent, position: Optional[Position] = None
+    ) -> ComponentUpdate:
+        """Add a component to the tree."""
+        node = ComponentNode(component=component)
+        self.flat_index[component.id] = node
+
+        if self.root is None:
+            self.root = node
+        else:
+            parent_node = self._find_parent(position)
+            if parent_node is not None:
+                node.parent_id = parent_node.component.id
+                parent_node.children.append(node)
+
+        return ComponentUpdate(
+            operation=UpdateOperation.CREATE,
+            target_id=component.id,
+            component=component,
+            position=position,
+        )
+
+    def update_component(
+        self, component_id: str, updates: Dict[str, Any]
+    ) -> Optional[ComponentUpdate]:
+        """Update a component's properties."""
+        node = self.flat_index.get(component_id)
+        if not node:
+            return None
+
+        # Create updated component
+        component_data = node.component.model_dump()
+        component_data.update(updates)
+        component_data["lifecycle"] = ComponentLifecycle.UPDATE
+        component_data["timestamp"] = datetime.utcnow().isoformat()
+
+        updated_component = node.component.__class__(**component_data)
+        node.component = updated_component
+
+        return ComponentUpdate(
+            operation=UpdateOperation.UPDATE,
+            target_id=component_id,
+            component=updated_component,
+            updates=updates,
+        )
+
+    def replace_component(
+        self, old_id: str, new_component: RichComponent
+    ) -> Optional[ComponentUpdate]:
+        """Replace one component with another."""
+        old_node = self.flat_index.get(old_id)
+        if not old_node:
+            return None
+
+        # Update the component in place
+        old_node.component = new_component
+
+        # Update index
+        del self.flat_index[old_id]
+        self.flat_index[new_component.id] = old_node
+
+        return ComponentUpdate(
+            operation=UpdateOperation.REPLACE, target_id=old_id, component=new_component
+        )
+
+    def remove_component(self, component_id: str) -> Optional[ComponentUpdate]:
+        """Remove a component and its children."""
+        node = self.flat_index.get(component_id)
+        if not node:
+            return None
+
+        # Remove from parent
+        if self.root and self.root.component.id == component_id:
+            self.root = None
+        else:
+            if self.root:
+                self.root.remove_child(component_id)
+
+        # Remove from flat index (including all children)
+        removed_ids = node.get_all_ids()
+        for removed_id in removed_ids:
+            self.flat_index.pop(removed_id, None)
+
+        return ComponentUpdate(operation=UpdateOperation.REMOVE, target_id=component_id)
+
+    def get_component(self, component_id: str) -> Optional[RichComponent]:
+        """Get a component by ID."""
+        node = self.flat_index.get(component_id)
+        return node.component if node else None
+
+    def _find_parent(self, position: Optional[Position]) -> Optional[ComponentNode]:
+        """Find the parent node for a new component."""
+        if not position or not position.anchor_id:
+            return self.root
+
+        anchor_node = self.flat_index.get(position.anchor_id)
+        if not anchor_node:
+            return self.root
+
+        if position.relation == "inside":
+            return anchor_node
+        elif position.relation in ["before", "after", "replace"]:
+            # Find the parent of the anchor
+            if anchor_node.parent_id:
+                parent_node = self.flat_index.get(anchor_node.parent_id)
+                return parent_node if parent_node else self.root
+            else:
+                return self.root
+        else:
+            return self.root
+
+
+class ComponentManager:
+    """Manages component lifecycle and state updates."""
+
+    def __init__(self) -> None:
+        self.components: Dict[str, RichComponent] = {}
+        self.component_tree = ComponentTree()
+        self.update_history: List[ComponentUpdate] = []
+        self.active_batch: Optional[str] = None
+
+    def emit(self, component: RichComponent) -> Optional[ComponentUpdate]:
+        """Emit a component with smart lifecycle management."""
+        if component.id in self.components:
+            # Existing component - determine if this is an update or replace
+            existing = self.components[component.id]
+
+            if component.lifecycle == ComponentLifecycle.UPDATE:
+                # Extract changes
+                old_data = existing.model_dump()
+                new_data = component.model_dump()
+                updates = {k: v for k, v in new_data.items() if old_data.get(k) != v}
+
+                update = self.component_tree.update_component(component.id, updates)
+            else:
+                # Replace
+                update = self.component_tree.replace_component(component.id, component)
+        else:
+            # New component - always append
+            update = self.component_tree.add_component(component, None)
+
+        if update:
+            self.components[component.id] = component
+            self.update_history.append(update)
+
+            if self.active_batch:
+                update.batch_id = self.active_batch
+
+        return update
+
+    def update_component(
+        self, component_id: str, **updates: Any
+    ) -> Optional[ComponentUpdate]:
+        """Update specific fields of an existing component."""
+        update = self.component_tree.update_component(component_id, updates)
+        if update and update.component:
+            self.components[component_id] = update.component
+            self.update_history.append(update)
+
+            if self.active_batch:
+                update.batch_id = self.active_batch
+
+        return update
+
+    def replace_component(
+        self, old_id: str, new_component: RichComponent
+    ) -> Optional[ComponentUpdate]:
+        """Replace one component with another."""
+        update = self.component_tree.replace_component(old_id, new_component)
+        if update:
+            self.components.pop(old_id, None)
+            self.components[new_component.id] = new_component
+            self.update_history.append(update)
+
+            if self.active_batch:
+                update.batch_id = self.active_batch
+
+        return update
+
+    def remove_component(self, component_id: str) -> Optional[ComponentUpdate]:
+        """Remove a component and handle cleanup."""
+        update = self.component_tree.remove_component(component_id)
+        if update:
+            self.components.pop(component_id, None)
+            self.update_history.append(update)
+
+            if self.active_batch:
+                update.batch_id = self.active_batch
+
+        return update
+
+    def get_component(self, component_id: str) -> Optional[RichComponent]:
+        """Get a component by ID."""
+        return self.components.get(component_id)
+
+    def get_all_components(self) -> List[RichComponent]:
+        """Get all components in the manager."""
+        return list(self.components.values())
+
+    def start_batch(self) -> str:
+        """Start a batch of related updates."""
+        self.active_batch = str(uuid.uuid4())
+        return self.active_batch
+
+    def end_batch(self) -> Optional[str]:
+        """End the current batch."""
+        batch_id = self.active_batch
+        self.active_batch = None
+        return batch_id
+
+    def get_updates_since(
+        self, timestamp: Optional[str] = None
+    ) -> List[ComponentUpdate]:
+        """Get all updates since a given timestamp."""
+        if not timestamp:
+            return self.update_history.copy()
+
+        try:
+            cutoff = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
+            return [
+                update
+                for update in self.update_history
+                if datetime.fromisoformat(update.timestamp.replace("Z", "+00:00"))
+                > cutoff
+            ]
+        except ValueError:
+            return self.update_history.copy()
+
+    def clear_history(self) -> None:
+        """Clear the update history."""
+        self.update_history.clear()

vanna/core/components.py

@@ -0,0 +1,53 @@
+"""
+UI component base class.
+
+This module defines the UiComponent class which is the return type for tool executions.
+It's placed in core/ because it's a fundamental type that tools return, not just a UI concern.
+"""
+
+from datetime import datetime
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class UiComponent(BaseModel):
+    """Base class for UI components streamed to client.
+
+    This wraps both rich and simple component representations,
+    allowing tools to return structured UI updates.
+
+    Note: We use Any for component types to avoid circular dependencies.
+    Type validation happens at runtime through validators.
+    """
+
+    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
+    rich_component: Any = Field(
+        ..., description="Rich component for advanced rendering"
+    )
+    simple_component: Optional[Any] = Field(
+        None, description="Simple component for basic rendering"
+    )
+
+    @model_validator(mode="after")
+    def validate_components(self) -> "UiComponent":
+        """Validate that components are the correct types at runtime."""
+        # Import from core - clean imports, no circular dependency
+        from .rich_component import RichComponent
+        from .simple_component import SimpleComponent
+
+        if not isinstance(self.rich_component, RichComponent):
+            raise ValueError(
+                f"rich_component must be a RichComponent, got {type(self.rich_component)}"
+            )
+
+        if self.simple_component is not None and not isinstance(
+            self.simple_component, SimpleComponent
+        ):
+            raise ValueError(
+                f"simple_component must be a SimpleComponent or None, got {type(self.simple_component)}"
+            )
+
+        return self
+
+    model_config = {"arbitrary_types_allowed": True}

vanna/core/enhancer/__init__.py

@@ -0,0 +1,11 @@
+"""
+LLM context enhancement system for adding context to prompts and messages.
+
+This module provides interfaces for enriching LLM system prompts and messages
+with additional context before LLM calls (e.g., from memory, RAG, documentation).
+"""
+
+from .base import LlmContextEnhancer
+from .default import DefaultLlmContextEnhancer
+
+__all__ = ["LlmContextEnhancer", "DefaultLlmContextEnhancer"]

vanna/core/enhancer/base.py

@@ -0,0 +1,94 @@
+"""
+LLM context enhancer interface.
+
+LLM context enhancers allow you to add additional context to the system prompt
+and user messages before LLM calls.
+"""
+
+from abc import ABC
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from ..user.models import User
+    from ..llm.models import LlmMessage
+
+
+class LlmContextEnhancer(ABC):
+    """Enhancer for adding context to LLM prompts and messages.
+
+    Subclass this to create custom enhancers that can:
+    - Add relevant context to the system prompt based on the user's initial message
+    - Enrich user messages with additional context (e.g., from memory/RAG)
+    - Inject relevant examples or documentation
+    - Add temporal or environmental context
+
+    Example:
+        class MemoryBasedEnhancer(LlmContextEnhancer):
+            def __init__(self, agent_memory):
+                self.agent_memory = agent_memory
+
+            async def enhance_system_prompt(
+                self,
+                system_prompt: str,
+                user_message: str,
+                user: User
+            ) -> str:
+                # Add relevant examples from memory based on user message
+                examples = await self.agent_memory.search_similar(user_message)
+                return system_prompt + "\\n\\nRelevant examples:\\n" + examples
+
+            async def enhance_user_messages(
+                self,
+                messages: list[LlmMessage],
+                user: User
+            ) -> list[LlmMessage]:
+                # Could modify or add to messages
+                return messages
+
+        agent = Agent(
+            llm_service=...,
+            llm_context_enhancer=MemoryBasedEnhancer(agent_memory)
+        )
+    """
+
+    async def enhance_system_prompt(
+        self, system_prompt: str, user_message: str, user: "User"
+    ) -> str:
+        """Enhance the system prompt with additional context.
+
+        This method is called before the first LLM request with the initial
+        user message, allowing you to add relevant context to the system prompt.
+
+        Args:
+            system_prompt: The original system prompt
+            user_message: The initial user message
+            user: The user making the request
+
+        Returns:
+            Enhanced system prompt with additional context
+
+        Note:
+            This is called once per conversation turn, before any tool calls.
+        """
+        return system_prompt
+
+    async def enhance_user_messages(
+        self, messages: list["LlmMessage"], user: "User"
+    ) -> list["LlmMessage"]:
+        """Enhance user messages with additional context.
+
+        This method is called to potentially modify or add context to user messages
+        before sending them to the LLM.
+
+        Args:
+            messages: The list of messages to enhance
+            user: The user making the request
+
+        Returns:
+            Enhanced list of messages
+
+        Note:
+            This is called before each LLM request, including after tool calls.
+            Be careful not to add context repeatedly on each iteration.
+        """
+        return messages

vanna/core/enhancer/default.py

@@ -0,0 +1,118 @@
+"""
+Default LLM context enhancer implementation using AgentMemory.
+
+This implementation enriches the system prompt with relevant memories
+based on the user's initial message.
+"""
+
+from typing import TYPE_CHECKING, List, Optional
+from .base import LlmContextEnhancer
+
+if TYPE_CHECKING:
+    from ..user.models import User
+    from ..llm.models import LlmMessage
+    from ...capabilities.agent_memory import AgentMemory, TextMemorySearchResult
+
+
+class DefaultLlmContextEnhancer(LlmContextEnhancer):
+    """Default enhancer that uses AgentMemory to add relevant context.
+
+    This enhancer searches the agent's memory for relevant examples and
+    tool use patterns based on the user's message, and adds them to the
+    system prompt.
+
+    Example:
+        agent = Agent(
+            llm_service=...,
+            agent_memory=agent_memory,
+            llm_context_enhancer=DefaultLlmContextEnhancer(agent_memory)
+        )
+    """
+
+    def __init__(self, agent_memory: Optional["AgentMemory"] = None):
+        """Initialize with optional agent memory.
+
+        Args:
+            agent_memory: Optional AgentMemory instance. If not provided,
+                         enhancement will be skipped.
+        """
+        self.agent_memory = agent_memory
+
+    async def enhance_system_prompt(
+        self, system_prompt: str, user_message: str, user: "User"
+    ) -> str:
+        """Enhance system prompt with relevant memories.
+
+        Searches agent memory for relevant text memories based on the
+        user's message and adds them to the system prompt.
+
+        Args:
+            system_prompt: The original system prompt
+            user_message: The initial user message
+            user: The user making the request
+
+        Returns:
+            Enhanced system prompt with relevant examples from memory
+        """
+        if not self.agent_memory:
+            return system_prompt
+
+        try:
+            # Import here to avoid circular dependency
+            from ..tool import ToolContext
+            import uuid
+
+            # Create a temporary context for memory search
+            context = ToolContext(
+                user=user,
+                conversation_id="temp",
+                request_id=str(uuid.uuid4()),
+                agent_memory=self.agent_memory,
+            )
+
+            # Search for relevant text memories based on user message
+            memories: List[
+                "TextMemorySearchResult"
+            ] = await self.agent_memory.search_text_memories(
+                query=user_message, context=context, limit=5
+            )
+
+            if not memories:
+                return system_prompt
+
+            # Format memories as context snippets to add to system prompt
+            examples_section = "\n\n## Relevant Context from Memory\n\n"
+            examples_section += "The following domain knowledge and context from prior interactions may be relevant:\n\n"
+
+            for result in memories:
+                memory = result.memory
+                examples_section += f"• {memory.content}\n"
+
+            # Append examples to system prompt
+            return system_prompt + examples_section
+
+        except Exception as e:
+            # If memory search fails, return original prompt
+            # Don't fail the entire request due to memory issues
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Failed to enhance system prompt with memories: {e}")
+            return system_prompt
+
+    async def enhance_user_messages(
+        self, messages: list["LlmMessage"], user: "User"
+    ) -> list["LlmMessage"]:
+        """Enhance user messages.
+
+        The default implementation doesn't modify user messages.
+        Override this to add context to user messages if needed.
+
+        Args:
+            messages: The list of messages
+            user: The user making the request
+
+        Returns:
+            Original list of messages (unmodified)
+        """
+        return messages

vanna/core/enricher/__init__.py

@@ -0,0 +1,10 @@
+"""
+Context enrichment system for adding data to tool execution context.
+
+This module provides interfaces for enriching ToolContext with additional
+data before tool execution.
+"""
+
+from .base import ToolContextEnricher
+
+__all__ = ["ToolContextEnricher"]

vanna/core/enricher/base.py

@@ -0,0 +1,59 @@
+"""
+Base context enricher interface.
+
+Context enrichers allow you to add additional data to the ToolContext
+before tools are executed.
+"""
+
+from abc import ABC
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..tool.models import ToolContext
+
+
+class ToolContextEnricher(ABC):
+    """Enricher for adding data to ToolContext.
+
+    Subclass this to create custom enrichers that can:
+    - Add user preferences from database
+    - Inject session state
+    - Add temporal context (timezone, current date)
+    - Include user history or profile data
+    - Add environment-specific configuration
+
+    Example:
+        class UserPreferencesEnricher(ToolContextEnricher):
+            def __init__(self, db):
+                self.db = db
+
+            async def enrich_context(self, context: ToolContext) -> ToolContext:
+                # Fetch user preferences
+                prefs = await self.db.get_user_preferences(context.user.id)
+
+                # Add to context metadata
+                context.metadata["preferences"] = prefs
+                context.metadata["timezone"] = prefs.get("timezone", "UTC")
+
+                return context
+
+        agent = AgentRunner(
+            llm_service=...,
+            context_enrichers=[UserPreferencesEnricher(db), SessionEnricher()]
+        )
+    """
+
+    async def enrich_context(self, context: "ToolContext") -> "ToolContext":
+        """Enrich the tool execution context with additional data.
+
+        Args:
+            context: The tool context to enrich
+
+        Returns:
+            Enriched context (typically modified in-place)
+
+        Note:
+            Enrichers typically modify the context.metadata dict to add
+            additional data that tools can access.
+        """
+        return context