1"""

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

3

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

5"""

6

7 from __future__ import annotations

8

9 import os

10 from abc import ABC, abstractmethod

11 from dataclasses import dataclass, field

12 from pathlib import Path

13 from typing import Any

14

15

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

17

18

19@dataclass

20 class SearchResult:

21 """检索结果。"""

22 content: str

23 score: float

24 metadata: dict = field(default_factory=dict)

25 source: str = ""

26

27

28 class VectorStore(ABC):

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

30

31 @abstractmethod

32 def add(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

39 @abstractmethod

40 def count(self) -> int:

41 ...

42

43 @abstractmethod

44 def clear(self):

45 ...

46

47

48 class ChromaStore(VectorStore):

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

50

51 Args:

52 collection_name: 集合名称

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

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

55 """

56

57 def __init__(

58 self,

59 collection_name: str = "default",

60 persist_dir: str | None = None,

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

62 ):

63 self._collection_name = collection_name

64 self._persist_dir = persist_dir

65 self._embedding_model = embedding_model

66 self._client = None

67 self._collection = None

68 self._initialized = False

69

70 def _get_embedding_function(self):

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

72 from chromadb.utils import embedding_functions

73 try:

74 import sentence_transformers # noqa: F401

75 return embedding_functions.SentenceTransformerEmbeddingFunction(

76 model_name=self._embedding_model,

77 )

78 except ImportError:

79 return embedding_functions.DefaultEmbeddingFunction()

80

81 def _ensure_init(self):

82 if self._initialized:

83 return

84 try:

85 import chromadb

86 from chromadb.utils import embedding_functions

87

88 if self._persist_dir:

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

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

91 else:

92 self._client = chromadb.Client()

93

94 self._ef = self._get_embedding_function()

95 self._collection = self._client.get_or_create_collection(

96 name=self._collection_name,

97 embedding_function=self._ef,

98 )

99 self._initialized = True

100 except ImportError:

101 raise ImportError(

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

103 )

104

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

106 self._ensure_init()

107 if ids is None:

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

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

110

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

112 self._ensure_init()

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

114 out = []

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

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

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

118 score = 0.0

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

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

121 meta = {}

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

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

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

125 return out

126

127 def count(self) -> int:

128 self._ensure_init()

129 return self._collection.count()

130

131 def clear(self):

132 self._ensure_init()

133 self._client.delete_collection(self._collection_name)

134 self._initialized = False