Coverage for astrocyte/integrations/llamaindex.py: 100%

1"""LlamaIndex integration — Astrocyte as a memory store for LlamaIndex agents. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.llamaindex import AstrocyteLlamaMemory 

6 

7 brain = Astrocyte.from_config("astrocyte.yaml") 

8 memory = AstrocyteLlamaMemory(brain, bank_id="user-123") 

9 

10 # Use with LlamaIndex chat engine 

11 chat_engine = index.as_chat_engine(memory=memory) 

12 

13 # Or standalone 

14 await memory.put("Calvin prefers dark mode") 

15 results = await memory.get("What does Calvin prefer?") 

16 

17Maps: 

18 - put → brain.retain() 

19 - get / search → brain.recall() 

20 - LlamaIndex expects memory.get() to return formatted string context 

21""" 

22 

23from __future__ import annotations 

24 

25from typing import TYPE_CHECKING, Any 

26 

27if TYPE_CHECKING: 

28 from astrocyte._astrocyte import Astrocyte 

29 

30from astrocyte.types import AstrocyteContext 

31 

32 

33class AstrocyteLlamaMemory: 

34 """Astrocyte-backed memory for LlamaIndex agents and chat engines. 

35 

36 Implements the interface pattern expected by LlamaIndex's memory system: 

37 put(), get(), get_all(), reset(). 

38 """ 

39 

40 def __init__( 

41 self, 

42 brain: Astrocyte, 

43 bank_id: str, 

44 *, 

45 context: AstrocyteContext | None = None, 

46 max_results: int = 10, 

47 ) -> None: 

48 self.brain = brain 

49 self.bank_id = bank_id 

50 self._context = context 

51 self.max_results = max_results 

52 

53 async def put( 

54 self, 

55 content: str, 

56 *, 

57 tags: list[str] | None = None, 

58 metadata: dict[str, Any] | None = None, 

59 ) -> str | None: 

60 """Store content into memory. Returns memory_id.""" 

61 result = await self.brain.retain( 

62 content, 

63 bank_id=self.bank_id, 

64 tags=tags or ["llamaindex"], 

65 metadata=metadata or {"source": "llamaindex"}, 

66 context=self._context, 

67 ) 

68 return result.memory_id if result.stored else None 

69 

70 async def get( 

71 self, 

72 query: str, 

73 *, 

74 max_results: int | None = None, 

75 ) -> str: 

76 """Retrieve relevant memories as formatted text (for prompt injection). 

77 

78 Returns a string of formatted memory hits, suitable for 

79 injection into system or user messages. 

80 """ 

81 result = await self.brain.recall( 

82 query, 

83 bank_id=self.bank_id, 

84 max_results=max_results or self.max_results, 

85 context=self._context, 

86 ) 

87 if not result.hits: 

88 return "" 

89 return "\n".join(f"- {h.text}" for h in result.hits) 

90 

91 async def get_all(self) -> list[dict[str, Any]]: 

92 """Retrieve all memories in the bank (for full context). 

93 

94 Returns list of dicts with text, score, metadata. 

95 """ 

96 result = await self.brain.recall( 

97 "*", 

98 bank_id=self.bank_id, 

99 max_results=1000, 

100 context=self._context, 

101 ) 

102 return [ 

103 { 

104 "text": h.text, 

105 "score": h.score, 

106 "metadata": h.metadata, 

107 "memory_id": h.memory_id, 

108 } 

109 for h in result.hits 

110 ] 

111 

112 async def search( 

113 self, 

114 query: str, 

115 *, 

116 max_results: int | None = None, 

117 tags: list[str] | None = None, 

118 ) -> list[dict[str, Any]]: 

119 """Search memory with structured results (for programmatic access).""" 

120 result = await self.brain.recall( 

121 query, 

122 bank_id=self.bank_id, 

123 max_results=max_results or self.max_results, 

124 tags=tags, 

125 context=self._context, 

126 ) 

127 return [ 

128 { 

129 "text": h.text, 

130 "score": h.score, 

131 "metadata": h.metadata, 

132 "memory_id": h.memory_id, 

133 } 

134 for h in result.hits 

135 ] 

136 

137 async def reset(self) -> None: 

138 """Clear all memories in the bank.""" 

139 await self.brain.clear_bank(self.bank_id, context=self._context)