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

1"""AutoGen / AG2 integration — Astrocyte memory for multi-agent conversations. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.autogen import AstrocyteAutoGenMemory 

6 

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

8 memory = AstrocyteAutoGenMemory(brain, bank_id="team-agents") 

9 

10 # Register with AutoGen agent 

11 agent = ConversableAgent("assistant", llm_config=llm_config) 

12 # Use memory in conversation hooks 

13 await memory.save("User prefers Python", agent_id="assistant") 

14 context = await memory.get_context("What does the user prefer?") 

15 

16Also provides tool registration for AutoGen's function-calling pattern: 

17 tools = memory.as_tools() 

18 agent.register_for_llm(tools) 

19""" 

20 

21from __future__ import annotations 

22 

23from typing import TYPE_CHECKING, Any 

24 

25if TYPE_CHECKING: 

26 from astrocyte._astrocyte import Astrocyte 

27 

28from astrocyte.types import AstrocyteContext 

29 

30 

31class AstrocyteAutoGenMemory: 

32 """Astrocyte-backed memory for AutoGen / AG2 agents. 

33 

34 Provides both a direct API (save/query/get_context) and tool registration 

35 for AutoGen's function-calling pattern. 

36 """ 

37 

38 def __init__( 

39 self, 

40 brain: Astrocyte, 

41 bank_id: str, 

42 *, 

43 context: AstrocyteContext | None = None, 

44 agent_banks: dict[str, str] | None = None, 

45 ) -> None: 

46 self.brain = brain 

47 self.bank_id = bank_id 

48 self._context = context 

49 self._agent_banks = agent_banks or {} 

50 

51 def _resolve_bank(self, agent_id: str | None = None) -> str: 

52 if agent_id and agent_id in self._agent_banks: 

53 return self._agent_banks[agent_id] 

54 return self.bank_id 

55 

56 async def save( 

57 self, 

58 content: str, 

59 *, 

60 agent_id: str | None = None, 

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

62 ) -> str | None: 

63 """Save content to memory. Returns memory_id.""" 

64 bank = self._resolve_bank(agent_id) 

65 result = await self.brain.retain( 

66 content, 

67 bank_id=bank, 

68 tags=tags or ["autogen"], 

69 metadata={"source": "autogen", "agent_id": agent_id or ""}, 

70 context=self._context, 

71 ) 

72 return result.memory_id if result.stored else None 

73 

74 async def query( 

75 self, 

76 query: str, 

77 *, 

78 agent_id: str | None = None, 

79 max_results: int = 5, 

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

81 """Query memory. Returns list of hit dicts.""" 

82 bank = self._resolve_bank(agent_id) 

83 result = await self.brain.recall(query, bank_id=bank, max_results=max_results, context=self._context) 

84 return [{"text": h.text, "score": h.score, "memory_id": h.memory_id} for h in result.hits] 

85 

86 async def get_context( 

87 self, 

88 query: str, 

89 *, 

90 agent_id: str | None = None, 

91 max_results: int = 5, 

92 ) -> str: 

93 """Get formatted memory context for injection into agent messages.""" 

94 hits = await self.query(query, agent_id=agent_id, max_results=max_results) 

95 if not hits: 

96 return "" 

97 return "\n".join(f"- {h['text']}" for h in hits) 

98 

99 def as_tools(self, *, include_reflect: bool = True) -> list[dict[str, Any]]: 

100 """Return OpenAI-format tool definitions for AutoGen function calling. 

101 

102 AutoGen uses OpenAI-compatible tool definitions internally. 

103 """ 

104 from astrocyte.integrations.openai_agents import astrocyte_tool_definitions 

105 

106 tools, _handlers = astrocyte_tool_definitions( 

107 self.brain, self.bank_id, context=self._context, include_reflect=include_reflect 

108 ) 

109 return tools 

110 

111 def get_handlers(self, *, include_reflect: bool = True) -> dict[str, Any]: 

112 """Return handler functions for dispatching tool calls.""" 

113 from astrocyte.integrations.openai_agents import astrocyte_tool_definitions 

114 

115 _tools, handlers = astrocyte_tool_definitions( 

116 self.brain, self.bank_id, context=self._context, include_reflect=include_reflect 

117 ) 

118 return handlers