86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
"""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)}
|