Coverage for agentos/memory/short_term.py: 34%

1""" 

2短期记忆 — 向量数据库存储，覆盖数天到数周的记忆。 

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass 

8 

9from agentos.memory.working import MemoryItem 

10 

11 

12@dataclass 

13class VectorMemory: 

14 """ 

15 短期记忆 — 基于ChromaDB的向量存储。 

16 存近期对话和重要上下文，按语义相似度检索。 

17 """ 

18 

19 def __init__(self, collection_name: str = "agentos_short_term"): 

20 self.collection_name = collection_name 

21 self._items: list[MemoryItem] = [] 

22 self._chroma_client = None 

23 

24 @property 

25 def chroma_client(self): 

26 """延迟加载ChromaDB客户端。""" 

27 if self._chroma_client is None: 

28 try: 

29 import chromadb 

30 

31 self._chroma_client = chromadb.PersistentClient(path="./.agentos/vector_db") 

32 except ImportError: 

33 self._chroma_client = None 

34 return self._chroma_client 

35 

36 async def add(self, item: MemoryItem): 

37 self._items.append(item) 

38 

39 async def search(self, query: str, limit: int = 5) -> list[MemoryItem]: 

40 """ 

41 向量语义搜索。 

42 如果ChromaDB不可用，降级为关键词匹配。 

43 """ 

44 if self._chroma_client: 

45 try: 

46 collection = self._chroma_client.get_or_create_collection(self.collection_name) 

47 results = collection.query(query_texts=[query], n_results=limit) 

48 ids = results.get("ids", [[]])[0] 

49 return [self._items[int(i)] for i in ids if int(i) < len(self._items)] 

50 except Exception: 

51 pass 

52 

53 # 降级：关键词匹配 

54 return [item for item in self._items[-limit:] if query.lower() in item.content.lower()] 

55 

56 def clear(self): 

57 self._items.clear()