Build a Native Tool Search Tool for Orchestrator Agents (LangGraph Example)

TOOL_CATALOG = [
    {
        "name": "github_create_pull_request",
        "description": "Create a GitHub pull request for a repository",
        "tags": ["github", "code", "pr", "version-control"],
        "risk_level": "medium",
    },
    {
        "name": "github_list_issues",
        "description": "List open issues for a GitHub repository",
        "tags": ["github", "issues", "code"],
        "risk_level": "low",
    },
    {
        "name": "slack_post_message",
        "description": "Post a message to a Slack channel",
        "tags": ["slack", "messaging", "notification"],
        "risk_level": "low",
    },
    {
        "name": "database_execute_query",
        "description": "Execute a SQL query against the production database",
        "tags": ["database", "sql", "data"],
        "risk_level": "high",
    },
    {
        "name": "email_send",
        "description": "Send an email to one or more recipients",
        "tags": ["email", "messaging", "notification"],
        "risk_level": "medium",
    },
    {
        "name": "jira_create_ticket",
        "description": "Create a new Jira ticket in a project",
        "tags": ["jira", "project-management", "tracking"],
        "risk_level": "low",
    },
]

def _search_catalog(query: str, top_k: int = 5) -> list[dict]:
    """Score and rank catalog entries by keyword overlap with the query."""
    q_tokens = query.lower().split()
 
    def score(entry: dict) -> int:
        haystack = " ".join([
            entry.get("name", ""),
            entry.get("description", ""),
            " ".join(entry.get("tags", [])),
        ]).lower()
        return sum(1 for token in q_tokens if token in haystack)
 
    scored = [(entry, score(entry)) for entry in TOOL_CATALOG]
    scored = [(e, s) for e, s in scored if s > 0]
    scored.sort(key=lambda x: x[1], reverse=True)
    return [e for e, _ in scored[:top_k]]

from langchain.tools import tool
 
 
@tool("tool_search")
def tool_search(query: str) -> str:
    """Search the tool catalog to discover available tools matching your intent.
 
    Use this tool BEFORE calling any specialized tool you are unsure about.
    It returns tool names, descriptions, and risk levels so you can make
    an informed choice.
 
    Args:
        query: Keywords describing the capability you need
               (e.g., 'create pull request', 'send message', 'query database').
    """
    matches = _search_catalog(query, top_k=5)
    if not matches:
        return "No matching tools found. Try different keywords."
 
    lines = [f"Found {len(matches)} matching tool(s):\n"]
    for m in matches:
        lines.append(
            f"- **{m['name']}**: {m['description']} "
            f"(risk: {m['risk_level']})"
        )
    return "\n".join(lines)

from langchain.tools import tool
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
 
 
# --- Always-loaded domain tools ---
 
@tool("get_weather")
def get_weather(location: str) -> str:
    """Get current weather for a location.
 
    Args:
        location: City name or coordinates.
    """
    return f"22 degrees C, partly cloudy in {location}."
 
 
# --- Orchestrator agent ---
 
SYSTEM_PROMPT = """\
You are an orchestrator agent with access to a large tool catalog.
 
You have a `tool_search` tool that lets you discover available
capabilities by keyword.
 
Workflow:
1. When the user's request involves a capability you are not sure
   about, call `tool_search` first with relevant keywords.
2. Review the returned tool names, descriptions, and risk levels.
3. If a matched tool is high-risk, inform the user and request
   confirmation before proceeding.
4. Call the appropriate domain tool once you have confirmed the
   right one.
 
Do NOT guess tool names. Always use `tool_search` to discover them.
"""
 
model = ChatOpenAI(model="gpt-4o")
 
agent = create_agent(
    model,
    tools=[tool_search, get_weather],
    system_prompt=SYSTEM_PROMPT,
)
 
# --- Run the agent ---
 
result = agent.invoke(
    {"messages": [{"role": "user", "content": "I need to open a PR on GitHub."}]}
)
 
# Expected agent behavior:
# 1. Call tool_search("create pull request github")
# 2. Receive: github_create_pull_request (risk: medium)
# 3. Inform the user about the available tool and its risk level
# 4. In a full system, the discovered tool would then be loaded and executed

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode, tools_condition
 
tools = [tool_search, get_weather]
 
model = ChatOpenAI(model="gpt-4o").bind_tools(tools)
 
 
def call_model(state: MessagesState):
    response = model.invoke(state["messages"])
    return {"messages": [response]}
 
 
builder = StateGraph(MessagesState)
builder.add_node("agent", call_model)
builder.add_node("tools", ToolNode(tools))
 
builder.add_edge(START, "agent")
builder.add_conditional_edges("agent", tools_condition)
builder.add_edge("tools", "agent")
 
graph = builder.compile()
 
result = graph.invoke(
    {"messages": [{"role": "user", "content": "I need to open a PR on GitHub."}]}
)