59 lines
2.3 KiB
Python
59 lines
2.3 KiB
Python
import asyncio
|
|
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 ResourceNotFoundError
|
|
from app.infrastructure.db import SessionFactory
|
|
from app.repository.conversation_repository import ConversationRepository
|
|
|
|
|
|
@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:
|
|
async with SessionFactory() as session:
|
|
rows = await ConversationRepository(session).run_result(run_id, int(context.user_id))
|
|
if rows is None:
|
|
raise ResourceNotFoundError("运行不存在")
|
|
run, message = rows
|
|
result = None
|
|
if run.status == "succeeded" and message is not None:
|
|
result = {"content": message.content, "tool_calls": message.tool_calls,
|
|
"intent": message.intent,
|
|
"confidence": str(message.confidence) if message.confidence else None,
|
|
"source_references": message.source_references or []}
|
|
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 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)
|