adding API docs

pydantic · samuelcolvin · Jan 15, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit 452a62f3ff760afc7ce788e607d6aa6aa9012d9a
diff --git a/docs/api/pydantic_graph/graph.md b/docs/api/pydantic_graph/graph.md
@@ -0,0 +1,3 @@
+# `pydantic_graph`
+
+::: pydantic_graph.graph
diff --git a/docs/api/pydantic_graph/mermaid.md b/docs/api/pydantic_graph/mermaid.md
@@ -0,0 +1,3 @@
+# `pydantic_graph.mermaid`
+
+::: pydantic_graph.mermaid
diff --git a/docs/api/pydantic_graph/nodes.md b/docs/api/pydantic_graph/nodes.md
@@ -0,0 +1,11 @@
+# `pydantic_graph.nodes`
+
+::: pydantic_graph.nodes
+    options:
+        members:
+            - GraphContext
+            - BaseNode
+            - End
+            - Edge
+            - RunEndT
+            - NodeRunEndT
diff --git a/docs/api/pydantic_graph/state.md b/docs/api/pydantic_graph/state.md
@@ -0,0 +1,3 @@
+# `pydantic_graph.state`
+
+::: pydantic_graph.state
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -55,6 +55,10 @@ nav:
     - api/models/ollama.md
     - api/models/test.md
     - api/models/function.md
+    - api/pydantic_graph/graph.md
+    - api/pydantic_graph/nodes.md
+    - api/pydantic_graph/state.md
+    - api/pydantic_graph/mermaid.md
 
 extra:
   # hide the "Made with Material for MkDocs" message
@@ -150,6 +154,7 @@ markdown_extensions:
 
 watch:
   - pydantic_ai_slim
+  - pydantic_graph
   - examples
 
 plugins:

diff --git a/pydantic_graph/pydantic_graph/__init__.py b/pydantic_graph/pydantic_graph/__init__.py
@@ -1,7 +1,7 @@
 from .exceptions import GraphRuntimeError, GraphSetupError
 from .graph import Graph
 from .nodes import BaseNode, Edge, End, GraphContext
-from .state import AbstractState, EndEvent, HistoryStep, NodeEvent
+from .state import AbstractState, EndStep, HistoryStep, NodeStep
 
 __all__ = (
     'Graph',
@@ -10,9 +10,9 @@
     'GraphContext',
     'Edge',
     'AbstractState',
-    'EndEvent',
+    'EndStep',
     'HistoryStep',
-    'NodeEvent',
+    'NodeStep',
     'GraphSetupError',
     'GraphRuntimeError',
 )
diff --git a/pydantic_graph/pydantic_graph/graph.py b/pydantic_graph/pydantic_graph/graph.py
@@ -14,7 +14,7 @@
 from . import _utils, exceptions, mermaid
 from ._utils import get_parent_namespace
 from .nodes import BaseNode, End, GraphContext, NodeDef
-from .state import EndEvent, HistoryStep, NodeEvent, StateT
+from .state import EndStep, HistoryStep, NodeStep, StateT
 
 __all__ = ('Graph',)
 
@@ -36,9 +36,16 @@ def __init__(
         self,
         *,
         nodes: Sequence[type[BaseNode[StateT, RunEndT]]],
-        state_type: type[StateT] | None = None,
         name: str | None = None,
     ):
+        """Create a graph from a sequence of nodes.
+
+        Args:
+            nodes: The nodes which make up the graph, nodes need to be unique and all be generic in the same
+                state type.
+            name: Optional name for the graph, if not provided the name will be inferred from the calling frame
+                on the first call to a graph method.
+        """
         self.name = name
 
         parent_namespace = get_parent_namespace(inspect.currentframe())
