import os # 在导入 chromadb 之前设置环境变量 os.environ["CHROMA_TELEMETRY"] = "false" os.environ["ANONYMIZED_TELEMETRY"] = "false" import chromadb from typing import List, Dict, Any import asyncio from chromadb.config import Settings from sentence_transformers import SentenceTransformer from langchain.text_splitter import RecursiveCharacterTextSplitter import uuid from datetime import datetime class AsyncVectorStore: """异步向量存储服务""" def __init__(self, persist_directory: str = "./chroma_db"): self.persist_directory = persist_directory self.client = chromadb.PersistentClient( path=persist_directory, settings=Settings(anonymized_telemetry=False) ) self.collection = self.client.get_or_create_collection( name="documents", metadata={"hnsw:space": "cosine"} ) # 尝试初始化向量编码器,如果网络失败则使用本地方案 try: print("正在加载向量编码模型...") self.encoder = SentenceTransformer("all-MiniLM-L6-v2") print("✓ 向量编码模型加载成功") except Exception as e: print(f"⚠️ 向量编码模型加载失败: {e}") print("使用简单的文本向量化方案(仅用于演示)") self.encoder = None self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=50, length_function=len ) async def add_document_async(self, content: str, filename: str) -> str: """异步添加文档到向量库""" doc_id = str(uuid.uuid4()) # 异步分割文本 chunks = await asyncio.to_thread(self.text_splitter.split_text, content) # 异步生成向量 embeddings = await asyncio.to_thread(self.encoder.encode, chunks) embeddings = embeddings.tolist() # 生成chunk IDs chunk_ids = [f"{doc_id}_{i}" for i in range(len(chunks))] # 准备元数据 metadatas = [ { "doc_id": doc_id, "filename": filename, "chunk_index": i, "upload_time": datetime.now().isoformat(), } for i in range(len(chunks)) ] # 异步添加到向量库 await asyncio.to_thread( self.collection.add, ids=chunk_ids, embeddings=embeddings, documents=chunks, metadatas=metadatas, ) return doc_id async def search_async(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]: """异步搜索相关文档""" # 异步生成查询向量 query_embedding = await asyncio.to_thread(self.encoder.encode, [query]) query_embedding = query_embedding.tolist()[0] # 异步查询 results = await asyncio.to_thread( self.collection.query, query_embeddings=[query_embedding], n_results=top_k, include=["documents", "metadatas", "distances"], ) formatted_results = [] if results["documents"] and results["documents"][0]: for i, doc in enumerate(results["documents"][0]): formatted_results.append( { "content": doc, "metadata": results["metadatas"][0][i], "distance": results["distances"][0][i], } ) return formatted_results async def get_documents_async(self) -> List[Dict[str, Any]]: """异步获取所有文档信息""" results = await asyncio.to_thread(self.collection.get, include=["metadatas"]) # 按文档ID分组 doc_info = {} for metadata in results["metadatas"]: doc_id = metadata["doc_id"] if doc_id not in doc_info: doc_info[doc_id] = { "id": doc_id, "filename": metadata["filename"], "upload_time": metadata["upload_time"], "chunks_count": 0, } doc_info[doc_id]["chunks_count"] += 1 return list(doc_info.values()) async def delete_document_async(self, doc_id: str) -> bool: """异步删除文档""" # 异步获取该文档的所有chunk IDs results = await asyncio.to_thread( self.collection.get, where={"doc_id": doc_id}, include=["metadatas"] ) if not results["ids"]: return False # 异步删除所有相关chunks await asyncio.to_thread(self.collection.delete, ids=results["ids"]) return True