126 lines
5.7 KiB
Python
126 lines
5.7 KiB
Python
import asyncio
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.contracts import RequestContext
|
|
from app.core.errors import RunNotFoundError
|
|
from app.infrastructure.db import SessionFactory
|
|
from app.model.conversation import ConversationMessage
|
|
from app.repository.conversation_repository import ConversationRepository
|
|
from app.service.financial_nl2sql_service import FinancialNL2SQLService
|
|
from app.core.nl2sql_contracts import FinancialNL2SQLInput
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunSnapshot:
|
|
run_id: str
|
|
trace_id: str
|
|
status: str
|
|
agent_type: str
|
|
session_id: str
|
|
result: dict[str, Any] | None
|
|
error_code: str | None
|
|
created_at: str
|
|
completed_at: str | None
|
|
|
|
|
|
class RunQueryService:
|
|
async def get(self, run_id: str, context: RequestContext) -> RunSnapshot:
|
|
recovery_question: str | None = None
|
|
async with SessionFactory() as session:
|
|
rows = await ConversationRepository(session).run_result(run_id, int(context.user_id))
|
|
if rows is None:
|
|
raise RunNotFoundError("运行不存在或不可见")
|
|
run, message = rows
|
|
result = None
|
|
if run.status == "succeeded" and message is not None:
|
|
# 「转人工标记」从 `tool_calls` 这个 JSON 列里取（与写入侧同一个位置）。
|
|
# 兼容两种历史形状：dict 里带 `transfer_required`（新），或 `calls` 裸列表（旧行）——
|
|
# 旧行取不到就按 False 处理，不猜、也不因为缺字段让整个响应失败。
|
|
transfer_required = False
|
|
transfer_reason = None
|
|
result_data: dict[str, Any] = {}
|
|
generated_sql: str | None = None
|
|
stored_calls = message.tool_calls
|
|
if isinstance(stored_calls, dict):
|
|
transfer_required = bool(stored_calls.get("transfer_required", False))
|
|
reason = stored_calls.get("transfer_reason")
|
|
transfer_reason = str(reason) if reason else None
|
|
stored_data = stored_calls.get("data")
|
|
if isinstance(stored_data, dict):
|
|
result_data = stored_data
|
|
stored_sql = stored_calls.get("sql")
|
|
if isinstance(stored_sql, str):
|
|
generated_sql = stored_sql
|
|
result = {"content": message.content, "tool_calls": stored_calls,
|
|
"intent": message.intent,
|
|
"confidence": str(message.confidence) if message.confidence else None,
|
|
"source_references": message.source_references or [],
|
|
# `docs/05` §6.3 规定 `result` 必须含这两个字段，此前未兑现。
|
|
# 前端据此判断"这轮要不要转人工"，不必再去猜兜底话术的开头。
|
|
"transfer_required": transfer_required,
|
|
"transfer_reason": transfer_reason}
|
|
if run.agent_type == "financial_nl2sql":
|
|
result["data"] = result_data
|
|
result["sql"] = generated_sql
|
|
if not result_data or not generated_sql:
|
|
request_message = await session.get(
|
|
ConversationMessage, run.request_message_id
|
|
)
|
|
if request_message is not None:
|
|
recovery_question = request_message.content
|
|
if recovery_question:
|
|
result = await self._recover_financial_result(
|
|
result, recovery_question, context
|
|
)
|
|
return RunSnapshot(
|
|
run.run_id, run.trace_id, run.status, run.agent_type, run.session_id, result,
|
|
run.error_code, run.created_at.isoformat() + "Z",
|
|
run.completed_at.isoformat() + "Z" if run.completed_at else None,
|
|
)
|
|
|
|
async def _recover_financial_result(
|
|
self, result: dict[str, Any], question: str, context: RequestContext
|
|
) -> dict[str, Any]:
|
|
"""兼容旧 Worker 只保存工具成功摘要、未保存查询载荷的历史运行。"""
|
|
try:
|
|
recovered = await FinancialNL2SQLService().query(
|
|
FinancialNL2SQLInput(question=question), context
|
|
)
|
|
except Exception:
|
|
logger.warning(
|
|
"历史 NL2SQL 结果补取失败，保留原运行摘要 trace_id=%s",
|
|
context.trace_id,
|
|
exc_info=True,
|
|
)
|
|
return result
|
|
if recovered.get("status") != "success":
|
|
return result
|
|
recovered_data = recovered.get("data")
|
|
recovered_sql = recovered.get("sql")
|
|
if isinstance(recovered_data, dict):
|
|
result["data"] = recovered_data
|
|
if isinstance(recovered_sql, str):
|
|
result["sql"] = recovered_sql
|
|
return result
|
|
|
|
async def watch(
|
|
self, initial: RunSnapshot, context: RequestContext
|
|
) -> AsyncIterator[RunSnapshot | None]:
|
|
settings = get_settings()
|
|
loop = asyncio.get_running_loop()
|
|
deadline = loop.time() + settings.sse_max_connection_seconds
|
|
snapshot = initial
|
|
while True:
|
|
yield snapshot
|
|
if snapshot.status in {"succeeded", "failed", "cancelled"} or loop.time() >= deadline:
|
|
return
|
|
yield None
|
|
await asyncio.sleep(max(0, min(settings.sse_heartbeat_seconds, deadline - loop.time())))
|
|
snapshot = await self.get(initial.run_id, context)
|