"""GEO 项目的 Custom Handler — 供 AgentKit Server 使用 所有 Handler 通过 HTTP 回调 GEO Backend 的 /internal/ 端点,不直接访问 DB。 """ import logging import os import httpx from agentkit.core.protocol import TaskMessage logger = logging.getLogger(__name__) GEO_BACKEND_URL = os.getenv("GEO_BACKEND_URL", "http://localhost:8000") INTERNAL_API_TOKEN = os.getenv("INTERNAL_API_TOKEN", "") def _internal_headers() -> dict: """获取内部 API 请求头""" headers = {"Content-Type": "application/json"} if INTERNAL_API_TOKEN: headers["X-Internal-Token"] = INTERNAL_API_TOKEN return headers async def handle_citation_task(task: TaskMessage) -> dict: """引用检测任务 — 通过 HTTP 回调 GEO Backend task_type 路由: - citation_detect: POST /internal/citation/detect - citation_detect_single: POST /internal/citation/detect-single """ if task.task_type == "citation_detect": return await _call_internal("/internal/citation/detect", task.input_data) elif task.task_type == "citation_detect_single": return await _call_internal("/internal/citation/detect-single", task.input_data) else: raise ValueError(f"Unsupported task type: {task.task_type}") async def handle_monitor_task(task: TaskMessage) -> dict: """效果追踪任务 — 通过 HTTP 回调 GEO Backend task_type 路由: - monitor_track: POST /internal/monitor/track - monitor_check_single: POST /internal/monitor/check-single """ if task.task_type == "monitor_track": return await _call_internal("/internal/monitor/track", task.input_data) elif task.task_type == "monitor_check_single": return await _call_internal("/internal/monitor/check-single", task.input_data) else: raise ValueError(f"Unsupported task type: {task.task_type}") async def handle_schema_task(task: TaskMessage) -> dict: """Schema 建议任务 — 通过 HTTP 回调 GEO Backend task_type 路由: - schema_advise: POST /internal/schema/advise """ if task.task_type == "schema_advise": return await _call_internal("/internal/schema/advise", task.input_data) else: raise ValueError(f"Unsupported task type: {task.task_type}") async def _call_internal(path: str, input_data: dict) -> dict: """调用 GEO Backend 内部 API""" try: async with httpx.AsyncClient(timeout=300.0) as client: resp = await client.post( f"{GEO_BACKEND_URL}{path}", json=input_data, headers=_internal_headers(), ) resp.raise_for_status() return resp.json() except httpx.HTTPStatusError as e: logger.error(f"HTTP error calling {path}: {e.response.status_code} {e.response.text[:500]}") return {"error": f"HTTP {e.response.status_code}", "detail": e.response.text[:500]} except Exception as e: logger.error(f"Error calling {path}: {e}") return {"error": str(e)}