1"""

2 Structured output validation for NexusAgent.

3

4 Provides Pydantic-style output validation for agents.

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

6 the output is automatically validated.

7"""

8

9 from __future__ import annotations

10

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

12 from dataclasses import dataclass

13

14 try:

15 from pydantic import BaseModel, ValidationError

16 PYDANTIC_AVAILABLE = True

17 except ImportError:

18 BaseModel = None

19 ValidationError = Exception

20 PYDANTIC_AVAILABLE = False

21

22 T = TypeVar("T")

23

24

25 if PYDANTIC_AVAILABLE:

26 from pydantic import ConfigDict

27

28

29 class 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

47 class 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

61 class 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

143 def 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

165 def 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