Coverage for agentos/rag/store.py: 30%

1""" 

2向量存储抽象层 — ChromaDB 封装。 

3 

4支持创建/加载 collection、添加文档、语义检索。 

5""" 

6 

7from __future__ import annotations 

8 

9import os 

10from abc import ABC, abstractmethod 

11from dataclasses import dataclass, field 

12from pathlib import Path 

13 

14 

15DEFAULT_PERSIST_DIR = Path.home() / ".agentos" / "chroma" 

16 

17 

18@dataclass 

19class SearchResult: 

20 """检索结果。""" 

21 content: str 

22 score: float 

23 metadata: dict = field(default_factory=dict) 

24 source: str = "" 

25 

26 

27class VectorStore(ABC): 

28 """向量存储抽象基类。""" 

29 

30 @abstractmethod 

31 def add(self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None): 

32 ... 

33 

34 @abstractmethod 

35 def search(self, query: str, top_k: int = 5) -> list[SearchResult]: 

36 ... 

37 

38 @abstractmethod 

39 def count(self) -> int: 

40 ... 

41 

42 @abstractmethod 

43 def clear(self): 

44 ... 

45 

46 

47class ChromaStore(VectorStore): 

48 """ChromaDB 向量存储实现。 

49 

50 Args: 

51 collection_name: 集合名称 

52 persist_dir: 持久化目录，None 则仅内存模式 

53 embedding_model: 嵌入模型名称（默认使用 sentence-transformers 轻量模型） 

54 """ 

55 

56 def __init__( 

57 self, 

58 collection_name: str = "default", 

59 persist_dir: str | None = None, 

60 embedding_model: str = "all-MiniLM-L6-v2", 

61 ): 

62 self._collection_name = collection_name 

63 self._persist_dir = persist_dir 

64 self._embedding_model = embedding_model 

65 self._client = None 

66 self._collection = None 

67 self._initialized = False 

68 

69 def _get_embedding_function(self): 

70 """获取 embedding 函数，优先 sentence-transformers，fallback 到 ONNX 内置模型。""" 

71 from chromadb.utils import embedding_functions 

72 try: 

73 import sentence_transformers # noqa: F401 

74 return embedding_functions.SentenceTransformerEmbeddingFunction( 

75 model_name=self._embedding_model, 

76 ) 

77 except ImportError: 

78 return embedding_functions.DefaultEmbeddingFunction() 

79 

80 def _ensure_init(self): 

81 if self._initialized: 

82 return 

83 try: 

84 import chromadb 

85 from chromadb.utils import embedding_functions 

86 

87 if self._persist_dir: 

88 os.makedirs(self._persist_dir, exist_ok=True) 

89 self._client = chromadb.PersistentClient(path=self._persist_dir) 

90 else: 

91 self._client = chromadb.Client() 

92 

93 self._ef = self._get_embedding_function() 

94 self._collection = self._client.get_or_create_collection( 

95 name=self._collection_name, 

96 embedding_function=self._ef, 

97 ) 

98 self._initialized = True 

99 except ImportError: 

100 raise ImportError( 

101 "chromadb 未安装。运行: pip install chromadb sentence-transformers" 

102 ) 

103 

104 def add(self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None): 

105 self._ensure_init() 

106 if ids is None: 

107 ids = [str(self.count() + i) for i in range(len(texts))] 

108 self._collection.add(documents=texts, metadatas=metadatas or None, ids=ids) 

109 

110 def search(self, query: str, top_k: int = 5) -> list[SearchResult]: 

111 self._ensure_init() 

112 results = self._collection.query(query_texts=[query], n_results=top_k) 

113 out = [] 

114 if results["documents"] and results["documents"][0]: 

115 for i in range(len(results["documents"][0])): 

116 doc = results["documents"][0][i] or "" 

117 score = 0.0 

118 if results.get("distances") and results["distances"][0]: 

119 score = 1.0 / (1.0 + float(results["distances"][0][i])) 

120 meta = {} 

121 if results.get("metadatas") and results["metadatas"][0]: 

122 meta = results["metadatas"][0][i] or {} 

123 out.append(SearchResult(content=doc, score=score, metadata=meta)) 

124 return out 

125 

126 def count(self) -> int: 

127 self._ensure_init() 

128 return self._collection.count() 

129 

130 def clear(self): 

131 self._ensure_init() 

132 self._client.delete_collection(self._collection_name) 

133 self._initialized = False