@@ -48,34 +55,6 @@ def __init__(
 
         self._validate_edges()
 
-    def _register_node(self, node: type[BaseNode[StateT, RunEndT]], parent_namespace: dict[str, Any] | None) -> None:
-        node_id = node.get_id()
-        if existing_node := self.node_defs.get(node_id):
-            raise exceptions.GraphSetupError(
-                f'Node ID `{node_id}` is not unique — found on {existing_node.node} and {node}'
-            )
-        else:
-            self.node_defs[node_id] = node.get_node_def(parent_namespace)
-
-    def _validate_edges(self):
-        known_node_ids = self.node_defs.keys()
-        bad_edges: dict[str, list[str]] = {}
-
-        for node_id, node_def in self.node_defs.items():
-            for edge in node_def.next_node_edges.keys():
-                if edge not in known_node_ids:
-                    bad_edges.setdefault(edge, []).append(f'`{node_id}`')
-
-        if bad_edges:
-            bad_edges_list = [f'`{k}` is referenced by {_utils.comma_and(v)}' for k, v in bad_edges.items()]
-            if len(bad_edges_list) == 1:
-                raise exceptions.GraphSetupError(f'{bad_edges_list[0]} but not included in the graph.')
-            else:
-                b = '\n'.join(f' {be}' for be in bad_edges_list)
-                raise exceptions.GraphSetupError(
-                    f'Nodes are referenced in the graph but not included in the graph:\n{b}'
-                )
-
     async def next(
         self,
         state: StateT,
@@ -84,13 +63,24 @@ async def next(
         *,
         infer_name: bool = True,
     ) -> BaseNode[StateT, Any] | End[RunEndT]:
+        """Run a node in the graph and return the next node to run.
+
+        Args:
+            state: The current state of the graph.
+            node: The node to run.
+            history: The history of the graph run so far. NOTE: this will be mutated to add the new step.
+            infer_name: Whether to infer the graph name from the calling frame.
+
+        Returns:
+            The next node to run or [`End`][pydantic_graph.nodes.End] if the graph has finished.
+        """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         node_id = node.get_id()
         if node_id not in self.node_defs:
             raise exceptions.GraphRuntimeError(f'Node `{node}` is not in the graph.')
 
-        history_step: NodeEvent[StateT, RunEndT] | None = NodeEvent(state, node)
+        history_step: NodeStep[StateT, RunEndT] | None = NodeStep(state, node)
         history.append(history_step)
 
         ctx = GraphContext(state)
@@ -107,6 +97,15 @@ async def run(
         *,
         infer_name: bool = True,
     ) -> tuple[RunEndT, list[HistoryStep[StateT, RunEndT]]]:
+        """Run the graph from a starting node until it ends.
+
+        Args:
+            state: The initial state of the graph.
+            start_node: the first node to run.
+            infer_name: Whether to infer the graph name from the calling frame.
+
+        Returns: The result type from ending the run and the history of the run.
+        """
         history: list[HistoryStep[StateT, RunEndT]] = []
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
@@ -119,7 +118,7 @@ async def run(
             while True:
                 next_node = await self.next(state, start_node, history, infer_name=False)
                 if isinstance(next_node, End):
-                    history.append(EndEvent(state, next_node))
+                    history.append(EndStep(state, next_node))
                     run_span.set_attribute('history', history)
                     return next_node.data, history
                 elif isinstance(next_node, BaseNode):
@@ -136,13 +135,29 @@ def mermaid_code(
         self,
         *,
         start_node: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
-        highlighted_nodes: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
-        highlight_css: str = mermaid.DEFAULT_HIGHLIGHT_CSS,
-        edge_labels: bool = True,
         title: str | None | Literal[False] = None,
+        edge_labels: bool = True,
         notes: bool = True,
+        highlighted_nodes: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
+        highlight_css: str = mermaid.DEFAULT_HIGHLIGHT_CSS,
         infer_name: bool = True,
     ) -> str:
+        """Generate a diagram representing the graph as [mermaid](https://mermaid.js.org/) chart.
+
+        This method calls [`pydantic_graph.mermaid.generate_code`][pydantic_graph.mermaid.generate_code].
+
+        Args:
+            start_node: The node or nodes to start the graph from.
+            title: The title of the diagram, use `False` to not include a title.
+            edge_labels: Whether to include edge labels.
+            notes: Whether to include notes on each node.
+            highlighted_nodes: Optional node or nodes to highlight.
+            highlight_css: The CSS to use for highlighting nodes.
+            infer_name: Whether to infer the graph name from the calling frame.
+
+        Returns:
+            The mermaid code for the graph, which can then be rendered as a diagram.
+        """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         if title is None and self.name:
@@ -158,6 +173,22 @@ def mermaid_code(
         )
 
     def mermaid_image(self, infer_name: bool = True, **kwargs: Unpack[mermaid.MermaidConfig]) -> bytes:
+        """Generate a diagram representing the graph as an image.
+
+        The format and diagram can be customized using `kwargs`,
+        see [`pydantic_graph.mermaid.MermaidConfig`][pydantic_graph.mermaid.MermaidConfig].
+
+        !!! note "Uses external service"
+            This method makes a request to [mermaid.ink](https://mermaid.ink) to render the image, `mermaid.ink`
+            is a free service not affiliated with Pydantic.
+
+        Args:
+            infer_name: Whether to infer the graph name from the calling frame.
+            **kwargs: Additional arguments to pass to `mermaid.request_image`.
+
+        Returns:
+            The image bytes.
+        """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         if 'title' not in kwargs and self.name:
@@ -167,12 +198,54 @@ def mermaid_image(self, infer_name: bool = True, **kwargs: Unpack[mermaid.Mermai
     def mermaid_save(
         self, path: Path | str, /, *, infer_name: bool = True, **kwargs: Unpack[mermaid.MermaidConfig]
     ) -> None:
+        """Generate a diagram representing the graph and save it as an image.
+
+        The format and diagram can be customized using `kwargs`,
+        see [`pydantic_graph.mermaid.MermaidConfig`][pydantic_graph.mermaid.MermaidConfig].
+
+        !!! note "Uses external service"
+            This method makes a request to [mermaid.ink](https://mermaid.ink) to render the image, `mermaid.ink`
+            is a free service not affiliated with Pydantic.
+
+        Args:
+            path: The path to save the image to.
+            infer_name: Whether to infer the graph name from the calling frame.
+            **kwargs: Additional arguments to pass to `mermaid.save_image`.
+        """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         if 'title' not in kwargs and self.name:
             kwargs['title'] = self.name
         mermaid.save_image(path, self, **kwargs)
 
+    def _register_node(self, node: type[BaseNode[StateT, RunEndT]], parent_namespace: dict[str, Any] | None) -> None:
+        node_id = node.get_id()
+        if existing_node := self.node_defs.get(node_id):
+            raise exceptions.GraphSetupError(
+                f'Node ID `{node_id}` is not unique — found on {existing_node.node} and {node}'
+            )
+        else:
+            self.node_defs[node_id] = node.get_node_def(parent_namespace)
+
+    def _validate_edges(self):
+        known_node_ids = self.node_defs.keys()
+        bad_edges: dict[str, list[str]] = {}
+
+        for node_id, node_def in self.node_defs.items():
+            for edge in node_def.next_node_edges.keys():
+                if edge not in known_node_ids:
+                    bad_edges.setdefault(edge, []).append(f'`{node_id}`')
+
+        if bad_edges:
+            bad_edges_list = [f'`{k}` is referenced by {_utils.comma_and(v)}' for k, v in bad_edges.items()]
+            if len(bad_edges_list) == 1:
+                raise exceptions.GraphSetupError(f'{bad_edges_list[0]} but not included in the graph.')
+            else:
+                b = '\n'.join(f' {be}' for be in bad_edges_list)
+                raise exceptions.GraphSetupError(
+                    f'Nodes are referenced in the graph but not included in the graph:\n{b}'
+                )
+
     def _infer_name(self, function_frame: FrameType | None) -> None:
         """Infer the agent name from the call frame.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# `pydantic_graph.mermaid`

		::: pydantic_graph.mermaid
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# `pydantic_graph.state`

		::: pydantic_graph.state