143 lines
5.1 KiB
Python
143 lines
5.1 KiB
Python
"""龙虾 Agent HTTP 客户端。
|
|
|
|
运行时配置来源优先级：DB (lobster_config 表) > .env (settings.LOBSTER_*)。
|
|
调用方通过 `build_lobster_client(db)` 或 `build_lobster_client_sync(api_base, api_key, timeout)`
|
|
获取实例；不再使用模块级单例，便于配置热更新。
|
|
|
|
- 当 api_base / api_key 任一为空时走 mock（便于本地自测）。
|
|
- 真调协议：POST {api_base}/runs 提交、GET {api_base}/runs/{run_id} 查状态。
|
|
实际字段由联调时与龙虾平台对齐。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import random
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from app.core.logging import logger
|
|
|
|
|
|
@dataclass
|
|
class LobsterRunResult:
|
|
run_id: str
|
|
status: str # QUEUED / RUNNING / SUCCEEDED / FAILED / TIMEOUT
|
|
output: dict[str, Any] | None = None
|
|
logs: list[dict[str, Any]] | None = None # [{"event_id": str, "level": str, "message": str}]
|
|
message: str | None = None
|
|
|
|
|
|
# 模块级共享 mock 状态机（同一 run_id 多次 poll 可继续推进）
|
|
_MOCK_RUNS: dict[str, dict[str, Any]] = {}
|
|
|
|
|
|
class LobsterClient:
|
|
"""按运行时配置构造；API_BASE / API_KEY 任一为空时退化为 mock。"""
|
|
|
|
def __init__(self, api_base: str = "", api_key: str = "", timeout: int = 30) -> None:
|
|
self.base = (api_base or "").rstrip("/")
|
|
self.key = api_key or ""
|
|
self.timeout = timeout or 30
|
|
|
|
@property
|
|
def is_mock(self) -> bool:
|
|
return not self.base or not self.key
|
|
|
|
async def submit(
|
|
self, template_id: str, params: dict[str, Any], task_no: str
|
|
) -> LobsterRunResult:
|
|
if self.is_mock:
|
|
return await self._mock_submit(template_id, params, task_no)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {self.key}",
|
|
"X-Loom-Task": task_no,
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {"template_id": template_id, "input": params, "task_no": task_no}
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
r = await client.post(f"{self.base}/runs", json=payload, headers=headers)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
return LobsterRunResult(
|
|
run_id=str(data.get("run_id") or data.get("id") or ""),
|
|
status=str(data.get("status") or "QUEUED"),
|
|
message=data.get("message"),
|
|
)
|
|
|
|
async def poll(self, run_id: str) -> LobsterRunResult:
|
|
if self.is_mock:
|
|
return await self._mock_poll(run_id)
|
|
|
|
headers = {"Authorization": f"Bearer {self.key}"}
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
r = await client.get(f"{self.base}/runs/{run_id}", headers=headers)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
return LobsterRunResult(
|
|
run_id=run_id,
|
|
status=str(data.get("status") or "RUNNING"),
|
|
output=data.get("output"),
|
|
logs=data.get("logs") or [],
|
|
message=data.get("message"),
|
|
)
|
|
|
|
# ---- mock ----
|
|
|
|
async def _mock_submit(
|
|
self, template_id: str, params: dict[str, Any], task_no: str
|
|
) -> LobsterRunResult:
|
|
run_id = f"mock-{task_no}-{int(time.time()*1000)}"
|
|
_MOCK_RUNS[run_id] = {
|
|
"created_at": time.time(),
|
|
"template_id": template_id,
|
|
"input": params,
|
|
"polls": 0,
|
|
}
|
|
logger.info("lobster_mock_submit", run_id=run_id, template_id=template_id)
|
|
await asyncio.sleep(0.05)
|
|
return LobsterRunResult(run_id=run_id, status="QUEUED")
|
|
|
|
async def _mock_poll(self, run_id: str) -> LobsterRunResult:
|
|
state = _MOCK_RUNS.get(run_id)
|
|
if state is None:
|
|
return LobsterRunResult(run_id=run_id, status="FAILED", message="mock run_id 未知")
|
|
state["polls"] += 1
|
|
if state["polls"] == 1:
|
|
return LobsterRunResult(
|
|
run_id=run_id,
|
|
status="RUNNING",
|
|
logs=[{"event_id": f"{run_id}-1", "level": "INFO", "message": "mock: 开始执行"}],
|
|
)
|
|
if state["polls"] >= 3:
|
|
output = {"mock": True, "input": state["input"], "random": random.randint(1, 99)}
|
|
return LobsterRunResult(
|
|
run_id=run_id,
|
|
status="SUCCEEDED",
|
|
output=output,
|
|
logs=[{"event_id": f"{run_id}-2", "level": "INFO", "message": "mock: 完成"}],
|
|
message="mock 成功",
|
|
)
|
|
return LobsterRunResult(
|
|
run_id=run_id,
|
|
status="RUNNING",
|
|
logs=[{"event_id": f"{run_id}-poll-{state['polls']}", "level": "INFO",
|
|
"message": f"mock: 第 {state['polls']} 次轮询"}],
|
|
)
|
|
|
|
|
|
async def build_lobster_client(db) -> LobsterClient:
|
|
"""根据当前 DB + .env 生效配置构造 client。调用方在执行任务前拿一次即可。"""
|
|
# 局部 import 避免模块循环依赖
|
|
from app.services import lobster_config_service
|
|
|
|
eff = await lobster_config_service.load_effective(db)
|
|
return LobsterClient(
|
|
api_base=eff.api_base,
|
|
api_key=eff.api_key,
|
|
timeout=eff.timeout_seconds,
|
|
)
|