Lab 04: Agent with Tool Calling
Goal
Build a multi-tool agent that can search, calculate, and reason across multiple steps. Implement approval gates, step logging, and loop detection.
Prerequisites
pip install openai httpx
Exercise 1: Single-Tool Agent
# lab04/01_single_tool.py
from openai import OpenAI
import json, math
client = OpenAI()
# Tool definition
CALCULATOR_TOOL = {
"type": "function",
"function": {
"name": "calculate",
"description": "Evaluate a mathematical expression. Supports basic arithmetic and math functions.",
"parameters": {
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "Math expression like '2 ** 10' or 'sqrt(144)'"
}
},
"required": ["expression"]
}
}
}
def safe_eval(expression: str) -> float:
safe_globals = {
"__builtins__": {},
"sqrt": math.sqrt, "log": math.log, "log2": math.log2,
"sin": math.sin, "cos": math.cos, "tan": math.tan,
"pi": math.pi, "e": math.e, "abs": abs, "round": round,
"pow": pow, "min": min, "max": max
}
result = eval(expression, safe_globals)
return result
def run_step(messages: list) -> tuple[str | None, list]:
"""Run one step of the agent loop. Returns (final_answer, updated_messages) or (None, updated_messages)."""
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
tools=[CALCULATOR_TOOL],
tool_choice="auto"
)
message = response.choices[0].message
messages = messages + [message] # Immutable pattern
if response.choices[0].finish_reason == "stop":
return message.content, messages
if message.tool_calls:
for tool_call in message.tool_calls:
args = json.loads(tool_call.function.arguments)
try:
result = safe_eval(args["expression"])
tool_result = str(result)
print(f" [TOOL] calculate({args['expression']}) = {result}")
except Exception as e:
tool_result = f"Error: {e}"
print(f" [TOOL] calculate({args['expression']}) = ERROR: {e}")
messages = messages + [{
"role": "tool",
"tool_call_id": tool_call.id,
"content": tool_result
}]
return None, messages
def run_agent(goal: str, max_steps: int = 10) -> str:
messages = [
{"role": "system", "content": "You are a math assistant. Use the calculator tool for all calculations."},
{"role": "user", "content": goal}
]
print(f"Goal: {goal}")
for step in range(max_steps):
print(f"\nStep {step + 1}:")
answer, messages = run_step(messages)
if answer:
print(f" [DONE] {answer}")
return answer
return "Max steps reached"
# Test
result = run_agent("What is 2^10 + sqrt(144) + the 7th Fibonacci number?")
print(f"\nFinal: {result}")
Exercise 2: Multi-Tool Agent
# lab04/02_multi_tool.py
from openai import OpenAI
import json, math, datetime
client = OpenAI()
# Tool registry
TOOLS = [
{
"type": "function",
"function": {
"name": "calculate",
"description": "Evaluate a mathematical expression.",
"parameters": {
"type": "object",
"properties": {
"expression": {"type": "string", "description": "Math expression to evaluate"}
},
"required": ["expression"]
}
}
},
{
"type": "function",
"function": {
"name": "get_current_date",
"description": "Get today's date and time.",
"parameters": {
"type": "object",
"properties": {}
}
}
},
{
"type": "function",
"function": {
"name": "search_knowledge",
"description": "Search a local knowledge base for factual information.",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query"}
},
"required": ["query"]
}
}
}
]
# Mock knowledge base
KNOWLEDGE_BASE = {
"llm": "A Large Language Model (LLM) is a type of AI model trained on large amounts of text data to generate and understand language.",
"transformer": "The transformer architecture uses self-attention mechanisms to process sequences of tokens in parallel.",
"kv cache": "The KV cache stores key-value pairs from previous tokens to avoid recomputing attention during generation.",
"rag": "RAG (Retrieval-Augmented Generation) retrieves relevant documents at query time and injects them into the LLM's context.",
"quantization": "Quantization reduces model precision (e.g., FP16 → INT4) to reduce memory usage and increase inference speed.",
}
def execute_tool(name: str, args: dict) -> str:
if name == "calculate":
try:
safe_globals = {"__builtins__": {}, "sqrt": math.sqrt, "pi": math.pi, "abs": abs}
result = eval(args["expression"], safe_globals)
return str(result)
except Exception as e:
return f"Error: {e}"
elif name == "get_current_date":
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
elif name == "search_knowledge":
query = args["query"].lower()
for key, value in KNOWLEDGE_BASE.items():
if key in query:
return value
return "No information found for this query."
return f"Unknown tool: {name}"
def run_agent(goal: str, max_steps: int = 15) -> str:
messages = [
{
"role": "system",
"content": "You are a helpful research assistant. Use tools to find information and perform calculations. Always use the search_knowledge tool before answering factual questions about AI/ML."
},
{"role": "user", "content": goal}
]
tool_call_counts = {}
for step in range(max_steps):
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
tools=TOOLS,
tool_choice="auto"
)
message = response.choices[0].message
messages.append(message)
if response.choices[0].finish_reason == "stop":
return message.content
if message.tool_calls:
for tool_call in message.tool_calls:
name = tool_call.function.name
args = json.loads(tool_call.function.arguments)
# Track for loop detection
call_key = f"{name}:{json.dumps(args, sort_keys=True)}"
tool_call_counts[call_key] = tool_call_counts.get(call_key, 0) + 1
if tool_call_counts[call_key] > 2:
return "Error: Detected repeated tool calls. Stopping."
result = execute_tool(name, args)
print(f"[Tool] {name}({args}) → {result[:80]}")
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": result
})
return "Max steps reached"
# Test
print("=" * 60)
result = run_agent("What is the KV cache, and how much memory does a 70B model need for it with a 4096-token context?")
print(f"\nAnswer:\n{result}")
print("\n" + "=" * 60)
result = run_agent("What is today's date? How many days until December 31st?")
print(f"\nAnswer:\n{result}")
Exercise 3: Agent with Approval Gate
# lab04/03_approval_gate.py
from openai import OpenAI
import json
client = OpenAI()
# Classify tools by risk level
LOW_RISK = {"search", "calculate", "read_file", "list_files"}
HIGH_RISK = {"write_file", "delete_file", "send_email", "make_api_call"}
class ApprovalAgent:
def __init__(self, auto_approve_low_risk: bool = True):
self.auto_approve_low_risk = auto_approve_low_risk
self.action_log = []
def request_approval(self, tool_name: str, args: dict) -> bool:
print(f"\n⚠️ APPROVAL REQUIRED")
print(f" Tool: {tool_name}")
print(f" Args: {json.dumps(args, indent=2)}")
response = input(" Approve? (yes/no): ").strip().lower()
return response in ["yes", "y"]
def execute_tool(self, tool_name: str, args: dict) -> str:
# Check risk level
if tool_name in HIGH_RISK:
approved = self.request_approval(tool_name, args)
if not approved:
return "Action denied by user."
# Log the action
self.action_log.append({"tool": tool_name, "args": args})
# Execute (mock implementations)
if tool_name == "search":
return f"Search results for: {args.get('query', '')}"
elif tool_name == "write_file":
return f"Written to {args.get('filename', 'file.txt')}"
elif tool_name == "send_email":
return f"Email sent to {args.get('to', '')}"
return f"Tool {tool_name} executed"
# Example: Run the approval agent
agent = ApprovalAgent(auto_approve_low_risk=True)
# Simulate a tool call to a high-risk tool
result = agent.execute_tool("send_email", {
"to": "team@company.com",
"subject": "Test",
"body": "Hello from agent"
})
print(f"\nResult: {result}")
print(f"\nAction log: {agent.action_log}")
Exercise 4: Observability — Log Every Step
# lab04/04_observability.py
import json, time
from dataclasses import dataclass, field
from openai import OpenAI
client = OpenAI()
@dataclass
class AgentTrace:
session_id: str
goal: str
steps: list = field(default_factory=list)
start_time: float = field(default_factory=time.time)
total_tokens: int = 0
def add_step(self, step_type: str, data: dict):
self.steps.append({
"step": len(self.steps) + 1,
"type": step_type,
"timestamp": time.time() - self.start_time,
**data
})
def summary(self) -> dict:
return {
"session_id": self.session_id,
"goal": self.goal[:80],
"total_steps": len(self.steps),
"total_duration_s": time.time() - self.start_time,
"total_tokens": self.total_tokens,
"tool_calls": [s for s in self.steps if s["type"] == "tool_call"]
}
def traced_agent(goal: str, tools: list, tool_handlers: dict) -> tuple[str, AgentTrace]:
import uuid
trace = AgentTrace(session_id=str(uuid.uuid4())[:8], goal=goal)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": goal}
]
for step in range(15):
step_start = time.time()
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
tools=tools if tools else None,
tool_choice="auto" if tools else None
)
message = response.choices[0].message
usage = response.usage
trace.total_tokens += usage.prompt_tokens + usage.completion_tokens
messages.append(message)
trace.add_step("llm_call", {
"finish_reason": response.choices[0].finish_reason,
"input_tokens": usage.prompt_tokens,
"output_tokens": usage.completion_tokens,
"duration_ms": (time.time() - step_start) * 1000
})
if response.choices[0].finish_reason == "stop":
return message.content, trace
if message.tool_calls:
for tc in message.tool_calls:
args = json.loads(tc.function.arguments)
handler = tool_handlers.get(tc.function.name)
result = handler(args) if handler else "Tool not found"
trace.add_step("tool_call", {
"tool": tc.function.name,
"args": args,
"result_preview": str(result)[:100]
})
messages.append({
"role": "tool",
"tool_call_id": tc.id,
"content": str(result)
})
return "Max steps reached", trace
# Test
import math
tools = [{
"type": "function",
"function": {
"name": "calc",
"description": "Evaluate a math expression",
"parameters": {
"type": "object",
"properties": {"expr": {"type": "string"}},
"required": ["expr"]
}
}
}]
handlers = {
"calc": lambda args: eval(args["expr"], {"__builtins__": {}, "sqrt": math.sqrt})
}
answer, trace = traced_agent("What is 100 * pi / 4, rounded to 2 decimal places?", tools, handlers)
print(f"Answer: {answer}")
print(f"\nTrace summary:")
print(json.dumps(trace.summary(), indent=2))
Deliverables
- Single-tool agent running (Exercise 1)
- Multi-tool agent with loop detection (Exercise 2)
- Approval gate tested with a "high risk" tool (Exercise 3)
- Full trace logged for a 3-step agent task (Exercise 4)
- Notes: what failure modes did you encounter?
Reflection Questions
- Why must the application (not the model) execute tool calls?
- An agent sent the same email 5 times. What safeguards would have prevented this?
- What is the confused deputy problem and how does tool whitelisting help?
- A tool call returns an error. How should the agent handle it?
- How would you build a memory system where the agent remembers information across sessions?