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

1""" 

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

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass, field 

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 self._chroma_client = chromadb.PersistentClient( 

31 path="./.agentos/vector_db" 

32 ) 

33 except ImportError: 

34 self._chroma_client = None 

35 return self._chroma_client 

36 

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

38 self._items.append(item) 

39 

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

41 """ 

42 向量语义搜索。 

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

44 """ 

45 if self._chroma_client: 

46 try: 

47 collection = self._chroma_client.get_or_create_collection( 

48 self.collection_name 

49 ) 

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

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

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

53 except Exception: 

54 pass 

55 

56 # 降级：关键词匹配 

57 return [ 

58 item 

59 for item in self._items[-limit:] 

60 if query.lower() in item.content.lower() 

61 ] 

62 

63 def clear(self): 

64 self._items.clear()