Merge branch 'main' into mistral-provider

Author: Kludex

.github/workflows/ci.yml

@@ -212,7 +212,7 @@ jobs:
           path: index.html
 
       - run: uv run coverage report --fail-under 95
-      - run: uv run diff-cover coverage.xml --fail-under 95
+      - run: uv run diff-cover coverage.xml --fail-under 100
 
   # https://github.com/marketplace/actions/alls-green#why used for branch protection checks
   check:

Makefile

@@ -49,11 +49,11 @@ test: ## Run tests and collect coverage data
 
 .PHONY: test-all-python
 test-all-python: ## Run tests on Python 3.9 to 3.13
-	UV_PROJECT_ENVIRONMENT=.venv39 uv run --python 3.9 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv39 uv run --python 3.9 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras --all-packages coverage run -p -m pytest
 	@uv run coverage combine
 	@uv run coverage report
 

docs/api/providers.md

@@ -11,3 +11,5 @@
 ::: pydantic_ai.providers.bedrock
 
 ::: pydantic_ai.providers.groq
+
+::: pydantic_ai.providers.azure

docs/api/pydantic_graph/nodes.md

@@ -3,6 +3,7 @@
 ::: pydantic_graph.nodes
     options:
         members:
+            - StateT
             - GraphRunContext
             - BaseNode
             - End

docs/api/pydantic_graph/persistence.md

@@ -0,0 +1,7 @@
+# `pydantic_graph.persistence`
+
+::: pydantic_graph.persistence
+
+::: pydantic_graph.persistence.in_mem
+
+::: pydantic_graph.persistence.file

docs/api/pydantic_graph/state.md

