Coverage for astrocyte/integrations/strands.py: 71%

1"""Strands Agents (AWS) integration — Astrocyte as agent tools. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.strands import astrocyte_strands_tools 

6 

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

8 tools = astrocyte_strands_tools(brain, bank_id="user-123") 

9 

10 # Register with Strands agent 

11 from strands import Agent 

12 agent = Agent(model=model, tools=tools) 

13 

14Strands Agents uses a tool specification pattern where each tool has: 

15- A spec dict (name, description, inputSchema in JSON Schema format) 

16- A handler function that receives the tool input 

17""" 

18 

19from __future__ import annotations 

20 

21import json 

22from typing import TYPE_CHECKING, Any 

23 

24if TYPE_CHECKING: 

25 from astrocyte._astrocyte import Astrocyte 

26 

27from astrocyte.types import AstrocyteContext 

28 

29 

30def astrocyte_strands_tools( 

31 brain: Astrocyte, 

32 bank_id: str, 

33 *, 

34 context: AstrocyteContext | None = None, 

35 include_reflect: bool = True, 

36 include_forget: bool = False, 

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

38 """Create Strands-compatible tool definitions backed by Astrocyte. 

39 

40 Returns a list of dicts, each with 'spec' (JSON Schema) and 'handler' (async callable). 

41 Strands Agents registers tools with this pattern. 

42 """ 

43 tools: list[dict[str, Any]] = [] 

44 

45 # ── memory_retain ── 

46 

47 async def retain_handler(tool_input: dict[str, Any]) -> str: 

48 content = tool_input["content"] 

49 tags = tool_input.get("tags") 

50 result = await brain.retain(content, bank_id=bank_id, tags=tags, context=context) 

51 return json.dumps({"stored": result.stored, "memory_id": result.memory_id}) 

52 

53 tools.append( 

54 { 

55 "spec": { 

56 "name": "memory_retain", 

57 "description": "Store content into long-term memory for future recall.", 

58 "inputSchema": { 

59 "json": { 

60 "type": "object", 

61 "properties": { 

62 "content": {"type": "string", "description": "The text to memorize."}, 

63 "tags": { 

64 "type": "array", 

65 "items": {"type": "string"}, 

66 "description": "Optional tags for filtering.", 

67 }, 

68 }, 

69 "required": ["content"], 

70 } 

71 }, 

72 }, 

73 "handler": retain_handler, 

74 } 

75 ) 

76 

77 # ── memory_recall ── 

78 

79 async def recall_handler(tool_input: dict[str, Any]) -> str: 

80 query = tool_input["query"] 

81 max_results = tool_input.get("max_results", 5) 

82 result = await brain.recall(query, bank_id=bank_id, max_results=max_results, context=context) 

83 hits = [{"text": h.text, "score": round(h.score, 4)} for h in result.hits] 

84 return json.dumps({"hits": hits, "total": result.total_available}) 

85 

86 tools.append( 

87 { 

88 "spec": { 

89 "name": "memory_recall", 

90 "description": "Search long-term memory for information relevant to a query.", 

91 "inputSchema": { 

92 "json": { 

93 "type": "object", 

94 "properties": { 

95 "query": {"type": "string", "description": "Natural language search query."}, 

96 "max_results": {"type": "integer", "description": "Maximum results.", "default": 5}, 

97 }, 

98 "required": ["query"], 

99 } 

100 }, 

101 }, 

102 "handler": recall_handler, 

103 } 

104 ) 

105 

106 # ── memory_reflect ── 

107 

108 if include_reflect: 

109 

110 async def reflect_handler(tool_input: dict[str, Any]) -> str: 

111 query = tool_input["query"] 

112 result = await brain.reflect(query, bank_id=bank_id, context=context) 

113 return json.dumps({"answer": result.answer}) 

114 

115 tools.append( 

116 { 

117 "spec": { 

118 "name": "memory_reflect", 

119 "description": "Synthesize a comprehensive answer from long-term memory.", 

120 "inputSchema": { 

121 "json": { 

122 "type": "object", 

123 "properties": { 

124 "query": {"type": "string", "description": "The question to answer."}, 

125 }, 

126 "required": ["query"], 

127 } 

128 }, 

129 }, 

130 "handler": reflect_handler, 

131 } 

132 ) 

133 

134 # ── memory_forget ── 

135 

136 if include_forget: 

137 

138 async def forget_handler(tool_input: dict[str, Any]) -> str: 

139 memory_ids = tool_input["memory_ids"] 

140 if isinstance(memory_ids, str): 

141 memory_ids = [mid.strip() for mid in memory_ids.split(",")] 

142 result = await brain.forget(bank_id, memory_ids=memory_ids, context=context) 

143 return json.dumps({"deleted_count": result.deleted_count}) 

144 

145 tools.append( 

146 { 

147 "spec": { 

148 "name": "memory_forget", 

149 "description": "Remove specific memories by their IDs.", 

150 "inputSchema": { 

151 "json": { 

152 "type": "object", 

153 "properties": { 

154 "memory_ids": { 

155 "type": "array", 

156 "items": {"type": "string"}, 

157 "description": "IDs of memories to delete.", 

158 }, 

159 }, 

160 "required": ["memory_ids"], 

161 } 

162 }, 

163 }, 

164 "handler": forget_handler, 

165 } 

166 ) 

167 

168 return tools