Custom RAG Retrieval Tool

Welcome! In this lesson, we’ll see how to implement a custom tool for Retrieval-Augmented Generation (RAG) retrieval in an AI agent. Previously, you learned how RAG agents use a knowledge base to answer questions. But what if you want your agent to use a specialized retrieval method or expose this as a reusable tool? That’s where custom tools come in.

Our goal: learn to design, implement, and integrate a custom retrieval tool that fetches the most relevant knowledge chunks for any user query. By the end, you’ll know how to build such a tool, validate its inputs, and connect it to your agent for context-aware responses.

Why do we need a custom tool for RAG retrieval? Imagine you’re building a personal learning assistant. You want it to answer questions using your own notes or learning plans, not just generic knowledge.

In AI agents, a “tool” is a callable function or service for a specific task. For RAG, this means a tool that searches your knowledge base and returns the most relevant information. For example, if a user asks, “What are my learning plans for SQL?” the agent should use the retrieval tool to find something like “Review different types of SQL joins — especially LEFT and FULL OUTER joins.”

A custom tool gives you control over how retrieval works, what data is returned, and how it’s formatted. This flexibility is key to building agents tailored to your needs.

Tools should validate their inputs. This ensures the tool gets the right data and handles errors well. Pydantic is a popular library for this in Python.

Define a Pydantic model for the tool’s arguments:

from pydantic import BaseModel

class FunctionArgs(BaseModel):
    user_query: str

This model enforces that user_query is a string. If the input is missing or incorrect, Pydantic raises an error.

Use this model in the tool’s function:

async def run_function(ctx, args: str) -> str:
    parsed = FunctionArgs.model_validate_json(args)
    chunks = retrieve_top_chunks(parsed.user_query)
    return "\n".join([c["chunk"] for c in chunks])
# Output (for a relevant query):
# Review different types of SQL joins — especially LEFT and FULL OUTER joins.

model_validate_json parses and validates the input. This step is crucial for robust tools.

With retrieval logic and input validation ready, let’s wrap it in a FunctionTool. This class describes the tool’s interface and connects it to the agent.

from agents import FunctionTool

rag_retrieval_tool = FunctionTool(
    name="rag_retrieval_tool",
    description="A tool to retrieve context from RAG based on user query",
    params_json_schema={
        "type": "object",
        "properties": {
            "user_query": {"type": "string"},
        },
        "required": ["user_query"],
        "additionalProperties": False
    },
    on_invoke_tool=run_function
)

Key parameters:

• name: The tool’s unique identifier.
• description: Explains what the tool does.
• params_json_schema: Describes the expected input.
• on_invoke_tool: The function to call (our validated retrieval function).

This makes the tool discoverable and usable by any agent that supports tool invocation.

Now, connect your custom tool to the agent. This lets the agent use the tool to retrieve information from the knowledge base.

from agents import Agent

AGENT = Agent(
    name="Learning Assistant",
    instructions=(
        "You are a personal learning assistant with access to rag tool. "
        "Whenever asked a question about learning plans, use the RAG retrieval tool to get context from the DB and answer user questions."
    ),
    tools=[rag_retrieval_tool]
)

Ask the agent a question, and it will use the retrieval tool as needed:

from agents import Runner

def ask_agent(prompt):
    result = Runner.run_sync(AGENT, prompt)
    return result.final_output

# Example usage
response = ask_agent("What are my learning plans for SQL?")
print(response)  # Review different types of SQL joins — especially LEFT and FULL OUTER joins.

The agent automatically invokes the retrieval tool, fetches relevant knowledge, and uses it to answer the user’s question.

You learned how to implement a custom tool for RAG retrieval and integrate it with an AI agent. We covered:

• Why custom tools are important for flexible, context-aware agents
• How to design a retrieval function that queries your knowledge base
• Using Pydantic for input validation
• Wrapping your logic in a FunctionTool
• Connecting the tool to an agent for relevant, knowledge-based answers

You now have the foundation to build and extend your own retrieval tools for any use case.

Now it’s your turn! Next, you’ll get hands-on experience building and testing your own custom RAG retrieval tool. You’ll practice designing retrieval logic, validating inputs, and integrating your tool with an agent. This will help solidify your understanding and prepare you for real-world applications.