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

vanna/examples/coding_agent_example.py

@@ -0,0 +1,300 @@
+"""
+Example coding agent using the vanna-agents framework.
+
+This example demonstrates building an agent that can edit code files,
+following the concepts from the "How to Build an Agent" article.
+The agent includes tools for file operations and uses an LLM service
+that can understand and modify code.
+
+Usage:
+  PYTHONPATH=. python vanna/examples/coding_agent_example.py
+"""
+
+import asyncio
+import uuid
+from typing import AsyncGenerator, List, Optional
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    ToolRegistry,
+    User,
+)
+from vanna.core.interfaces import LlmService
+from vanna.core.models import (
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+    ToolCall,
+    ToolSchema,
+)
+from vanna.tools.file_system import create_file_system_tools
+from vanna.tools.python import create_python_tools
+
+
+class CodingLlmService(LlmService):
+    """
+    LLM service that simulates a coding assistant.
+
+    This demonstrates the minimal implementation needed for an agent
+    as described in the article - just needs to understand tool calls
+    and respond appropriately.
+    """
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Handle non-streaming requests."""
+        await asyncio.sleep(0.1)  # Simulate thinking time
+        return self._build_response(request)
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Handle streaming requests."""
+        await asyncio.sleep(0.1)
+        response = self._build_response(request)
+
+        if response.tool_calls:
+            yield LlmStreamChunk(tool_calls=response.tool_calls)
+        if response.content:
+            # Simulate streaming by chunking the response
+            words = response.content.split()
+            for i, word in enumerate(words):
+                chunk = word if i == 0 else f" {word}"
+                await asyncio.sleep(0.05)  # Simulate streaming delay
+                yield LlmStreamChunk(content=chunk)
+
+        yield LlmStreamChunk(finish_reason=response.finish_reason)
+
+    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
+        """Validate tools - no errors for this simple implementation."""
+        return []
+
+    def _build_response(self, request: LlmRequest) -> LlmResponse:
+        """Build a response based on the conversation context."""
+        last_message = request.messages[-1] if request.messages else None
+
+        # If we just got a tool result, respond to it
+        if last_message and last_message.role == "tool":
+            tool_result = last_message.content or "Tool executed"
+            return LlmResponse(
+                content=f"I've completed the operation. {tool_result}",
+                finish_reason="stop",
+            )
+
+        # If user is asking for file operations, use tools
+        if last_message and last_message.role == "user":
+            user_message = last_message.content.lower()
+
+            if "list files" in user_message or "show files" in user_message:
+                return LlmResponse(
+                    content="I'll list the files for you.",
+                    tool_calls=[
+                        ToolCall(
+                            id=f"call_{uuid.uuid4().hex[:8]}",
+                            name="list_files",
+                            arguments={},
+                        )
+                    ],
+                    finish_reason="tool_calls",
+                )
+
+            elif "read" in user_message and (
+                "file" in user_message
+                or ".py" in user_message
+                or ".txt" in user_message
+            ):
+                filename = _extract_filename(user_message)
+
+                if filename:
+                    return LlmResponse(
+                        content=f"I'll read the file '{filename}' for you.",
+                        tool_calls=[
+                            ToolCall(
+                                id=f"call_{uuid.uuid4().hex[:8]}",
+                                name="read_file",
+                                arguments={"filename": filename},
+                            )
+                        ],
+                        finish_reason="tool_calls",
+                    )
+
+            elif "create" in user_message or "write" in user_message:
+                # Suggest creating a simple example file
+                return LlmResponse(
+                    content="I'll create an example Python file for you.",
+                    tool_calls=[
+                        ToolCall(
+                            id=f"call_{uuid.uuid4().hex[:8]}",
+                            name="write_file",
+                            arguments={
+                                "filename": "example.py",
+                                "content": "# Example Python file\nprint('Hello from the coding agent!')\n\ndef greet(name):\n    return f'Hello, {name}!'\n\nif __name__ == '__main__':\n    print(greet('World'))\n",
+                                "overwrite": True,
+                            },
+                        )
+                    ],
+                    finish_reason="tool_calls",
+                )
+
+            elif (
+                "run" in user_message or "execute" in user_message
+            ) and ".py" in user_message:
+                filename = _extract_filename(user_message)
+                if filename:
+                    return LlmResponse(
+                        content=f"I'll run the Python file '{filename}'.",
+                        tool_calls=[
+                            ToolCall(
+                                id=f"call_{uuid.uuid4().hex[:8]}",
+                                name="run_python_file",
+                                arguments={
+                                    "filename": filename,
+                                    "arguments": [],
+                                },
+                            )
+                        ],
+                        finish_reason="tool_calls",
+                    )
+
+            elif (
+                "edit" in user_message
+                or "update" in user_message
+                or "modify" in user_message
+            ):
+                return LlmResponse(
+                    content="I'll update the greet function to make it more descriptive.",
+                    tool_calls=[
+                        ToolCall(
+                            id=f"call_{uuid.uuid4().hex[:8]}",
+                            name="edit_file",
+                            arguments={
+                                "filename": "example.py",
+                                "edits": [
+                                    {
+                                        "start_line": 4,
+                                        "end_line": 5,
+                                        "new_content": (
+                                            "def greet(name):\n"
+                                            '    """Return a friendly greeting."""\n'
+                                            '    return f"Hello, {name}! Welcome to the coding agent."\n'
+                                        ),
+                                    }
+                                ],
+                            },
+                        )
+                    ],
+                    finish_reason="tool_calls",
+                )
+
+        # Default response
+        return LlmResponse(
+            content=(
+                "I'm a coding assistant. I can help you list, read, write, edit, and run Python files. "
+                "Try asking me to 'list files', 'read example.py', 'create a Python file', 'run example.py', or 'update example.py'."
+            ),
+            finish_reason="stop",
+        )
+
+
+def create_demo_agent() -> Agent:
+    """
+    Create a coding agent with file operation tools.
+
+    This follows the pattern from the article - minimal code
+    to create a powerful code-editing agent. Uses dependency injection
+    for file system operations with LocalFileSystem as default.
+    """
+    # Create tool registry and register file system tools
+    tool_registry = ToolRegistry()
+
+    # Use the convenience function to create tools with default LocalFileSystem
+    for tool in create_file_system_tools():
+        tool_registry.register(tool)
+
+    for tool in create_python_tools():
+        tool_registry.register(tool)
+
+    # Create LLM service
+    llm_service = CodingLlmService()
+
+    # Create agent with configuration
+    return Agent(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        config=AgentConfig(
+            stream_responses=True,
+            include_thinking_indicators=True,
+            max_tool_iterations=3,
+        ),
+    )
+
+
+async def main() -> None:
+    """
+    Demonstrate the coding agent in action.
+
+    As the article mentions: "300 lines of code and three tools and now
+    you're able to talk to an alien intelligence that edits your code."
+    """
+    print("🤖 Starting Coding Agent Demo")
+    print("This demonstrates the concepts from 'How to Build an Agent'")
+    print("-" * 50)
+
+    # Create the agent
+    agent = create_demo_agent()
+
+    # Create a test user
+    user = User(id="coder123", username="developer", permissions=[])
+
+    # Show available tools
+    tools = await agent.get_available_tools(user)
+    print(f"Available tools: {[tool.name for tool in tools]}")
+    print()
+
+    # Demo conversation
+    conversation_id = "coding-session"
+
+    demos = [
+        "Hello! Can you list the files in this directory?",
+        "Can you create a simple Python file for me?",
+        "Now read the example.py file you just created",
+        "Please update the greet function to include a docstring and a friendlier message.",
+        "Run example.py so I can see its output.",
+        "Great, read example.py again to confirm the changes.",
+    ]
+
+    for i, message in enumerate(demos, 1):
+        print(f"Demo {i}: {message}")
+        print("Agent response:")
+
+        async for component in agent.send_message(
+            user=user, message=message, conversation_id=conversation_id
+        ):
+            if (
+                hasattr(component.rich_component, "content")
+                and component.rich_component.content
+            ):
+                print(f"  📝 {component.rich_component.content}")
+            elif hasattr(component.rich_component, "message"):
+                print(f"  💬 {component.rich_component.message}")
+            elif component.simple_component and hasattr(
+                component.simple_component, "text"
+            ):
+                print(f"  📄 {component.simple_component.text}")
+
+        print("-" * 30)
+
+
+def _extract_filename(message: str) -> Optional[str]:
+    """Extract a likely filename token from a user message."""
+
+    for token in message.replace("\n", " ").split():
+        cleaned = token.strip("'\".,;!?")
+        if "." in cleaned and not cleaned.startswith("."):
+            return cleaned
+
+    return None
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

