141 lines
5.1 KiB
Python
141 lines
5.1 KiB
Python
"""A2A 网关接入客户端：Agent 服务启动后自动注册 + 心跳保活 + 注销。
|
||
|
||
用法（在 api_server.py 的 main() 中）:
|
||
from gateway_client import register_agent, start_heartbeat, unregister_agent
|
||
|
||
if register_agent(dev_app.name, f"http://127.0.0.1:{PORT}"):
|
||
start_heartbeat(dev_app.name)
|
||
|
||
环境变量:
|
||
GATEWAY_URL 网关地址，默认 http://127.0.0.1:8000
|
||
GATEWAY_AUTH 网关认证密码，需与网关 GATEWAY_AUTH 一致
|
||
AGENT_TAGS 能力标签，逗号分隔，默认 code
|
||
AGENT_MAX_CONCURRENT 最大并发数，默认 1
|
||
HEARTBEAT_INTERVAL 心跳间隔秒，默认 10
|
||
"""
|
||
import logging
|
||
import os
|
||
import threading
|
||
import time
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
GATEWAY_URL = os.getenv("GATEWAY_URL", "http://127.0.0.1:8000").rstrip("/")
|
||
GATEWAY_AUTH = os.getenv("GATEWAY_AUTH", "dev-gateway-auth")
|
||
HEARTBEAT_INTERVAL = int(os.getenv("HEARTBEAT_INTERVAL", "10"))
|
||
|
||
_client = httpx.Client(timeout=5.0)
|
||
|
||
|
||
def _parse_tags() -> list[str]:
|
||
raw = os.getenv("AGENT_TAGS", "code")
|
||
return [t.strip() for t in raw.split(",") if t.strip()]
|
||
|
||
|
||
def register_agent(agent_id: str, endpoint: str, agent_tags: list[str] | None = None,
|
||
max_concurrent: int | None = None) -> bool:
|
||
"""向网关注册本 Agent。成功返回 True，失败（网关未启动/认证失败）返回 False。"""
|
||
tags = agent_tags if agent_tags is not None else _parse_tags()
|
||
max_conc = max_concurrent or int(os.getenv("AGENT_MAX_CONCURRENT", "1"))
|
||
body = {
|
||
"auth": GATEWAY_AUTH,
|
||
"agent_id": agent_id,
|
||
"endpoint": endpoint,
|
||
"agent_tags": tags,
|
||
"max_concurrent": max_conc,
|
||
"current_load": 0,
|
||
}
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/register", json=body)
|
||
if resp.status_code == 200:
|
||
logger.info("registered to gateway agent=%s endpoint=%s tags=%s", agent_id, endpoint, tags)
|
||
return True
|
||
logger.error("register failed agent=%s status=%s body=%s", agent_id, resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("register network error agent=%s err=%s", agent_id, e)
|
||
return False
|
||
|
||
|
||
def _heartbeat_once(agent_id: str, current_load: int) -> bool:
|
||
try:
|
||
resp = _client.post(
|
||
f"{GATEWAY_URL}/api/agent/heartbeat",
|
||
json={"agent_id": agent_id, "current_load": current_load},
|
||
)
|
||
if resp.status_code == 200:
|
||
return True
|
||
logger.warning("heartbeat failed status=%s body=%s", resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.warning("heartbeat network error err=%s", e)
|
||
return False
|
||
|
||
|
||
def _heartbeat_loop(agent_id: str) -> None:
|
||
while True:
|
||
try:
|
||
load = get_current_load()
|
||
except Exception:
|
||
load = 0
|
||
_heartbeat_once(agent_id, load)
|
||
time.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
|
||
def start_heartbeat(agent_id: str) -> threading.Thread:
|
||
"""启动后台心跳线程（daemon，随进程退出）。"""
|
||
t = threading.Thread(target=_heartbeat_loop, args=(agent_id,), name=f"gateway-hb-{agent_id}", daemon=True)
|
||
t.start()
|
||
logger.info("heartbeat started agent=%s interval=%ss", agent_id, HEARTBEAT_INTERVAL)
|
||
return t
|
||
|
||
|
||
def unregister_agent(agent_id: str) -> bool:
|
||
"""向网关注销本 Agent。"""
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/unregister", json={"agent_id": agent_id})
|
||
if resp.status_code in (200, 404):
|
||
logger.info("unregistered from gateway agent=%s", agent_id)
|
||
return True
|
||
logger.warning("unregister failed status=%s body=%s", resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("unregister network error agent=%s err=%s", agent_id, e)
|
||
return False
|
||
|
||
|
||
def report_result(request_id: str, agent_id: str, status: str = "success",
|
||
progress: int = 100, result: dict | None = None,
|
||
error_info: str | None = None) -> bool:
|
||
"""任务执行完成后，将结果回传给网关（POST /api/agent/result）。
|
||
|
||
Args:
|
||
request_id: 网关分配的任务 ID
|
||
agent_id: 本 Agent ID
|
||
status: success / failed
|
||
progress: 进度百分比
|
||
result: 结果字典（可选）
|
||
error_info: 错误信息（失败时必填）
|
||
"""
|
||
body = {
|
||
"request_id": request_id,
|
||
"agent_id": agent_id,
|
||
"status": status,
|
||
"progress": progress,
|
||
"result": result,
|
||
"error_info": error_info,
|
||
}
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/result", json=body)
|
||
if resp.status_code == 200:
|
||
logger.info("result reported request=%s status=%s", request_id, status)
|
||
return True
|
||
logger.error("report result failed request=%s status=%s body=%s", request_id, resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("report result network error request=%s err=%s", request_id, e)
|
||
return False
|
||
|
||
|
||
def get_current_load() -> int:
|
||
"""Agent 当前负载。子类/调用方可覆写以报告真实并发数。"""
|
||
return 0
|