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 dataclasses import dataclass 

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

13 

14try: 

15 from pydantic import BaseModel, ValidationError 

16 

17 PYDANTIC_AVAILABLE = True 

18except ImportError: 

19 BaseModel = None 

20 ValidationError = Exception 

21 PYDANTIC_AVAILABLE = False 

22 

23T = TypeVar("T") 

24 

25 

26if PYDANTIC_AVAILABLE: 

27 from pydantic import ConfigDict 

28 

29 

30class StructuredOutput(BaseModel if PYDANTIC_AVAILABLE else object): 

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 

57 success: bool 

58 output: T | None = None 

59 error: str | None = None 

60 

61 

62class OutputValidator(Generic[T]): 

63 """ 

64 Validator for structured outputs. 

65 

66 Usage: 

67 validator = OutputValidator(MyOutput) 

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

69 if result.success: 

70 output = result.output # MyOutput instance 

71 """ 

72 

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

74 """ 

75 Initialize validator. 

76 

77 Args: 

78 output_type: Expected output type 

79 """ 

80 self.output_type = output_type 

81 self._is_pydantic = ( 

82 PYDANTIC_AVAILABLE 

83 and isinstance(output_type, type) 

84 and issubclass(output_type, BaseModel) 

85 ) 

86 

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

88 """ 

89 Validate data against output type. 

90 

91 Args: 

92 data: Data to validate 

93 

94 Returns: 

95 ValidationResult with success/error info 

96 """ 

97 # If not Pydantic, just check type 

98 if not self._is_pydantic: 

99 if isinstance(data, self.output_type): 

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

101 else: 

102 return ValidationResult( 

103 success=False, 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 

187 if issubclass(origin, Agent): 

188 args = get_args(base) 

189 if len(args) >= 2: 

190 return args[1] 

191 

192 return None