Use .iter() API to fully replace existing streaming implementation
#951
Conversation
Docs Preview
|
dmontagu
force-pushed
the
dmontagu/graph-run-streaming
branch
from
302b8bf to
55ab7f0
Compare
dmontagu
force-pushed
the
dmontagu/graph-run-streaming
branch
from
2cfa693 to
6586af3
Compare
|
I think this is basically ready to merge, but I have one question about API choice I'll call out in-line |
docs/agents.md
Outdated
| # Begin a node-by-node, streaming iteration | ||
| with weather_agent.iter(user_prompt, deps=deps) as run: | ||
| async for node in run: | ||
| if Agent.is_model_request_node(node): |
There was a problem hiding this comment.
I think the main question I have is whether we should have these node-narrowing typeguards as:
- static methods on
Agent, as I have here - static methods on
AgentRun - static methods on the various node types (i.e., an
is_type_ofstaticmethod; I'll note I tried this and felt it was worse/more confusing) - free-floating functions
so far I think static method on Agent is best, but I see merits to each of these, though I'll note that static methods on the node types was what I implemented first and was my least favorite of the alternatives for various reasons.
| @@ -107,8 +107,31 @@ class GraphAgentDeps(Generic[DepsT, ResultDataT]): | |||
| run_span: logfire_api.LogfireSpan | |||
|
|
|||
|
|
|||
| class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]): | |||
There was a problem hiding this comment.
I tried using a type-alias instead of a class here, and it caused some issues because I couldn't inherit from it.
I think may be possible to drop the AgentNode class and use a type alias (just inheriting from the non-aliased parametrized BaseNode), and either way I think it's not currently public, but I think having this type allows us to remove most references to BaseNode in agent.py, which I think ultimately makes a lot of types more readable.
Unless one of you really objects, I would prefer to keep it around for now rather than wrestle with the consequences of removing it (in terms of verbosity and/or type-checking challenges). Especially considering it isn't a public API anyway.
| # In case the user didn't manually consume the full stream, ensure it is fully consumed here, | ||
| # otherwise usage won't be properly counted: | ||
| async for _ in agent_stream: | ||
| pass |
There was a problem hiding this comment.
@Kludex not 100% sure if we should do this force-consume-the-stream (when there isn't an exception), but I think it's what we should do. (Unless we confirm that you don't get charged for streams you don't consume, but that seems unlikely.)
dmontagu
force-pushed
the
dmontagu/graph-run-streaming
branch
2 times, most recently
from
e736b47 to
54aff03
Compare
dmontagu
force-pushed
the
dmontagu/graph-run-streaming
branch
from
54aff03 to
e73b85b
Compare
This is a work in progress on top of #833. I wanted to see this through to make sure that we had a plan for using the node streaming from #833 to replace our current streaming implementation APIs, and with this branch in the state it is, I feel pretty convinced that it's going to work and we'll be able to use the implementation here (perhaps with some tweaks) to achieve full capability parity with the existing APIs, while also being significantly more flexible and powerful.
I think we should merge #833 first without it being capable of fully replacing the
Agent.run_streamAPI, etc., but once that is merged I think we should aim to merge this shortly after and use it to replace the streaming APIs more fully.