Tools Library

The tools library in Airtrain provides a comprehensive set of utilities that enable AI agents to interact with their environment. These tools are built on a flexible registry system that supports both stateless and stateful operation, allowing for a wide range of capabilities while maintaining clean abstractions.

Tool Registry System

Airtrain's tool registry system is designed to be:

• Extensible: Easily add custom tools to extend functionality
• Typed: Consistent interfaces make tools easy to use
• Validated: Tools must meet specific requirements to ensure reliability
• Stateful or Stateless: Support tools that maintain state between calls or are purely functional

Tool Registration

Tools are registered with the system using a decorator:

from airtrain.tools import register_tool, StatelessTool

@register_tool("my_tool")
class MyTool(StatelessTool):
    # Tool implementation
    ...

For stateful tools that need to maintain state:

from airtrain.tools import register_tool, StatefulTool

@register_tool("my_stateful_tool", tool_type="stateful")
class MyStatefulTool(StatefulTool):
    # Stateful tool implementation
    ...

Code Navigation Tools

Airtrain provides several tools for navigating and exploring codebases:

Directory Listing

The ListDirectoryTool allows agents to view the contents of directories:

from airtrain.tools import ToolFactory

# Get the tool
list_dir_tool = ToolFactory.get_tool("list_directory")

# List files in a directory
result = list_dir_tool(path="./src", show_hidden=False)

Directory Tree

The DirectoryTreeTool provides a hierarchical view of a directory structure:

dir_tree_tool = ToolFactory.get_tool("directory_tree")

# Show directory tree with depth limit
result = dir_tree_tool(
    path="./src",
    max_depth=3,
    show_hidden=False
)

Terminal Navigation

The TerminalNavigationTool lets agents navigate the file system with persistent state:

# Note: This is a stateful tool
nav_tool = ToolFactory.get_tool("terminal_navigation", "stateful")

# Show current directory
pwd_result = nav_tool(action="pwd")

# Change directory
cd_result = nav_tool(action="cd", directory="../other_dir")

# Push directory to stack
push_result = nav_tool(action="pushd", directory="src")

# Return to previous directory
pop_result = nav_tool(action="popd")

File and Content Search Tools

Find Files

The FindFilesTool locates files matching specific patterns:

find_tool = ToolFactory.get_tool("find_files")

# Find Python files
result = find_tool(
    directory="./src",
    pattern="**/*.py",
    max_results=50,
    show_hidden=False
)

Search Content

The SearchTermTool searches for specific terms within files:

search_tool = ToolFactory.get_tool("search_term")

# Search for a function name
result = search_tool(
    term="def process_data",
    directory="./src",
    file_pattern="*.py",
    case_sensitive=False,
    regex=False,
    max_results=20
)

Web Search

The WebSearchTool allows searching the web using the Exa search API:

web_search_tool = ToolFactory.get_tool("web_search")

# Search the web
result = web_search_tool(
    query="Latest advancements in AI safety research",
    num_results=5,
    include_domains=["arxiv.org", "openai.com"],
    exclude_domains=["reddit.com"],
    use_autoprompt=True
)

# Process the search results
for item in result["results"]:
    print(f"Title: {item['title']}")
    print(f"URL: {item['url']}")
    print(f"Content: {item['content'][:200]}...")

To use this tool, you must set the EXA_API_KEY environment variable with a valid API key from Exa.ai.

Command Execution Tools

Execute Commands

The ExecuteCommandTool allows running shell commands in a controlled environment:

cmd_tool = ToolFactory.get_tool("execute_command")

# Run a command
result = cmd_tool(
    command="git status",
    working_dir="./my_repo",
    timeout=10.0
)

Testing Tools

Run Python Tests

The RunPytestTool enables running pytest on specific test files:

pytest_tool = ToolFactory.get_tool("run_pytest")

# Run tests
result = pytest_tool(
    test_path="./tests/test_module.py",
    verbose=True,
    args=["-xvs"],
    timeout=30.0
)

Creating Custom Tools

Airtrain makes it easy to create custom tools for specific needs. Custom tools must implement:

1. An __init__ method defining the tool parameters
2. A __call__ method implementing the tool functionality
3. A to_dict method for LLM function calling

Example: Custom Stateless Tool

from airtrain.tools import register_tool, StatelessTool

