- Notifications
You must be signed in to change notification settings - Fork1k
Make it possible to stream structured while using.iter#2078
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.
Already on GitHub?Sign in to your account
base:main
Are you sure you want to change the base?
Conversation
| @@ -95,7 +95,7 @@ async def _validate_response( | |||
| match = self._output_schema.find_named_tool(message.parts, output_tool_name) | |||
| if match is None: | |||
| raise exceptions.UnexpectedModelBehavior( # pragma: no cover | |||
| f'Invalid response, unable to find tool: {self._output_schema.tool_names()}' | |||
| f'Invalid response, unable to find tool {output_tool_name!r}; expected one of {self._output_schema.tool_names()}' | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
this error message and the one below just seemed confusing; I suspect they are hard/impossible to hit in practice but they should still include the relevant information imo
9b55b99 to215b39fCompare| @@ -2052,6 +2052,69 @@ def usage(self) -> _usage.Usage: | |||
| """Get usage statistics for the run so far, including token usage, model requests, and so on.""" | |||
| return self._graph_run.state.usage | |||
| async def validate_structured_output( | |||
| self, response: _messages.ModelResponse, *, tool_name: str | None = None, allow_partial: bool = False | |||
| ) -> FinalResult[OutputDataT] | None: | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
My biggest concern with this implementation is that it is annotated as returning aFinalResult. That type is exactly what you want to be returning here, I think (well, it would be enough to just useOption, but that's not public, and why not include tool info?). But the nameFinalResult is kind of confusing when the purpose of adding this was to support cases where it literally is not the final result, i.e. partial streaming. On the other hand, it's pretty clear what's going on imo, and do we really want to introduce anotherOption-like result marker type? I'm open to others' opinions.
| tool_name: If provided, this should be the name of the tool that will produce the output. | ||
| This is only included so that you can skip the tool-lookup step if you already know which tool | ||
| produced the output, which may be the case if calling this method in a loop while streaming. | ||
| (You can get the `tool_name` from the return value of this method when it is not `None`.) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
Not sure if we really want to add the complexity associated with allowing the end user to specify the tool, or just always look up the tool here. I'd be okay either way, dropping this should be trivial though if we want to (just delete the argument and the firstif-statement referencing it below).
| if isinstance(output_schema, _output.ToolOutputSchema): | ||
| if tool_name is None: | ||
| # infer the tool name from the response parts | ||
| for part, _ in output_schema.find_tool(response.parts): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
This returns the sametuple[call, output_tool] we get fromfind_named_tool below, so if we're in this branch, we can skip that extra lookup
Docs Preview
|
| @@ -2052,6 +2052,69 @@ def usage(self) -> _usage.Usage: | |||
| """Get usage statistics for the run so far, including token usage, model requests, and so on.""" | |||
| return self._graph_run.state.usage | |||
| async def validate_structured_output( | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
This is basically a copy-paste of the other method in_agent_graph, I suggest moving it to a shared util.
Addresses the requesthere to support structured streaming while using
.iter.In particular, the analogous example to the whales streaming stuff would look like:
I have some concerns about the type names etc. but I think overall this seems like a reasonable thing to add.