@@ -151,7 +151,7 @@ agent = Agent('openai:gpt-4o', instrument=instrumentation_settings)
 Agent.instrument_all(instrumentation_settings)
 ```
 
-For now, this won't look as good in the Logfire UI, but we're working on it. **Once the UI supports it, `event_mode='logs'` will become the default.**
+For now, this won't look as good in the Logfire UI, but we're working on it.
 
 If you have very long conversations, the `events` span attribute may be truncated. Using `event_mode='logs'` will help avoid this issue.
 

docs/graph.md

@@ -828,6 +828,28 @@ Usage(requests=1, request_tokens=57, response_tokens=8, total_tokens=65, details
 1. The name of the model running on the remote server
 2. The url of the remote server
 
+### Azure AI Foundry
+
+If you want to use [Azure AI Foundry](https://ai.azure.com/) as your provider, you can do so by using the
+[`AzureProvider`][pydantic_ai.providers.azure.AzureProvider] class.
+
+```python {title="azure_provider_example.py"}
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.azure import AzureProvider
+
+model = OpenAIModel(
+    'gpt-4o',
+    provider=AzureProvider(
+        azure_endpoint='your-azure-endpoint',
+        api_version='your-api-version',
+        api_key='your-api-key',
+    ),
+)
+agent = Agent(model)
+...
+```
+
 ### OpenRouter
 
 To use [OpenRouter](https://openrouter.ai), first create an API key at [openrouter.ai/keys](https://openrouter.ai/keys).

docs/logfire.md

@@ -9,11 +9,16 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Annotated
 
 import logfire
-from devtools import debug
-from pydantic_graph import BaseNode, Edge, End, Graph, GraphRunContext, HistoryStep
+from groq import BaseModel
+from pydantic_graph import (
+    BaseNode,
+    End,
+    Graph,
+    GraphRunContext,
+)
+from pydantic_graph.persistence.file import FileStatePersistence
 
 from pydantic_ai import Agent
 from pydantic_ai.format_as_xml import format_as_xml
@@ -41,22 +46,23 @@ async def run(self, ctx: GraphRunContext[QuestionState]) -> Answer:
         )
         ctx.state.ask_agent_messages += result.all_messages()
         ctx.state.question = result.data
-        return Answer()
+        return Answer(result.data)
 
 
 @dataclass
 class Answer(BaseNode[QuestionState]):
-    answer: str | None = None
+    question: str
 
     async def run(self, ctx: GraphRunContext[QuestionState]) -> Evaluate:
-        assert self.answer is not None
-        return Evaluate(self.answer)
+        answer = input(f'{self.question}: ')
+        return Evaluate(answer)
 
 
-@dataclass
-class EvaluationResult:
+class EvaluationResult(BaseModel, use_attribute_docstrings=True):
     correct: bool
+    """Whether the answer is correct."""
     comment: str
+    """Comment on the answer, reprimand the user if the answer is wrong."""
 
 
 evaluate_agent = Agent(
@@ -67,101 +73,76 @@ class EvaluationResult:
 
 
 @dataclass
-class Evaluate(BaseNode[QuestionState]):
+class Evaluate(BaseNode[QuestionState, None, str]):
     answer: str
 
     async def run(
         self,
         ctx: GraphRunContext[QuestionState],
-    ) -> Congratulate | Reprimand:
+    ) -> End[str] | Reprimand:
         assert ctx.state.question is not None
         result = await evaluate_agent.run(
             format_as_xml({'question': ctx.state.question, 'answer': self.answer}),
             message_history=ctx.state.evaluate_agent_messages,
         )
         ctx.state.evaluate_agent_messages += result.all_messages()
         if result.data.correct:
-            return Congratulate(result.data.comment)
+            return End(result.data.comment)
         else:
             return Reprimand(result.data.comment)
 
 
-@dataclass
-class Congratulate(BaseNode[QuestionState, None, None]):
-    comment: str
-
-    async def run(
-        self, ctx: GraphRunContext[QuestionState]
-    ) -> Annotated[End, Edge(label='success')]:
-        print(f'Correct answer! {self.comment}')
-        return End(None)
-
-
 @dataclass
 class Reprimand(BaseNode[QuestionState]):
     comment: str
 
     async def run(self, ctx: GraphRunContext[QuestionState]) -> Ask:
         print(f'Comment: {self.comment}')
-        # > Comment: Vichy is no longer the capital of France.
         ctx.state.question = None
         return Ask()
 
 
 question_graph = Graph(
-    nodes=(Ask, Answer, Evaluate, Congratulate, Reprimand), state_type=QuestionState
+    nodes=(Ask, Answer, Evaluate, Reprimand), state_type=QuestionState
 )
 
 
 async def run_as_continuous():
     state = QuestionState()
     node = Ask()
-    history: list[HistoryStep[QuestionState, None]] = []
-    with logfire.span('run questions graph'):
-        while True:
-            node = await question_graph.next(node, history, state=state)
-            if isinstance(node, End):
-                debug([e.data_snapshot() for e in history])
-                break
-            elif isinstance(node, Answer):
-                assert state.question
-                node.answer = input(f'{state.question} ')
-            # otherwise just continue
+    end = await question_graph.run(node, state=state)
+    print('END:', end.output)
 
 
 async def run_as_cli(answer: str | None):
-    history_file = Path('question_graph_history.json')
-    history = (
-        question_graph.load_history(history_file.read_bytes())
-        if history_file.exists()
-        else []
-    )
-
-    if history:
-        last = history[-1]
-        assert last.kind == 'node', 'expected last step to be a node'
-        state = last.state
-        assert answer is not None, 'answer is required to continue from history'
-        node = Answer(answer)
+    persistence = FileStatePersistence(Path('question_graph.json'))
+    persistence.set_graph_types(question_graph)
+
+    if snapshot := await persistence.load_next():
+        state = snapshot.state
+        assert answer is not None, (
+            'answer required, usage "uv run -m pydantic_ai_examples.question_graph cli <answer>"'
+        )
+        node = Evaluate(answer)
     else:
         state = QuestionState()
         node = Ask()
-    debug(state, node)
+    # debug(state, node)
 
-    with logfire.span('run questions graph'):
+    async with question_graph.iter(node, state=state, persistence=persistence) as run:
         while True:
-            node = await question_graph.next(node, history, state=state)
+            node = await run.next()
             if isinstance(node, End):
-                debug([e.data_snapshot() for e in history])
+                print('END:', node.data)
+                history = await persistence.load_all()
+                print('history:', '\n'.join(str(e.node) for e in history), sep='\n')
                 print('Finished!')
                 break
             elif isinstance(node, Answer):
-                print(state.question)
+                print(node.question)
                 break
             # otherwise just continue
 
-    history_file.write_bytes(question_graph.dump_history(history, indent=2))
-
 
 if __name__ == '__main__':
     import asyncio

docs/models.md

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-examples"
-version = "0.0.39"
+version = "0.0.40"
 description = "Examples of how to use PydanticAI and what it can do."
 authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
 license = "MIT"
@@ -32,7 +32,7 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "pydantic-ai-slim[openai,vertexai,groq,anthropic]==0.0.39",
+    "pydantic-ai-slim[openai,vertexai,groq,anthropic]==0.0.40",
     "asyncpg>=0.30.0",
     "fastapi>=0.115.4",
     "logfire[asyncpg,fastapi,sqlite3]>=2.6",

examples/pydantic_ai_examples/question_graph.py

@@ -66,7 +66,7 @@ nav:
       - api/providers.md
       - api/pydantic_graph/graph.md
       - api/pydantic_graph/nodes.md
-      - api/pydantic_graph/state.md
+      - api/pydantic_graph/persistence.md
       - api/pydantic_graph/mermaid.md
       - api/pydantic_graph/exceptions.md
 

examples/pyproject.toml

@@ -475,8 +475,8 @@ async def main():
             start_node,
             state=state,
             deps=graph_deps,
-            infer_name=False,
             span=use_span(run_span, end_on_exit=True),
+            infer_name=False,
         ) as graph_run:
             yield AgentRun(graph_run)
 

mkdocs.yml

@@ -262,8 +262,14 @@ def model_name(self) -> str:
 
     @property
     @abstractmethod
-    def system(self) -> str | None:
-        """The system / model provider, ex: openai."""
+    def system(self) -> str:
+        """The system / model provider, ex: openai.
+
+        Use to populate the `gen_ai.system` OpenTelemetry semantic convention attribute,
+        so should use well-known values listed in
+        https://opentelemetry.io/docs/specs/semconv/attributes-registry/gen-ai/#gen-ai-system
+        when applicable.
+        """
         raise NotImplementedError()
 
     @property

pydantic_ai_slim/pydantic_ai/agent.py

@@ -33,13 +33,7 @@
 )
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import (
-    Model,
-    ModelRequestParameters,
-    StreamedResponse,
-    cached_async_http_client,
-    check_allow_model_requests,
-)
+from . import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client, check_allow_model_requests
 
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
@@ -115,7 +109,7 @@ class AnthropicModel(Model):
     client: AsyncAnthropic = field(repr=False)
 
     _model_name: AnthropicModelName = field(repr=False)
-    _system: str | None = field(default='anthropic', repr=False)
+    _system: str = field(default='anthropic', repr=False)
 
     def __init__(
         self,
@@ -183,7 +177,7 @@ def model_name(self) -> AnthropicModelName:
         return self._model_name
 
     @property
-    def system(self) -> str | None:
+    def system(self) -> str:
         """The system / model provider."""
         return self._system
 
@@ -355,8 +349,17 @@ async def _map_user_prompt(
                             source={'data': io.BytesIO(item.data), 'media_type': item.media_type, 'type': 'base64'},  # type: ignore
                             type='image',
                         )
+                    elif item.media_type == 'application/pdf':
+                        yield DocumentBlockParam(
+                            source=Base64PDFSourceParam(
+                                data=io.BytesIO(item.data),
+                                media_type='application/pdf',
+                                type='base64',
+                            ),
+                            type='document',
+                        )
                     else:
-                        raise RuntimeError('Only images are supported for binary content')
+                        raise RuntimeError('Only images and PDFs are supported for binary content')
                 elif isinstance(item, ImageUrl):
                     try:
                         response = await cached_async_http_client().get(item.url)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -119,15 +119,15 @@ class BedrockConverseModel(Model):
     client: BedrockRuntimeClient
 
     _model_name: BedrockModelName = field(repr=False)
-    _system: str | None = field(default='bedrock', repr=False)
+    _system: str = field(default='bedrock', repr=False)
 
     @property
     def model_name(self) -> str:
         """The model name."""
         return self._model_name
 
     @property
-    def system(self) -> str | None:
+    def system(self) -> str:
         """The system / model provider, ex: openai."""
         return self._system