Concept #39Hardsystem-design

How would you implement tool use in agents?

#gen-ai#agents#system-design

Answer

Implementing Tool Use in AI Agents

Tool use (function calling) allows an LLM to invoke external functions — APIs, databases, calculators — and incorporate the results into its reasoning.

OpenAI Function Calling

python
from openai import OpenAI
import json

client = OpenAI()

# Define tools
tools = [
    {
        "type": "function",
        "function": {
            "name": "search_knowledge_base",
            "description": "Search the internal knowledge base for relevant documents",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "The search query"},
                    "top_k": {"type": "integer", "description": "Number of results", "default": 5},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_order_status",
            "description": "Look up the status of a customer order",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string", "description": "The order ID"},
                },
                "required": ["order_id"],
            },
        },
    },
]

# Tool implementations
def search_knowledge_base(query: str, top_k: int = 5) -> list[dict]:
    # Real implementation would query vector store
    return [{"content": f"Result for: {query}", "score": 0.9}]

def get_order_status(order_id: str) -> dict:
    # Real implementation would query database
    return {"order_id": order_id, "status": "shipped", "eta": "2 days"}

TOOL_REGISTRY = {
    "search_knowledge_base": search_knowledge_base,
    "get_order_status": get_order_status,
}

def run_agent(user_message: str, max_iterations: int = 5) -> str:
    messages = [{"role": "user", "content": user_message}]

    for _ in range(max_iterations):
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=tools,
            tool_choice="auto",
        )

        message = response.choices[0].message
        messages.append(message)

        # No tool call — return final answer
        if not message.tool_calls:
            return message.content

        # Execute each requested tool
        for tool_call in message.tool_calls:
            fn_name = tool_call.function.name
            fn_args = json.loads(tool_call.function.arguments)

            tool_fn = TOOL_REGISTRY.get(fn_name)
            if tool_fn:
                result = tool_fn(**fn_args)
            else:
                result = {"error": f"Unknown tool: {fn_name}"}

            # Add tool result to conversation
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": json.dumps(result),
            })

    return "Max iterations reached without final answer."

# Run
answer = run_agent("What's the status of order #12345?")
print(answer)

LangChain Agent with Tools

python
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain import hub

@tool
def search_web(query: str) -> str:
    '''Search the web for current information.'''
    # Real implementation: use SerpAPI, Tavily, etc.
    return f"Web results for: {query}"

@tool
def run_python_code(code: str) -> str:
    '''Execute Python code and return the output. Use for calculations.'''
    import io, contextlib
    output = io.StringIO()
    with contextlib.redirect_stdout(output):
        exec(code)
    return output.getvalue()

llm = ChatOpenAI(model="gpt-4o")
tools = [search_web, run_python_code]
prompt = hub.pull("hwchase17/openai-tools-agent")

agent = create_tool_calling_agent(llm, tools, prompt)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True, max_iterations=5)

result = executor.invoke({"input": "What is 1234 × 5678? Show your work."})
print(result["output"])

Tool Design Best Practices

PrincipleExample
Clear description"Search internal KB for product info" not "search"
Precise parametersInclude type, format, valid values
Idempotent readsSearch/lookup tools are safe to retry
Guard writesConfirm before deleting, sending emails
Return structured dataJSON dict with status and result
Timeout all toolsPrevent infinite hangs

Safety: Never give agents direct database write access without confirmation. Use a "dry run" mode that shows the intended action and requires approval for destructive operations.