Coverage for core/langgraph/tools/registry.py: 100%

1""" 

2Tool resolution surface: a name-keyed store supporting eager registration and 

3lazy import-by-name. 

4 

5The registry holds only what it is explicitly given. It decides no policy: it 

6never reads manifests, never chooses an agent's tools, and never special-cases a 

7capability. One tool resolves to one module attribute (ADR-0038). 

8""" 

9from __future__ import annotations 

10 

11import importlib 

12from collections.abc import Callable, Sequence 

13from dataclasses import dataclass 

14from typing import Any 

15 

16from .contract import Tool, ToolContext, ToolResult 

17 

18 

19class ToolRegistryError(ValueError): 

20 """Raised for an unknown name or a duplicate registration.""" 

21 

22 

23@dataclass(frozen=True) 

24class _LazyLocator: 

25 """Deferred locator for a tool not yet imported.""" 

26 

27 module_path: str 

28 attr: str 

29 

30 

31class ToolRegistry: 

32 """Stores tools by name and resolves declarations in declaration order.""" 

33 

34 def __init__(self) -> None: 

35 self._entries: dict[str, Tool | _LazyLocator] = {} 

36 

37 def register(self, name: str, tool: Tool) -> None: 

38 """Eagerly register *tool* under *name*.""" 

39 self._reject_duplicate(name) 

40 self._entries[name] = tool 

41 

42 def register_lazy(self, name: str, module_path: str, attr: str = "TOOL") -> None: 

43 """Register a lazy locator; the module is imported on first resolution.""" 

44 self._reject_duplicate(name) 

45 self._entries[name] = _LazyLocator(module_path=module_path, attr=attr) 

46 

47 def has(self, name: str) -> bool: 

48 """Report registration without importing a lazy module.""" 

49 return name in self._entries 

50 

51 def registered_names(self) -> list[str]: 

52 """All registered names in registration order.""" 

53 return list(self._entries) 

54 

55 def resolve(self, names: Sequence[str]) -> list[Tool]: 

56 """Resolve *names* in the exact given order; duplicates are preserved.""" 

57 return [self._resolve_one(name) for name in names] 

58 

59 def definitions(self, names: Sequence[str]) -> list[dict[str, Any]]: 

60 """Request-level OpenAI function-tool definitions in declaration order.""" 

61 return [tool.definition() for tool in self.resolve(names)] 

62 

63 def dispatch_map( 

64 self, names: Sequence[str] 

65 ) -> dict[str, Callable[[ToolContext], ToolResult]]: 

66 """Map each declared name to a callable invoking that tool's execute.""" 

67 return { 

68 name: self._dispatcher(tool) 

69 for name, tool in zip(names, self.resolve(names)) 

70 } 

71 

72 def _reject_duplicate(self, name: str) -> None: 

73 if name in self._entries: 

74 raise ToolRegistryError(f"tool {name!r} is already registered") 

75 

76 def _resolve_one(self, name: str) -> Tool: 

77 entry = self._entries.get(name) 

78 if entry is None: 

79 raise ToolRegistryError(f"tool {name!r} is not registered") 

80 if isinstance(entry, _LazyLocator): 

81 entry = self._import_tool(name, entry) 

82 self._entries[name] = entry 

83 return entry 

84 

85 @staticmethod 

86 def _import_tool(name: str, locator: _LazyLocator) -> Tool: 

87 try: 

88 module = importlib.import_module(locator.module_path) 

89 except ImportError as exc: 

90 raise ToolRegistryError( 

91 f"tool {name!r} locator " 

92 f"{locator.module_path}:{locator.attr} could not be imported: {exc}" 

93 ) from exc 

94 tool = getattr(module, locator.attr, None) 

95 if tool is None: 

96 raise ToolRegistryError( 

97 f"tool {name!r} locator " 

98 f"{locator.module_path}:{locator.attr} has no such attribute" 

99 ) 

100 return tool 

101 

102 @staticmethod 

103 def _dispatcher(tool: Tool) -> Callable[[ToolContext], ToolResult]: 

104 def dispatch(context: ToolContext) -> ToolResult: 

105 return tool.execute(context) 

106 return dispatch