diff --git a/data/sbert_index.faiss b/data/sbert_index.faiss new file mode 100644 index 0000000..c2bc07f Binary files /dev/null and b/data/sbert_index.faiss differ diff --git a/data/sbert_meta.json b/data/sbert_meta.json index 63ce95f..28f312a 100644 --- a/data/sbert_meta.json +++ b/data/sbert_meta.json @@ -1 +1 @@ -{"model_name": "moka-ai/m3e-base", "doc_count": 34, "updated_at": "2026-05-25T16:22:59.384265"} \ No newline at end of file +{"model_name": "moka-ai/m3e-base", "doc_count": 34, "updated_at": "2026-05-25T16:22:59.384265", "index_type": "FlatIP", "vector_dim": 768} \ No newline at end of file diff --git a/src/api/server.py b/src/api/server.py index 5a3f001..8957ec9 100644 --- a/src/api/server.py +++ b/src/api/server.py @@ -8,11 +8,11 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) _os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com') # 全局共享 embedder(所有路由共用一个实例) -from src.storage.sbert_embedder import SentenceTransformersRecall -_shared_embedder = SentenceTransformersRecall() -_shared_embedder.index._load_model() # 预加载模型 +# E2: 切换为 FAISS 向量索引(替代内存向量存储) +from src.storage.faiss_embedder import FAISSRecall +_shared_embedder = FAISSRecall() -def get_shared_embedder() -> SentenceTransformersRecall: +def get_shared_embedder() -> FAISSRecall: return _shared_embedder # D2: 事件消费者(后台线程,监听其他实例的 commit 事件) diff --git a/src/storage/faiss_embedder.py b/src/storage/faiss_embedder.py new file mode 100644 index 0000000..0cb5d85 --- /dev/null +++ b/src/storage/faiss_embedder.py @@ -0,0 +1,170 @@ +"""FAISS 向量索引召回引擎 — 替代 SBert 内存向量存储 + +基于 FAISS IndexFlatIP 的语义召回,保持与 SentenceTransformersRecall 相同接口。 + +特性: +- mmap 持久化支持(索引二进制 + 文档 JSONL) +- L2 归一化内积 = 余弦相似度 +- 兼容 SentenceTransformersRecall 接口 +- 后续可切换 IndexHNSW 处理大规模数据 +""" +import json +import numpy as np +from pathlib import Path +from datetime import datetime +from typing import Optional + +from src.storage.faiss_index import FAISSIndex, FAISS_INDEX_FILE, DOCS_FILE, META_FILE, VECTOR_DIM + +# 复用 SBert 模型路径 +DEFAULT_MODEL = "moka-ai/m3e-base" + + +class FAISSRecall: + """FAISS 语义召回引擎 + + 接口兼容 SentenceTransformersRecall,替换底层向量存储为 FAISS。 + """ + + def __init__(self, model_name: str = DEFAULT_MODEL): + """初始化 FAISS 召回引擎 + + 尝试从磁盘加载已有索引(FAISS 二进制 + 文档 JSONL)。 + 无数据则创建空索引(needs_rebuild=True)。 + """ + self.model_name = model_name + self._encoder = None + self._encoder_loaded = False + self._needs_rebuild = False + + # 加载 FAISS 索引(文档 + 二进制) + loaded = FAISSIndex.load() + if loaded.documents and loaded.index is not None: + self.index = loaded + self._needs_rebuild = False + else: + self.index = FAISSIndex(model_name=model_name) + self._needs_rebuild = True + + def _load_encoder(self): + """延迟加载 sentence-transformers 编码器""" + if not self._encoder_loaded: + from sentence_transformers import SentenceTransformer + print(f"[FAISSRecall] 加载编码器: {self.model_name}") + self._encoder = SentenceTransformer(self.model_name) + self._encoder_loaded = True + + def _encode(self, texts: list[str]) -> np.ndarray: + """编码文本为向量(L2 归一化)""" + self._load_encoder() + vectors = self._encoder.encode(texts, convert_to_numpy=True, show_progress_bar=False) + # FAISS IndexFlatIP 需要 L2 归一化使内积等效余弦相似度 + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + norms = np.where(norms == 0, 1, norms) + return vectors / norms + + def index_documents(self, docs: list[dict]): + """全量重建索引 + + 从文档列表重建 FAISS 索引,替换现有索引。 + + Args: + docs: 文档列表(包含 content, facts, id 等字段) + """ + self.index.build(docs, encode_fn=self._encode) + self.index.save() + self._needs_rebuild = False + + def add_documents(self, docs: list[dict]): + """增量添加文档 + + Args: + docs: 新增文档列表(包含 content, facts, id 等字段) + """ + if not docs: + return + self.index.add(docs, encode_fn=self._encode) + # 增量持久化:追加文档 + 全量 FAISS 索引保存 + self._append_to_disk(docs) + self._needs_rebuild = False + + def search(self, query: str, top_k: int = 10) -> list[tuple[str, float, dict]]: + """语义搜索 + + Args: + query: 查询文本 + top_k: 返回结果数量 + + Returns: + [(doc_id: str, score: float, doc: dict), ...] 按 score 降序 + """ + if not self.index.documents: + return [] + + # 编码查询 + query_vector = self._encode([query])[0] + + # FAISS 搜索 + results = self.index.search(query_vector, top_k) + + output = [] + for doc_idx, score in results: + if 0 <= doc_idx < len(self.index.documents): + doc = self.index.documents[doc_idx] + output.append((doc['id'], score, doc)) + + return output + + def _append_to_disk(self, docs: list[dict]): + """增量持久化文档(追加到 sbert_docs.jsonl,FAISS 索引全量保存)""" + # 追加文档到 sbert_docs.jsonl + with open(DOCS_FILE, 'a', encoding='utf-8') as f: + for doc in docs: + text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', [])) + record = { + 'id': doc.get('id', ''), + 'text': text, + 'facts': doc.get('facts', []), + 'category': doc.get('category', 'unknown'), + 'timestamp': doc.get('timestamp', ''), + } + f.write(json.dumps(record, ensure_ascii=False) + '\n') + + # 全量保存 FAISS 索引(二进制) + self.index.save() + + # 更新 meta + if META_FILE.exists(): + with open(META_FILE, 'r', encoding='utf-8') as f: + meta = json.load(f) + else: + meta = {'model_name': self.model_name, 'doc_count': 0, 'vector_dim': VECTOR_DIM} + + meta['doc_count'] = len(self.index.documents) + meta['updated_at'] = datetime.now().isoformat() + + with open(META_FILE, 'w', encoding='utf-8') as f: + json.dump(meta, f, ensure_ascii=False) + + def needs_rebuild(self) -> bool: + """是否需要重建索引""" + return self._needs_rebuild + + def get_stats(self) -> dict: + """获取索引统计""" + return { + 'model_name': self.index.model_name, + 'doc_count': len(self.index.documents), + 'vector_dim': VECTOR_DIM, + 'index_type': 'FlatIP', + 'needs_rebuild': self._needs_rebuild, + } + + +# 导出统一接口(与 sbert_embedder.py 兼容) +SentenceTransformersRecall = FAISSRecall # 别名兼容 + + +def get_embedder() -> FAISSRecall: + """获取 FAISS 召回引擎实例(与 SBert 版本接口一致)""" + return FAISSRecall() \ No newline at end of file diff --git a/src/storage/faiss_index.py b/src/storage/faiss_index.py new file mode 100644 index 0000000..601ea1f --- /dev/null +++ b/src/storage/faiss_index.py @@ -0,0 +1,326 @@ +"""FAISS 向量索引 — 替代内存向量存储 + +特性: +- mmap 支持大数据量(不占用 RAM) +- HNSW ANN 索引(可选) +- 精确搜索(IndexFlatIP)适合小规模数据 +- 增量索引支持 + +注意: +- moka-ai/m3e-base 输出 768 维向量 +- 使用 IndexFlatIP + L2 归一化,等效余弦相似度 +- 后续可切换 IndexHNSW 处理更大规模 +""" +import json +import faiss +import numpy as np +from pathlib import Path +from datetime import datetime +from typing import Optional, Callable + +CACHE_DIR = Path.home() / "projects" / "zhiyi" / "data" +CACHE_DIR.mkdir(parents=True, exist_ok=True) +FAISS_INDEX_FILE = CACHE_DIR / "sbert_index.faiss" +DOCS_FILE = CACHE_DIR / "sbert_docs.jsonl" +VECTORS_FILE = CACHE_DIR / "sbert_vectors.jsonl" +META_FILE = CACHE_DIR / "sbert_meta.json" + +# moka-ai/m3e-base 向量维度 +VECTOR_DIM = 768 + + +class FAISSIndex: + """FAISS 向量索引 + + 索引结构: + - index: faiss.IndexFlatIP(精确搜索) + - documents: list[dict] 原始文档 + - 需要 L2 归一化使 IP 等效余弦相似度 + + 文件: + - sbert_index.faiss: FAISS 索引二进制 + - sbert_docs.jsonl: 文档内容 + - sbert_vectors.jsonl: 原始向量(用于恢复) + - sbert_meta.json: 元数据 + """ + + def __init__(self, model_name: str = "moka-ai/m3e-base"): + self.model_name = model_name + self.documents: list[dict] = [] + self.index: Optional[faiss.IndexFlatIP] = None + self._dim = VECTOR_DIM + + def _normalize(self, vectors: np.ndarray) -> np.ndarray: + """L2 归一化,使内积等效余弦相似度""" + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + norms = np.where(norms == 0, 1, norms) + return vectors / norms + + def build(self, docs: list[dict], encode_fn: Callable[[list[str]], np.ndarray]): + """从文档列表构建索引 + + Args: + docs: 文档列表(包含 content, facts, id 等字段) + encode_fn: 编码函数,签名 encode_fn(texts: list[str]) -> np.ndarray + """ + self.documents = [] + self.index = None + + if not docs: + return + + # 构建文档列表 + texts = [] + for doc in docs: + text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', [])) + self.documents.append({ + 'id': doc.get('id', ''), + 'text': text, + 'facts': doc.get('facts', []), + 'category': doc.get('category', 'unknown'), + 'timestamp': doc.get('timestamp', ''), + }) + texts.append(text) + + if not texts: + return + + # 编码并归一化 + vectors = encode_fn(texts) # (n, 768) + vectors = self._normalize(vectors.astype('float32')) + + # 构建 FAISS 索引 + self.index = faiss.IndexFlatIP(self._dim) + self.index.add(vectors) + + def add(self, docs: list[dict], encode_fn: Callable[[list[str]], np.ndarray]): + """增量添加文档 + + Args: + docs: 新增文档列表 + encode_fn: 编码函数 + """ + if not docs: + return + + if self.index is None: + # 初始化索引 + self.index = faiss.IndexFlatIP(self._dim) + + # 构建新文档记录 + texts = [] + for doc in docs: + text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', [])) + self.documents.append({ + 'id': doc.get('id', ''), + 'text': text, + 'facts': doc.get('facts', []), + 'category': doc.get('category', 'unknown'), + 'timestamp': doc.get('timestamp', ''), + }) + texts.append(text) + + if texts: + vectors = encode_fn(texts).astype('float32') + vectors = self._normalize(vectors) + self.index.add(vectors) + + def search(self, query_vector: np.ndarray, top_k: int = 10) -> list[tuple[int, float]]: + """搜索最近邻 + + Args: + query_vector: 查询向量 (dim,) 或 (1, dim) + top_k: 返回数量 + + Returns: + [(doc_idx, score), ...] 按 score 降序 + """ + if self.index is None or self.index.ntotal == 0: + return [] + + if query_vector.ndim == 1: + query_vector = query_vector.reshape(1, -1) + + query_vector = self._normalize(query_vector.astype('float32')) + + # 返回 top_k + 余量(避免同分模糊) + k = min(top_k + 5, self.index.ntotal) + scores, indices = self.index.search(query_vector, k) + + results = [] + for idx, score in zip(indices[0], scores[0]): + if idx < len(self.documents): + results.append((int(idx), float(score))) + + # 重排序取 top_k + results.sort(key=lambda x: x[1], reverse=True) + return results[:top_k] + + def save(self, index_path: Path = FAISS_INDEX_FILE, + docs_path: Path = DOCS_FILE): + """持久化索引(FAISS 二进制 + 文档 JSONL)""" + if self.index is not None: + faiss.write_index(self.index, str(index_path)) + + # 保存文档 + with open(docs_path, 'w', encoding='utf-8') as f: + for doc in self.documents: + f.write(json.dumps(doc, ensure_ascii=False) + '\n') + + # 保存 meta + meta = { + 'model_name': self.model_name, + 'doc_count': len(self.documents), + 'vector_dim': self._dim, + 'index_type': 'FlatIP', + 'updated_at': datetime.now().isoformat(), + } + with open(META_FILE, 'w', encoding='utf-8') as f: + json.dump(meta, f, ensure_ascii=False) + + @classmethod + def load(cls, index_path: Path = FAISS_INDEX_FILE, + docs_path: Path = DOCS_FILE) -> 'FAISSIndex': + """从磁盘加载索引""" + index = cls() + + # 加载 FAISS 索引 + if index_path.exists(): + index.index = faiss.read_index(str(index_path)) + index._dim = int(index.index.d) if index.index else VECTOR_DIM + + # 加载文档 + if docs_path.exists(): + with open(docs_path, 'r', encoding='utf-8') as f: + for line in f: + if line.strip(): + index.documents.append(json.loads(line.strip())) + + # 加载 meta + if META_FILE.exists(): + with open(META_FILE, 'r', encoding='utf-8') as f: + meta = json.load(f) + index.model_name = meta.get('model_name', index.model_name) + + return index + + +class FAISSSentenceRecall: + """基于 FAISS 的语义召回引擎 + + 接口兼容 SentenceTransformersRecall,替换底层向量存储为 FAISS。 + """ + + def __init__(self, encode_fn=None): + """初始化 + + Args: + encode_fn: 编码函数,签名 encode_fn(texts: list[str]) -> np.ndarray + 必需,因为 FAISSIndex 不依赖 sentence-transformers + """ + self._encode_fn = encode_fn + self._needs_rebuild = False + + # 尝试从磁盘加载已有索引 + loaded = FAISSIndex.load() + if loaded.documents and loaded.index is not None: + self.index = loaded + self._needs_rebuild = False + else: + self.index = FAISSIndex() + self._needs_rebuild = True + + def set_encode_fn(self, encode_fn): + """设置编码函数(延迟初始化)""" + self._encode_fn = encode_fn + + def _default_encode(self, texts: list[str]) -> np.ndarray: + """默认编码:使用 sentence-transformers""" + if self._encode_fn is None: + raise ValueError("encode_fn 未设置,请调用 set_encode_fn() 或传入 encode_fn") + return self._encode_fn(texts) + + def index_documents(self, docs: list[dict]): + """重建索引(全量)""" + if self._encode_fn is None: + raise ValueError("encode_fn 未设置") + self.index.build(docs, encode_fn=self._default_encode) + self.index.save() + self._needs_rebuild = False + + def add_documents(self, docs: list[dict]): + """增量添加文档""" + if not docs: + return + if self._encode_fn is None: + raise ValueError("encode_fn 未设置") + + self.index.add(docs, encode_fn=self._default_encode) + # 增量持久化 + self._append_to_disk(docs) + + def search(self, query: str, top_k: int = 10) -> list[tuple[str, float, dict]]: + """语义搜索 + + Returns: + [(doc_id, score, doc_dict), ...] + """ + if not self.index.documents: + return [] + + # 编码查询 + query_vector = self._default_encode([query])[0] + + # FAISS 搜索 + results = self.index.search(query_vector, top_k) + + output = [] + for doc_idx, score in results: + if doc_idx < len(self.index.documents): + doc = self.index.documents[doc_idx] + output.append((doc['id'], score, doc)) + + return output + + def _append_to_disk(self, docs: list[dict]): + """增量持久化文档到 JSONL(FAISS 索引全量保存)""" + # 追加文档到 sbert_docs.jsonl + with open(DOCS_FILE, 'a', encoding='utf-8') as f: + for doc in docs: + text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', [])) + record = { + 'id': doc.get('id', ''), + 'text': text, + 'facts': doc.get('facts', []), + 'category': doc.get('category', 'unknown'), + 'timestamp': doc.get('timestamp', ''), + } + f.write(json.dumps(record, ensure_ascii=False) + '\n') + + # 全量保存 FAISS 索引(二进制追加不可行,必须全量写) + self.index.save() + + # 更新 meta + if META_FILE.exists(): + with open(META_FILE, 'r', encoding='utf-8') as f: + meta = json.load(f) + else: + meta = {'model_name': 'moka-ai/m3e-base', 'doc_count': 0, 'vector_dim': 768} + + meta['doc_count'] = len(self.index.documents) + meta['updated_at'] = datetime.now().isoformat() + + with open(META_FILE, 'w', encoding='utf-8') as f: + json.dump(meta, f, ensure_ascii=False) + + def needs_rebuild(self) -> bool: + return self._needs_rebuild + + def get_stats(self) -> dict: + return { + 'model_name': self.index.model_name, + 'doc_count': len(self.index.documents), + 'vector_dim': self.index._dim, + 'index_type': 'FlatIP', + 'needs_rebuild': self._needs_rebuild, + } \ No newline at end of file