Various minor fixes (#834)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/agent.py

@@ -282,7 +282,7 @@ async def main():
         result_schema: _result.ResultSchema[RunResultDataT] | None = self._prepare_result_schema(result_type)
 
         # Build the graph
-        graph = _agent_graph.build_agent_graph(self.name, self._deps_type, result_type or self.result_type)
+        graph = self._build_graph(result_type)
 
         # Build the initial state
         state = _agent_graph.GraphAgentState(

pydantic_graph/pydantic_graph/graph.py

@@ -181,22 +181,22 @@ async def main():
                         start=start_node,
                     )
                 )
-        while True:
-            next_node = await self.next(start_node, history, state=state, deps=deps, infer_name=False)
-            if isinstance(next_node, End):
-                history.append(EndStep(result=next_node))
-                if run_span is not None:
-                    run_span.set_attribute('history', history)
-                return next_node.data, history
-            elif isinstance(next_node, BaseNode):
-                start_node = next_node
-            else:
-                if TYPE_CHECKING:
-                    typing_extensions.assert_never(next_node)
-                else:
-                    raise exceptions.GraphRuntimeError(
-                        f'Invalid node return type: `{type(next_node).__name__}`. Expected `BaseNode` or `End`.'
-                    )
+
+            next_node = start_node
+            while True:
+                next_node = await self.next(next_node, history, state=state, deps=deps, infer_name=False)
+                if isinstance(next_node, End):
+                    history.append(EndStep(result=next_node))
+                    if run_span is not None:
+                        run_span.set_attribute('history', history)
+                    return next_node.data, history
+                elif not isinstance(next_node, BaseNode):
+                    if TYPE_CHECKING:
+                        typing_extensions.assert_never(next_node)
+                    else:
+                        raise exceptions.GraphRuntimeError(
+                            f'Invalid node return type: `{type(next_node).__name__}`. Expected `BaseNode` or `End`.'
+                        )
 
     def run_sync(
         self: Graph[StateT, DepsT, T],