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 

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

15 

16 

17@dataclass 

18class SearchResult: 

19 """检索结果。""" 

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( 

32 self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None 

33 ): ... 

34 

35 @abstractmethod 

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

37 

38 @abstractmethod 

39 def count(self) -> int: ... 

40 

41 @abstractmethod 

42 def clear(self): ... 

43 

44 

45class ChromaStore(VectorStore): 

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

47 

48 Args: 

49 collection_name: 集合名称 

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

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

52 """ 

53 

54 def __init__( 

55 self, 

56 collection_name: str = "default", 

57 persist_dir: str | None = None, 

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

59 ): 

60 self._collection_name = collection_name 

61 self._persist_dir = persist_dir 

62 self._embedding_model = embedding_model 

63 self._client = None 

64 self._collection = None 

65 self._initialized = False 

66 

67 def _get_embedding_function(self): 

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

69 from chromadb.utils import embedding_functions 

70 

71 try: 

72 import sentence_transformers # noqa: F401 

73 

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 

86 if self._persist_dir: 

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

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

89 else: 

90 self._client = chromadb.Client() 

91 

92 self._ef = self._get_embedding_function() 

93 self._collection = self._client.get_or_create_collection( 

94 name=self._collection_name, 

95 embedding_function=self._ef, 

96 ) 

97 self._initialized = True 

98 except ImportError: 

99 raise ImportError("chromadb 未安装。运行: pip install chromadb sentence-transformers") 

100 

101 def add( 

102 self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None 

103 ): 

104 self._ensure_init() 

105 if ids is None: 

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

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

108 

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

110 self._ensure_init() 

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

112 out = [] 

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

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

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

116 score = 0.0 

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

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

119 meta = {} 

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

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

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

123 return out 

124 

125 def count(self) -> int: 

126 self._ensure_init() 

127 return self._collection.count() 

128 

129 def clear(self): 

130 self._ensure_init() 

131 self._client.delete_collection(self._collection_name) 

132 self._initialized = False