# How to Add Custom Tools to an Agent's Toolset in Bettafish

> Learn to add custom tools to Bettafish agents. Wrap APIs, define response classes, add retry logic, and register your tool for enhanced agent capabilities.

- Repository: [BaiFu/bettafish](https://github.com/666ghj/bettafish)
- Tags: how-to-guide
- Published: 2026-02-23

---

**To add custom tools to an agent's toolset in Bettafish, create a new agency class that wraps your external API, define typed response dataclasses, decorate the core execution method with `with_graceful_retry`, and register the tool name in the agent's `execute_search_tool` dispatcher.**

In the Bettafish framework, agents like `DeepSearchAgent` interact with external services through modular toolsets that expose atomic search methods. Adding a custom tool requires following the established agency pattern used by existing implementations such as `TavilyNewsAgency` in [`QueryEngine/tools/search.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/tools/search.py). This guide walks through the exact implementation steps using the actual source architecture from the `666ghj/bettafish` repository.

## Understanding the Toolset Architecture

Bettafish employs a **dispatcher pattern** where the agent routes LLM tool requests to specific API wrappers. The architecture consists of four core components:

- **Tool Client**: A class like `TavilyNewsAgency` that wraps the external API and exposes atomic methods (e.g., `basic_search_news`)
- **Response Dataclasses**: Typed containers such as `TavilyResponse` and `SearchResult` that standardize data flow between the tool and agent
- **Agent Dispatcher**: The `DeepSearchAgent.execute_search_tool` method in [`QueryEngine/agent.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/agent.py) that maps string tool names to agency methods
- **Retry Infrastructure**: The `with_graceful_retry` decorator from [`utils/retry_helper.py`](https://github.com/666ghj/bettafish/blob/main/utils/retry_helper.py) that provides exponential backoff for transient failures

Every new tool must follow the signature `def tool_name(self, query: str, ...) -> ResponseType` to integrate seamlessly with this pipeline.

## Step 1: Create the API Client and Response Models

First, create a new Python file for your custom toolset. Define dataclasses for type-safe responses and a client wrapper for authentication.

```python

# file: QueryEngine/tools/custom_search.py

import os
import sys
from dataclasses import dataclass, field
from typing import List, Optional

# Add utils path for the retry decorator

root_dir = os.path.abspath(os.path.join(__file__, "..", ".."))
sys.path.append(os.path.join(root_dir, "utils"))
from retry_helper import with_graceful_retry, SEARCH_API_RETRY_CONFIG


class CustomClient:
    """Lightweight wrapper around the external service."""
    def __init__(self, api_key: str | None = None):
        self.api_key = api_key or os.getenv("CUSTOM_API_KEY")
        if not self.api_key:
            raise ValueError("CUSTOM_API_KEY environment variable is required")
        # Initialize third-party SDK here


@dataclass
class CustomResult:
    title: str
    url: str
    snippet: str
    published_date: Optional[str] = None


@dataclass
class CustomResponse:
    query: str
    results: List[CustomResult] = field(default_factory=list)
    response_time: Optional[float] = None

```

## Step 2: Implement the Tool Method with Retry Logic

Create an agency class that exposes public methods the LLM can invoke. Decorate the internal execution method with `with_graceful_retry` to match the resilience pattern used by existing tools in [`QueryEngine/tools/search.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/tools/search.py).

```python

# file: QueryEngine/tools/custom_search.py (continuation)

class CustomSearchAgency:
    """Collection of custom search tools for the agent."""
    
    def __init__(self, api_key: str | None = None):
        self._client = CustomClient(api_key)

    @with_graceful_retry(SEARCH_API_RETRY_CONFIG, default_return=CustomResponse(query="failed"))
    def _execute(self, *, query: str, **kwargs) -> CustomResponse:
        """Internal method handling the actual API call."""
        raw = self._client.search(query, **kwargs)
        results = [
            CustomResult(
                title=r["title"],
                url=r["url"],
                snippet=r["snippet"],
                published_date=r.get("published_date")
            )
            for r in raw.get("hits", [])
        ]
        return CustomResponse(
            query=query, 
            results=results, 
            response_time=raw.get("time")
        )

    # Public tool that the LLM can request by name

    def curated_search(self, query: str, max_results: int = 8) -> CustomResponse:
        """Custom curated search returning the top max_results hits."""
        print(f"--- TOOL: Custom curated search (query={query}) ---")
        return self._execute(query=query, limit=max_results)

```

The `with_graceful_retry` decorator ensures your tool handles transient HTTP errors consistently with the existing **Tavily** and **Bocha** integrations.

## Step 3: Register the Custom Tool in the Agent Dispatcher

Modify [`QueryEngine/agent.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/agent.py) to import your new agency and add a dispatch case in the `execute_search_tool` method.

```python

# QueryEngine/agent.py

from .tools.custom_search import CustomSearchAgency, CustomResponse

class DeepSearchAgent:
    def __init__(self, config: Optional[Settings] = None):
        # ... existing initialization ...

        self.custom_agency = CustomSearchAgency(
            api_key=self.config.CUSTOM_API_KEY
        )
        # ... rest of initialization ...

    def execute_search_tool(self, tool_name: str, query: str, **kwargs):
        """Dispatcher mapping tool names to agency methods."""
        if tool_name == "tavily_search":
            return self.tavily_agency.basic_search_news(query, **kwargs)
        elif tool_name == "curated_search":  # <-- New tool registration

            return self.custom_agency.curated_search(query, **kwargs)
        # ... other tools ...

```

Ensure the return type remains compatible with downstream processing. If the agent expects `TavilyResponse` specifically, adapt your `CustomResponse` to match that schema or convert it within the dispatcher.

## Step 4: Update Configuration and Environment Variables

Expose any new API keys through the central `Settings` class in [`config.py`](https://github.com/666ghj/bettafish/blob/main/config.py) to maintain consistent configuration management.

```python

# config.py

class Settings(BaseSettings):
    # ... existing settings ...

    CUSTOM_API_KEY: str | None = None
    MAX_REFLECTIONS: int = 3

```

Agents instantiated via the factory pattern will automatically pick up `CUSTOM_API_KEY` from the environment, enabling immediate usage without code changes to the agent initialization logic.

## Example: Calling the Custom Tool from a Prompt

Once registered, the LLM can invoke your tool through the standard action schema:

```python

# Example prompt or test script

tool_request = {
    "action": "use_tool",
    "tool_name": "curated_search",
    "arguments": {
        "query": "latest quantum computing breakthroughs",
        "max_results": 5
    }
}

# The agent routes this to CustomSearchAgency.curated_search

response = agent.execute_search_tool(
    tool_request["tool_name"],
    tool_request["arguments"]["query"],
    max_results=tool_request["arguments"]["max_results"]
)

```

The agent receives a `CustomResponse` object, converts it to the common search-result dictionary format, and passes it to downstream summarization nodes transparently.

## Summary

- **Create** a new agency class in `QueryEngine/tools/` that wraps your external API client and exposes atomic methods
- **Define** typed dataclasses for responses (e.g., `CustomResponse`, `CustomResult`) to ensure type safety across the agent pipeline
- **Decorate** low-level API calls with `@with_graceful_retry` from [`utils/retry_helper.py`](https://github.com/666ghj/bettafish/blob/main/utils/retry_helper.py) to inherit exponential backoff behavior
- **Register** the tool in [`QueryEngine/agent.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/agent.py) by adding an `elif` branch in `execute_search_tool` that maps a string name to your agency method
- **Configure** API keys in [`config.py`](https://github.com/666ghj/bettafish/blob/main/config.py) using the `Settings` class to support environment-based configuration
- **Document** the tool method with docstrings specifying parameters and return types for future maintainers

## Frequently Asked Questions

### What is the role of the `execute_search_tool` dispatcher in Bettafish?

The `execute_search_tool` method in [`QueryEngine/agent.py`](https://github.com/666ghj/bettafish/blob/main/QueryEngine/agent.py) acts as a central router that translates string tool names requested by the LLM into concrete method calls on agency instances. This dispatcher pattern allows you to add new tools without modifying the agent's core logic—simply add a new conditional branch mapping the tool name to your custom agency method.

### Why should custom tools use the `with_graceful_retry` decorator?

All production tools in Bettafish use `with_graceful_retry` (defined in [`utils/retry_helper.py`](https://github.com/666ghj/bettafish/blob/main/utils/retry_helper.py)) to handle transient network failures, rate limiting, and timeouts with exponential backoff. Applying this decorator to your custom tool's internal execution method ensures consistent error handling and prevents agent failures due to temporary API outages, matching the resilience of built-in Tavily and Bocha integrations.

### How do I handle APIs that return different data formats than Tavily?

If your custom API returns a schema incompatible with `TavilyResponse`, you have two options: either define your response dataclass to match the expected downstream format in the agent, or perform transformation logic inside the `execute_search_tool` dispatcher before returning the result. The Bettafish architecture does not enforce a single response type, but downstream nodes may expect specific fields like `results` and `response_time`.

### Can I add custom tools to MediaEngine agents as well?

Yes, the pattern is identical across engines. The [`MediaEngine/agent.py`](https://github.com/666ghj/bettafish/blob/main/MediaEngine/agent.py) implementation also uses an `execute_search_tool` dispatcher and follows the same agency pattern seen in `QueryEngine`. Create your tool class, import it into the relevant agent file, instantiate it in `__init__`, and add the dispatch case to make it available to that specific engine's workflow.