AIRegulation-DocAnalysis/backend/app/api/routes/compliance.py

"""Define API routes for compliance."""

from __future__ import annotations

import asyncio
import json
from pathlib import Path
from typing import AsyncGenerator

from fastapi import APIRouter, File, UploadFile
from fastapi.responses import StreamingResponse

from app.schemas.compliance import (
    AnalyzeResponse,
    ComplianceChatRequest,
)
from app.services.mock_data import generate_task_id, get_mock_compliance_result
from app.shared.bootstrap import get_agent_conversation_service


router = APIRouter(prefix="/compliance", tags=["合规分析"])

tasks_store: dict[str, dict] = {}

RAW_DATA_DIR = Path(__file__).resolve().parents[3] / "data" / "raw"


@router.post("/analyze", response_model=AnalyzeResponse)
async def analyze_document(file: UploadFile = File(...)):
    """Handle analyze document."""
    task_id = generate_task_id()

    RAW_DATA_DIR.mkdir(parents=True, exist_ok=True)
    file_path = RAW_DATA_DIR / f"compliance_{task_id}_{file.filename}"

    content = await file.read()
    with file_path.open("wb") as f:
        f.write(content)

    tasks_store[task_id] = {
        "task_id": task_id,
        "file_path": str(file_path),
        "status": "processing",
        "result": None,
    }

    tasks_store[task_id]["status"] = "completed"
    tasks_store[task_id]["result"] = get_mock_compliance_result(task_id)

    return AnalyzeResponse(task_id=task_id)


@router.get("/result/{task_id}")
async def get_result(task_id: str):
    """Return result."""
    if task_id not in tasks_store:
        return get_mock_compliance_result(task_id)

    task = tasks_store[task_id]
    if task["status"] == "processing":
        return {"status": "processing", "message": "分析进行中"}
    return task["result"]


@router.post("/chat/{segment_id}")
async def compliance_chat(segment_id: int, request: ComplianceChatRequest):
    """Stream compliance Q&A grounded in real vector retrieval."""
    query = request.query
    if request.segment_context:
        query = f"[段落分析上下文]\n{request.segment_context}\n\n用户问题：{request.query}"

    _, event_stream = get_agent_conversation_service().stream_chat(
        query=query,
        top_k=5,
        prompt_template="compliance_qa",
    )

    async def generate() -> AsyncGenerator[str, None]:
        """Translate agent SSE events to compliance chunk/done format."""
        for event in event_stream:
            event_type = event.get("event", "")
            if event_type == "content":
                text = event.get("data", "")
                if text:
                    yield (
                        "event: message\n"
                        f"data: {json.dumps({'type': 'chunk', 'text': text}, ensure_ascii=False)}\n\n"
                    )
            elif event_type == "done":
                yield (
                    "event: message\n"
                    f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
                )
            await asyncio.sleep(0)

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
    )