Coverage for astrocyte/integrations/beeai.py: 79%

1"""BeeAI / IBM Bee Agent Framework integration. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.beeai import astrocyte_bee_tools 

6 

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

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

9 

10 # Register with BeeAI agent 

11 agent = BeeAgent(llm=llm, tools=tools) 

12 

13BeeAI uses a tool pattern where each tool has a name, description, 

14input schema, and an async handler function. Similar to OpenAI tools 

15but with BeeAI-specific tool class wrappers. 

16""" 

17 

18from __future__ import annotations 

19 

20import json 

21from typing import TYPE_CHECKING, Any 

22 

23if TYPE_CHECKING: 

24 from astrocyte._astrocyte import Astrocyte 

25 

26from astrocyte.types import AstrocyteContext 

27 

28 

29class AstrocyteBeeTool: 

30 """A single Astrocyte tool compatible with BeeAI's Tool interface. 

31 

32 BeeAI expects tools with: name, description, input_schema, and a run() method. 

33 """ 

34 

35 def __init__( 

36 self, 

37 name: str, 

38 description: str, 

39 input_schema: dict[str, Any], 

40 handler: Any, 

41 ) -> None: 

42 self.name = name 

43 self.description = description 

44 self.input_schema = input_schema 

45 self._handler = handler 

46 

47 async def run(self, input_data: dict[str, Any]) -> str: 

48 """Execute the tool. Returns a string result.""" 

49 return await self._handler(input_data) 

50 

51 

52def astrocyte_bee_tools( 

53 brain: Astrocyte, 

54 bank_id: str, 

55 *, 

56 context: AstrocyteContext | None = None, 

57 include_reflect: bool = True, 

58 include_forget: bool = False, 

59) -> list[AstrocyteBeeTool]: 

60 """Create BeeAI-compatible tools backed by Astrocyte.""" 

61 tools: list[AstrocyteBeeTool] = [] 

62 

63 async def _retain(input_data: dict[str, Any]) -> str: 

64 content = input_data["content"] 

65 tags = input_data.get("tags") 

66 tag_list = [t.strip() for t in tags.split(",")] if isinstance(tags, str) and tags else tags 

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

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

69 

70 tools.append( 

71 AstrocyteBeeTool( 

72 name="memory_retain", 

73 description="Store content into long-term memory.", 

74 input_schema={"type": "object", "properties": {"content": {"type": "string"}}, "required": ["content"]}, 

75 handler=_retain, 

76 ) 

77 ) 

78 

79 async def _recall(input_data: dict[str, Any]) -> str: 

80 query = input_data["query"] 

81 max_results = input_data.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 AstrocyteBeeTool( 

88 name="memory_recall", 

89 description="Search long-term memory for relevant information.", 

90 input_schema={"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, 

91 handler=_recall, 

92 ) 

93 ) 

94 

95 if include_reflect: 

96 

97 async def _reflect(input_data: dict[str, Any]) -> str: 

98 result = await brain.reflect(input_data["query"], bank_id=bank_id, context=context) 

99 return result.answer 

100 

101 tools.append( 

102 AstrocyteBeeTool( 

103 name="memory_reflect", 

104 description="Synthesize an answer from long-term memory.", 

105 input_schema={"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, 

106 handler=_reflect, 

107 ) 

108 ) 

109 

110 if include_forget: 

111 

112 async def _forget(input_data: dict[str, Any]) -> str: 

113 ids = input_data["memory_ids"] 

114 if isinstance(ids, str): 

115 ids = [mid.strip() for mid in ids.split(",")] 

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

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

118 

119 tools.append( 

120 AstrocyteBeeTool( 

121 name="memory_forget", 

122 description="Remove specific memories by their IDs.", 

123 input_schema={ 

124 "type": "object", 

125 "properties": {"memory_ids": {"type": "array", "items": {"type": "string"}}}, 

126 "required": ["memory_ids"], 

127 }, 

128 handler=_forget, 

129 ) 

130 ) 

131 

132 return tools