Tools as Class

Learn how to create and use class-based tools with AI agents for enhanced functionality.

from praisonaiagents import Agent
from pydantic import BaseModel

class CustomTool(BaseModel):
    def run(self, query: str) -> str:
        return f"Processed: {query}"

agent = Agent(name="CustomAgent", tools=[CustomTool])
agent.start("Run the custom tool on sample data")

The user wraps behaviour in a class-based tool; the agent invokes it like any other capability.

A workflow demonstrating how to create and use class-based tools that can be integrated with AI agents to extend their capabilities with custom functionality.

Quick Start

1. Install Package

   First, install the PraisonAI Agents package:

   pip install praisonaiagents
   
2. Set API Key

   Set your OpenAI API key and EXA API key as environment variables in your terminal:

   export OPENAI_API_KEY=your_api_key_here
   export EXA_API_KEY=your_exa_api_key_here
   
3. Create a file

   Create a new file app.py with the basic setup:

   from praisonaiagents import Agent, Task, AgentTeam
   import os
   import requests
   from typing import Any, Dict, List, Optional
   from pydantic import BaseModel, Field
   
   class EXASearchTool(BaseModel):
       """Wrapper for EXA Search API."""
       search_url: str = "https://api.exa.ai/search"
       headers: Dict = {
           "accept": "application/json",
           "content-type": "application/json",
       }
       max_results: Optional[int] = None
   
       def run(self, query: str) -> str:
           """Run query through EXA and return concatenated results."""
           payload = {
               "query": query,
               "type": "magic",
           }
   
           headers = self.headers.copy()
           headers["x-api-key"] = os.environ['EXA_API_KEY']
   
           response = requests.post(self.search_url, json=payload, headers=headers)
           results = response.json()
           
           if 'results' in results:
               return self._parse_results(results['results'])
           return ""
   
       def results(self, query: str, max_results: Optional[int] = None) -> List[Dict[str, Any]]:
           """Run query through EXA and return metadata."""
           payload = {
               "query": query,
               "type": "magic",
           }
   
           headers = self.headers.copy()
           headers["x-api-key"] = os.environ['EXA_API_KEY']
   
           response = requests.post(self.search_url, json=payload, headers=headers)
           results = response.json()
           
           if 'results' in results:
               return results['results'][:max_results] if max_results else results['results']
           return []
   
       def _parse_results(self, results: List[Dict[str, Any]]) -> str:
           """Parse results into a readable string format."""
           strings = []
           for result in results:
               try:
                   strings.append('\n'.join([
                       f"Title: {result['title']}",
                       f"Score: {result['score']}",
                       f"Url: {result['url']}",
                       f"ID: {result['id']}",
                       "---"
                   ]))
               except KeyError:
                   continue
   
           content = '\n'.join(strings)
           return f"\nSearch results: {content}\n"
   
   # Create an agent with the tool
   agent = Agent(
       name="SearchAgent",
       role="Research Assistant",
       goal="Search for information about 'AI Agents Framework'",
       backstory="I am an AI assistant that can search GitHub.",
       tools=[EXASearchTool],
       reflection=False
   )
   
   # Create task to demonstrate the tool
   task = Task(
       name="search_task",
       description="Search for information about 'AI Agents Framework'",
       expected_output="Information about AI Agents Framework",
       agent=agent
   )
   
   # Create and start the workflow
   agents = AgentTeam(
       agents=[agent],
       tasks=[task]
   )
   
   agents.start()
   
4. Start Agents

   Type this in your terminal to run your agents:

   python app.py
   

Requirements

• Python 3.10 or higher
• OpenAI API key. Generate OpenAI API key here
• EXA API key for search functionality
• Basic understanding of Python and Pydantic

Understanding Tools as Class

What are Class-based Tools?

Class-based tools enable:

• Custom functionality encapsulation
• Reusable tool components
• Type-safe tool interfaces
• Complex API integrations

Features

Pydantic Integration

Built-in validation and type safety with Pydantic models.

API Wrapping

Easily wrap external APIs as agent tools.

Method Flexibility

Support for multiple methods within a single tool.

Type Hints

Strong typing for better code reliability.

Configuration Options

# Create a custom tool class
class CustomTool(BaseModel):
    """Custom tool with configuration options."""
    api_url: str = Field(default="https://api.example.com")
    headers: Dict[str, str] = Field(default_factory=dict)
    max_retries: int = Field(default=3)

    def run(self, input_data: str) -> str:
        """Main execution method."""
        # Tool implementation
        return "Result"

    def configure(self, **kwargs):
        """Update tool configuration."""
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)

# Use the tool with an agent
agent = Agent(
    name="CustomAgent",
    role="Tool User",
    goal="Use custom tool functionality",
    tools=[CustomTool]
)

Troubleshooting

Tool Issues

If tool execution fails:

• Check API credentials
• Verify network connectivity
• Enable verbose logging

Type Errors

If type validation fails:

• Review input types
• Check Pydantic model
• Verify method signatures

Next Steps

Function Tools

Learn about function-based tools

API Tools

Explore API integration tools

For optimal results, ensure your tool classes are well-documented and follow Pydantic best practices for model definition.

Best Practices

Keep tools focused on one task

Each tool should do one thing well - this makes agent behavior more predictable.

Return descriptive strings

Return human-readable strings from tools so the agent can understand and relay results.

Include error messages in returns

On failure, return a descriptive error string rather than raising an exception.

Add docstrings to tools

The agent uses tool docstrings to decide when to call a tool - write clear, specific descriptions.

Related

Custom Tools

Build your own agent tools

Tools Overview

Browse PraisonAI tool documentation