440 lines
21 KiB
Python
440 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import re
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from deepagents import create_deep_agent
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_core.tools import StructuredTool
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import BaseModel, Field
|
|
|
|
from a2a_pack import (
|
|
A2AAgent,
|
|
ConsumerSetup,
|
|
ConsumerSetupField,
|
|
ConsumerSetupMissing,
|
|
EgressPolicy,
|
|
LLMProvisioning,
|
|
Pricing,
|
|
Resources,
|
|
RunContext,
|
|
skill,
|
|
)
|
|
|
|
|
|
DEFAULT_BASE_URL = "https://tasks.a2acloud.io"
|
|
OPERATIONS = json.loads("{\n \"create_task\": {\n \"description\": \"Create Task\",\n \"destructive\": true,\n \"method\": \"POST\",\n \"operation_id\": \"create_task\",\n \"parameters\": [],\n \"path\": \"/api/tasks\",\n \"request_body\": {\n \"$ref\": \"#/components/schemas/TaskCreate\"\n },\n \"security\": [\n {\n \"bearerAuth\": []\n },\n {\n \"apiKeyAuth\": []\n }\n ],\n \"skill_name\": \"create_task\",\n \"summary\": \"Create Task\",\n \"tags\": [\n \"Tasks\"\n ]\n },\n \"delete_task\": {\n \"description\": \"Delete Task\",\n \"destructive\": true,\n \"method\": \"DELETE\",\n \"operation_id\": \"delete_task\",\n \"parameters\": [\n {\n \"description\": \"\",\n \"in\": \"path\",\n \"name\": \"task_id\",\n \"required\": true,\n \"schema\": {\n \"title\": \"Task Id\",\n \"type\": \"string\"\n }\n }\n ],\n \"path\": \"/api/tasks/{task_id}\",\n \"request_body\": null,\n \"security\": [\n {\n \"bearerAuth\": []\n },\n {\n \"apiKeyAuth\": []\n }\n ],\n \"skill_name\": \"delete_task\",\n \"summary\": \"Delete Task\",\n \"tags\": [\n \"Tasks\"\n ]\n },\n \"get_task\": {\n \"description\": \"Get Task\",\n \"destructive\": false,\n \"method\": \"GET\",\n \"operation_id\": \"get_task\",\n \"parameters\": [\n {\n \"description\": \"\",\n \"in\": \"path\",\n \"name\": \"task_id\",\n \"required\": true,\n \"schema\": {\n \"title\": \"Task Id\",\n \"type\": \"string\"\n }\n }\n ],\n \"path\": \"/api/tasks/{task_id}\",\n \"request_body\": null,\n \"security\": [\n {\n \"bearerAuth\": []\n },\n {\n \"apiKeyAuth\": []\n }\n ],\n \"skill_name\": \"get_task\",\n \"summary\": \"Get Task\",\n \"tags\": [\n \"Tasks\"\n ]\n },\n \"get_tasks_health\": {\n \"description\": \"Healthz\",\n \"destructive\": false,\n \"method\": \"GET\",\n \"operation_id\": \"get_tasks_health\",\n \"parameters\": [],\n \"path\": \"/healthz\",\n \"request_body\": null,\n \"security\": null,\n \"skill_name\": \"get_tasks_health\",\n \"summary\": \"Healthz\",\n \"tags\": [\n \"Health\"\n ]\n },\n \"list_tasks\": {\n \"description\": \"List Tasks\",\n \"destructive\": false,\n \"method\": \"GET\",\n \"operation_id\": \"list_tasks\",\n \"parameters\": [\n {\n \"description\": \"\",\n \"in\": \"query\",\n \"name\": \"status\",\n \"required\": false,\n \"schema\": {\n \"anyOf\": [\n {\n \"enum\": [\n \"todo\",\n \"in_progress\",\n \"done\",\n \"canceled\"\n ],\n \"type\": \"string\"\n },\n {\n \"type\": \"null\"\n }\n ],\n \"title\": \"Status\"\n }\n },\n {\n \"description\": \"\",\n \"in\": \"query\",\n \"name\": \"q\",\n \"required\": false,\n \"schema\": {\n \"anyOf\": [\n {\n \"maxLength\": 120,\n \"minLength\": 1,\n \"type\": \"string\"\n },\n {\n \"type\": \"null\"\n }\n ],\n \"title\": \"Q\"\n }\n },\n {\n \"description\": \"\",\n \"in\": \"query\",\n \"name\": \"due_before\",\n \"required\": false,\n \"schema\": {\n \"anyOf\": [\n {\n \"format\": \"date-time\",\n \"type\": \"string\"\n },\n {\n \"type\": \"null\"\n }\n ],\n \"title\": \"Due Before\"\n }\n },\n {\n \"description\": \"\",\n \"in\": \"query\",\n \"name\": \"limit\",\n \"required\": false,\n \"schema\": {\n \"default\": 50,\n \"maximum\": 200,\n \"minimum\": 1,\n \"title\": \"Limit\",\n \"type\": \"integer\"\n }\n },\n {\n \"description\": \"\",\n \"in\": \"query\",\n \"name\": \"offset\",\n \"required\": false,\n \"schema\": {\n \"default\": 0,\n \"minimum\": 0,\n \"title\": \"Offset\",\n \"type\": \"integer\"\n }\n }\n ],\n \"path\": \"/api/tasks\",\n \"request_body\": null,\n \"security\": [\n {\n \"bearerAuth\": []\n },\n {\n \"apiKeyAuth\": []\n }\n ],\n \"skill_name\": \"list_tasks\",\n \"summary\": \"List Tasks\",\n \"tags\": [\n \"Tasks\"\n ]\n },\n \"update_task\": {\n \"description\": \"Update Task\",\n \"destructive\": true,\n \"method\": \"PATCH\",\n \"operation_id\": \"update_task\",\n \"parameters\": [\n {\n \"description\": \"\",\n \"in\": \"path\",\n \"name\": \"task_id\",\n \"required\": true,\n \"schema\": {\n \"title\": \"Task Id\",\n \"type\": \"string\"\n }\n }\n ],\n \"path\": \"/api/tasks/{task_id}\",\n \"request_body\": {\n \"$ref\": \"#/components/schemas/TaskUpdate\"\n },\n \"security\": [\n {\n \"bearerAuth\": []\n },\n {\n \"apiKeyAuth\": []\n }\n ],\n \"skill_name\": \"update_task\",\n \"summary\": \"Update Task\",\n \"tags\": [\n \"Tasks\"\n ]\n }\n}")
|
|
ROOT_SECURITY = json.loads("[]")
|
|
SECURITY_SCHEMES = json.loads("{\n \"apiKeyAuth\": {\n \"description\": \"Use the configured `TASKS_API_KEY` value.\",\n \"in\": \"header\",\n \"name\": \"x-api-key\",\n \"type\": \"apiKey\"\n },\n \"bearerAuth\": {\n \"description\": \"Use `Authorization: Bearer $TASKS_API_KEY`.\",\n \"scheme\": \"bearer\",\n \"type\": \"http\"\n }\n}")
|
|
SECURITY_FIELDS = json.loads("{\n \"apiKeyAuth\": {\n \"field\": \"X_API_KEY\",\n \"kind\": \"apiKey\",\n \"location\": \"header\",\n \"name\": \"x-api-key\"\n },\n \"bearerAuth\": {\n \"field\": \"BEARERAUTH_TOKEN\",\n \"kind\": \"http\",\n \"scheme\": \"bearer\"\n }\n}")
|
|
PATH_PARAMETER_RE = re.compile(r"{([^}/]+)}")
|
|
|
|
|
|
class OperationInput(BaseModel):
|
|
parameters: dict[str, Any] = Field(
|
|
default_factory=dict,
|
|
description="Path, query, header, and cookie parameters keyed by OpenAPI parameter name.",
|
|
)
|
|
body: Any | None = Field(default=None, description="JSON request body, when the operation accepts one.")
|
|
|
|
|
|
class TasksApi(A2AAgent):
|
|
name = "tasks-api"
|
|
description = "API-only service for tracking tasks in Postgres."
|
|
version = "0.1.0"
|
|
consumer_setup = ConsumerSetup.from_fields(
|
|
ConsumerSetupField.config("OPENAPI_BASE_URL", label="API base URL", description="Override the default API server (https://tasks.a2acloud.io).", required=False, input_type="url"),
|
|
ConsumerSetupField.secret("BEARERAUTH_TOKEN", label="bearerAuth bearer credential", description="Use `Authorization: Bearer $TASKS_API_KEY`.", required=True),
|
|
ConsumerSetupField.secret("X_API_KEY", label="apiKeyAuth API key", description="Use the configured `TASKS_API_KEY` value.", required=True),
|
|
)
|
|
llm_provisioning = LLMProvisioning.PLATFORM
|
|
pricing = Pricing(
|
|
price_per_call_usd=0.0,
|
|
caller_pays_llm=False,
|
|
notes="Uses a scoped platform LLM grant through ctx.llm.",
|
|
)
|
|
resources = Resources(cpu="200m", memory="512Mi")
|
|
egress = EgressPolicy(
|
|
allow_hosts=('tasks.a2acloud.io',),
|
|
deny_internet_by_default=True,
|
|
)
|
|
tools_used = ("openapi", "deepagents")
|
|
capabilities = {
|
|
"openapi_auto_agent": {
|
|
"operation_count": len(OPERATIONS),
|
|
"default_base_url": DEFAULT_BASE_URL,
|
|
"security_schemes": list(SECURITY_SCHEMES),
|
|
}
|
|
}
|
|
|
|
@skill(
|
|
name="auto",
|
|
description="Use the OpenAPI service to complete a natural-language goal.",
|
|
tags=("openapi", "auto"),
|
|
timeout_seconds=180,
|
|
)
|
|
async def auto(self, ctx: RunContext, goal: str) -> dict[str, Any]:
|
|
creds = ctx.llm
|
|
if not creds.api_key:
|
|
return {
|
|
"error": "llm_credentials_missing",
|
|
"final": (
|
|
"LLM credentials were not available for this run. Hosted "
|
|
"generated OpenAPI agents should receive a platform LLM grant."
|
|
),
|
|
"messages": [],
|
|
}
|
|
model_kwargs: dict[str, Any] = {
|
|
"model": creds.model,
|
|
"base_url": creds.base_url,
|
|
"api_key": creds.api_key,
|
|
}
|
|
if creds.temperature_mode != "omit" and creds.temperature is not None:
|
|
model_kwargs["temperature"] = creds.temperature
|
|
if creds.extra_body:
|
|
model_kwargs["extra_body"] = dict(creds.extra_body)
|
|
model = ChatOpenAI(**model_kwargs)
|
|
graph = create_deep_agent(
|
|
model=model,
|
|
tools=self._operation_tools(ctx),
|
|
system_prompt=self._system_prompt(),
|
|
)
|
|
result = await graph.ainvoke({"messages": [{"role": "user", "content": goal}]})
|
|
messages = result.get("messages", []) if isinstance(result, dict) else []
|
|
final = _message_text(messages[-1]) if messages else result
|
|
return {
|
|
"final": final,
|
|
"messages": [_message_to_dict(message) for message in messages[-8:]],
|
|
}
|
|
|
|
@skill(
|
|
name="get_tasks_health",
|
|
description="GET /healthz - Healthz",
|
|
tags=('Health',),
|
|
timeout_seconds=60,
|
|
)
|
|
async def get_tasks_health(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"get_tasks_health",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
@skill(
|
|
name="list_tasks",
|
|
description="GET /api/tasks - List Tasks",
|
|
tags=('Tasks',),
|
|
timeout_seconds=60,
|
|
)
|
|
async def list_tasks(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"list_tasks",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
@skill(
|
|
name="create_task",
|
|
description="POST /api/tasks - Create Task",
|
|
tags=('Tasks',),
|
|
timeout_seconds=120,
|
|
)
|
|
async def create_task(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"create_task",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
@skill(
|
|
name="get_task",
|
|
description="GET /api/tasks/{task_id} - Get Task",
|
|
tags=('Tasks',),
|
|
timeout_seconds=60,
|
|
)
|
|
async def get_task(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"get_task",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
@skill(
|
|
name="update_task",
|
|
description="PATCH /api/tasks/{task_id} - Update Task",
|
|
tags=('Tasks',),
|
|
timeout_seconds=120,
|
|
)
|
|
async def update_task(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"update_task",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
@skill(
|
|
name="delete_task",
|
|
description="DELETE /api/tasks/{task_id} - Delete Task",
|
|
tags=('Tasks',),
|
|
timeout_seconds=120,
|
|
)
|
|
async def delete_task(
|
|
self,
|
|
ctx: RunContext,
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
"delete_task",
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
def _operation_tools(self, ctx: RunContext) -> list[StructuredTool]:
|
|
tools: list[StructuredTool] = []
|
|
for operation_id, operation in OPERATIONS.items():
|
|
async def call(
|
|
parameters: dict[str, Any] | None = None,
|
|
body: Any | None = None,
|
|
*,
|
|
_operation_id: str = operation_id,
|
|
) -> dict[str, Any]:
|
|
return await self._request(
|
|
ctx,
|
|
_operation_id,
|
|
parameters=parameters or {},
|
|
body=body,
|
|
)
|
|
|
|
tools.append(
|
|
StructuredTool.from_function(
|
|
coroutine=call,
|
|
name=operation["skill_name"],
|
|
description=self._tool_description(operation),
|
|
args_schema=OperationInput,
|
|
)
|
|
)
|
|
return tools
|
|
|
|
def _system_prompt(self) -> str:
|
|
lines = [
|
|
"You operate an API through generated OpenAPI tools.",
|
|
"Call tools to get real results. Do not invent API responses.",
|
|
"For write, update, or delete operations, explain the intended action before calling the tool.",
|
|
"If an operation reports missing consumer setup, tell the user which setup field is required.",
|
|
"",
|
|
"Available operations:",
|
|
]
|
|
for operation in OPERATIONS.values():
|
|
marker = "WRITE" if operation.get("destructive") else "READ"
|
|
lines.append(
|
|
f"- {operation['skill_name']}: {marker} {operation['method']} {operation['path']} — {operation['summary']}"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
def _tool_description(self, operation: dict[str, Any]) -> str:
|
|
parts = [
|
|
f"{operation['method']} {operation['path']}",
|
|
str(operation.get("description") or operation.get("summary") or ""),
|
|
]
|
|
if operation.get("parameters"):
|
|
parts.append(
|
|
"Parameters: "
|
|
+ ", ".join(
|
|
f"{p['name']} in {p['in']}{' required' if p.get('required') else ''}"
|
|
for p in operation["parameters"]
|
|
)
|
|
)
|
|
if operation.get("request_body") is not None:
|
|
parts.append("Accepts a JSON request body.")
|
|
return "\n".join(part for part in parts if part)
|
|
|
|
async def _request(
|
|
self,
|
|
ctx: RunContext,
|
|
operation_id: str,
|
|
*,
|
|
parameters: dict[str, Any],
|
|
body: Any | None,
|
|
) -> dict[str, Any]:
|
|
operation = OPERATIONS[operation_id]
|
|
base_url = str(ctx.consumer_config("OPENAPI_BASE_URL", DEFAULT_BASE_URL) or DEFAULT_BASE_URL).rstrip("/")
|
|
url, query, headers = self._request_parts(ctx, operation, parameters)
|
|
request_kwargs: dict[str, Any] = {
|
|
"params": query,
|
|
"headers": headers,
|
|
}
|
|
if body is not None:
|
|
request_kwargs["json"] = body
|
|
async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
|
|
response = await client.request(
|
|
operation["method"],
|
|
f"{base_url}{url}",
|
|
**request_kwargs,
|
|
)
|
|
content_type = response.headers.get("content-type", "")
|
|
try:
|
|
payload: Any = response.json() if "json" in content_type else response.text
|
|
except ValueError:
|
|
payload = response.text
|
|
if response.status_code >= 400:
|
|
return {
|
|
"ok": False,
|
|
"status_code": response.status_code,
|
|
"operation_id": operation_id,
|
|
"error": payload,
|
|
}
|
|
return {
|
|
"ok": True,
|
|
"status_code": response.status_code,
|
|
"operation_id": operation_id,
|
|
"result": payload,
|
|
}
|
|
|
|
def _request_parts(
|
|
self,
|
|
ctx: RunContext,
|
|
operation: dict[str, Any],
|
|
parameters: dict[str, Any],
|
|
) -> tuple[str, dict[str, Any], dict[str, str]]:
|
|
path = operation["path"]
|
|
query: dict[str, Any] = {}
|
|
headers: dict[str, str] = {}
|
|
cookies: dict[str, Any] = {}
|
|
for param in operation.get("parameters") or []:
|
|
name = param["name"]
|
|
if name not in parameters:
|
|
if param.get("required"):
|
|
raise ValueError(f"missing required parameter {name!r}")
|
|
continue
|
|
value = parameters[name]
|
|
location = param["in"]
|
|
if location == "path":
|
|
path = path.replace("{" + name + "}", str(value))
|
|
elif location == "query":
|
|
query[name] = value
|
|
elif location == "header":
|
|
headers[name] = str(value)
|
|
elif location == "cookie":
|
|
cookies[name] = value
|
|
unresolved = PATH_PARAMETER_RE.findall(path)
|
|
if unresolved:
|
|
missing = ", ".join(sorted(set(unresolved)))
|
|
raise ValueError("missing required path parameter(s): " + missing)
|
|
if cookies:
|
|
headers["cookie"] = "; ".join(f"{key}={value}" for key, value in cookies.items())
|
|
auth_headers, auth_query = self._auth_for_operation(ctx, operation)
|
|
headers.update(auth_headers)
|
|
query.update(auth_query)
|
|
return path, query, headers
|
|
|
|
def _auth_for_operation(self, ctx: RunContext, operation: dict[str, Any]) -> tuple[dict[str, str], dict[str, Any]]:
|
|
requirements = operation.get("security")
|
|
if requirements is None:
|
|
requirements = ROOT_SECURITY
|
|
if requirements == []:
|
|
return {}, {}
|
|
missing: list[str] = []
|
|
for requirement in requirements or []:
|
|
if not isinstance(requirement, dict) or not requirement:
|
|
return {}, {}
|
|
headers: dict[str, str] = {}
|
|
query: dict[str, Any] = {}
|
|
ok = True
|
|
for scheme_name in requirement:
|
|
mapping = SECURITY_FIELDS.get(scheme_name) or {}
|
|
secret_name = mapping.get("field")
|
|
value = _consumer_secret_optional(ctx, secret_name) if secret_name else ""
|
|
if not value:
|
|
ok = False
|
|
if secret_name:
|
|
missing.append(secret_name)
|
|
continue
|
|
kind = mapping.get("kind")
|
|
if kind == "apiKey":
|
|
prefix = mapping.get("prefix") or ""
|
|
sent_value = f"{prefix}{value}"
|
|
if mapping.get("location") == "query":
|
|
query[mapping.get("name") or scheme_name] = sent_value
|
|
else:
|
|
headers[mapping.get("name") or scheme_name] = sent_value
|
|
elif kind in {"http", "oauth2"}:
|
|
scheme = mapping.get("scheme") or "Bearer"
|
|
if scheme.lower() == "basic":
|
|
token = base64.b64encode(value.encode("utf-8")).decode("ascii")
|
|
headers["authorization"] = f"Basic {token}"
|
|
else:
|
|
headers["authorization"] = f"{scheme} {value}"
|
|
else:
|
|
ok = False
|
|
if ok:
|
|
return headers, query
|
|
if missing:
|
|
raise ConsumerSetupMissing(
|
|
"operation requires consumer setup secret(s): "
|
|
+ ", ".join(sorted(set(missing)))
|
|
)
|
|
return {}, {}
|
|
|
|
|
|
def _consumer_secret_optional(ctx: RunContext, name: str | None) -> str:
|
|
if not name:
|
|
return ""
|
|
try:
|
|
return ctx.consumer_secret(name)
|
|
except ConsumerSetupMissing:
|
|
return ""
|
|
|
|
|
|
def _message_text(message: Any) -> Any:
|
|
content = getattr(message, "content", message)
|
|
return content
|
|
|
|
|
|
def _message_to_dict(message: Any) -> dict[str, Any]:
|
|
if isinstance(message, BaseMessage):
|
|
return message.model_dump(mode="json")
|
|
if hasattr(message, "model_dump"):
|
|
return message.model_dump(mode="json")
|
|
if isinstance(message, dict):
|
|
return message
|
|
return {"content": str(message)}
|
|
|
|
|
|
agent = TasksApi()
|