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, "llm_provider": llm["provider"] 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")