hgzero/rag/src/api/term_routes.py

"""
용어 관련 API 엔드포인트
"""
from fastapi import APIRouter, HTTPException
from sse_starlette.sse import EventSourceResponse
import asyncio
import json
import logging

from ..services.sse_manager import sse_manager

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/rag/terms", tags=["terms"])


@router.get("/stream/{session_id}")
async def stream_terms(session_id: str):
    """
    용어 검색 결과 SSE 스트림

    Args:
        session_id: 회의 세션 ID

    Returns:
        SSE 스트림
    """
    try:
        # SSE 연결 등록
        queue = sse_manager.register(session_id)
        logger.info(f"용어 스트림 시작: {session_id}")

        async def event_generator():
            """SSE 이벤트 생성기"""
            try:
                # 연결 확인 메시지
                yield {
                    "event": "connected",
                    "data": json.dumps({"session_id": session_id, "status": "connected"})
                }

                # 메시지 수신 및 전송
                while True:
                    try:
                        # Timeout을 두어 주기적으로 heartbeat 전송
                        message = await asyncio.wait_for(queue.get(), timeout=30.0)

                        yield {
                            "event": message["event"],
                            "data": json.dumps(message["data"])
                        }

                    except asyncio.TimeoutError:
                        # Heartbeat 전송
                        yield {
                            "event": "heartbeat",
                            "data": json.dumps({"type": "heartbeat"})
                        }

            except asyncio.CancelledError:
                logger.info(f"용어 스트림 취소됨: {session_id}")
            except Exception as e:
                logger.error(f"이벤트 생성 중 에러: {str(e)}")
            finally:
                # 연결 정리
                sse_manager.unregister(session_id)
                logger.info(f"용어 스트림 종료: {session_id}")

        return EventSourceResponse(event_generator())

    except ValueError as e:
        raise HTTPException(status_code=429, detail=str(e))
    except Exception as e:
        logger.error(f"스트림 시작 실패: {str(e)}")
        raise HTTPException(status_code=500, detail="스트림 시작 실패")


@router.get("/stream/{session_id}/status")
async def get_stream_status(session_id: str):
    """
    스트림 연결 상태 확인

    Args:
        session_id: 회의 세션 ID

    Returns:
        연결 상태
    """
    is_connected = sse_manager.is_connected(session_id)

    return {
        "session_id": session_id,
        "connected": is_connected
    }