Coverage for astrocyte/integrations/pydantic_ai.py: 93%

1"""Pydantic AI integration — Astrocyte as agent tools. 

2 

3Usage: 

4 from astrocyte import Astrocyte 

5 from astrocyte.integrations.pydantic_ai import astrocyte_tools 

6 

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

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

9 

10 agent = Agent(model="claude-sonnet-4-20250514", tools=tools) 

11 

12Exposes retain, recall, reflect as tool functions that 

13Pydantic AI agents can call during execution. 

14""" 

15 

16from __future__ import annotations 

17 

18from typing import TYPE_CHECKING, Any 

19 

20if TYPE_CHECKING: 

21 from astrocyte._astrocyte import Astrocyte 

22 

23from astrocyte.types import AstrocyteContext 

24 

25 

26def astrocyte_tools( 

27 brain: Astrocyte, 

28 bank_id: str, 

29 *, 

30 context: AstrocyteContext | None = None, 

31 include_reflect: bool = True, 

32 include_forget: bool = False, 

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

34 """Create Pydantic AI-compatible tool definitions backed by Astrocyte. 

35 

36 Returns a list of tool definition dicts that Pydantic AI can register. 

37 Each tool is a dict with 'name', 'description', 'function', and 'parameters'. 

38 """ 

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

40 

41 async def memory_retain(content: str, tags: list[str] | None = None) -> str: 

42 """Store content into long-term memory.""" 

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

44 if result.stored: 

45 return f"Stored memory (id: {result.memory_id})" 

46 return f"Failed to store: {result.error}" 

47 

48 tools.append( 

49 { 

50 "name": "memory_retain", 

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

52 "function": memory_retain, 

53 } 

54 ) 

55 

56 async def memory_recall(query: str, max_results: int = 5) -> str: 

57 """Search long-term memory for relevant information.""" 

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

59 if not result.hits: 

60 return "No relevant memories found." 

61 lines = [f"- [{h.score:.2f}] {h.text}" for h in result.hits] 

62 return f"Found {len(result.hits)} memories:\n" + "\n".join(lines) 

63 

64 tools.append( 

65 { 

66 "name": "memory_recall", 

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

68 "function": memory_recall, 

69 } 

70 ) 

71 

72 if include_reflect: 

73 

74 async def memory_reflect(query: str) -> str: 

75 """Synthesize an answer from long-term memory.""" 

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

77 return result.answer 

78 

79 tools.append( 

80 { 

81 "name": "memory_reflect", 

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

83 "function": memory_reflect, 

84 } 

85 ) 

86 

87 if include_forget: 

88 

89 async def memory_forget(memory_ids: list[str]) -> str: 

90 """Remove specific memories by their IDs.""" 

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

92 return f"Deleted {result.deleted_count} memories." 

93 

94 tools.append( 

95 { 

96 "name": "memory_forget", 

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

98 "function": memory_forget, 

99 } 

100 ) 

101 

102 return tools