Attempted improvements to "fix streamed response messages" #278
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
dmontagu
commented
Dec 16, 2024
Comment on lines
1018
to
1030
| async for item in stream_structured_parts(model_response): | ||
| handled_any_parts = True | ||
| if self._result_schema and (match := self._result_schema.find_tool([item])): | ||
| call, _ = match | ||
| return _MarkFinalResult(model_response, call.tool_name) | ||
| elif isinstance(item, _messages.ToolCallPart): | ||
| call = item | ||
| if tool := self._function_tools.get(call.tool_name): | ||
| # NOTE: The next line starts the tool running, which is at odds with the | ||
| # logfire span created below saying that's where the tool is run. | ||
| tasks.append(asyncio.create_task(tool.run(deps, call, conv_messages), name=call.tool_name)) | ||
| else: | ||
| parts.append(self._unknown_tool(call.tool_name)) |
There was a problem hiding this comment.
this is one of the interesting bits of this PR
dmontagu
commented
Dec 16, 2024
Comment on lines
1127
to
1152
| async def stream_structured_parts(model_response: StreamStructuredResponse) -> AsyncIterator[ModelResponsePart]: | ||
| last_part_index = 0 | ||
| last_part: ModelResponsePart | None = None | ||
|
|
||
| def new_or_updated_parts() -> Iterator[ModelResponsePart]: | ||
| nonlocal last_part_index, last_part | ||
|
|
||
| structured_msg = model_response.get() | ||
| new_last_part_index = last_part_index | ||
| for i, part in enumerate(structured_msg.parts[last_part_index:]): | ||
| if i == 0 and part == last_part: | ||
| continue # this part was not updated | ||
| yield part | ||
| last_part = part | ||
| new_last_part_index = last_part_index + i | ||
| last_part_index = new_last_part_index | ||
|
|
||
| while True: | ||
| for p in new_or_updated_parts(): | ||
| yield p | ||
| try: | ||
| await model_response.__anext__() | ||
| except StopAsyncIteration: | ||
| break | ||
| for p in new_or_updated_parts(): | ||
| yield p |
There was a problem hiding this comment.
This is the other interesting bit of this PR — the idea is that I've added a function that streams the changes to parts coming out of the model response. Once we have this, we can just handle them one at a time above.
Deploying pydantic-ai with
|
| Latest commit: |
a78fb53
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://b6791f89.pydantic-ai.pages.dev |
| Branch Preview URL: | https://dmontagu-streamed-response-m.pydantic-ai.pages.dev |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This gets one of the new tests passing, and I think improves the state of that PR conceptually, but now a test is failing and I'm not sure what the intended behavior of that test was or whether the new output makes sense.
In particular,
tests.test_streaming.test_call_toolfails in a way that I can't tell what's going wrong.