vanna/examples/custom_system_prompt_example.py

@@ -0,0 +1,174 @@
+"""
+Example demonstrating custom system prompt builder with dependency injection.
+
+This example shows how to create a custom SystemPromptBuilder that dynamically
+generates system prompts based on user context and available tools.
+
+Usage:
+  python -m vanna.examples.custom_system_prompt_example
+"""
+
+from typing import List, Optional
+
+from vanna.core.interfaces import SystemPromptBuilder
+from vanna.core.models import ToolSchema, User
+
+
+class CustomSystemPromptBuilder(SystemPromptBuilder):
+    """Custom system prompt builder that personalizes prompts based on user."""
+
+    async def build_system_prompt(
+        self, user: User, tools: List[ToolSchema]
+    ) -> Optional[str]:
+        """Build a personalized system prompt.
+
+        Args:
+            user: The user making the request
+            tools: List of tools available to the user
+
+        Returns:
+            Personalized system prompt
+        """
+        # Build personalized greeting
+        username = user.username or user.id
+        greeting = f"Hello {username}! I'm your AI assistant."
+
+        # Add role-specific instructions based on user permissions
+        role_instructions = []
+        if "admin" in user.permissions:
+            role_instructions.append(
+                "As an admin user, you have access to all tools and capabilities."
+            )
+        elif "analyst" in user.permissions:
+            role_instructions.append(
+                "You're working as an analyst. I'll help you query and visualize data."
+            )
+        else:
+            role_instructions.append("I'm here to help you with your tasks.")
+
+        # List available tools
+        tool_info = []
+        if tools:
+            tool_info.append("\nAvailable tools:")
+            for tool in tools:
+                tool_info.append(f"- {tool.name}: {tool.description}")
+
+        # Combine all parts
+        parts = [greeting] + role_instructions + tool_info
+        return "\n".join(parts)
+
+
+class SQLAssistantSystemPromptBuilder(SystemPromptBuilder):
+    """System prompt builder specifically for SQL database assistants."""
+
+    def __init__(self, database_name: str = "database"):
+        """Initialize with database context.
+
+        Args:
+            database_name: Name of the database being queried
+        """
+        self.database_name = database_name
+
+    async def build_system_prompt(
+        self, user: User, tools: List[ToolSchema]
+    ) -> Optional[str]:
+        """Build a SQL-focused system prompt.
+
+        Args:
+            user: The user making the request
+            tools: List of tools available to the user
+
+        Returns:
+            SQL-focused system prompt
+        """
+        prompt = f"""You are an expert SQL database assistant for the {self.database_name} database.
+
+Your primary responsibilities:
+1. Write efficient, correct SQL queries
+2. Explain query results clearly
+3. Suggest optimizations when relevant
+4. Visualize data when appropriate
+
+Guidelines:
+- Always validate SQL syntax before execution
+- Use appropriate JOINs and avoid Cartesian products
+- Limit result sets to reasonable sizes by default
+- Format numbers and dates for readability
+"""
+
+        # Add tool-specific instructions
+        has_viz_tool = any(tool.name == "visualize_data" for tool in tools)
+        if has_viz_tool:
+            prompt += "\n- Create visualizations for numerical data when it helps understanding"
+
+        return prompt
+
+
+async def demo() -> None:
+    """Demonstrate custom system prompt builders."""
+    from vanna import Agent, User
+    from vanna.core.registry import ToolRegistry
+    from vanna.integrations.anthropic.mock import MockLlmService
+
+    # Example 1: Custom personalized system prompt
+    print("=" * 60)
+    print("Example 1: Custom Personalized System Prompt")
+    print("=" * 60)
+
+    custom_builder = CustomSystemPromptBuilder()
+    admin_user = User(id="user-1", username="Alice", permissions=["admin"])
+
+    # Simulate some tools
+    mock_tools = [
+        ToolSchema(
+            name="query_database", description="Query the SQL database", parameters={}
+        ),
+        ToolSchema(
+            name="visualize_data",
+            description="Create data visualizations",
+            parameters={},
+        ),
+    ]
+
+    prompt = await custom_builder.build_system_prompt(admin_user, mock_tools)
+    print("\nGenerated system prompt for admin user:")
+    print("-" * 60)
+    print(prompt)
+    print("-" * 60)
+
+    # Example 2: SQL-specific system prompt
+    print("\n" + "=" * 60)
+    print("Example 2: SQL Assistant System Prompt")
+    print("=" * 60)
+
+    sql_builder = SQLAssistantSystemPromptBuilder(database_name="Chinook")
+    analyst_user = User(id="user-2", username="Bob", permissions=["analyst"])
+
+    prompt = await sql_builder.build_system_prompt(analyst_user, mock_tools)
+    print("\nGenerated system prompt for SQL assistant:")
+    print("-" * 60)
+    print(prompt)
+    print("-" * 60)
+
+    # Example 3: Using custom builder with Agent
+    print("\n" + "=" * 60)
+    print("Example 3: Using Custom Builder with Agent")
+    print("=" * 60)
+
+    mock_llm = MockLlmService()
+    tool_registry = ToolRegistry()
+
+    agent = Agent(
+        llm_service=mock_llm,
+        tool_registry=tool_registry,
+        system_prompt_builder=sql_builder,  # Inject custom builder here
+    )
+
+    print("\nAgent created with custom SQL system prompt builder!")
+    print("The agent will now use the SQL-focused system prompt for all interactions.")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(demo())

