161718192021222324252627282930313233343536373839404142434445464748495051

classAgentOutputSchemaBase(abc.ABC):"""An object that captures the JSON schema of the output, as well as validating/parsing JSON    produced by the LLM into the output type.    """@abc.abstractmethoddefis_plain_text(self)->bool:"""Whether the output type is plain text (versus a JSON object)."""pass@abc.abstractmethoddefname(self)->str:"""The name of the output type."""pass@abc.abstractmethoddefjson_schema(self)->dict[str,Any]:"""Returns the JSON schema of the output. Will only be called if the output type is not        plain text.        """pass@abc.abstractmethoddefis_strict_json_schema(self)->bool:"""Whether the JSON schema is in strict mode. Strict mode constrains the JSON schema        features, but guarantees valid JSON. See here for details:        https://platform.openai.com/docs/guides/structured-outputs#supported-schemas        """pass@abc.abstractmethoddefvalidate_json(self,json_str:str)->Any:"""Validate a JSON string against the output type. You must return the validated object,        or raise a `ModelBehaviorError` if the JSON is invalid.        """pass

 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168

@dataclass(init=False)classAgentOutputSchema(AgentOutputSchemaBase):"""An object that captures the JSON schema of the output, as well as validating/parsing JSON    produced by the LLM into the output type.    """output_type:type[Any]"""The type of the output."""_type_adapter:TypeAdapter[Any]"""A type adapter that wraps the output type, so that we can validate JSON."""_is_wrapped:bool"""Whether the output type is wrapped in a dictionary. This is generally done if the base    output type cannot be represented as a JSON Schema object.    """_output_schema:dict[str,Any]"""The JSON schema of the output."""_strict_json_schema:bool"""Whether the JSON schema is in strict mode. We **strongly** recommend setting this to True,    as it increases the likelihood of correct JSON input.    """def__init__(self,output_type:type[Any],strict_json_schema:bool=True):"""        Args:            output_type: The type of the output.            strict_json_schema: Whether the JSON schema is in strict mode. We **strongly** recommend                setting this to True, as it increases the likelihood of correct JSON input.        """self.output_type=output_typeself._strict_json_schema=strict_json_schemaifoutput_typeisNoneoroutput_typeisstr:self._is_wrapped=Falseself._type_adapter=TypeAdapter(output_type)self._output_schema=self._type_adapter.json_schema()return# We should wrap for things that are not plain text, and for things that would definitely# not be a JSON Schema object.self._is_wrapped=not_is_subclass_of_base_model_or_dict(output_type)ifself._is_wrapped:OutputType=TypedDict("OutputType",{_WRAPPER_DICT_KEY:output_type,# type: ignore},)self._type_adapter=TypeAdapter(OutputType)self._output_schema=self._type_adapter.json_schema()else:self._type_adapter=TypeAdapter(output_type)self._output_schema=self._type_adapter.json_schema()ifself._strict_json_schema:try:self._output_schema=ensure_strict_json_schema(self._output_schema)exceptUserErrorase:raiseUserError("Strict JSON schema is enabled, but the output type is not valid. ""Either make the output type strict, ""or wrap your type with AgentOutputSchema(YourType, strict_json_schema=False)")fromedefis_plain_text(self)->bool:"""Whether the output type is plain text (versus a JSON object)."""returnself.output_typeisNoneorself.output_typeisstrdefis_strict_json_schema(self)->bool:"""Whether the JSON schema is in strict mode."""returnself._strict_json_schemadefjson_schema(self)->dict[str,Any]:"""The JSON schema of the output type."""ifself.is_plain_text():raiseUserError("Output type is plain text, so no JSON schema is available")returnself._output_schemadefvalidate_json(self,json_str:str)->Any:"""Validate a JSON string against the output type. Returns the validated object, or raises        a `ModelBehaviorError` if the JSON is invalid.        """validated=_json.validate_json(json_str,self._type_adapter,partial=False)ifself._is_wrapped:ifnotisinstance(validated,dict):_error_tracing.attach_error_to_current_span(SpanError(message="Invalid JSON",data={"details":f"Expected a dict, got{type(validated)}"},))raiseModelBehaviorError(f"Expected a dict, got{type(validated)} for JSON:{json_str}")if_WRAPPER_DICT_KEYnotinvalidated:_error_tracing.attach_error_to_current_span(SpanError(message="Invalid JSON",data={"details":f"Could not find key{_WRAPPER_DICT_KEY} in JSON"},))raiseModelBehaviorError(f"Could not find key{_WRAPPER_DICT_KEY} in JSON:{json_str}")returnvalidated[_WRAPPER_DICT_KEY]returnvalidateddefname(self)->str:"""The name of the output type."""return_type_to_str(self.output_type)