1"""

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

3"""

4

5 from __future__ import annotations

6

7 from dataclasses import dataclass, field

8

9 from agentos.memory.working import MemoryItem

10

11

12@dataclass

13 class 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()