Commit593ff0e

authored

Yield initial node during Graph (and therefore Agent) iteration (pydantic#1412)

1 parent8ba6234 commit593ff0eCopy full SHA for 593ff0e

File tree

+31

-1

lines changed

+31

-1

lines changed

Lines changed: 7 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -121,6 +121,12 @@ async def main():`
`121`	`121`	`print(nodes)`
`122`	`122`	`"""`
`123`	`123`	`[`
	`124`	`+ UserPromptNode(`
	`125`	`+ user_prompt='What is the capital of France?',`
	`126`	`+ system_prompts=(),`
	`127`	`+ system_prompt_functions=[],`
	`128`	`+ system_prompt_dynamic_functions={},`
	`129`	`+ ),`
`124`	`130`	`ModelRequestNode(`
`125`	`131`	`request=ModelRequest(`
`126`	`132`	`parts=[`
`@@ -338,6 +344,7 @@ if __name__ == '__main__':`
`338`	`344`	`print(output_messages)`
`339`	`345`	`"""`
`340`	`346`	`[`
	`347`	`+ '=== UserPromptNode: What will the weather be like in Paris on Tuesday? ===',`
`341`	`348`	`'=== ModelRequestNode: streaming partial request tokens ===',`
`342`	`349`	`'[Request] Starting part 0: ToolCallPart(tool_name=\'weather_forecast\', args=\'{"location":"Pa\', tool_call_id=\'0001\', part_kind=\'tool-call\')',`
`343`	`350`	`'[Request] Part 0 args_delta=ris","forecast_',`

Lines changed: 1 addition & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -510,6 +510,7 @@ async def main():`
`510`	`510`	`#> Node: CountDown()`
`511`	`511`	`#> Node: CountDown()`
`512`	`512`	`#> Node: CountDown()`
	`513`	`+#> Node: CountDown()`
`513`	`514`	`#> Node: End(data=0)`
`514`	`515`	`print('Final result:', run.result.output)# (3)!`
`515`	`516`	`#> Final result: 0`

Lines changed: 12 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -372,6 +372,12 @@ async def main():`
`372`	`372`	`print(nodes)`
`373`	`373`	`'''`
`374`	`374`	`[`
	`375`	`+ UserPromptNode(`
	`376`	`+ user_prompt='What is the capital of France?',`
	`377`	`+ system_prompts=(),`
	`378`	`+ system_prompt_functions=[],`
	`379`	`+ system_prompt_dynamic_functions={},`
	`380`	`+ ),`
`375`	`381`	`ModelRequestNode(`
`376`	`382`	`request=ModelRequest(`
`377`	`383`	`parts=[`
`@@ -1355,6 +1361,12 @@ async def main():`
`1355`	`1361`	`print(nodes)`
`1356`	`1362`	`'''`
`1357`	`1363`	`[`
	`1364`	`+ UserPromptNode(`
	`1365`	`+ user_prompt='What is the capital of France?',`
	`1366`	`+ system_prompts=(),`
	`1367`	`+ system_prompt_functions=[],`
	`1368`	`+ system_prompt_dynamic_functions={},`
	`1369`	`+ ),`
`1358`	`1370`	`ModelRequestNode(`
`1359`	`1371`	`request=ModelRequest(`
`1360`	`1372`	`parts=[`

Lines changed: 8 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -607,6 +607,7 @@ async def main():`
`607`	`607`	`print(node_states)`
`608`	`608`	`'''`
`609`	`609`	`[`
	`610`	`+ (Increment(), MyState(number=1)),`
`610`	`611`	`(Increment(), MyState(number=1)),`
`611`	`612`	`(Check42(), MyState(number=2)),`
`612`	`613`	`(End(data=2), MyState(number=2)),`
`@@ -621,6 +622,7 @@ async def main():`
`621`	`622`	`print(node_states)`
`622`	`623`	`'''`
`623`	`624`	`[`
	`625`	`+ (Increment(), MyState(number=41)),`
`624`	`626`	`(Increment(), MyState(number=41)),`
`625`	`627`	`(Check42(), MyState(number=42)),`
`626`	`628`	`(Increment(), MyState(number=42)),`
`@@ -665,6 +667,7 @@ def __init__(`
`665`	`667`	`self.deps=deps`
`666`	`668`
`667`	`669`	`self._next_node:BaseNode[StateT,DepsT,RunEndT]\|End[RunEndT]=start_node`
	`670`	`+self._is_started:bool=False`
`668`	`671`
`669`	`672`	`@property`
`670`	`673`	`defnext_node(self)->BaseNode[StateT,DepsT,RunEndT]\|End[RunEndT]:`
`@@ -777,8 +780,13 @@ def __aiter__(self) -> AsyncIterator[BaseNode[StateT, DepsT, RunEndT] \| End[RunE`
`777`	`780`
`778`	`781`	`asyncdef__anext__(self)->BaseNode[StateT,DepsT,RunEndT]\|End[RunEndT]:`
`779`	`782`	"""Use the last returned node as the input to `Graph.next`."""
	`783`	`+ifnotself._is_started:`
	`784`	`+self._is_started=True`
	`785`	`+returnself._next_node`
	`786`	`+`
`780`	`787`	`ifisinstance(self._next_node,End):`
`781`	`788`	`raiseStopAsyncIteration`
	`789`	`+`
`782`	`790`	`returnawaitself.next(self._next_node)`
`783`	`791`
`784`	`792`	`def__repr__(self)->str:`

Lines changed: 3 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -312,7 +312,9 @@ async def test_iter():`
`312`	`312`	`assertgraph_iter.result`
`313`	`313`	`assertgraph_iter.result.output==8`
`314`	`314`
`315`		`-assertnode_reprs==snapshot(["String2Length(input_data='3.14')",'Double(input_data=4)','End(data=8)'])`
	`315`	`+assertnode_reprs==snapshot(`
	`316`	`+ ['Float2String(input_data=3.14)',"String2Length(input_data='3.14')",'Double(input_data=4)','End(data=8)']`
	`317`	`+ )`
`316`	`318`
`317`	`319`
`318`	`320`	`asyncdeftest_iter_next(mock_snapshot_id:object):`

Comments

(0)