1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
|
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
API Server for custom user agents
"""
import argparse
import importlib.util
import os
import sys
from src.models import AgentRequest, AgentResponse
from src.registry import agentRegistry
from agents import Runner, Agent
import fastapi
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
import uvicorn
class ArgParser:
"""
Argument parser for command line arguments.
"""
def __init__(self):
parser = argparse.ArgumentParser()
parser.add_argument(
"--plugin-dir",
type=str,
help="Directory containing plugins",
default="/useragent/src/custom_agents",
)
parser.add_argument(
"--port", type=int, help="Port to run the server on", default=443
)
parser.add_argument(
"--address", type=str, help="which address to bind to", default="0.0.0.0"
)
self.args = parser.parse_args()
def load_plugins(plugin_dir: str):
"""
Load all Python files in the specified directory as modules.
Args:
plugin_dir (str): Directory containing the plugin files.
"""
for filename in os.listdir(plugin_dir):
if filename.endswith(".py") and filename != "__init__.py":
filepath = os.path.join(plugin_dir, filename)
print(filepath)
module_name = filename[:-3]
spec = importlib.util.spec_from_file_location(module_name, filepath)
if spec and spec.loader:
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
print(f"Loaded module: {module_name} from {filepath}")
else:
print(f"Warning: Could not load module {module_name} from {filepath}")
def get_agent_generator(agent_name: str) -> Agent:
"""
Find and return the agent generator function from the custom agent registry
"""
for k, v in agentRegistry.registry.items():
print(k, v)
try:
agent = agentRegistry.registry[agent_name]
except KeyError:
raise ValueError(f"Agent {agent_name} not found in registry.")
return agent
class AgentRunner:
"""
Class to run the agent with the provided request.
"""
def __init__(self, agent_request):
self.agent_request = agent_request
async def run(self) -> str:
print(
f"Running agent {self.agent_request.agent_name} with instructions: {self.agent_request.instructions} with query: {self.agent_request.query}"
)
try:
agent = get_agent_generator(self.agent_request.agent_name)(
self.agent_request
)
except ValueError as e:
print(f"Error: {e}")
return str(e)
result = await Runner.run(agent, self.agent_request.query)
print(result.final_output_as(str))
return result.final_output_as(str)
class APIServer:
"""
The API server
"""
def __init__(self):
self.app = fastapi.FastAPI()
self.router = fastapi.APIRouter()
self.router.add_api_route(
"/api/v1/agent", self.agent_handler, methods=["POST"], tags=[]
)
self.agent_registry = {}
self.app.include_router(self.router)
async def agent_handler(self, agent_request: AgentRequest) -> fastapi.Response:
print(f"Received request: {agent_request}")
response = await AgentRunner(agent_request).run()
result = AgentResponse(agent_name=agent_request.agent_name, response=response)
print(f"Response: {result}")
return JSONResponse(
content=jsonable_encoder(result),
status_code=200,
)
def main():
argparser = ArgParser()
load_plugins(argparser.args.plugin_dir)
app = APIServer().app
uvicorn.run(
app,
host=argparser.args.address,
port=argparser.args.port,
ssl_keyfile="./server.key",
ssl_certfile="./server.cert",
)
if __name__ == "__main__":
main()
|