@register_tool("word_count")
class WordCountTool(StatelessTool):
    """Tool for counting words in text."""
    
    def __init__(self):
        self.name = "word_count"
        self.description = "Count words in a text string"
        self.parameters = {
            "type": "object",
            "properties": {
                "text": {
                    "type": "string", 
                    "description": "Text to analyze"
                },
                "ignore_case": {
                    "type": "boolean",
                    "description": "Whether to ignore case when counting"
                }
            },
            "required": ["text"]
        }
    
    def __call__(self, text: str, ignore_case: bool = True) -> dict:
        """Count words in the text."""
        if ignore_case:
            text = text.lower()
            
        words = text.split()
        word_counts = {}
        
        for word in words:
            # Strip punctuation
            word = word.strip('.,!?:;()[]{}\'\"')
            if word:
                word_counts[word] = word_counts.get(word, 0) + 1
        
        return {
            "total_words": len(words),
            "unique_words": len(word_counts),
            "word_counts": word_counts
        }
    
    def to_dict(self):
        """Convert tool to dictionary format for LLM function calling."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters
            }
        }

Example: Custom Stateful Tool

from airtrain.tools import register_tool, StatefulTool

@register_tool("conversation_memory", tool_type="stateful")
class ConversationMemoryTool(StatefulTool):
    """Tool for storing conversation history with memory."""
    
    def __init__(self):
        self.name = "conversation_memory"
        self.description = "Store and retrieve messages from conversation history"
        self.parameters = {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["add", "get", "clear"],
                    "description": "Action to perform on the conversation memory"
                },
                "message": {
                    "type": "string",
                    "description": "Message to add when action is 'add'"
                },
                "role": {
                    "type": "string",
                    "enum": ["user", "assistant", "system"],
                    "description": "Role of the message when action is 'add'"
                }
            },
            "required": ["action"]
        }
        self.reset()
    
    @classmethod
    def create_instance(cls):
        """Create a new instance with fresh state."""
        return cls()
    
    def reset(self):
        """Reset the conversation memory."""
        self.messages = []
    
    def __call__(self, action: str, message: str = None, role: str = "user"):
        """Execute the conversation memory tool."""
        if action == "add" and message:
            self.messages.append({
                "role": role,
                "content": message
            })
            return {"status": "success", "message_count": len(self.messages)}
        
        elif action == "get":
            return {"status": "success", "messages": self.messages}
        
        elif action == "clear":
            self.reset()
            return {"status": "success", "message_count": 0}
        
        return {"status": "error", "message": "Invalid action"}
    
    def to_dict(self):
        """Convert tool to dictionary format for LLM function calling."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters
            }
        }

Using Tools with Agents

Airtrain's tools are designed to work seamlessly with agents. To provide tools to an agent:

from airtrain.tools import ToolFactory
from airtrain.agents import Agent

# Create an agent with tools
agent = Agent(
    name="coding_assistant",
    tools=[
        ToolFactory.get_tool("list_directory"),
        ToolFactory.get_tool("search_term"),
        ToolFactory.get_tool("execute_command"),
        ToolFactory.get_tool("web_search")  # Web search capability
    ]
)

# The agent can now use these tools in its processing

Advanced Tool Combinations

The tools can be combined in powerful ways to create sophisticated workflows:

Web Search with LLM Reasoning

from airtrain.tools import ToolFactory
from airtrain.integrations.openai import OpenAIChatSkill, OpenAICredentials, OpenAIChatInput

# Get web search tool
web_search = ToolFactory.get_tool("web_search")

# Set up OpenAI integration
openai_credentials = OpenAICredentials(api_key="your-api-key")
chat_skill = OpenAIChatSkill(credentials=openai_credentials)

# Use the web search tool
search_results = web_search(
    query="Latest advancements in quantum computing",
    num_results=3
)

# Extract useful information from results
context = ""
for i, result in enumerate(search_results["results"], 1):
    context += f"[{i}] {result['title']}\n"
    context += f"URL: {result['url']}\n"
    context += f"Content: {result['content'][:200]}...\n\n"

# Use GPT-4o to reason about the search results
chat_input = OpenAIChatInput(
    messages=[
        {"role": "system", "content": "You are a helpful AI research assistant."},
        {"role": "user", "content": f"Based on these search results about quantum computing, provide a summary of the latest advancements:\n\n{context}"}
    ],
    model="gpt-4o",
    temperature=0.2
)

# Get the enhanced response
response = chat_skill.process(chat_input)

Best Practices

When working with tools:

1. Security: Always validate inputs and limit command execution to safe operations
2. Error Handling: Return clear error messages when tools encounter problems
3. Resource Management: Set appropriate timeouts and limits for resource-intensive operations
4. State Management: Be careful about when to use stateful vs. stateless tools
5. Tool Composition: Build complex workflows by chaining multiple simple tools

Summary

The tools library provides a flexible foundation for Airtrain agents to interact with their environment. By combining navigation, search, execution, and testing capabilities, agents can perform complex tasks in a structured and controlled manner. Extending the library with custom tools allows for domain-specific functionality while maintaining a consistent interface.