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

vanna/examples/mock_custom_tool.py

@@ -0,0 +1,311 @@
+"""
+Mock example showing how to create and use custom tools.
+
+This example demonstrates creating a simple calculator tool
+and registering it with an agent that uses a mock LLM service.
+It now includes a `MockCalculatorLlmService` that automatically
+invokes the calculator tool with random numbers before echoing
+back the computed answer.
+
+Usage:
+  Template: Copy this file and modify for your custom tools
+  Interactive: python -m vanna.examples.mock_custom_tool
+  REPL: from vanna.examples.mock_custom_tool import create_demo_agent
+  Server: python -m vanna.servers --example mock_custom_tool
+"""
+
+import asyncio
+import random
+import uuid
+from typing import AsyncGenerator, Dict, List, Optional, Tuple, Type
+
+from pydantic import BaseModel, Field
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    Tool,
+    ToolContext,
+    ToolRegistry,
+    ToolResult,
+    User,
+    UiComponent,
+)
+from vanna.core.interfaces import LlmService
+from vanna.core.models import (
+    LlmRequest,
+    LlmResponse,
+    LlmStreamChunk,
+    ToolCall,
+    ToolSchema,
+)
+from vanna.core.rich_components import (
+    CardComponent,
+    NotificationComponent,
+    ComponentType,
+)
+from vanna.core.simple_components import (
+    SimpleTextComponent,
+)
+
+
+class CalculatorArgs(BaseModel):
+    """Arguments for the calculator tool."""
+
+    operation: str = Field(
+        description="The operation to perform: add, subtract, multiply, divide"
+    )
+    a: float = Field(description="First number")
+    b: float = Field(description="Second number")
+
+
+class CalculatorTool(Tool[CalculatorArgs]):
+    """A simple calculator tool."""
+
+    @property
+    def name(self) -> str:
+        return "calculator"
+
+    @property
+    def description(self) -> str:
+        return "Perform basic arithmetic operations (add, subtract, multiply, divide)"
+
+    def get_args_schema(self) -> Type[CalculatorArgs]:
+        return CalculatorArgs
+
+    async def execute(self, context: ToolContext, args: CalculatorArgs) -> ToolResult:
+        """Execute the calculator operation."""
+        symbol_map = {"add": "+", "subtract": "-", "multiply": "×", "divide": "÷"}
+
+        try:
+            if args.operation == "add":
+                result = args.a + args.b
+            elif args.operation == "subtract":
+                result = args.a - args.b
+            elif args.operation == "multiply":
+                result = args.a * args.b
+            elif args.operation == "divide":
+                if args.b == 0:
+                    message = "Cannot divide by zero"
+                    await asyncio.sleep(3)
+                    return ToolResult(
+                        success=False,
+                        result_for_llm=message,
+                        ui_component=UiComponent(
+                            rich_component=NotificationComponent(
+                                type=ComponentType.NOTIFICATION,
+                                level="error",
+                                message=message,
+                            ),
+                            simple_component=SimpleTextComponent(text=message),
+                        ),
+                        error=message,
+                    )
+                result = args.a / args.b
+            else:
+                message = f"Unknown operation: {args.operation}"
+                await asyncio.sleep(3)
+                return ToolResult(
+                    success=False,
+                    result_for_llm=message,
+                    ui_component=UiComponent(
+                        rich_component=NotificationComponent(
+                            type=ComponentType.NOTIFICATION,
+                            level="warning",
+                            message=message,
+                        ),
+                        simple_component=SimpleTextComponent(text=message),
+                    ),
+                    error=message,
+                )
+
+            await asyncio.sleep(3)
+
+            symbol = symbol_map.get(args.operation, args.operation)
+            expression = f"{args.a:g} {symbol} {args.b:g} = {result:g}"
+            return ToolResult(
+                success=True,
+                result_for_llm=str(result),
+                ui_component=UiComponent(
+                    rich_component=CardComponent(
+                        type=ComponentType.CARD,
+                        title="Calculator Result",
+                        content=expression,
+                    ),
+                    simple_component=SimpleTextComponent(text=expression),
+                ),
+                error=None,
+            )
+
+        except Exception as e:
+            message = str(e)
+            await asyncio.sleep(3)
+            return ToolResult(
+                success=False,
+                result_for_llm=message,
+                ui_component=UiComponent(
+                    rich_component=NotificationComponent(
+                        type=ComponentType.NOTIFICATION,
+                        level="error",
+                        message=message,
+                    ),
+                    simple_component=SimpleTextComponent(text=message),
+                ),
+                error=message,
+            )
+
+
+class MockCalculatorLlmService(LlmService):
+    """LLM service that exercises the calculator tool before echoing the result."""
+
+    def __init__(self, seed: Optional[int] = None):
+        self._random = random.Random(seed)
+
+    async def send_request(self, request: LlmRequest) -> LlmResponse:
+        """Handle non-streaming calculator interactions."""
+        await asyncio.sleep(0.05)
+        return self._build_response(request)
+
+    async def stream_request(
+        self, request: LlmRequest
+    ) -> AsyncGenerator[LlmStreamChunk, None]:
+        """Provide streaming compatibility by yielding a single chunk."""
+        await asyncio.sleep(0.05)
+        response = self._build_response(request)
+
+        if response.tool_calls:
+            yield LlmStreamChunk(tool_calls=response.tool_calls)
+        if response.content is not None:
+            yield LlmStreamChunk(
+                content=response.content, finish_reason=response.finish_reason
+            )
+        else:
+            yield LlmStreamChunk(finish_reason=response.finish_reason)
+
+    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
+        """Mock validation - no errors."""
+        return []
+
+    def _build_response(self, request: LlmRequest) -> LlmResponse:
+        """Create a response that either calls the tool or echoes its result."""
+        last_message = request.messages[-1] if request.messages else None
+
+        if last_message and last_message.role == "tool":
+            answer = last_message.content or "No result provided"
+            return LlmResponse(
+                content=answer,
+                finish_reason="stop",
+                usage={
+                    "prompt_tokens": 30,
+                    "completion_tokens": 10,
+                    "total_tokens": 40,
+                },
+            )
+
+        operation, a, b = self._random_operands()
+        tool_call = ToolCall(
+            id=f"call_{uuid.uuid4().hex[:8]}",
+            name="calculator",
+            arguments={"operation": operation, "a": a, "b": b},
+        )
+
+        return LlmResponse(
+            content="Let me ask my calculator friend for help...",
+            tool_calls=[tool_call],
+            finish_reason="tool_calls",
+            usage={"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35},
+        )
+
+    def _random_operands(self) -> Tuple[str, float, float]:
+        """Generate operation and operands suited for the calculator tool."""
+        operation = self._random.choice(["add", "subtract", "multiply", "divide"])
+
+        if operation == "divide":
+            b = float(self._random.randint(1, 10))
+            multiplier = self._random.randint(1, 10)
+            a = float(b * multiplier)
+        elif operation == "subtract":
+            b = float(self._random.randint(1, 10))
+            a = b + float(self._random.randint(0, 10))
+        else:
+            a = float(self._random.randint(1, 12))
+            b = float(self._random.randint(1, 12))
+
+        return operation, a, b
+
+
+def create_demo_agent() -> Agent:
+    """Create a demo agent with custom calculator tool.
+
+    Returns:
+        Configured Agent with calculator tool and mock calculator LLM
+    """
+    tool_registry = ToolRegistry()
+    calculator_tool = CalculatorTool()
+    tool_registry.register(calculator_tool)
+
+    llm_service = MockCalculatorLlmService()
+
+    return Agent(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        config=AgentConfig(
+            stream_responses=False,
+            include_thinking_indicators=False,
+        ),
+    )
+
+
+async def main() -> None:
+    """Run the mock custom tool example."""
+
+    # Create agent using factory function
+    agent = create_demo_agent()
+    tool_registry = agent.tool_registry
+
+    # Create a test user
+    user = User(id="user123", username="testuser", permissions=[])
+
+    # Test the tool directly
+    print("Testing calculator tool directly:")
+    tool_call = ToolCall(
+        id="test123", name="calculator", arguments={"operation": "add", "a": 5, "b": 3}
+    )
+
+    context = ToolContext(user=user, conversation_id="test", request_id="test")
+
+    result = await tool_registry.execute(tool_call, context)
+    print(f"5 + 3 = {result.result_for_llm if result.success else result.error}")
+
+    # Show available tools
+    schemas = await tool_registry.get_schemas(user)
+    print(f"\nAvailable tools for user: {[schema.name for schema in schemas]}")
+
+    # Demonstrate the mock LLM triggering a tool call
+    print("\nAgent conversation demo:")
+    conversation_id = "calc-demo"
+    async for component in agent.send_message(
+        user=user,
+        message="Can you compute something for me?",
+        conversation_id=conversation_id,
+    ):
+        print(f"- Component type: {component.rich_component.type}")
+        if (
+            hasattr(component.rich_component, "content")
+            and component.rich_component.content
+        ):
+            print(f"Assistant: {component.rich_component.content}")
+        elif component.simple_component and hasattr(component.simple_component, "text"):
+            print(f"Assistant: {component.simple_component.text}")
+        else:
+            print(f"- Component data: {component.rich_component.data}")
+
+
+def run_interactive() -> None:
+    """Entry point for interactive usage."""
+    print("Starting mock custom tool example...")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()

vanna/examples/mock_quickstart.py

@@ -0,0 +1,79 @@
+"""
+Mock quickstart example for the Vanna Agents framework.
+
+This example shows how to create a basic agent with a mock LLM service
+and have a simple conversation.
+
+Usage:
+  Template: Copy this file and modify for your needs
+  Interactive: python -m vanna.examples.mock_quickstart
+  REPL: from vanna.examples.mock_quickstart import create_demo_agent
+  Server: python -m vanna.servers --example mock_quickstart
+"""
+
+import asyncio
+
+from vanna import (
+    AgentConfig,
+    Agent,
+    MemoryConversationStore,
+    MockLlmService,
+    User,
+)
+
+
+def create_demo_agent() -> Agent:
+    """Create a demo agent for REPL and server usage.
+
+    Returns:
+        Configured Agent instance
+    """
+    llm_service = MockLlmService(
+        response_content="Hello! I'm a helpful AI assistant created using the Vanna Agents framework."
+    )
+
+    return Agent(
+        llm_service=llm_service,
+        config=AgentConfig(
+            stream_responses=True,  # Enable streaming for better server experience
+            include_thinking_indicators=True,
+        ),
+    )
+
+
+async def main() -> None:
+    """Run the mock quickstart example."""
+
+    # Create agent using factory function
+    agent = create_demo_agent()
+
+    # Create a test user
+    user = User(
+        id="user123", username="testuser", email="test@example.com", permissions=[]
+    )
+
+    # Start a conversation
+    conversation_id = "conversation123"
+    user_message = "Hello! Can you introduce yourself?"
+
+    print(f"User: {user_message}")
+    print("Agent: ", end="")
+
+    # Send message and collect response
+    async for component in agent.send_message(
+        user=user, message=user_message, conversation_id=conversation_id
+    ):
+        if hasattr(component, "content"):
+            print(component.content, end="")
+
+    print()
+
+
+def run_interactive() -> None:
+    """Entry point for interactive usage."""
+    print("Starting Vanna Agents mock quickstart demo...")
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    run_interactive()

vanna/examples/mock_quota_example.py

@@ -0,0 +1,145 @@
+"""
+Mock quota-based agent example using Mock LLM service.
+
+This example demonstrates how to create a custom agent runner that
+enforces user-based message quotas. It shows:
+- Custom agent runner subclass
+- Quota management and enforcement
+- Error handling for quota exceeded cases
+- Multiple users with different quotas
+
+Run:
+  PYTHONPATH=. python vanna/examples/mock_quota_example.py
+"""
+
+import asyncio
+
+from vanna import (
+    AgentConfig,
+    MemoryConversationStore,
+    MockLlmService,
+    User,
+)
+from vanna.core.registry import ToolRegistry
+from vanna.tools import ListFilesTool
+from vanna.examples.quota_agent import QuotaAgentRunner, QuotaExceededError
+
+
+async def demonstrate_quota_system() -> None:
+    """Demonstrate the quota-based agent system."""
+    print("🚀 Starting Mock Quota-based Agent Example\n")
+
+    # Create a mock LLM service
+    llm_service = MockLlmService(
+        response_content="Hello! I'm here to help you with your questions."
+    )
+
+    # Create tool registry with list_files tool
+    tool_registry = ToolRegistry()
+    list_files_tool = ListFilesTool()
+    tool_registry.register(list_files_tool)
+
+    # Create conversation store
+    conversation_store = MemoryConversationStore()
+
+    # Create the quota-based agent
+    agent = QuotaAgentRunner(
+        llm_service=llm_service,
+        tool_registry=tool_registry,
+        conversation_store=conversation_store,
+        config=AgentConfig(
+            stream_responses=False,
+            include_thinking_indicators=False,
+        ),
+    )
+
+    # Create users with different quota settings
+    regular_user = User(
+        id="user1", username="alice", email="alice@example.com", permissions=[]
+    )
+
+    premium_user = User(
+        id="user2", username="bob", email="bob@example.com", permissions=["premium"]
+    )
+
+    # Set custom quotas
+    agent.set_user_quota(regular_user.id, 3)  # Alice gets 3 messages
+    agent.set_user_quota(premium_user.id, 5)  # Bob gets 5 messages (premium)
+
+    print("📋 User Quotas:")
+    print(
+        f"  • {regular_user.username}: {agent.get_user_quota(regular_user.id)} messages"
+    )
+    print(
+        f"  • {premium_user.username}: {agent.get_user_quota(premium_user.id)} messages"
+    )
+    print()
+
+    # Test regular user within quota
+    print("💬 Testing regular user (Alice) within quota:")
+    for i in range(1, 4):  # Send 3 messages (within quota)
+        print(f"  Message {i}/3:")
+        async for component in agent.send_message(
+            user=regular_user,
+            message=f"Hello, this is message {i}",
+            conversation_id="alice-conv",
+        ):
+            if hasattr(component, "content") and component.content:
+                print(f"    Agent: {component.content}")
+        print()
+
+    # Test regular user exceeding quota
+    print("⚠️  Testing regular user (Alice) exceeding quota:")
+    async for component in agent.send_message(
+        user=regular_user,
+        message="This message should be blocked",
+        conversation_id="alice-conv",
+    ):
+        if hasattr(component, "content") and component.content:
+            print(f"    Agent: {component.content}")
+    print()
+
+    # Test premium user with higher quota
+    print("⭐ Testing premium user (Bob) with higher quota:")
+    for i in range(1, 4):  # Send 3 messages
+        print(f"  Message {i}/5:")
+        async for component in agent.send_message(
+            user=premium_user,
+            message=f"Premium user message {i}",
+            conversation_id="bob-conv",
+        ):
+            if hasattr(component, "content") and component.content:
+                print(f"    Agent: {component.content}")
+    print()
+
+    # Demonstrate quota reset
+    print("🔄 Resetting Alice's usage:")
+    agent.reset_user_usage(regular_user.id)
+    print(f"  Alice's remaining messages: {agent.get_user_remaining(regular_user.id)}")
+    print()
+
+    print("✅ After reset, Alice can send messages again:")
+    async for component in agent.send_message(
+        user=regular_user,
+        message="This should work after reset",
+        conversation_id="alice-conv2",
+    ):
+        if hasattr(component, "content") and component.content:
+            print(f"    Agent: {component.content}")
+
+    print("\n📊 Final Usage Summary:")
+    print(
+        f"  • Alice: {agent.get_user_usage(regular_user.id)}/{agent.get_user_quota(regular_user.id)} used"
+    )
+    print(
+        f"  • Bob: {agent.get_user_usage(premium_user.id)}/{agent.get_user_quota(premium_user.id)} used"
+    )
+
+
+async def main() -> None:
+    """Run the mock quota example."""
+    await demonstrate_quota_system()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())