Coverage for astrocyte/documents/retrieval/tools.py: 33%

1"""make_retrieval_tools() — agent framework integration. 

2 

3Returns the three PageIndex-parity tools as plain async callables. 

4No framework-specific schema is baked in — callers add their own 

5tool descriptor (OpenAI function schema, Anthropic tool definition, 

6etc.) around these callables. 

7""" 

8 

9from __future__ import annotations 

10 

11from typing import Callable 

12 

13from astrocyte.documents.retrieval.retriever import DocumentRetriever 

14from astrocyte.documents.retrieval.types import DocumentInfo, NodeContent, TreeSkeleton 

15 

16 

17def make_retrieval_tools(retriever: DocumentRetriever) -> list[Callable]: 

18 """Return the three PageIndex-parity async tool callables. 

19 

20 Tools: 

21 get_document_info(doc_id: str) → DocumentInfo 

22 get_document_structure(doc_id: str) → TreeSkeleton 

23 get_node_content(doc_id: str, node_id: str) → NodeContent 

24 """ 

25 

26 async def get_document_info(doc_id: str) -> DocumentInfo: 

27 return await retriever.get_document_info(doc_id) 

28 

29 async def get_document_structure(doc_id: str) -> TreeSkeleton: 

30 return await retriever.get_document_structure(doc_id) 

31 

32 async def get_node_content(doc_id: str, node_id: str) -> NodeContent: 

33 return await retriever.get_node_content(doc_id, node_id) 

34 

35 get_document_info.__name__ = "get_document_info" 

36 get_document_structure.__name__ = "get_document_structure" 

37 get_node_content.__name__ = "get_node_content" 

38 

39 return [get_document_info, get_document_structure, get_node_content]