Graph Support (pydantic#528)

Co-authored-by: David Montague <35119617+dmontagu@users.noreply.github.com>
Co-authored-by: Israel Ekpo <44282278+izzyacademy@users.noreply.github.com>

Author: 3 people

.github/workflows/ci.yml

@@ -10,7 +10,7 @@ on:
 
 env:
   CI: true
-  COLUMNS: 120
+  COLUMNS: 150
   UV_PYTHON: 3.12
   UV_FROZEN: '1'
 
@@ -129,8 +129,8 @@ jobs:
 
       - run: mkdir coverage
 
-      # run tests with just `pydantic-ai-slim` dependencies
-      - run: uv run --package pydantic-ai-slim coverage run -m pytest
+      # run tests with just `pydantic-ai-slim` and `pydantic-graph` dependencies
+      - run: uv run --package pydantic-ai-slim --extra graph coverage run -m pytest
         env:
           COVERAGE_FILE: coverage/.coverage.${{ runner.os }}-py${{ matrix.python-version }}-slim
 

.gitignore

@@ -13,3 +13,4 @@ env*/
 examples/pydantic_ai_examples/.chat_app_messages.sqlite
 .cache/
 .vscode/
+/question_graph_history.json

Makefile

@@ -34,7 +34,7 @@ typecheck-pyright:
 
 .PHONY: typecheck-mypy
 typecheck-mypy:
-	uv run mypy --strict tests/typed_agent.py
+	uv run mypy
 
 .PHONY: typecheck
 typecheck: typecheck-pyright ## Run static type checking

docs/.hooks/main.py

@@ -19,11 +19,28 @@ def on_page_markdown(markdown: str, page: Page, config: Config, files: Files) ->
     return markdown
 
 
+# path to the main mkdocs material bundle file, found during `on_env`
+bundle_path: Path | None = None
+
+
 def on_env(env: Environment, config: Config, files: Files) -> Environment:
+    global bundle_path
+    for file in files:
+        if re.match('assets/javascripts/bundle.[a-z0-9]+.min.js', file.src_uri):
+            bundle_path = Path(file.dest_dir) / file.src_uri
+
     env.globals['build_timestamp'] = str(int(time.time()))
     return env
 
 
+def on_post_build(config: Config) -> None:
+    """Inject extra CSS into mermaid styles to avoid titles being the same color as the background in dark mode."""
+    if bundle_path.exists():
+        content = bundle_path.read_text()
+        content, _ = re.subn(r'}(\.statediagram)', '}.statediagramTitleText{fill:#888}\1', content, count=1)
+        bundle_path.write_text(content)
+
+
 def replace_uv_python_run(markdown: str) -> str:
     return re.sub(r'```bash\n(.*?)(python/uv[\- ]run|pip/uv[\- ]add|py-cli)(.+?)\n```', sub_run, markdown)
 

docs/api/models/function.md

@@ -9,7 +9,7 @@ Its primary use case is for more advanced unit testing than is possible with `Te
 
 Here's a minimal example:
 
-```py {title="function_model_usage.py" call_name="test_my_agent" lint="not-imports"}
+```py {title="function_model_usage.py" call_name="test_my_agent" noqa="I001"}
 from pydantic_ai import Agent
 from pydantic_ai.messages import ModelMessage, ModelResponse
 from pydantic_ai.models.function import FunctionModel, AgentInfo

docs/api/models/test.md

@@ -4,7 +4,7 @@ Utility model for quickly testing apps built with PydanticAI.
 
 Here's a minimal example:
 
-```py {title="test_model_usage.py" call_name="test_my_agent" lint="not-imports"}
+```py {title="test_model_usage.py" call_name="test_my_agent" noqa="I001"}
 from pydantic_ai import Agent
 from pydantic_ai.models.test import TestModel
 

docs/api/pydantic_graph/exceptions.md

@@ -0,0 +1,3 @@
+# `pydantic_graph.exceptions`
+
+::: pydantic_graph.exceptions

docs/api/pydantic_graph/graph.md

@@ -0,0 +1,3 @@
+# `pydantic_graph`
+
+::: pydantic_graph.graph

docs/api/pydantic_graph/mermaid.md

@@ -0,0 +1,3 @@
+# `pydantic_graph.mermaid`
+
+::: pydantic_graph.mermaid

docs/api/pydantic_graph/nodes.md

@@ -0,0 +1,12 @@
+# `pydantic_graph.nodes`
+
+::: pydantic_graph.nodes
+    options:
+        members:
+            - GraphRunContext
+            - BaseNode
+            - End
+            - Edge
+            - DepsT
+            - RunEndT
+            - NodeRunEndT

docs/api/pydantic_graph/state.md

@@ -0,0 +1,3 @@
+# `pydantic_graph.state`
+
+::: pydantic_graph.state

docs/examples/question-graph.md

@@ -0,0 +1,37 @@
+# Question Graph
+
+Example of a graph for asking and evaluating questions.
+
+Demonstrates:
+
+* [`pydantic_graph`](../graph.md)
+
+## Running the Example
+
+With [dependencies installed and environment variables set](./index.md#usage), run:
+
+```bash
+python/uv-run -m pydantic_ai_examples.question_graph
+```
+
+## Example Code
+
+```python {title="question_graph.py"}
+#! examples/pydantic_ai_examples/question_graph.py
+```
+
+The mermaid diagram generated in this example looks like this:
+
+```mermaid
+---
+title: question_graph
+---
+stateDiagram-v2
+  [*] --> Ask
+  Ask --> Answer: ask the question
+  Answer --> Evaluate: answer the question
+  Evaluate --> Congratulate
+  Evaluate --> Castigate
+  Congratulate --> [*]: success
+  Castigate --> Ask: try again
+```