Coverage for agentos/protocols/output.py: 38%

1""" 

2Structured output validation for NexusAgent. 

3 

4Provides Pydantic-style output validation for agents. 

5When Agent[Deps, Out] has Out as a Pydantic BaseModel, 

6the output is automatically validated. 

7""" 

8 

9from __future__ import annotations 

10 

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

12from dataclasses import dataclass 

13 

14try: 

15 from pydantic import BaseModel, ValidationError 

16 PYDANTIC_AVAILABLE = True 

17except ImportError: 

18 BaseModel = None 

19 ValidationError = Exception 

20 PYDANTIC_AVAILABLE = False 

21 

22T = TypeVar("T") 

23 

24 

25if PYDANTIC_AVAILABLE: 

26 from pydantic import ConfigDict 

27 

28 

29class StructuredOutput(BaseModel if PYDANTIC_AVAILABLE else object): 

30 

31 """Agent 结构化输出。""" 

32 

33 """ 

34 Base class for structured outputs. 

35 

36 Usage: 

37 class MyOutput(StructuredOutput): 

38 answer: str 

39 confidence: float 

40 sources: list[str] 

41 """ 

42 if PYDANTIC_AVAILABLE: 

43 model_config = ConfigDict(extra="forbid") 

44 

45 

46@dataclass 

47class ValidationResult(Generic[T]): 

48 """ 

49 Result of output validation. 

50 

51 Attributes: 

52 success: Whether validation passed 

53 output: Validated output (if success) 

54 error: Validation error (if failed) 

55 """ 

56 success: bool 

57 output: T | None = None 

58 error: str | None = None 

59 

60 

61class OutputValidator(Generic[T]): 

62 """ 

63 Validator for structured outputs. 

64 

65 Usage: 

66 validator = OutputValidator(MyOutput) 

67 result = validator.validate({"answer": "42", "confidence": 0.9}) 

68 if result.success: 

69 output = result.output # MyOutput instance 

70 """ 

71 

72 def __init__(self, output_type: type[T]): 

73 """ 

74 Initialize validator. 

75 

76 Args: 

77 output_type: Expected output type 

78 """ 

79 self.output_type = output_type 

80 self._is_pydantic = ( 

81 PYDANTIC_AVAILABLE and 

82 isinstance(output_type, type) and 

83 issubclass(output_type, BaseModel) 

84 ) 

85 

86 def validate(self, data: Any) -> ValidationResult[T]: 

87 """ 

88 Validate data against output type. 

89 

90 Args: 

91 data: Data to validate 

92 

93 Returns: 

94 ValidationResult with success/error info 

95 """ 

96 # If not Pydantic, just check type 

97 if not self._is_pydantic: 

98 if isinstance(data, self.output_type): 

99 return ValidationResult(success=True, output=data) 

100 else: 

101 return ValidationResult( 

102 success=False, 

103 error=f"Expected {self.output_type}, got {type(data)}" 

104 ) 

105 

106 # Pydantic validation 

107 try: 

108 if isinstance(data, self.output_type): 

109 # Already correct type 

110 return ValidationResult(success=True, output=data) 

111 elif isinstance(data, dict): 

112 # Try to construct from dict 

113 output = self.output_type(**data) 

114 return ValidationResult(success=True, output=output) 

115 else: 

116 # Try model_validate 

117 output = self.output_type.model_validate(data) 

118 return ValidationResult(success=True, output=output) 

119 except ValidationError as e: 

120 return ValidationResult(success=False, error=str(e)) 

121 except Exception as e: 

122 return ValidationResult(success=False, error=str(e)) 

123 

124 def validate_or_raise(self, data: Any) -> T: 

125 """ 

126 Validate data, raise on failure. 

127 

128 Args: 

129 data: Data to validate 

130 

131 Returns: 

132 Validated output 

133 

134 Raises: 

135 ValueError: If validation fails 

136 """ 

137 result = self.validate(data) 

138 if not result.success: 

139 raise ValueError(f"Output validation failed: {result.error}") 

140 return result.output 

141 

142 

143def validate_output(output_type: type[T], data: Any) -> ValidationResult[T]: 

144 """ 

145 Validate data against output type. 

146 

147 Convenience function wrapping OutputValidator. 

148 

149 Args: 

150 output_type: Expected output type 

151 data: Data to validate 

152 

153 Returns: 

154 ValidationResult with success/error info 

155 

156 Usage: 

157 result = validate_output(MyOutput, {"answer": "42"}) 

158 if result.success: 

159 output = result.output 

160 """ 

161 validator = OutputValidator(output_type) 

162 return validator.validate(data) 

163 

164 

165def get_output_type(agent_class: type) -> type | None: 

166 """ 

167 Extract output type from Agent class. 

168 

169 Args: 

170 agent_class: Agent subclass 

171 

172 Returns: 

173 Output type if found, None otherwise 

174 """ 

175 # Check type hints 

176 hints = get_type_hints(agent_class) 

177 if 'Out' in hints: 

178 return hints['Out'] 

179 

180 # Check generic base 

181 for base in agent_class.__mro__: 

182 origin = get_origin(base) 

183 if origin is not None: 

184 # Check if it's Agent 

185 from agentos.core.di import Agent 

186 if issubclass(origin, Agent): 

187 args = get_args(base) 

188 if len(args) >= 2: 

189 return args[1] 

190 

191 return None