Coverage for agentos/core/di.py: 34%

1""" 

2Dependency Injection system for NexusAgent. 

3 

4Provides type-safe Agent[Deps, Out] generic base class, 

5RunContext for dependency injection, and Depends() for 

6automatic dependency resolution. 

7""" 

8 

9from __future__ import annotations 

10 

11import uuid 

12from collections.abc import Callable 

13from dataclasses import dataclass, field 

14from typing import Any, Generic, TypeVar, get_args, get_origin, get_type_hints 

15 

16# Type variables for Agent generic 

17Deps = TypeVar("Deps") 

18Out = TypeVar("Out") 

19 

20 

21@dataclass 

22class RunContext(Generic[Deps]): 

23 """ 

24 Runtime context passed to Agent.run(). 

25 

26 Contains: 

27 - deps: The dependencies for this agent 

28 - agent_name: Name of the agent 

29 - run_id: Unique ID for this run 

30 - metadata: Additional metadata 

31 """ 

32 

33 deps: Deps 

34 agent_name: str = "" 

35 run_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) 

36 metadata: dict[str, Any] = field(default_factory=dict) 

37 

38 def get(self, key: str, default: Any = None) -> Any: 

39 """Get metadata value.""" 

40 return self.metadata.get(key, default) 

41 

42 def set(self, key: str, value: Any) -> None: 

43 """Set metadata value.""" 

44 self.metadata[key] = value 

45 

46 

47class Depends: 

48 """ 

49 Dependency marker for automatic injection. 

50 

51 Usage: 

52 def get_db() -> Database: 

53 return Database() 

54 

55 class MyAgent(Agent[Depends(get_db), str]): 

56 async def run(self, ctx): 

57 db = ctx.deps # Database instance 

58 """ 

59 

60 def __init__(self, callable: Callable[..., Any]): 

61 self.callable = callable 

62 

63 def resolve(self) -> Any: 

64 """Resolve the dependency.""" 

65 return self.callable() 

66 

67 

68def inject_tool(tool: Callable[..., Any]) -> Callable[..., Any]: 

69 """ 

70 Decorator to inject a tool into an agent. 

71 

72 Usage: 

73 @inject_tool(search_tool) 

74 class MyAgent(Agent): 

75 ... 

76 """ 

77 

78 def decorator(cls): 

79 if not hasattr(cls, "_tools"): 

80 cls._tools = [] 

81 cls._tools.append(tool) 

82 return cls 

83 

84 return decorator 

85 

86 

87def requires_context(*fields: str) -> Callable[..., Any]: 

88 """ 

89 Decorator to declare required context fields. 

90 

91 Usage: 

92 @requires_context("user_id", "session_id") 

93 class MyAgent(Agent): 

94 ... 

95 """ 

96 

97 def decorator(cls): 

98 if not hasattr(cls, "_required_context"): 

99 cls._required_context = [] 

100 cls._required_context.extend(fields) 

101 return cls 

102 

103 return decorator 

104 

105 

106class Agent(Generic[Deps, Out]): 

107 """ 

108 Base class for all agents. 

109 

110 Type-safe generic: Agent[Deps, Out] 

111 - Deps: Type of dependencies 

112 - Out: Type of output 

113 

114 Usage: 

115 class MyAgent(Agent[str, str]): 

116 async def run(self, ctx: RunContext[str]) -> str: 

117 return f"Hello, {ctx.deps}!" 

118 

119 agent = MyAgent() 

120 result = await agent.invoke("World") 

121 """ 

122 

123 def __init__(self, name: str = ""): 

124 self.name = name or self.__class__.__name__ 

125 self._tools: list[Callable[..., Any]] = getattr(self.__class__, "_tools", []) 

126 self._required_context: list[str] = getattr(self.__class__, "_required_context", []) 

127 

128 async def run(self, ctx: RunContext[Deps]) -> Out: 

129 """ 

130 Main agent logic. Override in subclass. 

131 

132 Args: 

133 ctx: Runtime context with dependencies 

134 

135 Returns: 

136 Agent output (type-checked against Out) 

137 """ 

138 raise NotImplementedError("Subclass must implement run()") 

139 

140 async def invoke(self, deps: Deps, **metadata) -> Out: 

141 """ 

142 Invoke the agent with dependencies. 

143 

144 Args: 

145 deps: Dependencies to inject 

146 **metadata: Additional metadata 

147 

148 Returns: 

149 Agent output 

150 """ 

151 # Resolve Depends if needed 

152 if isinstance(deps, Depends): 

153 deps = deps.resolve() 

154 

155 # Create context 

156 ctx = RunContext[Deps]( 

157 deps=deps, 

158 agent_name=self.name, 

159 metadata=metadata, 

160 ) 

161 

162 # Validate required context 

163 for f in self._required_context: 

164 if f not in ctx.metadata: 

165 raise ValueError(f"Required context field missing: {f}") 

166 

167 # Run agent 

168 result = await self.run(ctx) 

169 

170 # Validate output type (if type hints available) 

171 result = self._validate_output(result) 

172 

173 return result 

174 

175 def _validate_output(self, result: Any) -> Out: 

176 """ 

177 Validate output against declared type. 

178 

179 Uses Pydantic validation if Out is a BaseModel, 

180 otherwise basic type checking. 

181 """ 

182 # Get type hints 

183 hints = get_type_hints(self.__class__) 

184 out_type = hints.get("Out") 

185 

186 if out_type is None: 

187 # Try to get from generic base 

188 for base in self.__class__.__mro__: 

189 origin = get_origin(base) 

190 if origin is Agent: 

191 args = get_args(base) 

192 if len(args) >= 2: 

193 out_type = args[1] 

194 break 

195 

196 if out_type is None: 

197 return result 

198 

199 # Check if it's a Pydantic model 

200 try: 

201 from pydantic import BaseModel 

202 

203 if isinstance(out_type, type) and issubclass(out_type, BaseModel): 

204 if not isinstance(result, out_type): 

205 # Try to validate/convert 

206 if isinstance(result, dict): 

207 result = out_type(**result) 

208 else: 

209 result = out_type.model_validate(result) 

210 except ImportError: 

211 pass 

212 

213 return result 

214 

215 def get_tools(self) -> list[Callable[..., Any]]: 

216 """Get registered tools.""" 

217 return self._tools.copy() 

218 

219 def __repr__(self) -> str: 

220 return f"{self.__class__.__name__}(name={self.name!r})"