Init: 第一个版本
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""知迹：本地 embedding + 向量数据库的 RAG 问答。"""
|
||||
@@ -0,0 +1,4 @@
|
||||
from app.main import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,46 @@
|
||||
import re
|
||||
|
||||
from app.config import CHUNK_OVERLAP, CHUNK_SIZE
|
||||
|
||||
|
||||
def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
|
||||
"""按中文标点和段落边界切分，相邻切片保留一小段重叠。"""
|
||||
text = re.sub(r"\r\n?", "\n", text).replace("\x00", "").strip()
|
||||
if not text:
|
||||
return []
|
||||
if len(text) <= size:
|
||||
return [text]
|
||||
|
||||
units = [unit for unit in re.split(r"(?<=[。！？；\n])", text) if unit]
|
||||
pieces: list[str] = []
|
||||
step = max(1, size - overlap)
|
||||
for unit in units:
|
||||
if len(unit) <= size:
|
||||
pieces.append(unit)
|
||||
continue
|
||||
for start in range(0, len(unit), step):
|
||||
pieces.append(unit[start : start + size])
|
||||
|
||||
chunks: list[str] = []
|
||||
buffer = ""
|
||||
for piece in pieces:
|
||||
if not buffer:
|
||||
buffer = piece
|
||||
continue
|
||||
if len(buffer) + len(piece) <= size:
|
||||
buffer += piece
|
||||
continue
|
||||
chunks.append(buffer.strip())
|
||||
tail = buffer[-overlap:] if overlap else ""
|
||||
candidate = f"{tail}{piece}"
|
||||
buffer = candidate if len(candidate) <= size else piece
|
||||
if buffer.strip():
|
||||
chunks.append(buffer.strip())
|
||||
|
||||
result: list[str] = []
|
||||
for chunk in chunks:
|
||||
chunk = re.sub(r"\n{3,}", "\n\n", chunk).strip()
|
||||
if len(chunk) < 8 or (result and chunk == result[-1]):
|
||||
continue
|
||||
result.append(chunk)
|
||||
return result or [text]
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
DATA_DIR = ROOT / "data"
|
||||
STATIC_DIR = ROOT / "static"
|
||||
|
||||
EMBEDDING_MODEL = "BAAI/bge-small-zh-v1.5"
|
||||
COLLECTION_NAME = "knowledge_bge_small_zh_v15"
|
||||
CHUNK_SIZE = 360
|
||||
CHUNK_OVERLAP = 60
|
||||
DEFAULT_TOP_K = 4
|
||||
LOW_SCORE = 0.45
|
||||
|
||||
os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")
|
||||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||||
|
||||
|
||||
def load_dotenv() -> None:
|
||||
path = ROOT / ".env"
|
||||
if not path.exists():
|
||||
return
|
||||
for raw in path.read_text(encoding="utf-8").splitlines():
|
||||
line = raw.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
key, value = line.split("=", 1)
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
if key and key not in os.environ:
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
def llm_config() -> dict | None:
|
||||
key = os.getenv("OPENAI_API_KEY", "").strip()
|
||||
if not key:
|
||||
return None
|
||||
base = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1").strip()
|
||||
model = os.getenv("OPENAI_MODEL", "gpt-4o-mini").strip()
|
||||
return {
|
||||
"key": key,
|
||||
"base": (base or "https://api.openai.com/v1").rstrip("/"),
|
||||
"model": model or "gpt-4o-mini",
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
import logging
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
from app.config import EMBEDDING_MODEL
|
||||
|
||||
logger = logging.getLogger("learn_rag")
|
||||
QUERY_PREFIX = "为这个句子生成表示以用于检索相关文章："
|
||||
|
||||
|
||||
def configure_download() -> None:
|
||||
"""Hugging Face 不可达时改用国内镜像，避免首次下载卡住。"""
|
||||
if os.getenv("HF_ENDPOINT"):
|
||||
return
|
||||
try:
|
||||
with urllib.request.urlopen("https://huggingface.co", timeout=4) as response:
|
||||
if response.status < 400:
|
||||
return
|
||||
except Exception:
|
||||
logger.info("无法连接 huggingface.co，嵌入模型改从 hf-mirror.com 下载")
|
||||
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
|
||||
|
||||
|
||||
class Embedder:
|
||||
def __init__(self) -> None:
|
||||
configure_download()
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
logger.info("正在加载嵌入模型 %s", EMBEDDING_MODEL)
|
||||
self.model = TextEmbedding(model_name=EMBEDDING_MODEL)
|
||||
probe = next(iter(self.model.embed(["书店"])))
|
||||
self.dimension = len(probe)
|
||||
logger.info("嵌入模型就绪，维度 %s", self.dimension)
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
return [vector.tolist() for vector in self.model.embed(texts)]
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
query_embed = getattr(self.model, "query_embed", None)
|
||||
if query_embed is not None:
|
||||
try:
|
||||
return next(iter(query_embed([text]))).tolist()
|
||||
except Exception:
|
||||
logger.info("query_embed 不可用，改用查询前缀")
|
||||
prefixed = QUERY_PREFIX + text
|
||||
return next(iter(self.model.embed([prefixed]))).tolist()
|
||||
@@ -0,0 +1,85 @@
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import LOW_SCORE, llm_config
|
||||
|
||||
logger = logging.getLogger("learn_rag")
|
||||
|
||||
SYSTEM_PROMPT = """你是「知迹」的知识库问答助手。只能根据给出的知识切片回答，不要补充切片之外的信息，也不要编造。
|
||||
如果切片不足以回答，就说明知识库里没有这项信息。
|
||||
用自然的中文直接回答问题，不要使用 Markdown，不要写标题。
|
||||
引用事实时在相关句末标注编号，格式必须是〔来源1〕这种全角括号。编号只能使用给定切片的编号。"""
|
||||
|
||||
|
||||
def build_context(chunks: list[dict]) -> str:
|
||||
blocks = []
|
||||
for chunk in chunks:
|
||||
blocks.append(
|
||||
f"〔来源{chunk['n']}〕《{chunk['title']}》相关度 {chunk['score']:.0%}\n{chunk['text']}"
|
||||
)
|
||||
return "\n\n".join(blocks)
|
||||
|
||||
|
||||
def extractive_answer(chunks: list[dict], low_confidence: bool) -> str:
|
||||
if not chunks:
|
||||
return "知识库里还没有可以检索的内容。"
|
||||
lead = ""
|
||||
if low_confidence:
|
||||
lead = "最接近的切片和问题只是部分相关，下面按原文摘出，请以右侧切片为准。\n\n"
|
||||
lines = [lead + chunks[0]["text"].strip() + f"〔来源{chunks[0]['n']}〕"]
|
||||
for chunk in chunks[1:2]:
|
||||
if chunks[0]["score"] - chunk["score"] <= 0.08:
|
||||
lines.append("\n" + chunk["text"].strip() + f"〔来源{chunk['n']}〕")
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def generate(question: str, chunks: list[dict]) -> tuple[str, str, str | None, bool]:
|
||||
low_confidence = not chunks or chunks[0]["score"] < LOW_SCORE
|
||||
if not chunks:
|
||||
return "知识库里还没有可以检索的内容。", "extractive", None, True
|
||||
|
||||
settings = llm_config()
|
||||
if settings is None:
|
||||
notice = "还没有配置大模型，回答直接摘自相关度最高的切片。在项目目录的 .env 中填入 OPENAI_API_KEY 后重启，即可生成连贯回答。"
|
||||
return extractive_answer(chunks, low_confidence), "extractive", notice, low_confidence
|
||||
|
||||
caution = ""
|
||||
if low_confidence:
|
||||
caution = "注意：最高相关度偏低。如果切片不能直接支持答案，请明确说知识库里没有这项信息。\n\n"
|
||||
user_prompt = f"{caution}知识切片如下：\n\n{build_context(chunks)}\n\n问题：{question}"
|
||||
try:
|
||||
answer = call_llm(settings, user_prompt)
|
||||
except Exception:
|
||||
logger.exception("大模型调用失败")
|
||||
notice = "大模型暂时不可用，已改为直接摘录检索到的切片。"
|
||||
return extractive_answer(chunks, low_confidence), "extractive", notice, low_confidence
|
||||
return answer, "llm", None, low_confidence
|
||||
|
||||
|
||||
def call_llm(settings: dict, user_prompt: str) -> str:
|
||||
response = httpx.post(
|
||||
f"{settings['base']}/chat/completions",
|
||||
headers={"Authorization": f"Bearer {settings['key']}"},
|
||||
json={
|
||||
"model": settings["model"],
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 800,
|
||||
"messages": [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
},
|
||||
timeout=httpx.Timeout(60.0, connect=10.0),
|
||||
)
|
||||
response.raise_for_status()
|
||||
message = response.json()["choices"][0]["message"]
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
content = "".join(
|
||||
part.get("text", "") if isinstance(part, dict) else str(part) for part in content
|
||||
)
|
||||
text = (content or "").strip()
|
||||
if not text:
|
||||
raise ValueError("模型返回了空回答")
|
||||
return text
|
||||
+215
@@ -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")
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
"""预置的虚构书店规章。首次启动、知识库为空时写入。"""
|
||||
|
||||
SAMPLES: list[dict[str, str]] = [
|
||||
{
|
||||
"title": "门店与营业时间",
|
||||
"content": """云栖书店位于杭州市西湖区北山街 128 号，出门即是白堤，步行到断桥大约八分钟。门店电话是 0571-8800-2218，邮箱是 hello@yunqi-books.example。
|
||||
|
||||
平日周一至周五的营业时间是上午 10:00 到晚上 21:30。周六、周日和法定节假日的营业时间是上午 9:30 到晚上 22:00。元旦当天以及春节除夕提前到 18:00 结束营业，正月初一全天闭店。
|
||||
|
||||
每周二上午 10:00 到 12:00 是闭架整理时段。这段时间不接待散客购书，只有松卡和柏卡会员可以入内取书、借阅和参加已经预约的活动。12:00 起恢复正常开放。
|
||||
|
||||
书店一共三层。一楼出售新书、人文社科和旅行读物，服务台也在一楼，负责会员、退换和停车券。二楼是文学、艺术、童书和绘本角。三楼是阅读区，以及可容纳 40 人的活动厅。地下一层是仓库，不对顾客开放。
|
||||
|
||||
雨天大堂提供一次性伞套。店内不接受大件行李寄存，折叠伞和手提包可以带入。""",
|
||||
},
|
||||
{
|
||||
"title": "会员等级与权益",
|
||||
"content": """云栖书店的会员分为竹卡、松卡、柏卡三级，都在一楼服务台办理，当天生效。
|
||||
|
||||
竹卡免费。消费按 1 元积 1 分，不足 1 元的部分不计分。积分满 200 分可以兑换一张 10 元购书券，兑换后扣掉相应积分。购书券自发放之日起 90 天内有效，过期作废，不能兑现。竹卡购书没有折扣，也不能在周二闭架整理时段进店。
|
||||
|
||||
松卡年费 128 元。持卡购书享受 9 折，特价区除外。松卡每个月可以免费参加一场书店自己举办的读书会，名额要提前在服务台预约。松卡可以在周二闭架整理时段入店。
|
||||
|
||||
柏卡年费 368 元。持卡购书享受 8.5 折，特价区除外。柏卡可以同时借阅 5 本店内藏书，借期 21 天，并且每次进入三楼阅读区时可以免费带一位同行友人。友人没有独立的座位延时资格。
|
||||
|
||||
升级按自然年累计实付金额计算，用掉折扣和购书券之后的金额才计入。当年实付满 800 元，竹卡可以升为松卡；当年实付满 2400 元，可以升为柏卡。升级时年费按剩余月份折算，也可以用超出的消费直接抵扣。积分跨年不清零。会员到期前 15 天，服务台会用短信提醒续费。""",
|
||||
},
|
||||
{
|
||||
"title": "购书退换与发票",
|
||||
"content": """未拆封、无折痕、无污渍、不影响二次销售的图书，自购买之日起 7 天内可以凭购物小票退货或换货。已经盖了云栖书店店章的书、有明显折角或笔记的书，以及特价区图书，不支持退换。
|
||||
|
||||
期刊、报纸、文创商品和饮品一经售出，不退不换。如果图书有缺页、倒装、严重破损等质量问题，不受 7 天限制，凭小票在 30 天内免费更换同一书号的新书。没有同款库存时，按原支付方式退款。
|
||||
|
||||
退货走原支付路径。现金支付的订单退成购书余额，不退现金。购书券抵扣的部分退回原券，沿用原来的有效期，不重新计算 90 天。
|
||||
|
||||
发票可以在购买当天到服务台开具，也可以在成交后 30 天内凭小票补开电子发票。需要税号时请当场提供。书店不把图书发票改成其他报销品类。""",
|
||||
},
|
||||
{
|
||||
"title": "藏书借阅",
|
||||
"content": """只有柏卡会员可以借阅店内藏书。店内藏书和正在出售的新书不是同一批，书脊贴有「藏」字蓝标，放在二楼西墙。可以外借的书，书目卡上会写明「可借」。
|
||||
|
||||
每次最多同时持有 5 本。单本借期 21 天，到期前可以续借 1 次，续借再延长 14 天。续借要把书带到服务台，不能打电话续借。
|
||||
|
||||
借阅押金是该书定价的 30%，四舍五入到元。还书检查无损后，押金退回原支付方式。逾期滞纳金是每本每天 1 元，从逾期的次日开始计算。逾期超过 30 天视为买断，按图书定价扣款，押金可以冲抵一部分，不够的需要补齐。买断之后的书不必归还。
|
||||
|
||||
外借图书不能批注、折角，也不能撕掉书标。丢失或严重污损时，按定价的 1.2 倍赔偿，押金优先抵扣。赔偿完成后，原书如果后来找到，书店不再收回，也不退还赔偿金。
|
||||
|
||||
借出的藏书仍然属于书店，不能转借给非会员。转借造成的污损由办卡人承担责任。""",
|
||||
},
|
||||
{
|
||||
"title": "阅读区与餐饮",
|
||||
"content": """三楼阅读区的入场条件是：当天在店内消费满 38 元，或者持有松卡、柏卡。竹卡和没有办卡的读者，如果没有达到消费金额，不能只为了占座进入阅读区。柏卡可以免费带一位同行友人，但友人不能单独延长座位时间。
|
||||
|
||||
普通日子每个座位限时 3 小时。周六、周日和法定节假日的 13:00 到 17:00 是高峰，限时改为 2 小时。服务台不会到点强制清场，不过有人排队时，超时的座位需要让出来。
|
||||
|
||||
阅读区供应手冲咖啡、桂花乌龙和巴斯克蛋糕。可以自带有盖的杯装饮品。不可以携带外来热食、盒饭、炸物，以及气味浓重的食物。一楼门口有饮水机，杯子需要自备，或者购买店内饮品。
|
||||
|
||||
阅读区可以使用笔记本和平板电脑，请调成静音。打电话请到楼梯间。儿童进入阅读区需要成人陪同。绘本角在二楼东侧，不属于阅读区，不设 3 小时限制；周末请不要长时间占用绘本角的矮桌。
|
||||
|
||||
每周六下午 15:00 有一小时亲子共读，地点在二楼绘本角，不需要会员，也不收费。需要提前在一楼服务台预约，每次限额 12 组家庭。迟到超过 10 分钟，名额让给现场等候的家庭。""",
|
||||
},
|
||||
{
|
||||
"title": "活动停车与其他服务",
|
||||
"content": """2026 年 10 月的读书会主题是「河流与城市」，嘉宾是作家林晚晴。时间是 10 月 18 日周六晚上 19:30，地点在三楼活动厅。松卡和柏卡会员免费，但要提前预约。非会员票价 49 元，包含一杯桂花乌龙。活动大约 90 分钟，19:20 开始签到，开始 15 分钟后停止入场。
|
||||
|
||||
11 月 2 日周日下午 13:00 到 17:00，书店门口举办旧书交换市集。读者可以带不超过 5 本品相良好的旧书来交换，不收摊位费，也不保证一定能换到想要的书。教材、破损书和期刊不接受。
|
||||
|
||||
书店不提供打印和复印。需要复印的读者，可以步行前往附近的浙江图书馆。店内无线网络名称是 Yunqi-Guest，密码写在服务台立牌上，只供应阅读区，不保证视频会议的带宽。
|
||||
|
||||
停车可以使用北山街地下停车场。当天在店内消费满 68 元，可以在一楼服务台领取 2 小时停车券。每车每天限领 1 张，需要出示停车票根。书店门口没有非机动车棚，共享单车不要停在台阶上。""",
|
||||
},
|
||||
]
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
from app.chunker import chunk_text
|
||||
from app.config import COLLECTION_NAME, DATA_DIR, EMBEDDING_MODEL
|
||||
from app.embeddings import Embedder
|
||||
from app.seed import SAMPLES
|
||||
|
||||
logger = logging.getLogger("learn_rag")
|
||||
|
||||
|
||||
class DisabledEmbedding:
|
||||
"""占位嵌入函数。向量全部由 fastembed 预先算好，避免 Chroma 再下载一套模型。"""
|
||||
|
||||
@staticmethod
|
||||
def name() -> str:
|
||||
return "precomputed"
|
||||
|
||||
def __call__(self, input):
|
||||
raise RuntimeError("向量由 fastembed 预先计算")
|
||||
|
||||
def embed_documents(self, input):
|
||||
raise RuntimeError("向量由 fastembed 预先计算")
|
||||
|
||||
def embed_query(self, input):
|
||||
raise RuntimeError("向量由 fastembed 预先计算")
|
||||
|
||||
|
||||
def to_score(distance: float) -> float:
|
||||
if distance <= 0:
|
||||
return 1.0
|
||||
if distance <= 1:
|
||||
return 1 - distance
|
||||
return 1 / (1 + distance)
|
||||
|
||||
|
||||
class KnowledgeBase:
|
||||
def __init__(self, embedder: Embedder) -> None:
|
||||
self.embedder = embedder
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
self.path = DATA_DIR / "documents.json"
|
||||
self._lock = threading.Lock()
|
||||
self.client = chromadb.PersistentClient(
|
||||
path=str(DATA_DIR / "chroma"),
|
||||
settings=Settings(anonymized_telemetry=False),
|
||||
)
|
||||
self.collection = self._open_collection()
|
||||
self.docs: list[dict] = []
|
||||
|
||||
def _open_collection(self):
|
||||
return self.client.get_or_create_collection(
|
||||
name=COLLECTION_NAME,
|
||||
embedding_function=DisabledEmbedding(),
|
||||
metadata={"hnsw:space": "cosine"},
|
||||
)
|
||||
|
||||
def _reset_collection(self) -> None:
|
||||
try:
|
||||
self.client.delete_collection(COLLECTION_NAME)
|
||||
except Exception:
|
||||
logger.info("向量集合尚不存在，直接创建")
|
||||
self.collection = self._open_collection()
|
||||
|
||||
def startup(self) -> None:
|
||||
needs_seed = not self.path.exists()
|
||||
payload = {"model": EMBEDDING_MODEL, "documents": []}
|
||||
if self.path.exists():
|
||||
payload = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
self.docs = payload.get("documents", [])
|
||||
if needs_seed:
|
||||
logger.info("写入示例知识库：云栖书店")
|
||||
self._reset_collection()
|
||||
for sample in SAMPLES:
|
||||
self._index(sample["title"], sample["content"])
|
||||
self._save()
|
||||
return
|
||||
if payload.get("model") != EMBEDDING_MODEL:
|
||||
logger.info("嵌入模型已更换，重新构建向量库")
|
||||
previous = list(self.docs)
|
||||
self.docs = []
|
||||
self._reset_collection()
|
||||
for doc in previous:
|
||||
self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
|
||||
self._save()
|
||||
return
|
||||
self._reconcile()
|
||||
|
||||
def _save(self) -> None:
|
||||
payload = {"model": EMBEDDING_MODEL, "documents": self.docs}
|
||||
temporary = self.path.with_suffix(".json.tmp")
|
||||
temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
temporary.replace(self.path)
|
||||
|
||||
def _summary(self, doc: dict) -> dict:
|
||||
preview = re.sub(r"\s+", " ", doc["content"]).strip()[:80]
|
||||
return {
|
||||
"id": doc["id"],
|
||||
"title": doc["title"],
|
||||
"chunk_count": doc["chunk_count"],
|
||||
"created_at": doc["created_at"],
|
||||
"preview": preview,
|
||||
}
|
||||
|
||||
def _index(
|
||||
self,
|
||||
title: str,
|
||||
content: str,
|
||||
doc_id: str | None = None,
|
||||
created_at: str | None = None,
|
||||
) -> dict:
|
||||
chunks = chunk_text(content)
|
||||
if not chunks:
|
||||
raise ValueError("正文无法切分")
|
||||
doc_id = doc_id or uuid.uuid4().hex
|
||||
created_at = created_at or datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
vectors = self.embedder.embed_documents(chunks)
|
||||
ids = [f"{doc_id}:{index}" for index in range(len(chunks))]
|
||||
metadatas = [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"title": title,
|
||||
"chunk_index": index,
|
||||
"chunk_total": len(chunks),
|
||||
}
|
||||
for index in range(len(chunks))
|
||||
]
|
||||
self.collection.add(
|
||||
ids=ids,
|
||||
documents=chunks,
|
||||
embeddings=vectors,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
record = {
|
||||
"id": doc_id,
|
||||
"title": title,
|
||||
"content": content,
|
||||
"created_at": created_at,
|
||||
"chunk_count": len(chunks),
|
||||
}
|
||||
self.docs.append(record)
|
||||
return record
|
||||
|
||||
def _reconcile(self) -> None:
|
||||
stored = self.collection.get(include=["metadatas"])
|
||||
present = {
|
||||
meta["doc_id"]
|
||||
for meta in (stored.get("metadatas") or [])
|
||||
if meta and meta.get("doc_id")
|
||||
}
|
||||
wanted = {doc["id"] for doc in self.docs}
|
||||
for doc_id in present - wanted:
|
||||
self.collection.delete(where={"doc_id": doc_id})
|
||||
for doc in self.docs:
|
||||
if doc["id"] not in present:
|
||||
logger.info("重新索引《%s》", doc["title"])
|
||||
self.docs = [item for item in self.docs if item["id"] != doc["id"]]
|
||||
self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
|
||||
self._save()
|
||||
|
||||
def list_documents(self) -> list[dict]:
|
||||
with self._lock:
|
||||
return [self._summary(doc) for doc in self.docs]
|
||||
|
||||
def counts(self) -> tuple[int, int]:
|
||||
with self._lock:
|
||||
return len(self.docs), int(self.collection.count())
|
||||
|
||||
def add(self, title: str, content: str) -> dict:
|
||||
with self._lock:
|
||||
record = self._index(title, content)
|
||||
try:
|
||||
self._save()
|
||||
except Exception:
|
||||
self.docs.pop()
|
||||
self.collection.delete(where={"doc_id": record["id"]})
|
||||
raise
|
||||
logger.info("新增《%s》，%s 个切片", title, record["chunk_count"])
|
||||
return self._summary(record)
|
||||
|
||||
def delete(self, doc_id: str) -> bool:
|
||||
with self._lock:
|
||||
if not any(doc["id"] == doc_id for doc in self.docs):
|
||||
return False
|
||||
self.collection.delete(where={"doc_id": doc_id})
|
||||
self.docs = [doc for doc in self.docs if doc["id"] != doc_id]
|
||||
self._save()
|
||||
return True
|
||||
|
||||
def search(self, question: str, top_k: int) -> list[dict]:
|
||||
with self._lock:
|
||||
count = int(self.collection.count())
|
||||
if count == 0:
|
||||
return []
|
||||
vector = self.embedder.embed_query(question)
|
||||
found = self.collection.query(
|
||||
query_embeddings=[vector],
|
||||
n_results=min(top_k, count),
|
||||
include=["documents", "metadatas", "distances"],
|
||||
)
|
||||
documents = (found.get("documents") or [[]])[0] or []
|
||||
metadatas = (found.get("metadatas") or [[]])[0] or []
|
||||
distances = (found.get("distances") or [[]])[0] or []
|
||||
chunks = []
|
||||
for text, meta, distance in zip(documents, metadatas, distances):
|
||||
meta = meta or {}
|
||||
chunks.append(
|
||||
{
|
||||
"id": f"{meta.get('doc_id', '')}:{meta.get('chunk_index', 0)}",
|
||||
"title": meta.get("title") or "未命名资料",
|
||||
"text": text or "",
|
||||
"score": round(to_score(float(distance)), 4),
|
||||
"chunk_index": int(meta.get("chunk_index") or 0) + 1,
|
||||
"chunk_total": int(meta.get("chunk_total") or 1),
|
||||
}
|
||||
)
|
||||
chunks.sort(key=lambda item: item["score"], reverse=True)
|
||||
for number, chunk in enumerate(chunks, start=1):
|
||||
chunk["n"] = number
|
||||
return chunks
|
||||
Reference in New Issue
Block a user