Repository navigation
Expand file tree
/
Copy pathlambda_function.py
More file actions
229 lines (195 loc) · 7.19 KB
/
Copy pathlambda_function.py
File metadata and controls
229 lines (195 loc) · 7.19 KB
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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
import asyncio
import base64
import json
import logging
import os
import re
from typing import Any
from a2a.server.apps.jsonrpc import A2ALambdaApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from agent_cards import load_agent_card
from agent_executor import AgentExecutorImpl
from agent_schemas import load_agent_schemas
from agent_environment import agent_env_routes
from agent_mcp_servers import agent_mcp_routes
from agent_cards import agent_card_routes
from agent_schemas import agent_schema_routes
from agent_vectorizer import agent_vectorizer_routes
from agent_prompts_generator import agent_prompts_route
from agent_schema_generator import agent_schemageneration_route
from agent_capabilities_auto import agent_capability_routes
from test_payload_generator import test_payloads_route
from agent.starlette_shim import (
Response,
Request,
JSONResponse,
PlainTextResponse,
FileResponse,
)
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
_cached_card = None
_cached_schemas = None
_lambda_app = None
_compiled_routes: list[tuple[re.Pattern, Any]] = []
def _load_card_and_schemas(agent_name: str) -> tuple[Any, Any]:
global _cached_card, _cached_schemas
if _cached_card is None or _cached_schemas is None:
_cached_card = load_agent_card(agent_name)
_cached_schemas = load_agent_schemas(agent_name)
return _cached_card, _cached_schemas
def _build_lambda_app(agent_name: str) -> A2ALambdaApplication:
global _lambda_app
if _lambda_app:
return _lambda_app
card, schemas = _load_card_and_schemas(agent_name)
handler = DefaultRequestHandler(
agent_executor=AgentExecutorImpl(
schemas["input_doc"], schemas["output_doc"]
),
task_store=InMemoryTaskStore(),
)
_lambda_app = A2ALambdaApplication(
agent_card=card,
http_handler=handler,
rpc_url="/",
)
return _lambda_app
def _response(body: dict, status: int = 200, headers: dict | None = None):
return {
"statusCode": status,
"headers": {"Content-Type": "application/json", **(headers or {})},
"body": json.dumps(body),
"isBase64Encoded": False,
}
def _handle_well_known(path: str, base_url: str, agent_name: str):
card, schemas = _load_card_and_schemas(agent_name)
if path == "/.well-known/a2a/agent.json":
data = card.model_dump(exclude_none=True)
data["url"] = base_url
data.setdefault("extensions", {})
data["extensions"]["schemas_url"] = f"{base_url}/.well-known/a2a/schemas.json"
data["extensions"]["schemas"] = {
"index": agent_name + "_index",
"version": schemas["version"],
"input": {
"id": f"{agent_name}-input-v1",
"schema": schemas["input_doc"]["schema"],
},
"output": {
"id": f"{agent_name}-output-v1",
"schema": schemas["output_doc"]["schema"],
},
}
return _response(data)
if path == "/.well-known/a2a/schemas.json":
data = {
"index": agent_name + "_index",
"version": schemas["version"],
"input": {
"id": f"{agent_name}-input-v1",
"schema": schemas["input_doc"]["schema"],
},
"output": {
"id": f"{agent_name}-output-v1",
"schema": schemas["output_doc"]["schema"],
},
}
return _response(data)
return None
def _compile_routes():
global _compiled_routes
if _compiled_routes:
return _compiled_routes
all_routes = (
agent_env_routes
+ agent_mcp_routes
+ agent_card_routes
+ agent_schema_routes
+ agent_vectorizer_routes
+ agent_prompts_route
+ agent_schemageneration_route
+ agent_capability_routes
+ test_payloads_route
)
compiled = []
for rt in all_routes:
pattern = re.sub(r"{(\w+)(?::[^}]+)?}", r"(?P<\1>[^/]+)", rt.path)
regex = re.compile(f"^{pattern}$")
compiled.append((regex, rt))
_compiled_routes = compiled
return _compiled_routes
def _match_route(method: str, path: str):
for regex, rt in _compile_routes():
if method.upper() not in rt.methods:
continue
m = regex.match(path)
if m:
return rt, m.groupdict()
return None, None
def _serve_ui(path: str):
# Simple static serving for packaged UI
if path == "/ui" or path == "/ui/":
return FileResponse("ui/dist/index.html")
if path.startswith("/ui/assets"):
rel = path.replace("/ui/assets/", "")
return FileResponse(os.path.join("ui/dist/assets", rel))
if path.startswith("/ui/"):
return FileResponse("ui/dist/index.html")
return None
def _to_lambda_response(resp: Any):
if isinstance(resp, Response):
body = resp.body
is_binary = isinstance(body, (bytes, bytearray))
return {
"statusCode": resp.status_code,
"headers": resp.headers,
"body": base64.b64encode(body).decode("utf-8") if is_binary else (body if isinstance(body, str) else json.dumps(body)),
"isBase64Encoded": is_binary,
}
if isinstance(resp, dict) and {"statusCode", "body"} <= set(resp.keys()):
return resp
if isinstance(resp, (dict, list)):
return _response(resp)
return _response(str(resp))
async def _async_handler(event, context, response_stream=None):
agent_name = context.function_name if context else "agent"
raw_path = event.get("rawPath") or event.get("path") or "/"
path = raw_path
if path.startswith("/$default"):
path = path[len("/$default") :] or "/"
base_url = (
f"{event.get('headers', {}).get('x-forwarded-proto', 'https')}://"
f"{event.get('headers', {}).get('host', '')}"
)
method = event.get("requestContext", {}).get("http", {}).get("method", "").upper()
if method == "GET":
meta_resp = _handle_well_known(path, base_url, agent_name)
if meta_resp:
return meta_resp
ui_resp = _serve_ui(path)
if ui_resp:
return _to_lambda_response(ui_resp)
rt, params = _match_route(method, path)
if rt:
headers = event.get("headers") or {}
body_raw = event.get("body")
if event.get("isBase64Encoded"):
body_raw = base64.b64decode(body_raw or b"")
try:
parsed_body = json.loads(body_raw) if isinstance(body_raw, (str, bytes, bytearray)) else body_raw
except Exception:
parsed_body = body_raw
req = Request(method=method, url=base_url + path, headers=headers, path_params=params or {}, body=parsed_body)
resp = await rt.endpoint(req)
return _to_lambda_response(resp)
app = _build_lambda_app(agent_name)
result = await app.handle_event(event, context, response_stream=response_stream)
if result is not None:
return result
return _response({}, status=200)
def handler(event, context, response_stream=None):
return asyncio.run(
_async_handler(event, context, response_stream=response_stream)
)