Add Agent.output_json_schema() method #3454
Conversation
| @DouweM is there a source of truth I can reference for what the schema output should be for tests? I think I will get stuck figuring this out so any guidance will be appreciated. Edit: I added snapshots of expected output to the tests. Please review them carefully as I'm guessing what they are supposed to be. |
g-eoj changed the title 
| @DouweM thanks for the reviews so far. Fixed: allows_text to True. This doesn't seem right but I'll wait for confirmation. |
| @DouweM ready for review |
| json_schemas: list[JsonSchema] = [] | ||
| for output_spec in flat_output_type: | ||
| json_schema = TypeAdapter(output_spec).json_schema(mode='serialization') | ||
| if json_schema not in json_schemas: |
There was a problem hiding this comment.
Maybe json_schemas can be a set so that it's automatically de-duped
There was a problem hiding this comment.
I don't think I can do that because dictionaries are mutable.
| elif isinstance(output_spec, _output.ToolOutput): | ||
| flat_output_type += _flatten_output_spec(output_spec.output) | ||
| elif inspect.isfunction(output_spec) or inspect.ismethod(output_spec): | ||
| return_annotation = inspect.signature(output_spec).return_annotation |
There was a problem hiding this comment.
What if there is no return annotation on the output function?
There was a problem hiding this comment.
In function_schema we use a helper method for this:
I think its return value will have a 'return' key. Can we use it here as well?
There was a problem hiding this comment.
Done - if there is no key, I set the return type to string.
| raise ValueError('`BinaryImage` must be have a media type that starts with "image/"') # pragma: no cover | ||
| | ||
| @classmethod | ||
| def __get_pydantic_json_schema__( |
There was a problem hiding this comment.
On second thought, let's leave these changes out for now. Once we get to #3122, we'll also need to think about the schemas for BinaryContent, FileUrl etc, so I'd rather tackle this in one go than instead of giving BinaryImage special treatment now.
There was a problem hiding this comment.
I reverted all changes to BinaryImage, hopefully that is okay.
| @DouweM ready for review |
| | ||
| def flatten_output_type( | ||
| output_type: OutputSpec[OutputDataT], | ||
| ) -> list[OutputSpec[OutputDataT] | type[str]]: |
There was a problem hiding this comment.
Would this be list[OutputDataT], without the OutputSpec? 🤔
There was a problem hiding this comment.
_flatten_output_spec returns OutputSpec sequences, so out of the box it doesn't work (I gave it a try, type checkers not pleased). So I think I would have to cast flat_output_type or write my own version of _flatten_output_spec to achieve the effect of having list[OutputDataT] as the return.
| return outputs_flat | ||
| | ||
| | ||
| def flatten_output_type( |
| @g-eoj Thanks Joe, looks great! 🤞🏻 that CI goes green |
2 participants
This PR adds a method to get the json schema for agent output.
It aims to support all output types, including special ones:
BinaryImageDeferredToolRequestsSimple example:
{'type': 'string'}Override agent output types:
{'type': 'bool'}Closes #3225