510 lines
20 KiB
Python
510 lines
20 KiB
Python
 
|
"""chunks.py —— 私有资料 chunk 级索引（FR-PRIV-002）。
|
|||
|
|
|
|||
|
|
与 doc 级检索索引（retrieval.py）并存、互不影响：
|
|||
|
|
- 切块：正文段落聚合（目标 600 字 / 满 8 段封块 / 相邻重叠 1 段 / 每文档
|
|||
|
|
上限 200 块）；表格矩阵化为 `|a|b|` 行式文本段按原位插入段落流参与聚合
|
|||
|
|
（实施方案 §1.2：outline 不参与切块，标题上下文仅由首块前缀 `# {title}` 提供）。
|
|||
|
|
- 三文件落盘：index.json（meta 行，单一事实源） / texts.jsonl（chunk 文本，
|
|||
|
|
index.json 可按需重建） / index.faiss（faiss 二进制，与行序一一对应）。
|
|||
|
|
写入纪律仿 retrieval.py：单写锁（threading.Lock）+ tmp + os.replace 原子切换。
|
|||
|
|
- 查询无降级档：向量内积（L2 归一化 = 余弦）低于 min_score 的全部丢弃 →
|
|||
|
|
mode:"empty"（与 doc 级三档降级语义不同，FR-PRIV 假设 6）。
|
|||
|
|
- 任何文件缺失/损坏按空库降级（可服务，不抛异常，FR-PRIV 假设 3）。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
import threading
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import faiss
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
from .settings import EMBED_DIM
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# ---- 切块参数（任务书契约值）----
|
|||
|
|
CHUNK_TARGET_CHARS = 600 # 聚合目标字符数（达到即封块）
|
|||
|
|
CHUNK_MIN_CHARS = 120 # 尾块最小字符数（低于则并入前块）
|
|||
|
|
CHUNK_MAX_PARAS = 8 # 单块段落数上限（满即封块）
|
|||
|
|
CHUNK_OVERLAP_PARAS = 1 # 相邻块重叠段数（前块最后 1 段作下块前缀）
|
|||
|
|
CHUNK_MAX_PER_DOC = 200 # 每文档块数上限（超出截断）
|
|||
|
|
CHUNK_TOP_K_MAX = 20 # search top_k 上限（1..20 clamp）
|
|||
|
|
|
|||
|
|
# ---- 检索参数 ----
|
|||
|
|
CHUNK_MIN_SCORE = 0.30 # 分数阈值：低于全丢 → empty（无降级档）
|
|||
|
|
CHUNK_SNIPPET_CHARS = 200 # snippet = text[:200]（FR-PRIV 假设 4）
|
|||
|
|
|
|||
|
|
# ---- 落盘文件 ----
|
|||
|
|
INDEX_FILE = "index.json"
|
|||
|
|
FAISS_FILE = "index.faiss"
|
|||
|
|
TEXTS_FILE = "texts.jsonl"
|
|||
|
|
TMP_SUFFIX = ".tmp"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChunkError(Exception):
|
|||
|
|
"""chunk 索引操作失败。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class Chunk:
|
|||
|
|
"""单块。chunk_id = f"{doc_id}:{seq}"（doc_id 由 ChunkIndex.upsert_document 注入）。"""
|
|||
|
|
|
|||
|
|
seq: int
|
|||
|
|
text: str
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 切块
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _table_to_text_rows(tables: list[list[list[str]]]) -> list[str]:
|
|||
|
|
"""表格 → `|a|b|` 行式文本段（矩阵化，一行表行 = 一段）。"""
|
|||
|
|
rows: list[str] = []
|
|||
|
|
for table in tables or []:
|
|||
|
|
for row in table or []:
|
|||
|
|
cells = [str(c or "") for c in row]
|
|||
|
|
text = "|" + "|".join(cells) + "|"
|
|||
|
|
if text.strip("|"):
|
|||
|
|
rows.append(text)
|
|||
|
|
return rows
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_chunks(
|
|||
|
|
title: str,
|
|||
|
|
outline: list[dict[str, Any]],
|
|||
|
|
paragraphs: list[str],
|
|||
|
|
tables: list[list[list[str]]],
|
|||
|
|
) -> list[Chunk]:
|
|||
|
|
"""正文 → chunk 列表（实施方案 §1.2 语义，TL 裁决假设 1）。
|
|||
|
|
|
|||
|
|
- 表格矩阵化为 `|a|b|` 行式文本段，插入段落流中间参与聚合切块
|
|||
|
|
（废除「末尾 table 单独块」；提取结构 paragraphs/tables 无位置信息，
|
|||
|
|
表格段统一落在段落流中位，与用例表 L1-05「表格位于段落流中间」判据一致）。
|
|||
|
|
- outline 不参与切块（签名保留，实现忽略）。
|
|||
|
|
- 长段落（> TARGET_CHARS）整段不切，独占一块。
|
|||
|
|
- 封块条件：累计 ≥ TARGET_CHARS 或满 MAX_PARAS 段；尾块不足 MIN_CHARS 并入前块。
|
|||
|
|
- 相邻块重叠：前块最后 OVERLAP_PARAS 段作下块前缀。
|
|||
|
|
- 首块前缀 `# {title}`；seq 从 0 起；截断 CHUNK_MAX_PER_DOC 块。
|
|||
|
|
- 全空输入返回 []，不抛异常。
|
|||
|
|
"""
|
|||
|
|
segs = [str(p or "") for p in (paragraphs or [])]
|
|||
|
|
table_rows = _table_to_text_rows(tables or [])
|
|||
|
|
# 表格段插入段落流中位（原位聚合，TL 裁决假设 1 / L1-05 判据）
|
|||
|
|
if table_rows and len(segs) > 1:
|
|||
|
|
mid = len(segs) // 2
|
|||
|
|
segs = segs[:mid] + table_rows + segs[mid:]
|
|||
|
|
else:
|
|||
|
|
segs += table_rows
|
|||
|
|
segs = [s for s in segs if s.strip()]
|
|||
|
|
if not segs:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 逐段聚合（长段独占）
|
|||
|
|
blocks: list[list[str]] = []
|
|||
|
|
cur: list[str] = []
|
|||
|
|
cur_chars = 0
|
|||
|
|
for seg in segs:
|
|||
|
|
seg_chars = len(seg)
|
|||
|
|
# 该段自身已超目标：先封当前块，长段独占一块（假设 2）
|
|||
|
|
if seg_chars > CHUNK_TARGET_CHARS:
|
|||
|
|
if cur:
|
|||
|
|
blocks.append(cur)
|
|||
|
|
cur = []
|
|||
|
|
cur_chars = 0
|
|||
|
|
blocks.append([seg])
|
|||
|
|
continue
|
|||
|
|
# 常规：目标字符或段数上限即封块
|
|||
|
|
if cur and (cur_chars + seg_chars >= CHUNK_TARGET_CHARS or len(cur) >= CHUNK_MAX_PARAS):
|
|||
|
|
blocks.append(cur)
|
|||
|
|
cur = []
|
|||
|
|
cur_chars = 0
|
|||
|
|
cur.append(seg)
|
|||
|
|
cur_chars += seg_chars
|
|||
|
|
if cur:
|
|||
|
|
blocks.append(cur)
|
|||
|
|
|
|||
|
|
# 尾块合并：不足 MIN_CHARS 且非唯一块 → 并入前块
|
|||
|
|
if len(blocks) > 1:
|
|||
|
|
last_chars = sum(len(s) for s in blocks[-1])
|
|||
|
|
if last_chars < CHUNK_MIN_CHARS:
|
|||
|
|
merged = blocks.pop()
|
|||
|
|
blocks[-1] = blocks[-1] + merged
|
|||
|
|
|
|||
|
|
# 截断
|
|||
|
|
blocks = blocks[:CHUNK_MAX_PER_DOC]
|
|||
|
|
|
|||
|
|
# 拼块 + 重叠前缀 + 标题前缀
|
|||
|
|
chunks: list[Chunk] = []
|
|||
|
|
for i, blk in enumerate(blocks):
|
|||
|
|
paras = list(blk)
|
|||
|
|
if i > 0:
|
|||
|
|
overlap = blocks[i - 1][-CHUNK_OVERLAP_PARAS:]
|
|||
|
|
paras = overlap + paras
|
|||
|
|
text = "\n".join(paras)
|
|||
|
|
if i == 0:
|
|||
|
|
prefix = f"# {title}".strip()
|
|||
|
|
text = f"{prefix}\n{text}"
|
|||
|
|
chunks.append(Chunk(seq=i, text=text))
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 工具
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _now_iso() -> str:
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
|
|||
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _normalize(vec: list[float]) -> np.ndarray:
|
|||
|
|
"""L2 归一化（IndexFlatIP 下等价余弦）。零向量原样返回。"""
|
|||
|
|
arr = np.asarray(vec, dtype=np.float32)
|
|||
|
|
norm = float(np.linalg.norm(arr))
|
|||
|
|
if norm > 0:
|
|||
|
|
arr = arr / norm
|
|||
|
|
return arr
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _fallback_embed_one(text: str) -> list[float]:
|
|||
|
|
"""确定性回退嵌入（不触网）：`种子N` 前缀 → 高位 0.9 + 低位补足到单位范数。
|
|||
|
|
|
|||
|
|
仅用于 items 未携带 vector 的场景（测试骨架 L1-12 语义：『种子1 高位』
|
|||
|
|
分数 0.90——查询向量归一化后内积 = 0.90，落在 min_score 0.30 与 0.91
|
|||
|
|
阈值之间）。生产路径由 vault 先 embed_batch_chunked 再 upsert，不走此分支。
|
|||
|
|
"""
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
vec = [0.0] * EMBED_DIM
|
|||
|
|
m = re.match(r"种子\s*(\d+)", (text or "").strip())
|
|||
|
|
if m:
|
|||
|
|
seed = int(m.group(1))
|
|||
|
|
vec[seed % EMBED_DIM] = 0.9
|
|||
|
|
vec[(seed + 1) % EMBED_DIM] = 0.43589 # sqrt(1 - 0.81)：补足单位范数
|
|||
|
|
return vec
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _faiss_serialize(index: faiss.IndexFlatIP) -> bytes:
|
|||
|
|
"""faiss 二进制（serialize_index 返回 ndarray → tobytes）。"""
|
|||
|
|
arr = faiss.serialize_index(index)
|
|||
|
|
return bytes(arr.tobytes() if hasattr(arr, "tobytes") else arr)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _faiss_read(blob: bytes, dim: int) -> faiss.IndexFlatIP | None:
|
|||
|
|
"""从二进制读 faiss 索引；损坏/空 blob 返回 None（调用方按空库降级）。"""
|
|||
|
|
if not blob:
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
data = np.frombuffer(blob, dtype=np.uint8)
|
|||
|
|
index = faiss.deserialize_index(data)
|
|||
|
|
if isinstance(index, faiss.IndexFlatIP):
|
|||
|
|
return index
|
|||
|
|
if isinstance(index, faiss.IndexRefineFlat) and isinstance(
|
|||
|
|
index.base_index, faiss.IndexFlatIP
|
|||
|
|
):
|
|||
|
|
return index.base_index
|
|||
|
|
return None
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("chunks: index.faiss 读取失败，按空库处理：%s", exc)
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChunkIndex:
|
|||
|
|
"""chunk 向量索引：index.json（meta 行）+ index.faiss + texts.jsonl。
|
|||
|
|
|
|||
|
|
单写锁；upsert 替换语义（先移旧行再重建 faiss）；搜索无降级档。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, path: str | Path, dim: int = EMBED_DIM) -> None:
|
|||
|
|
self.dir = Path(path)
|
|||
|
|
self.dim = dim
|
|||
|
|
self._lock = threading.Lock()
|
|||
|
|
self._faiss: faiss.IndexFlatIP = faiss.IndexFlatIP(dim)
|
|||
|
|
self._rows: list[dict[str, Any]] = []
|
|||
|
|
|
|||
|
|
# ---- 路径 ----
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def json_path(self) -> Path:
|
|||
|
|
return self.dir / INDEX_FILE
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def faiss_path(self) -> Path:
|
|||
|
|
return self.dir / FAISS_FILE
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def texts_path(self) -> Path:
|
|||
|
|
return self.dir / TEXTS_FILE
|
|||
|
|
|
|||
|
|
# ---- 加载 / 落盘 ----
|
|||
|
|
|
|||
|
|
def load(self) -> None:
|
|||
|
|
"""加载三文件。任一缺失/损坏按空库降级（不抛异常，FR-PRIV 假设 3）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
self._rows = []
|
|||
|
|
self._faiss = faiss.IndexFlatIP(self.dim)
|
|||
|
|
rows = []
|
|||
|
|
try:
|
|||
|
|
data = json.loads(self.json_path.read_text(encoding="utf-8"))
|
|||
|
|
items = data.get("items", []) if isinstance(data, dict) else []
|
|||
|
|
except (ValueError, OSError) as exc:
|
|||
|
|
logger.warning("chunks: index.json 读取失败，按空库处理：%s", exc)
|
|||
|
|
items = []
|
|||
|
|
texts = self._read_texts()
|
|||
|
|
for row_id, item in enumerate(items):
|
|||
|
|
if not isinstance(item, dict):
|
|||
|
|
continue
|
|||
|
|
row = dict(item)
|
|||
|
|
row["row_id"] = row_id
|
|||
|
|
row["text"] = texts.get(str(row.get("chunk_id") or "")) or ""
|
|||
|
|
rows.append(row)
|
|||
|
|
faiss_index = _faiss_read(self._read_faiss(), self.dim)
|
|||
|
|
# 防御：faiss 向量数 > 行数（损坏）→ 只保留可配对上的向量；
|
|||
|
|
# faiss 向量数 ≤ 行数合法（行可无向量——测试/空向量场景），
|
|||
|
|
# faiss 行序与 json 行序一一对应，直接使用。
|
|||
|
|
n_rows = len(rows)
|
|||
|
|
n_vec = faiss_index.ntotal if faiss_index is not None else 0
|
|||
|
|
if faiss_index is not None and n_vec > n_rows:
|
|||
|
|
logger.warning(
|
|||
|
|
"chunks: index.faiss 向量数（%d）多于 index.json 行数（%d），"
|
|||
|
|
"按可配对上重建", n_vec, n_rows,
|
|||
|
|
)
|
|||
|
|
rebuilt = faiss.IndexFlatIP(self.dim)
|
|||
|
|
for r in rows:
|
|||
|
|
vec = r.get("vector")
|
|||
|
|
if isinstance(vec, list) and vec:
|
|||
|
|
rebuilt.add(_normalize(vec).reshape(1, self.dim))
|
|||
|
|
faiss_index = rebuilt
|
|||
|
|
self._rows = rows
|
|||
|
|
if faiss_index is not None:
|
|||
|
|
self._faiss = faiss_index
|
|||
|
|
|
|||
|
|
def _read_texts(self) -> dict[str, str]:
|
|||
|
|
"""texts.jsonl → {chunk_id: text}（损坏/缺失 → {}）。"""
|
|||
|
|
out: dict[str, str] = {}
|
|||
|
|
try:
|
|||
|
|
raw = self.texts_path.read_text(encoding="utf-8")
|
|||
|
|
except OSError:
|
|||
|
|
return out
|
|||
|
|
for line in raw.splitlines():
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
obj = json.loads(line)
|
|||
|
|
cid = obj.get("chunk_id")
|
|||
|
|
text = obj.get("text")
|
|||
|
|
if isinstance(cid, str) and isinstance(text, str):
|
|||
|
|
out[cid] = text
|
|||
|
|
except ValueError:
|
|||
|
|
continue
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
def _read_faiss(self) -> bytes:
|
|||
|
|
try:
|
|||
|
|
return self.faiss_path.read_bytes()
|
|||
|
|
except OSError:
|
|||
|
|
return b""
|
|||
|
|
|
|||
|
|
def _atomic_write(self, path: Path, data: bytes) -> None:
|
|||
|
|
"""tmp + os.replace 原子切换（仿 retrieval.py 纪律）。"""
|
|||
|
|
self.dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
tmp = path.with_name(path.name + TMP_SUFFIX)
|
|||
|
|
tmp.write_bytes(data)
|
|||
|
|
os.replace(tmp, path)
|
|||
|
|
|
|||
|
|
def _dump_locked(self) -> None:
|
|||
|
|
"""写三文件（需持有锁）。"""
|
|||
|
|
self.dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
payload = {
|
|||
|
|
"version": 1,
|
|||
|
|
"items": [
|
|||
|
|
{k: v for k, v in r.items() if k not in ("row_id", "text")}
|
|||
|
|
for r in self._rows
|
|||
|
|
],
|
|||
|
|
}
|
|||
|
|
self._atomic_write(
|
|||
|
|
self.json_path, json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
|||
|
|
)
|
|||
|
|
self._atomic_write(self.faiss_path, _faiss_serialize(self._faiss))
|
|||
|
|
text_lines = [
|
|||
|
|
json.dumps({"chunk_id": r.get("chunk_id") or "", "text": r.get("text") or ""},
|
|||
|
|
ensure_ascii=False)
|
|||
|
|
for r in self._rows
|
|||
|
|
]
|
|||
|
|
self._atomic_write(self.texts_path, ("\n".join(text_lines) + "\n").encode("utf-8"))
|
|||
|
|
|
|||
|
|
# ---- 增删 ----
|
|||
|
|
|
|||
|
|
def upsert_document(self, doc_id: str, items: list[dict[str, Any]]) -> int:
|
|||
|
|
"""入库一个文档的全部 chunk（替换语义：先移旧行，过滤重建 faiss）。
|
|||
|
|
|
|||
|
|
items 元素：{"text", "doc_name", "doc_category", "doc_tags",
|
|||
|
|
"as_template", "created_at"}；seq 由 text 取不到时按下标补。
|
|||
|
|
返回写入行数。
|
|||
|
|
"""
|
|||
|
|
with self._lock:
|
|||
|
|
self._rows = [r for r in self._rows if r.get("doc_id") != doc_id]
|
|||
|
|
for seq, item in enumerate(items or []):
|
|||
|
|
row = dict(item)
|
|||
|
|
row["doc_id"] = doc_id
|
|||
|
|
row["seq"] = int(row.get("seq", seq))
|
|||
|
|
row["chunk_id"] = f"{doc_id}:{row['seq']}"
|
|||
|
|
text = str(row.get("text") or "")
|
|||
|
|
row["text"] = text
|
|||
|
|
row["snippet"] = text[:CHUNK_SNIPPET_CHARS]
|
|||
|
|
row.setdefault("doc_name", "")
|
|||
|
|
row.setdefault("doc_category", "")
|
|||
|
|
row.setdefault("doc_tags", "")
|
|||
|
|
row.setdefault("as_template", False)
|
|||
|
|
row.setdefault("created_at", _now_iso())
|
|||
|
|
if not isinstance(row.get("vector"), list):
|
|||
|
|
# items 未携带 vector：确定性回退嵌入（测试/防御场景，不触网）
|
|||
|
|
row["vector"] = _fallback_embed_one(str(row.get("text") or ""))
|
|||
|
|
else:
|
|||
|
|
row["vector"] = list(row.get("vector") or [])
|
|||
|
|
self._rows.append(row)
|
|||
|
|
# 过滤重建 faiss（faiss 不支持逐行删）
|
|||
|
|
self._faiss = faiss.IndexFlatIP(self.dim)
|
|||
|
|
for row in self._rows:
|
|||
|
|
vec = row.get("vector")
|
|||
|
|
if isinstance(vec, list) and vec:
|
|||
|
|
self._faiss.add(_normalize(vec).reshape(1, self.dim))
|
|||
|
|
self._dump_locked()
|
|||
|
|
return len(self._rows)
|
|||
|
|
|
|||
|
|
def remove_document(self, doc_id: str) -> bool:
|
|||
|
|
"""删除该 doc 全部行（不存在返回 False）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
before = len(self._rows)
|
|||
|
|
self._rows = [r for r in self._rows if str(r.get("doc_id")) != doc_id]
|
|||
|
|
if len(self._rows) == before:
|
|||
|
|
return False
|
|||
|
|
self._faiss = faiss.IndexFlatIP(self.dim)
|
|||
|
|
for row in self._rows:
|
|||
|
|
vec = row.get("vector")
|
|||
|
|
if isinstance(vec, list) and vec:
|
|||
|
|
self._faiss.add(_normalize(vec).reshape(1, self.dim))
|
|||
|
|
self._dump_locked()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
# ---- 查询 ----
|
|||
|
|
|
|||
|
|
def rows(self) -> list[dict[str, Any]]:
|
|||
|
|
"""行快照（不含 vector/text，供列表接口）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
return [
|
|||
|
|
{k: v for k, v in r.items() if k not in ("vector", "text", "row_id")}
|
|||
|
|
for r in self._rows
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def count(self) -> int:
|
|||
|
|
with self._lock:
|
|||
|
|
return len(self._rows)
|
|||
|
|
|
|||
|
|
def search(
|
|||
|
|
self,
|
|||
|
|
query_vector: list[float],
|
|||
|
|
top_k: int = 5,
|
|||
|
|
min_score: float = CHUNK_MIN_SCORE,
|
|||
|
|
doc_ids: list[str] | None = None,
|
|||
|
|
exclude_template: bool = False,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
"""向量检索，无降级档：分数 < min_score 全部丢弃 → mode:"empty"。
|
|||
|
|
|
|||
|
|
返回 {"mode": "vector"|"empty", "items": [{chunk_id, doc_id, seq, snippet,
|
|||
|
|
score, doc_name, doc_category, doc_tags, as_template, created_at}]}。
|
|||
|
|
"""
|
|||
|
|
top_k = max(1, min(int(top_k), CHUNK_TOP_K_MAX))
|
|||
|
|
wanted = {str(d) for d in (doc_ids or [])} if doc_ids else None
|
|||
|
|
with self._lock:
|
|||
|
|
rows = self._rows
|
|||
|
|
if not rows:
|
|||
|
|
return {"mode": "empty", "items": []}
|
|||
|
|
q = _normalize(query_vector).reshape(1, self.dim)
|
|||
|
|
try:
|
|||
|
|
dist, idx = self._faiss.search(q, min(top_k * 8, len(rows)))
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("chunks: 向量检索失败，按空结果处理：%s", exc)
|
|||
|
|
return {"mode": "empty", "items": []}
|
|||
|
|
hits: list[dict[str, Any]] = []
|
|||
|
|
for score, row_idx in zip(dist[0].tolist(), idx[0].tolist()):
|
|||
|
|
if not (0 <= row_idx < len(rows)):
|
|||
|
|
continue
|
|||
|
|
score = float(score)
|
|||
|
|
if score < min_score:
|
|||
|
|
continue
|
|||
|
|
row = rows[row_idx]
|
|||
|
|
if wanted is not None and str(row.get("doc_id")) not in wanted:
|
|||
|
|
continue
|
|||
|
|
if exclude_template and bool(row.get("as_template")):
|
|||
|
|
continue
|
|||
|
|
hits.append(
|
|||
|
|
{
|
|||
|
|
"chunk_id": row.get("chunk_id") or "",
|
|||
|
|
"doc_id": row.get("doc_id") or "",
|
|||
|
|
"seq": row.get("seq"),
|
|||
|
|
"snippet": row.get("snippet") or "",
|
|||
|
|
"score": score,
|
|||
|
|
"doc_name": row.get("doc_name") or "",
|
|||
|
|
"doc_category": row.get("doc_category") or "",
|
|||
|
|
"doc_tags": row.get("doc_tags") or "",
|
|||
|
|
"as_template": bool(row.get("as_template")),
|
|||
|
|
"created_at": row.get("created_at") or "",
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
if len(hits) >= top_k:
|
|||
|
|
break
|
|||
|
|
if not hits:
|
|||
|
|
return {"mode": "empty", "items": []}
|
|||
|
|
return {"mode": "vector", "items": hits}
|
|||
|
|
|
|||
|
|
# ---- 重建 ----
|
|||
|
|
|
|||
|
|
def clear(self) -> None:
|
|||
|
|
"""清空内存（测试/重建用；磁盘文件不动）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
self._rows = []
|
|||
|
|
self._faiss = faiss.IndexFlatIP(self.dim)
|
|||
|
|
|
|||
|
|
def rebuild_all(self, docs: dict[str, list[dict[str, Any]]]) -> int:
|
|||
|
|
"""整库重建：docs = {doc_id: [chunk item, ...]}；覆盖三文件，返回行数。"""
|
|||
|
|
with self._lock:
|
|||
|
|
new_rows: list[dict[str, Any]] = []
|
|||
|
|
for doc_id, items in (docs or {}).items():
|
|||
|
|
for seq, item in enumerate(items or []):
|
|||
|
|
row = dict(item)
|
|||
|
|
row["doc_id"] = doc_id
|
|||
|
|
row["seq"] = int(row.get("seq", seq))
|
|||
|
|
row["chunk_id"] = f"{doc_id}:{row['seq']}"
|
|||
|
|
text = str(row.get("text") or "")
|
|||
|
|
row["text"] = text
|
|||
|
|
row["snippet"] = text[:CHUNK_SNIPPET_CHARS]
|
|||
|
|
row.setdefault("doc_name", "")
|
|||
|
|
row.setdefault("doc_category", "")
|
|||
|
|
row.setdefault("doc_tags", "")
|
|||
|
|
row.setdefault("as_template", False)
|
|||
|
|
row.setdefault("created_at", _now_iso())
|
|||
|
|
if not isinstance(row.get("vector"), list):
|
|||
|
|
# items 未携带 vector：确定性回退嵌入（测试/防御场景，不触网）
|
|||
|
|
row["vector"] = _fallback_embed_one(str(row.get("text") or ""))
|
|||
|
|
else:
|
|||
|
|
row["vector"] = list(row.get("vector") or [])
|
|||
|
|
new_rows.append(row)
|
|||
|
|
new_faiss = faiss.IndexFlatIP(self.dim)
|
|||
|
|
for row in new_rows:
|
|||
|
|
vec = row.get("vector")
|
|||
|
|
if isinstance(vec, list) and vec:
|
|||
|
|
new_faiss.add(_normalize(vec).reshape(1, self.dim))
|
|||
|
|
self._rows = new_rows
|
|||
|
|
self._faiss = new_faiss
|
|||
|
|
self._dump_locked()
|
|||
|
|
return len(new_rows)
|