""" 向量数据库服务 (ChromaDB) """ import os from typing import List, Dict, Optional import chromadb from chromadb.config import Settings # 配置 VECTOR_STORE_PATH = os.getenv("VECTOR_STORE_PATH", "./data/vector_store") COLLECTION_NAME_PREFIX = "kb_" class VectorStore: """向量存储服务""" def __init__(self): os.makedirs(VECTOR_STORE_PATH, exist_ok=True) self.client = chromadb.PersistentClient( path=VECTOR_STORE_PATH, settings=Settings(anonymized_telemetry=False) ) def get_collection(self, kb_id: str): """获取知识库集合""" collection_name = f"{COLLECTION_NAME_PREFIX}{kb_id}" try: return self.client.get_collection(name=collection_name) except (ValueError, chromadb.errors.NotFoundError): return None def create_collection(self, kb_id: str, embedding_model: str = "text-embedding-3-small"): """创建知识库向量集合""" collection_name = f"{COLLECTION_NAME_PREFIX}{kb_id}" try: self.client.get_collection(name=collection_name) return collection_name except (ValueError, chromadb.errors.NotFoundError): self.client.create_collection( name=collection_name, metadata={ "kb_id": kb_id, "embedding_model": embedding_model } ) return collection_name def delete_collection(self, kb_id: str): """删除知识库向量集合""" collection_name = f"{COLLECTION_NAME_PREFIX}{kb_id}" try: self.client.delete_collection(name=collection_name) return True except (ValueError, chromadb.errors.NotFoundError): return False def add_documents( self, kb_id: str, documents: List[str], embeddings: Optional[List[List[float]]] = None, ids: Optional[List[str]] = None, metadatas: Optional[List[Dict]] = None ): """添加文档片段到向量库""" collection = self.get_collection(kb_id) if ids is None: ids = [f"chunk-{i}" for i in range(len(documents))] if embeddings is not None: collection.add( documents=documents, embeddings=embeddings, ids=ids, metadatas=metadatas ) else: collection.add( documents=documents, ids=ids, metadatas=metadatas ) return len(documents) def search( self, kb_id: str, query: str, n_results: int = 5, where: Optional[Dict] = None ) -> Dict: """检索相似文档""" collection = self.get_collection(kb_id) # 生成查询向量 query_embedding = embedding_service.embed_query(query) results = collection.query( query_embeddings=[query_embedding], n_results=n_results, where=where ) return results def get_stats(self, kb_id: str) -> Dict: """获取向量库统计""" collection = self.get_collection(kb_id) return { "count": collection.count(), "kb_id": kb_id } def delete_documents(self, kb_id: str, ids: List[str]): """删除指定文档片段""" collection = self.get_collection(kb_id) collection.delete(ids=ids) def delete_by_metadata(self, kb_id: str, document_id: str): """根据文档 ID 删除所有片段""" collection = self.get_collection(kb_id) results = collection.get(where={"document_id": document_id}) if results["ids"]: collection.delete(ids=results["ids"]) class EmbeddingService: """ embedding 服务（支持多种模型）""" def __init__(self, model: str = "text-embedding-3-small"): self.model = model self._client = None def _get_client(self): """获取 OpenAI 客户端""" if self._client is None: try: from openai import OpenAI api_key = os.getenv("OPENAI_API_KEY") if api_key: self._client = OpenAI(api_key=api_key) except ImportError: pass return self._client def embed(self, texts: List[str]) -> List[List[float]]: """生成 embedding 向量""" client = self._get_client() if client is None: # 返回随机向量（仅用于测试） import random import math dim = 1536 if "3-small" in self.model else 1024 return [[random.uniform(-1, 1) for _ in range(dim)] for _ in texts] response = client.embeddings.create( model=self.model, input=texts ) return [data.embedding for data in response.data] def embed_query(self, query: str) -> List[float]: """生成查询向量""" return self.embed([query])[0] class DocumentProcessor: """文档处理服务""" def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap def chunk_text(self, text: str, document_id: str = "") -> List[Dict]: """将文本分块""" # 简单分块（按句子/段落） import re # 按句子分割 sentences = re.split(r'[。！？\n]', text) chunks = [] current_chunk = "" current_size = 0 for i, sentence in enumerate(sentences): sentence = sentence.strip() if not sentence: continue sentence_len = len(sentence) if current_size + sentence_len > self.chunk_size and current_chunk: # 保存当前块 chunks.append({ "content": current_chunk.strip(), "document_id": document_id, "chunk_index": len(chunks), "metadata": { "source": "text" } }) # 处理重叠 if self.chunk_overlap > 0: # 保留末尾部分 overlap_chars = current_chunk[-self.chunk_overlap:] current_chunk = overlap_chars + " " + sentence current_size = len(overlap_chars) + sentence_len + 1 else: current_chunk = sentence current_size = sentence_len else: if current_chunk: current_chunk += " " current_chunk += sentence current_size += sentence_len + 1 # 保存最后一个块 if current_chunk.strip(): chunks.append({ "content": current_chunk.strip(), "document_id": document_id, "chunk_index": len(chunks), "metadata": { "source": "text" } }) return chunks def process_document(self, text: str, document_id: str = "") -> List[Dict]: """完整处理文档""" return self.chunk_text(text, document_id) # 全局实例 vector_store = VectorStore() embedding_service = EmbeddingService() def search_knowledge(kb_id: str, query: str, n_results: int = 5) -> Dict: """知识库检索""" # 生成查询向量 query_vector = embedding_service.embed_query(query) # 检索 results = vector_store.search( kb_id=kb_id, query=query, n_results=n_results ) return { "query": query, "results": [ { "content": doc, "metadata": meta, "distance": dist } for doc, meta, dist in zip( results.get("documents", [[]])[0] if results.get("documents") else [], results.get("metadatas", [[]])[0] if results.get("metadatas") else [], results.get("distances", [[]])[0] if results.get("distances") else [] ) ] } def index_document(kb_id: str, document_id: str, text: str) -> int: """索引文档到向量库""" # 分块 processor = DocumentProcessor() chunks = processor.process_document(text, document_id) if not chunks: return 0 # 生成向量 contents = [c["content"] for c in chunks] embeddings = embedding_service.embed(contents) # 添加到向量库 ids = [f"{document_id}-{c['chunk_index']}" for c in chunks] metadatas = [ { "document_id": c["document_id"], "chunk_index": c["chunk_index"], "kb_id": kb_id } for c in chunks ] vector_store.add_documents( kb_id=kb_id, documents=contents, embeddings=embeddings, ids=ids, metadatas=metadatas ) return len(chunks) def delete_document_from_vector(kb_id: str, document_id: str): """从向量库删除文档""" vector_store.delete_by_metadata(kb_id, document_id)