Coverage for agentos/vectorstore/db.py: 22%
153 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1"""
2AgentOS v0.30 向量数据库集成 — Chroma + FAISS。
3语义记忆检索、知识库索引。
4"""
6import os
7import pickle
8import uuid
9from dataclasses import dataclass, field
12@dataclass
13class VectorEntry:
14 """向量条目。"""
16 id: str
17 text: str
18 metadata: dict = field(default_factory=dict)
19 score: float = 0.0
22class BaseVectorStore:
23 """向量存储基类。"""
25 def add(
26 self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None
27 ) -> list[str]: ...
28 def search(self, query: str, top_k: int = 5) -> list[VectorEntry]: ...
29 def delete(self, ids: list[str]): ...
30 def count(self) -> int: ...
33class FAISSVectorStore(BaseVectorStore):
34 """基于 FAISS 的轻量向量存储。"""
36 def __init__(self, dim: int = 768, index_path: str = ""):
37 self.dim = dim
38 self.index_path = index_path
39 self._index = None
40 self._store: dict[str, tuple[list[float], str, dict]] = {}
41 self._next_id = 0
42 if index_path and os.path.exists(index_path):
43 self._load()
45 def _init_index(self):
46 try:
47 import faiss
49 self._index = faiss.IndexFlatIP(self.dim)
50 except ImportError:
51 self._index = None
53 def add(
54 self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None
55 ) -> list[str]:
56 embeddings = self._embed(texts)
57 if not self._index:
58 self._init_index()
59 if self._index:
60 import numpy as np
62 vecs = np.array(embeddings, dtype=np.float32)
63 self._index.add(vecs)
65 res_ids = []
66 for i, text in enumerate(texts):
67 rid = ids[i] if ids else f"v{self._next_id}"
68 self._next_id += 1
69 self._store[rid] = (embeddings[i], text, metadatas[i] if metadatas else {})
70 res_ids.append(rid)
71 return res_ids
73 def _fallback_search(self, q_vec, top_k):
74 """Fallback余弦相似度搜索（无faiss时使用）。"""
75 import math
77 scores = []
78 for rid, (vec, text, meta) in self._store.items():
79 dot = sum(a * b for a, b in zip(q_vec, vec))
80 na = math.sqrt(sum(a * a for a in q_vec))
81 nb = math.sqrt(sum(b * b for b in vec))
82 sim = dot / (na * nb) if na * nb > 0 else 0.0
83 scores.append((sim, rid, text, meta))
84 scores.sort(key=lambda x: x[0], reverse=True)
85 return [
86 VectorEntry(id=rid, text=text, metadata=meta, score=float(s))
87 for s, rid, text, meta in scores[:top_k]
88 ]
90 def search(self, query: str, top_k: int = 5) -> list[VectorEntry]:
91 if not self._store:
92 return []
93 q_vec = self._embed([query])[0]
94 if not self._index:
95 return self._fallback_search(q_vec, top_k)
96 q_vec = self._embed([query])[0]
97 import numpy as np
99 distances, indices = self._index.search(np.array([q_vec], dtype=np.float32), min(top_k, self.count()))
100 results = []
101 for score, idx in zip(distances[0], indices[0]):
102 if idx < 0:
103 continue
104 rid = f"v{idx}"
105 if rid in self._store:
106 _, text, meta = self._store[rid]
107 results.append(VectorEntry(id=rid, text=text, metadata=meta, score=float(score)))
108 return results
110 def delete(self, ids: list[str]):
111 for rid in ids:
112 self._store.pop(rid, None)
114 def count(self) -> int:
115 return len(self._store)
117 def _embed(self, texts: list[str]) -> list[list[float]]:
118 """轻量嵌入：使用 all-MiniLM-L6-v2 或回退到 TF-IDF。"""
119 try:
120 from sentence_transformers import SentenceTransformer
122 model = SentenceTransformer("all-MiniLM-L6-v2")
123 embeddings = model.encode(texts, normalize_embeddings=True)
124 return embeddings.tolist()
125 except ImportError:
126 return self._tfidf_embed(texts)
128 def _tfidf_embed(self, texts: list[str]) -> list[list[float]]:
129 """TF-IDF 回退，仅作占位。"""
130 import hashlib
132 dim = self.dim
133 result = []
134 for t in texts:
135 h = hashlib.sha256(t.encode()).digest()
136 vec = [(h[i] / 255.0) for i in range(min(len(h), dim))]
137 vec += [0.0] * (dim - len(vec))
138 result.append(vec)
139 return result
141 def _save(self):
142 if self.index_path:
143 os.makedirs(os.path.dirname(self.index_path) or ".", exist_ok=True)
144 with open(self.index_path, "wb") as f:
145 pickle.dump({"store": self._store, "next_id": self._next_id}, f)
147 def _load(self):
148 with open(self.index_path, "rb") as f:
149 data = pickle.load(f)
150 self._store = data["store"]
151 self._next_id = data["next_id"]
153 def __del__(self):
154 if self.index_path:
155 self._save()
158class ChromaVectorStore(BaseVectorStore):
159 """Chroma 向量存储。"""
161 def __init__(self, collection_name: str = "agentos", persist_dir: str = "./chroma_data"):
162 self.collection_name = collection_name
163 self.persist_dir = persist_dir
164 self._client = None
165 self._collection = None
166 self._init()
168 def _init(self):
169 try:
170 import chromadb
172 self._client = chromadb.PersistentClient(path=self.persist_dir)
173 self._collection = self._client.get_or_create_collection(self.collection_name)
174 except ImportError:
175 self._collection = None
177 def add(
178 self, texts: list[str], metadatas: list[dict] | None = None, ids: list[str] | None = None
179 ) -> list[str]:
180 if not self._collection:
181 ids = ids or [f"v{len(self._fallback_store)}-{i}" for i in range(len(texts))]
182 for i, t in enumerate(texts):
183 self._fallback_store[ids[i]] = {
184 "text": t,
185 "metadata": metadatas[i] if metadatas else {},
186 }
187 return ids
189 ids = ids or [str(uuid.uuid4())[:8] for _ in texts]
190 self._collection.add(documents=texts, metadatas=metadatas or [{}] * len(texts), ids=ids)
191 return ids
193 def search(self, query: str, top_k: int = 5) -> list[VectorEntry]:
194 if not self._collection:
195 if self._fallback_store:
196 return [
197 VectorEntry(id=k, text=v["text"], metadata=v["metadata"], score=0.5)
198 for k, v in list(self._fallback_store.items())[:top_k]
199 ]
200 return []
201 results = self._collection.query(query_texts=[query], n_results=top_k)
202 entries = []
203 for i, rid in enumerate(results.get("ids", [[]])[0]):
204 entries.append(
205 VectorEntry(
206 id=rid,
207 text=results["documents"][0][i] if results.get("documents") else "",
208 metadata=results["metadatas"][0][i] if results.get("metadatas") else {},
209 score=1.0 - results["distances"][0][i] if results.get("distances") else 0.0,
210 )
211 )
212 return entries
214 def delete(self, ids: list[str]):
215 if self._collection:
216 self._collection.delete(ids=ids)
217 else:
218 for rid in ids:
219 self._fallback_store.pop(rid, None)
221 def count(self) -> int:
222 if self._collection:
223 return self._collection.count()
224 return len(self._fallback_store)
226 @property
227 def _fallback_store(self) -> dict:
228 if not hasattr(self, "_fb"):
229 self._fb = {}
230 return self._fb