group_fqcd_jr/app/service/run_query_service.py

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 RunNotFoundError
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 RunNotFoundError("运行不存在或不可见")
            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)