vanna/examples/default_workflow_handler_example.py

@@ -0,0 +1,208 @@
+"""
+Example demonstrating the DefaultWorkflowHandler with setup health checking.
+
+This example shows how the DefaultWorkflowHandler provides intelligent starter UI
+that adapts based on available tools and helps users understand their setup status.
+
+Run:
+  PYTHONPATH=. python vanna/examples/default_workflow_handler_example.py
+"""
+
+import asyncio
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    MemoryConversationStore,
+    MockLlmService,
+    User,
+    DefaultWorkflowHandler,
+)
+from vanna.core.registry import ToolRegistry
+from vanna.core.user.resolver import SimpleUserResolver
+from vanna.tools import ListFilesTool
+
+
+async def demonstrate_setup_scenarios():
+    """Demonstrate different setup scenarios with DefaultWorkflowHandler."""
+    print("🚀 Starting DefaultWorkflowHandler Setup Health Check Demo\n")
+
+    # Create basic components
+    llm_service = MockLlmService(response_content="I'm ready to help!")
+    conversation_store = MemoryConversationStore()
+    user_resolver = SimpleUserResolver()
+
+    # Create test user
+    user = User(
+        id="user1",
+        username="alice",
+        email="alice@example.com",
+        group_memberships=["user"],
+    )
+
+    print("=" * 60)
+    print("SCENARIO 1: Empty Setup (No Tools)")
+    print("=" * 60)
+
+    # Empty tool registry
+    empty_registry = ToolRegistry()
+
+    agent_empty = Agent(
+        llm_service=llm_service,
+        tool_registry=empty_registry,
+        user_resolver=user_resolver,
+        conversation_store=conversation_store,
+        config=AgentConfig(stream_responses=False),
+        workflow_handler=DefaultWorkflowHandler(),
+    )
+
+    print("📋 Starter UI for empty setup:")
+    async for component in agent_empty.send_message(
+        request_context=user_resolver.create_request_context(
+            metadata={"starter_ui_request": True}
+        ),
+        message="",
+        conversation_id="empty-setup",
+    ):
+        if hasattr(component, "simple_component") and component.simple_component:
+            print(f"  📄 {component.simple_component.text[:100]}...")
+        elif hasattr(component, "rich_component"):
+            comp = component.rich_component
+            if hasattr(comp, "title"):
+                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
+            elif hasattr(comp, "content"):
+                print(f"  📝 {comp.content[:100]}...")
+
+    print("\n" + "=" * 60)
+    print("SCENARIO 2: Functional Setup (SQL + Basic Tools)")
+    print("=" * 60)
+
+    # Tool registry with SQL tool (simulated)
+    functional_registry = ToolRegistry()
+
+    # Register a mock SQL tool (we'll simulate by tool name)
+    list_tool = ListFilesTool()
+    list_tool.name = "run_sql"  # Simulate SQL tool
+    functional_registry.register(list_tool)
+
+    agent_functional = Agent(
+        llm_service=llm_service,
+        tool_registry=functional_registry,
+        user_resolver=user_resolver,
+        conversation_store=conversation_store,
+        config=AgentConfig(stream_responses=False),
+        workflow_handler=DefaultWorkflowHandler(),
+    )
+
+    print("📋 Starter UI for functional setup:")
+    async for component in agent_functional.send_message(
+        request_context=user_resolver.create_request_context(
+            metadata={"starter_ui_request": True}
+        ),
+        message="",
+        conversation_id="functional-setup",
+    ):
+        if hasattr(component, "simple_component") and component.simple_component:
+            print(f"  📄 {component.simple_component.text[:100]}...")
+        elif hasattr(component, "rich_component"):
+            comp = component.rich_component
+            if hasattr(comp, "title"):
+                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
+            elif hasattr(comp, "content"):
+                print(f"  📝 {comp.content[:100]}...")
+
+    print("\n" + "=" * 60)
+    print("SCENARIO 3: Complete Setup (SQL + Memory + Visualization)")
+    print("=" * 60)
+
+    # Complete tool registry
+    complete_registry = ToolRegistry()
+
+    # Mock SQL tool
+    sql_tool = ListFilesTool()
+    sql_tool.name = "run_sql"
+    complete_registry.register(sql_tool)
+
+    # Mock memory tools
+    search_tool = ListFilesTool()
+    search_tool.name = "search_saved_correct_tool_uses"
+    complete_registry.register(search_tool)
+
+    save_tool = ListFilesTool()
+    save_tool.name = "save_question_tool_args"
+    complete_registry.register(save_tool)
+
+    # Mock visualization tool
+    viz_tool = ListFilesTool()
+    viz_tool.name = "visualize_data"
+    complete_registry.register(viz_tool)
+
+    agent_complete = Agent(
+        llm_service=llm_service,
+        tool_registry=complete_registry,
+        user_resolver=user_resolver,
+        conversation_store=conversation_store,
+        config=AgentConfig(stream_responses=False),
+        workflow_handler=DefaultWorkflowHandler(),
+    )
+
+    print("📋 Starter UI for complete setup:")
+    async for component in agent_complete.send_message(
+        request_context=user_resolver.create_request_context(
+            metadata={"starter_ui_request": True}
+        ),
+        message="",
+        conversation_id="complete-setup",
+    ):
+        if hasattr(component, "simple_component") and component.simple_component:
+            print(f"  📄 {component.simple_component.text[:100]}...")
+        elif hasattr(component, "rich_component"):
+            comp = component.rich_component
+            if hasattr(comp, "title"):
+                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
+            elif hasattr(comp, "content"):
+                print(f"  📝 {comp.content[:100]}...")
+
+    print("\n" + "=" * 60)
+    print("SCENARIO 4: Testing Commands")
+    print("=" * 60)
+
+    print("📋 Testing /help command:")
+    async for component in agent_complete.send_message(
+        request_context=user_resolver.create_request_context(),
+        message="/help",
+        conversation_id="help-test",
+    ):
+        if hasattr(component, "rich_component") and hasattr(
+            component.rich_component, "content"
+        ):
+            print(f"  📝 Help: {component.rich_component.content[:200]}...")
+
+    print("\n📋 Testing /status command:")
+    async for component in agent_complete.send_message(
+        request_context=user_resolver.create_request_context(),
+        message="/status",
+        conversation_id="status-test",
+    ):
+        if hasattr(component, "rich_component"):
+            comp = component.rich_component
+            if hasattr(comp, "title"):
+                print(f"  📊 {comp.title}: {comp.status}")
+            elif hasattr(comp, "content"):
+                print(f"  📝 Status: {comp.content[:150]}...")
+
+    print("\n✅ Demo complete! The DefaultWorkflowHandler provides:")
+    print("  • Smart setup health checking")
+    print("  • Contextual starter UI based on available tools")
+    print("  • Helpful error messages and setup guidance")
+    print("  • Built-in command handling (/help, /status)")
+    print("  • Automatic tool analysis and recommendations")
+
+
+async def main():
+    """Run the DefaultWorkflowHandler demonstration."""
+    await demonstrate_setup_scenarios()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())