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

1"""OpenAI-compatible tool definitions for Astrocyte memory. 

2 

3Works with OpenAI Agents SDK, Claude Agent SDK, or any system 

4that accepts OpenAI-format function/tool definitions. 

5 

6Usage: 

7 from astrocyte import Astrocyte 

8 from astrocyte.integrations.openai_agents import astrocyte_tool_definitions 

9 

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

11 tools, handlers = astrocyte_tool_definitions(brain, bank_id="user-123") 

12 

13 # tools: list of OpenAI-format tool dicts (for the API) 

14 # handlers: dict[name → async callable] (for dispatching tool calls) 

15""" 

16 

17from __future__ import annotations 

18 

19import json 

20from collections.abc import Awaitable, Callable 

21from typing import TYPE_CHECKING, Any 

22 

23if TYPE_CHECKING: 

24 from astrocyte._astrocyte import Astrocyte 

25 

26from astrocyte.types import AstrocyteContext 

27 

28# Type alias for handler functions 

29ToolHandler = Callable[..., Awaitable[str]] 

30 

31 

32def astrocyte_tool_definitions( 

33 brain: Astrocyte, 

34 bank_id: str, 

35 *, 

36 context: AstrocyteContext | None = None, 

37 include_reflect: bool = True, 

38 include_forget: bool = False, 

39) -> tuple[list[dict[str, Any]], dict[str, ToolHandler]]: 

40 """Create OpenAI-format tool definitions backed by Astrocyte. 

41 

42 Returns: 

43 (tools, handlers) where: 

44 - tools: list of dicts in OpenAI function-calling format 

45 - handlers: dict mapping tool name → async handler function 

46 

47 The tools list can be passed directly to OpenAI's `tools` parameter. 

48 When the model calls a tool, look up the handler by name and call it. 

49 """ 

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

51 handlers: dict[str, ToolHandler] = {} 

52 

53 # ── memory_retain ── 

54 

55 tools.append( 

56 { 

57 "type": "function", 

58 "function": { 

59 "name": "memory_retain", 

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

61 "parameters": { 

62 "type": "object", 

63 "properties": { 

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

65 "tags": { 

66 "type": "array", 

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

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

69 }, 

70 }, 

71 "required": ["content"], 

72 }, 

73 }, 

74 } 

75 ) 

76 

77 async def _retain(content: str, tags: list[str] | None = None, **kwargs: Any) -> str: 

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

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

80 

81 handlers["memory_retain"] = _retain 

82 

83 # ── memory_recall ── 

84 

85 tools.append( 

86 { 

87 "type": "function", 

88 "function": { 

89 "name": "memory_recall", 

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

91 "parameters": { 

92 "type": "object", 

93 "properties": { 

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

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

96 }, 

97 "required": ["query"], 

98 }, 

99 }, 

100 } 

101 ) 

102 

103 async def _recall(query: str, max_results: int = 5, **kwargs: Any) -> str: 

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

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

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

107 

108 handlers["memory_recall"] = _recall 

109 

110 # ── memory_reflect ── 

111 

112 if include_reflect: 

113 tools.append( 

114 { 

115 "type": "function", 

116 "function": { 

117 "name": "memory_reflect", 

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

119 "parameters": { 

120 "type": "object", 

121 "properties": { 

122 "query": {"type": "string", "description": "The question to answer from memory."}, 

123 }, 

124 "required": ["query"], 

125 }, 

126 }, 

127 } 

128 ) 

129 

130 async def _reflect(query: str, **kwargs: Any) -> str: 

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

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

133 

134 handlers["memory_reflect"] = _reflect 

135 

136 # ── memory_forget ── 

137 

138 if include_forget: 

139 tools.append( 

140 { 

141 "type": "function", 

142 "function": { 

143 "name": "memory_forget", 

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

145 "parameters": { 

146 "type": "object", 

147 "properties": { 

148 "memory_ids": { 

149 "type": "array", 

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

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

152 }, 

153 }, 

154 "required": ["memory_ids"], 

155 }, 

156 }, 

157 } 

158 ) 

159 

160 async def _forget(memory_ids: list[str], **kwargs: Any) -> str: 

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

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

163 

164 handlers["memory_forget"] = _forget 

165 

166 return tools, handlers