217 lines
6.7 KiB
Python
217 lines
6.7 KiB
Python
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")
|