Coverage for astrocyte/integrations/google_adk.py: 81%

1"""Google ADK (Agent Development Kit) integration. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.google_adk import astrocyte_adk_tools 

6 

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

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

9 

10 # Register with ADK agent 

11 agent = Agent(model="gemini-2.0-flash", tools=tools) 

12 

13Google ADK uses a callable-based tool pattern. Each tool is an async function 

14with a docstring (used as the tool description) and type-annotated parameters. 

15""" 

16 

17from __future__ import annotations 

18 

19from typing import TYPE_CHECKING, Any 

20 

21if TYPE_CHECKING: 

22 from astrocyte._astrocyte import Astrocyte 

23 

24from astrocyte.types import AstrocyteContext 

25 

26 

27def astrocyte_adk_tools( 

28 brain: Astrocyte, 

29 bank_id: str, 

30 *, 

31 context: AstrocyteContext | None = None, 

32 include_reflect: bool = True, 

33 include_forget: bool = False, 

34) -> list: 

35 """Create Google ADK-compatible tool functions backed by Astrocyte. 

36 

37 Returns a list of async callables that ADK can register as tools. 

38 Each function has proper type annotations and docstrings for ADK schema generation. 

39 """ 

40 tools: list = [] 

41 

42 async def memory_retain(content: str, tags: str | None = None) -> dict[str, Any]: 

43 """Store content into long-term memory for future recall. 

44 

45 Args: 

46 content: The text to memorize. 

47 tags: Comma-separated tags for filtering (optional). 

48 """ 

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

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

51 return {"stored": result.stored, "memory_id": result.memory_id} 

52 

53 tools.append(memory_retain) 

54 

55 async def memory_recall(query: str, max_results: int = 5) -> dict[str, Any]: 

56 """Search long-term memory for information relevant to a query. 

57 

58 Args: 

59 query: Natural language search query. 

60 max_results: Maximum number of results to return. 

61 """ 

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

63 return { 

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

65 "total": result.total_available, 

66 } 

67 

68 tools.append(memory_recall) 

69 

70 if include_reflect: 

71 

72 async def memory_reflect(query: str) -> dict[str, str]: 

73 """Synthesize a comprehensive answer from long-term memory. 

74 

75 Args: 

76 query: The question to answer from memory. 

77 """ 

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

79 return {"answer": result.answer} 

80 

81 tools.append(memory_reflect) 

82 

83 if include_forget: 

84 

85 async def memory_forget(memory_ids: str) -> dict[str, int]: 

86 """Remove specific memories by their IDs. 

87 

88 Args: 

89 memory_ids: Comma-separated memory IDs to delete. 

90 """ 

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

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

93 return {"deleted_count": result.deleted_count} 

94 

95 tools.append(memory_forget) 

96 

97 return tools