1from openai import OpenAI
2import json
3
4client = OpenAI()
5
6TOOLS = [
7 {
8 "type": "function",
9 "function": {
10 "name": "search_docs",
11 "description": "Search documentation for relevant content",
12 "parameters": {
13 "type": "object",
14 "properties": {"query": {"type": "string"}},
15 "required": ["query"],
16 },
17 },
18 },
19 {
20 "type": "function",
21 "function": {
22 "name": "write_file",
23 "description": "Write content to a file",
24 "parameters": {
25 "type": "object",
26 "properties": {"filename": {"type": "string"}, "content": {"type": "string"}},
27 "required": ["filename", "content"],
28 },
29 },
30 },
31]
32
33def run_agent(task: str) -> str:
34 messages = [{"role": "system", "content": "You are a helpful agent. Use tools to complete tasks."},
35 {"role": "user", "content": task}]
36 while True:
37 response = client.chat.completions.create(
38 model="gpt-4o", messages=messages, tools=TOOLS, tool_choice="auto"
39 )
40 msg = response.choices[0].message
41 if msg.tool_calls:
42 messages.append(msg)
43 for tc in msg.tool_calls:
44 result = dispatch_tool(tc.function.name, json.loads(tc.function.arguments))
45 messages.append({"role": "tool", "tool_call_id": tc.id, "content": str(result)})
46 else:
47 return msg.content
48
49def dispatch_tool(name: str, args: dict) -> str:
50 if name == "search_docs":
51 return f"Found 3 documents about: {args['query']}"
52 if name == "write_file":
53 return f"Wrote {len(args['content'])} bytes to {args['filename']}"
54 return "Unknown tool"