pydantic-ai-graph - simplify public generics (#539)

pydantic · samuelcolvin · Jan 15, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit b03e7bc64107016e28456e270c95ea1c00819499
diff --git a/examples/pydantic_ai_examples/email_extract_graph.py b/examples/pydantic_ai_examples/email_extract_graph.py
@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import asyncio
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import logfire
@@ -48,7 +49,8 @@ def extract_system_prompt(ctx: RunContext[list[str]]):
     return prompt
 
 
-class ExtractEvent(BaseNode[State, None]):
+@dataclass
+class ExtractEvent(BaseNode[State]):
     async def run(self, ctx: GraphContext[State]) -> CleanEvent:
         event = await extract_agent.run(
             ctx.state.email_content, deps=ctx.state.skip_events
@@ -73,7 +75,10 @@ def timestamp_system_prompt():
 )
 
 
-class CleanEvent(BaseNode[State, RawEventDetails]):
+@dataclass
+class CleanEvent(BaseNode[State]):
+    input_data: RawEventDetails
+
     async def run(self, ctx: GraphContext[State]) -> InspectEvent:
         start_ts, duration = await asyncio.gather(
             timestamp_agent.run(self.input_data.start_ts),
@@ -89,7 +94,10 @@ async def run(self, ctx: GraphContext[State]) -> InspectEvent:
         )
 
 
-class InspectEvent(BaseNode[State, EventDetails, EventDetails | None]):
+@dataclass
+class InspectEvent(BaseNode[State, EventDetails | None]):
+    input_data: EventDetails
+
     async def run(
         self, ctx: GraphContext[State]
     ) -> ExtractEvent | End[EventDetails | None]:
@@ -104,15 +112,18 @@ async def run(
             return End(None)
         else:
             ctx.state.skip_events.append(self.input_data.title)
-            return ExtractEvent(None)
+            return ExtractEvent()
 
 
-graph = Graph[State, None, EventDetails | None](
-    ExtractEvent,
-    CleanEvent,
-    InspectEvent,
+graph = Graph[State, EventDetails | None](
+    nodes=(
+        ExtractEvent,
+        CleanEvent,
+        InspectEvent,
+    )
 )
-print(graph.mermaid_code())
+graph_runner = graph.get_runner(ExtractEvent)
+print(graph_runner.mermaid_code())
 
 email = """
 Hi Samuel,
@@ -137,7 +148,7 @@ async def run(
 
 async def main():
     state = State(email_content=email)
-    result, history = await graph.run(None, state)
+    result, history = await graph_runner.run(state, None)
     debug(result, history)
 
 

diff --git a/pydantic_ai_graph/pydantic_ai_graph/__init__.py b/pydantic_ai_graph/pydantic_ai_graph/__init__.py
@@ -1,5 +1,16 @@
-from .graph import Graph
+from .graph import Graph, GraphRun, GraphRunner
 from .nodes import BaseNode, End, GraphContext
-from .state import AbstractState, Snapshot
+from .state import AbstractState, EndEvent, Step, StepOrEnd
 
-__all__ = 'BaseNode', 'End', 'GraphContext', 'Graph', 'Snapshot', 'AbstractState'
+__all__ = (
+    'Graph',
+    'GraphRunner',
+    'GraphRun',
+    'BaseNode',
+    'End',
+    'GraphContext',
+    'AbstractState',
+    'EndEvent',
+    'StepOrEnd',
+    'Step',
+)
diff --git a/pydantic_ai_graph/pydantic_ai_graph/_utils.py b/pydantic_ai_graph/pydantic_ai_graph/_utils.py
@@ -49,3 +49,17 @@ def type_arg_name(arg: Any) -> str:
         return 'None'
     else:
         return arg.__name__
+
+
+def get_parent_namespace(frame: types.FrameType | None) -> dict[str, Any] | None:
+    """Attempt to get the namespace where the graph was defined.
+
+    If the graph is defined with generics `Graph[a, b]` then another frame is inserted, and we have to skip that
+    to get the correct namespace.
+    """
+    if frame is not None:
+        if back := frame.f_back:
+            if back.f_code.co_filename.endswith('/typing.py'):
+                return get_parent_namespace(back)
+            else:
+                return back.f_locals