Learn-LLM-RAG/app/main.py

import asyncio
import logging
import os
import time
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from app.config import (
    CHUNK_OVERLAP,
    CHUNK_SIZE,
    DEFAULT_TOP_K,
    EMBEDDING_MODEL,
    STATIC_DIR,
    llm_config,
)
from app.embeddings import Embedder
from app.generator import generate
from app.store import KnowledgeBase

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger("learn_rag")


class AskRequest(BaseModel):
    question: str = Field(min_length=1, max_length=2000)
    top_k: int = Field(default=DEFAULT_TOP_K, ge=1, le=8)


class DocumentIn(BaseModel):
    title: str = Field(min_length=1, max_length=80)
    content: str = Field(min_length=1, max_length=200_000)


def decode_text(raw: bytes) -> str:
    if len(raw) > 1_000_000:
        raise HTTPException(status_code=413, detail="文件超过 1MB")
    for encoding in ("utf-8-sig", "utf-8", "gb18030"):
        try:
            return raw.decode(encoding)
        except UnicodeDecodeError:
            continue
    raise HTTPException(status_code=400, detail="无法识别文件编码，请使用 UTF-8 或 GB18030")


def clean_document(title: str, content: str) -> tuple[str, str]:
    title = " ".join(title.replace("\x00", "").split())
    content = content.replace("\x00", "").strip()
    if not title:
        raise HTTPException(status_code=400, detail="标题不能为空")
    if len(title) > 80:
        title = title[:80]
    if len(content) < 10:
        raise HTTPException(status_code=400, detail="正文至少需要 10 个字")
    return title, content


def build_runtime() -> KnowledgeBase:
    embedder = Embedder()
    knowledge = KnowledgeBase(embedder)
    knowledge.startup()
    documents, chunks = knowledge.counts()
    logger.info("知识库就绪：%s 篇文档，%s 个切片", documents, chunks)
    return knowledge


async def boot(application: FastAPI) -> None:
    if application.state.booting or application.state.ready:
        return
    application.state.booting = True
    application.state.error = None
    try:
        application.state.kb = await asyncio.to_thread(build_runtime)
        application.state.ready = True
    except Exception as exc:
        logger.exception("初始化失败")
        application.state.error = f"初始化失败：{exc}"
    finally:
        application.state.booting = False


@asynccontextmanager
async def lifespan(application: FastAPI):
    application.state.ready = False
    application.state.booting = False
    application.state.error = None
    application.state.kb = None
    asyncio.create_task(boot(application))
    yield


app = FastAPI(title="知迹", lifespan=lifespan)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


def require_kb(request: Request) -> KnowledgeBase:
    if not request.app.state.ready or request.app.state.kb is None:
        detail = request.app.state.error or "嵌入模型还在准备，请稍后再试"
        raise HTTPException(status_code=503, detail=detail)
    return request.app.state.kb


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/status")
def status(request: Request) -> dict:
    knowledge = request.app.state.kb
    documents, chunks = knowledge.counts() if knowledge is not None else (0, 0)
    llm = llm_config()
    return {
        "ready": bool(request.app.state.ready),
        "booting": bool(request.app.state.booting),
        "error": request.app.state.error,
        "embedding_model": EMBEDDING_MODEL,
        "vector_db": "Chroma",
        "chunk_size": CHUNK_SIZE,
        "chunk_overlap": CHUNK_OVERLAP,
        "llm_ready": llm is not None,
        "llm_model": llm["model"] if llm else None,
        "documents": documents,
        "chunks": chunks,
    }


@app.post("/api/boot")
async def boot_again(request: Request) -> dict:
    if request.app.state.ready:
        return {"ok": True}
    asyncio.create_task(boot(request.app))
    return {"ok": True, "booting": True}


@app.get("/api/documents")
def list_documents(request: Request) -> dict:
    knowledge = request.app.state.kb
    if knowledge is None:
        return {"documents": []}
    return {"documents": knowledge.list_documents()}


@app.post("/api/documents")
def create_document(body: DocumentIn, request: Request) -> dict:
    knowledge = require_kb(request)
    title, content = clean_document(body.title, body.content)
    try:
        return knowledge.add(title, content)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.post("/api/documents/upload")
async def upload_document(request: Request, file: UploadFile = File(...)) -> dict:
    knowledge = require_kb(request)
    raw = await file.read()
    content = decode_text(raw)
    filename = (file.filename or "未命名资料").rsplit("/", 1)[-1]
    title = filename.rsplit(".", 1)[0]
    title, content = clean_document(title, content)
    try:
        return knowledge.add(title, content)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.delete("/api/documents/{doc_id}")
def delete_document(doc_id: str, request: Request) -> dict:
    knowledge = require_kb(request)
    if not knowledge.delete(doc_id):
        raise HTTPException(status_code=404, detail="没有找到这篇资料")
    return {"ok": True}


@app.post("/api/ask")
async def ask(body: AskRequest, request: Request) -> dict:
    knowledge = require_kb(request)
    question = body.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="请输入问题")
    started = time.perf_counter()

    def run() -> dict:
        chunks = knowledge.search(question, body.top_k)
        logger.info(
            "检索「%s」→ %s",
            question,
            [(chunk["title"], chunk["score"]) for chunk in chunks],
        )
        answer, mode, notice, low_confidence = generate(question, chunks)
        return {
            "question": question,
            "answer": answer,
            "mode": mode,
            "notice": notice,
            "low_confidence": low_confidence,
            "elapsed_ms": int((time.perf_counter() - started) * 1000),
            "chunks": chunks,
        }

    return await asyncio.to_thread(run)


def main() -> None:
    port = int(os.getenv("PORT", "8000"))
    logger.info("知迹已启动：http://127.0.0.1:%s", port)
    uvicorn.run(app, host="127.0.0.1", port=port, log_level="info")