PlantAI-rag-new/vstore/api.py

"""api.py —— vstore FastAPI 服务（8094）。

端点（docs/team/brief_dev1_p1.md）：
- POST   /api/vstore/ingest           multipart: file + name/category/description/tags/as_template
- GET    /api/vstore/search           ?q&top_k&category&as_template
- GET    /api/vstore/documents/{id}   明文全文 + outline + meta（仅生成时取）
- GET    /api/vstore/documents/{id}/style
- DELETE /api/vstore/documents/{id}
- POST   /api/vstore/rebuild-index
- GET    /api/vstore/health
FR-PRIV 新增（chunk 级索引，A1）：
- GET    /api/vstore/search-chunks    ?q&top_k&min_score&doc_ids&exclude_template
- POST   /api/vstore/rebuild-chunks
"""
from __future__ import annotations

import asyncio
import logging
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse

from . import vault as _vault
from .backend import BackendError, handshake, parse_backend
from .chunks import CHUNK_MIN_SCORE, CHUNK_TOP_K_MAX
from .settings import (
    DEFAULT_TOP_K,
    MAX_TOP_K,
    get_paths,
)
from .vault import VaultError

logger = logging.getLogger(__name__)

app = FastAPI(title="PlantAI vstore", version="0.1.0")


def _error(code: str, message: str, status: int, degraded: bool = False) -> JSONResponse:
    """错误信封（与 BFF 现有格式一致：{ok, error:{code, message, degraded}}）。"""
    return JSONResponse(
        status_code=status,
        content={"ok": False, "error": {"code": code, "message": message, "degraded": degraded}},
    )


# ---- 启动钩子：解析后端 + 握手校验（失败抛异常拒绝启动）----

def _init_service() -> None:
    paths = get_paths()
    backend = parse_backend(paths.vault_path, paths.index_path)
    try:
        handshake(backend)
    except BackendError as exc:
        logger.error("vstore: 启动握手失败，拒绝启动：%s", exc)
        raise
    _vault.set_backend(backend)
    _vault.init_vault()


@app.on_event("startup")
def _startup() -> None:
    _init_service()


# ---- 端点 ----

@app.get("/api/vstore/health")
async def health() -> dict[str, Any]:
    """健康检查：后端 ok + 库存量 + chunk 库存量（FR-PRIV-007）。"""
    vault = _vault.get_vault()
    count = vault.index.count() if hasattr(vault, "index") else len(vault.list_rows())
    chunks = getattr(vault, "chunks", None)
    chunk_count = chunks.count() if chunks is not None else 0
    return {"ok": True, "count": count, "backend": "local", "chunk_count": chunk_count}


@app.post("/api/vstore/ingest")
async def ingest(
    file: UploadFile = File(...),
    name: str = Form(""),
    category: str = Form(""),
    description: str = Form(""),
    tags: str = Form(""),
    as_template: bool = Form(False),
) -> dict[str, Any]:
    """multipart 入库。文件解析/嵌入/加密落盘/索引由 vault.add 完成。"""
    content = await file.read()
    try:
        item = _vault.get_vault().add(
            content,
            filename=file.filename or "document",
            meta={"name": name, "category": category, "description": description, "tags": tags},
            as_template=bool(as_template),
        )
    except VaultError as exc:
        return _error("VALIDATION", str(exc), 422)
    return {"ok": True, "item": item}


@app.get("/api/vstore/search")
async def search(
    q: str = "",
    top_k: int = DEFAULT_TOP_K,
    category: str = "",
    as_template: bool | None = None,
) -> dict[str, Any]:
    """三档降级检索。"""
    try:
        result = _vault.get_vault().search(
            q,
            top_k=max(1, min(top_k, MAX_TOP_K)),
            category=category or None,
            as_template=as_template,
        )
        return {"ok": True, **result}
    except VaultError as exc:
        return _error("INTERNAL", str(exc), 500, degraded=True)


@app.get("/api/vstore/documents/{doc_id}")
async def get_document(doc_id: str) -> dict[str, Any]:
    """明文全文 + outline + meta（仅生成时取）。"""
    doc = _vault.get_vault().get_document_json(doc_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="document not found")
    return {"ok": True, "document": doc}


@app.get("/api/vstore/documents/{doc_id}/style")
async def get_style(doc_id: str) -> dict[str, Any]:
    """docx 样式提取（导出对齐用）。"""
    style = _vault.get_vault().get_style(doc_id)
    if style is None:
        raise HTTPException(status_code=404, detail="document not found")
    return {"ok": True, "style_hints": style}


@app.delete("/api/vstore/documents/{doc_id}")
async def delete_document(doc_id: str) -> dict[str, Any]:
    """删除：索引行 + 密文 secure_delete。"""
    removed = _vault.get_vault().delete(doc_id)
    return {"ok": True, "removed": removed}


@app.post("/api/vstore/rebuild-index")
async def rebuild_index() -> dict[str, Any]:
    """重建索引：全部密文重新解密 → 嵌入 → 原子切换。

    明文只在内存流转；任何一条失败（密钥不符/嵌入失败）整库拒绝切换。
    """
    vault = _vault.get_vault()
    rebuilt = vault.rebuild_index()
    if rebuilt["status"] == "error":
        return _error("INTERNAL", rebuilt["message"], 500, degraded=True)
    return {"ok": True, "count": rebuilt["count"]}


# ---------------------------------------------------------------------------
# FR-PRIV（A1）：chunk 级索引端点
# ---------------------------------------------------------------------------

@app.get("/api/vstore/search-chunks")
async def search_chunks(
    q: str = "",
    top_k: int = 5,
    min_score: float = CHUNK_MIN_SCORE,
    doc_ids: str = "",
    exclude_template: bool = False,
) -> dict[str, Any]:
    """chunk 级检索（FR-PRIV-002）。无降级档：低于 min_score 全部丢弃 → empty。

    top_k clamp 1..20；doc_ids 为逗号分隔白名单（空 = 不过滤）。
    """
    vault = _vault.get_vault()
    chunks = getattr(vault, "chunks", None)
    if chunks is None:
        return {"ok": True, "mode": "empty", "query": q, "count": 0, "items": []}
    top_k = max(1, min(int(top_k), CHUNK_TOP_K_MAX))
    wanted = [d.strip() for d in doc_ids.split(",") if d.strip()] or None
    try:
        # 同步阻塞段（查询嵌入 + faiss 检索）一律 to_thread，不阻塞事件循环
        vector = await asyncio.to_thread(_vault.get_embed_fn(), [q or ""])
        result = await asyncio.to_thread(
            chunks.search,
            vector[0] if vector else [0.0] * _vault.EMBED_DIM,
            top_k,
            min_score,
            wanted,
            bool(exclude_template),
        )
        return {
            "ok": True,
            "mode": result.get("mode", "empty"),
            "query": q,
            "count": len(result.get("items") or []),
            "items": result.get("items") or [],
        }
    except VaultError as exc:
        return _error("INTERNAL", str(exc), 500, degraded=True)


@app.post("/api/vstore/rebuild-chunks")
async def rebuild_chunks() -> dict[str, Any]:
    """重建 chunk 索引（FR-PRIV-007）：全部密文解密 → 切块 → 批量嵌入 →
    原子切换。同步段 to_thread。
    """
    vault = _vault.get_vault()
    try:
        result = await asyncio.to_thread(vault.rebuild_chunks)
    except Exception as exc:
        logger.exception("vstore: rebuild-chunks 失败")
        return _error("INTERNAL", f"chunk 重建失败：{exc}", 500, degraded=True)
    return {"ok": True, **result}