Coverage for src / lexigram / contracts / ai / tools.py: 76%

1"""Tools + LCEL contracts for G-03 parity. 

2 

3Defines tool decorators and classes analogous to LangChain's tool system. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from collections.abc import Callable 

10from dataclasses import dataclass 

11from functools import wraps 

12from typing import Any 

13 

14 

15@dataclass(frozen=True) 

16class Tool: 

17 """Base tool class (like LangChain's BaseTool). 

18 

19 Attributes: 

20 name: Tool name. 

21 description: Tool description. 

22 func: The underlying function. 

23 """ 

24 

25 name: str 

26 description: str = "" 

27 func: Callable[..., Any] | None = None 

28 

29 def invoke(self, *args: Any, **kwargs: Any) -> Any: 

30 """Synchronously invoke the tool.""" 

31 if self.func is None: 

32 raise ValueError("Tool function not set") 

33 return self.func(*args, **kwargs) 

34 

35 async def ainvoke(self, *args: Any, **kwargs: Any) -> Any: 

36 """Asynchronously invoke the tool.""" 

37 if self.func is None: 

38 raise ValueError("Tool function not set") 

39 if asyncio.iscoroutinefunction(self.func): 

40 return await self.func(*args, **kwargs) 

41 return self.func(*args, **kwargs) 

42 

43 

44class StructuredTool: 

45 """Tool with structured input/output (like LangChain's StructuredTool). 

46 

47 Analogous to LangChain's StructuredTool for tools with typed parameters. 

48 """ 

49 

50 def __init__( 

51 self, 

52 name: str, 

53 description: str, 

54 func: Callable[..., Any], 

55 ) -> None: 

56 self.name = name 

57 self.description = description 

58 self.func = func 

59 

60 def invoke(self, *args: Any, **kwargs: Any) -> Any: 

61 """Synchronously invoke the tool.""" 

62 return self.func(*args, **kwargs) 

63 

64 async def ainvoke(self, *args: Any, **kwargs: Any) -> Any: 

65 """Asynchronously invoke the tool.""" 

66 if asyncio.iscoroutinefunction(self.func): 

67 return await self.func(*args, **kwargs) 

68 return self.func(*args, **kwargs) 

69 

70 

71def tool(name: str | None = None) -> Callable[[Callable[..., Any]], Tool]: 

72 """Decorator to create a tool from a function (like @tool). 

73 

74 Args: 

75 name: Optional tool name. Defaults to function name. 

76 

77 Returns: 

78 A decorator that creates a Tool from a function. 

79 

80 Example: 

81 @tool("add") 

82 def add(a: int, b: int) -> int: 

83 '''Add two numbers.''' 

84 return a + b 

85 """ 

86 

87 def decorator(func: Callable[..., Any]) -> Tool: 

88 tool_name = name or func.__name__ 

89 tool_description = func.__doc__ or "" 

90 

91 @wraps(func) 

92 def wrapper(*args: Any, **kwargs: Any) -> Any: 

93 return func(*args, **kwargs) 

94 

95 return Tool( 

96 name=tool_name, 

97 description=tool_description.strip(), 

98 func=func, 

99 ) 

100 

101 return decorator 

102 

103 

104__all__ = [ 

105 "StructuredTool", 

106 "Tool", 

107 "tool", 

108]