finrep-report/algo/src/finrep_algo_agent/api/routers/rag.py

from __future__ import annotations

from fastapi import APIRouter, File, Form, HTTPException, UploadFile

from finrep_algo_agent.api.deps import RagServiceDep, SettingsDep
from finrep_algo_agent.rag.ingestion import extract_text_from_upload
from finrep_algo_agent.schemas.rag import (
    RagDeleteResponse,
    RagDocumentIn,
    RagFileProcessResult,
    RagIngestFilesResponse,
    RagIngestRequest,
    RagIngestResponse,
    RagRetrieveRequest,
    RagRetrieveResponse,
)

router = APIRouter()


@router.post("/ingest-files", response_model=RagIngestFilesResponse)
async def rag_ingest_files(
    settings: SettingsDep,
    rag: RagServiceDep,
    task_id: str = Form(..., description="报告任务 ID，材料向量按 task 隔离"),
    replace: bool = Form(True),
    files: list[UploadFile] = File(..., description="业务上传材料，服务端解析文本后自动分块入库"),
) -> RagIngestFilesResponse:
    """文件直达 Python 时走本接口：解析 → 向量化 → 写入本任务 RAG 索引（无需 Java 先抽正文）。"""
    if not (settings.embedding_api_key or settings.llm_api_key):
        raise HTTPException(
            status_code=400,
            detail="RAG 入库需配置 FINREP_EMBEDDING_API_KEY 或 FINREP_LLM_API_KEY",
        )
    if not files:
        raise HTTPException(status_code=422, detail="files 不能为空")

    file_results: list[RagFileProcessResult] = []
    documents: list[RagDocumentIn] = []
    for uf in files:
        raw = await uf.read()
        ex = extract_text_from_upload(filename=uf.filename or "upload.bin", data=raw)
        fr = RagFileProcessResult(
            filename=ex.source_label,
            doc_id=ex.doc_id,
            characters=len(ex.text),
            skipped=not bool(ex.text),
            warning=ex.warning,
        )
        file_results.append(fr)
        if ex.text:
            documents.append(
                RagDocumentIn(
                    doc_id=ex.doc_id,
                    title="",
                    text=ex.text,
                    source_label=ex.source_label,
                )
            )

    if not documents:
        raise HTTPException(
            status_code=422,
            detail="所有文件均未解析出有效文本，未写入索引",
        )

    try:
        ing = await rag.ingest(task_id, documents, replace=replace)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"向量化服务异常: {e}") from e

    return RagIngestFilesResponse(
        task_id=ing.task_id,
        document_count=ing.document_count,
        chunk_count=ing.chunk_count,
        files=file_results,
    )


@router.post("/ingest", response_model=RagIngestResponse)
async def rag_ingest(
    body: RagIngestRequest,
    settings: SettingsDep,
    rag: RagServiceDep,
) -> RagIngestResponse:
    if not (settings.embedding_api_key or settings.llm_api_key):
        raise HTTPException(
            status_code=400,
            detail="RAG 入库需配置 FINREP_EMBEDDING_API_KEY 或 FINREP_LLM_API_KEY",
        )
    try:
        return await rag.ingest(
            body.task_id,
            body.documents,
            replace=body.replace,
        )
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"向量化服务异常: {e}") from e


@router.post("/retrieve", response_model=RagRetrieveResponse)
async def rag_retrieve(
    body: RagRetrieveRequest,
    settings: SettingsDep,
    rag: RagServiceDep,
) -> RagRetrieveResponse:
    if not (settings.embedding_api_key or settings.llm_api_key):
        raise HTTPException(
            status_code=400,
            detail="RAG 检索需配置 FINREP_EMBEDDING_API_KEY 或 FINREP_LLM_API_KEY",
        )
    try:
        return await rag.retrieve(
            body.task_id,
            body.query,
            top_k=body.top_k,
            min_score=body.min_score,
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"检索服务异常: {e}") from e


@router.delete("/{task_id}", response_model=RagDeleteResponse)
async def rag_delete_index(task_id: str, rag: RagServiceDep) -> RagDeleteResponse:
    deleted = rag.delete_index(task_id)
    return RagDeleteResponse(task_id=task_id, deleted=deleted)