Custom Tools

Create custom tools to extend your agent with any functionality you need.

Creating a Custom Tool

Every tool extends the Tool base class and implements four key components:

from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.components import UiComponent, SimpleTextComponent
from pydantic import BaseModel, Field
from typing import Type

# 1. Define argument schema
class MyToolArgs(BaseModel):
    query: str = Field(description="Search query")
    limit: int = Field(default=10, description="Max results")

# 2. Implement the tool
class MyCustomTool(Tool[MyToolArgs]):
    @property
    def name(self) -> str:
        return "search_knowledge_base"
    
    @property
    def description(self) -> str:
        return "Search the company knowledge base for information"
    
    @property
    def access_groups(self) -> List[str]:
        return []  # Empty = available to all users
    
    def get_args_schema(self) -> Type[MyToolArgs]:
        return MyToolArgs
    
    async def execute(self, context: ToolContext, args: MyToolArgs) -> ToolResult:
        # 3. Implement your logic
        results = await self.search_kb(args.query, limit=args.limit)
        
        # 4. Return ToolResult with UI component
        return ToolResult(
            success=True,
            result_for_llm=f"Found {len(results)} results:\n{results}",
            ui_component=UiComponent(
                rich_component=...,  # Rich UI representation
                simple_component=SimpleTextComponent(text=str(results))
            ),
            metadata={"result_count": len(results)}
        )

Tool Components

1. Argument Schema (Pydantic Model)

Define what arguments your tool accepts:

from pydantic import BaseModel, Field

class EmailToolArgs(BaseModel):
    to: str = Field(description="Recipient email address")
    subject: str = Field(description="Email subject")
    body: str = Field(description="Email body")
    cc: List[str] = Field(default=[], description="CC recipients")

Pydantic automatically:

• Validates argument types
• Generates JSON schema for the LLM
• Provides helpful error messages

2. Tool Properties

class MyTool(Tool):
    @property
    def name(self) -> str:
        return "my_tool"  # Must be unique
    
    @property
    def description(self) -> str:
        return "Clear description helps the LLM know when to use this tool"
    
    @property
    def access_groups(self) -> List[str]:
        return ['admin', 'power_users']  # Or [] for all users

3. Execute Method

The execute method receives:

• context: User info, conversation ID, request ID
• args: Validated arguments

async def execute(self, context: ToolContext, args: MyToolArgs) -> ToolResult:
    # Access user information
    user_id = context.user.id
    user_email = context.user.email
    user_groups = context.user.group_memberships
    
    # Access conversation metadata
    conversation_id = context.conversation_id
    request_id = context.request_id
    
    # Your tool logic here
    result = await do_something(args, user_id)
    
    return ToolResult(
        success=True,
        result_for_llm="Human-readable result for the LLM",
        ui_component=create_ui_component(result),
        metadata={"custom": "data"}
    )

4. ToolResult

Return a structured result:

ToolResult(
    success=True,  # or False for errors
    result_for_llm="Text the LLM sees and can reason about",
    ui_component=UiComponent(...),  # Optional: UI to show user
    error="Error message if success=False",  # Optional
    metadata={"key": "value"}  # Optional: structured data
)

Real-World Examples

Example 1: External API Tool

import httpx
from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.components import UiComponent, NotificationComponent, ComponentType
from pydantic import BaseModel

class WeatherArgs(BaseModel):
    city: str
    country_code: str = "US"

class WeatherTool(Tool[WeatherArgs]):
    def __init__(self, api_key: str):
        self.api_key = api_key
    
    @property
    def name(self) -> str:
        return "get_weather"
    
    @property
    def description(self) -> str:
        return "Get current weather information for a city"
    
    def get_args_schema(self) -> Type[WeatherArgs]:
        return WeatherArgs
    
    async def execute(self, context: ToolContext, args: WeatherArgs) -> ToolResult:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"https://api.weather.com/v1/current",
                params={
                    "city": args.city,
                    "country": args.country_code,
                    "apikey": self.api_key
                }
            )
            
            if response.status_code != 200:
                return ToolResult(
                    success=False,
                    result_for_llm=f"Failed to fetch weather data",
                    error=response.text
                )
            
            data = response.json()
            temp = data['temperature']
            conditions = data['conditions']
            
            result_text = f"Weather in {args.city}: {temp}°F, {conditions}"
            
            return ToolResult(
                success=True,
                result_for_llm=result_text,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="info",
                        message=result_text
                    ),
                    simple_component=SimpleTextComponent(text=result_text)
                ),
                metadata={"temperature": temp, "conditions": conditions}
            )

Example 2: Database Tool with User Permissions

class UserAnalyticsTool(Tool[AnalyticsArgs]):
    def __init__(self, db_connection):
        self.db = db_connection
    
    @property
    def access_groups(self) -> List[str]:
        return ['analytics', 'admin']  # Restricted access
    
    async def execute(self, context: ToolContext, args: AnalyticsArgs) -> ToolResult:
        # Check fine-grained permissions
        if 'pii_access' not in context.user.permissions:
            return ToolResult(
                success=False,
                result_for_llm="You don't have permission to access PII data",
                error="Permission denied"
            )
        
        # Execute query with user-specific filters
        results = await self.db.query(
            args.metric,
            user_id=context.user.id  # Row-level security
        )
        
        return ToolResult(success=True, result_for_llm=str(results))

Example 3: Tool with State/Caching

class CachedSearchTool(Tool[SearchArgs]):
    def __init__(self):
        self.cache = {}
    
    async def execute(self, context: ToolContext, args: SearchArgs) -> ToolResult:
        cache_key = f"{args.query}:{context.user.id}"
        
        # Check cache
        if cache_key in self.cache:
            cached_result = self.cache[cache_key]
            return ToolResult(
                success=True,
                result_for_llm=f"(Cached) {cached_result}",
                metadata={"cached": True}
            )
        
        # Perform search
        results = await search_engine.search(args.query)
        
        # Cache result
        self.cache[cache_key] = results
        
        return ToolResult(
            success=True,
            result_for_llm=results,
            metadata={"cached": False}
        )

Registering Tools with Agent

from vanna import Agent

agent = Agent(
    llm_service=llm,
    sql_runner=sql_runner
)

# Register your custom tools
agent.register_tool(WeatherTool(api_key="..."))
agent.register_tool(UserAnalyticsTool(db_connection))
agent.register_tool(CachedSearchTool())

Best Practices

1. Clear naming - Use descriptive, action-oriented names (send_email, not email)
2. Detailed descriptions - Help the LLM understand when to use the tool
3. Validate inputs - Use Pydantic Field with descriptions and constraints
4. Handle errors gracefully - Return ToolResult with success=False
5. User-aware - Use context.user for permissions and personalization
6. Provide UI components - Give users visual feedback
7. Use async - All tool methods are async
8. Type hints - Use proper typing for better IDE support

Dependency Injection

Tools support constructor injection for testability and flexibility:

class EmailTool(Tool):
    def __init__(self, smtp_client, from_address: str):
        self.smtp = smtp_client
        self.from_address = from_address

# Easy to test with mocks
tool = EmailTool(smtp_client=MockSMTP(), from_address="test@example.com")

# Easy to configure for different environments
prod_tool = EmailTool(
    smtp_client=ProductionSMTP(host="smtp.gmail.com"),
    from_address="noreply@company.com"
)

See Also

• Built-in Tools
• Tool Memory
• Custom Charts
• Authentication & Permissions