Init: 第一个版本

2026-10-10 23:11:55 +08:00
commit c1a3461c49
17 changed files with 3806 additions and 0 deletions
@@ -0,0 +1,215 @@
+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")