Merge branch 'main' into update-versins

Author: medaminezghal

.gitignore

@@ -15,7 +15,6 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 .vscode/
 /question_graph_history.json
 /docs-site/.wrangler/
-/CLAUDE.md
 node_modules/
 **.idea/
 .coverage*

CLAUDE.md

@@ -0,0 +1,127 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Development Commands
+
+### Core Development Tasks
+- **Install dependencies**: `make install` (requires uv, pre-commit, and deno)
+- **Run all checks**: `make` (format, lint, typecheck, test with coverage)
+- **Format code**: `make format`
+- **Lint code**: `make lint`
+- **Type checking**: `make typecheck` (uses pyright) or `make typecheck-both` (pyright + mypy)
+- **Run tests**: `make test` (with coverage) or `make test-fast` (parallel, no coverage)
+- **Build docs**: `make docs` or `make docs-serve` (local development)
+
+### Single Test Commands
+- **Run specific test**: `uv run pytest tests/test_agent.py::test_function_name -v`
+- **Run test file**: `uv run pytest tests/test_agent.py -v`
+- **Run with debug**: `uv run pytest tests/test_agent.py -v -s`
+
+### Multi-Python Testing
+- **Install all Python versions**: `make install-all-python`
+- **Test all Python versions**: `make test-all-python`
+
+## Project Architecture
+
+### Core Components
+
+**Agent System (`pydantic_ai_slim/pydantic_ai/agent.py`)**
+- `Agent[AgentDepsT, OutputDataT]`: Main orchestrator class with generic types for dependency injection and output validation
+- Entry points: `run()`, `run_sync()`, `run_stream()` methods
+- Handles tool management, system prompts, and model interaction
+
+**Model Integration (`pydantic_ai_slim/pydantic_ai/models/`)**
+- Unified interface across providers: OpenAI, Anthropic, Google, Groq, Cohere, Mistral, Bedrock, HuggingFace
+- Model strings: `"openai:gpt-4o"`, `"anthropic:claude-3-5-sonnet"`, `"google:gemini-1.5-pro"`
+- `ModelRequestParameters` for configuration, `StreamedResponse` for streaming
+
+**Graph-based Execution (`pydantic_graph/` + `_agent_graph.py`)**
+- State machine execution through: `UserPromptNode` → `ModelRequestNode` → `CallToolsNode`
+- `GraphAgentState` maintains message history and usage tracking
+- `GraphRunContext` provides execution context
+
+**Tool System (`tools.py`, `toolsets/`)**
+- `@agent.tool` decorator for function registration
+- `RunContext[AgentDepsT]` provides dependency injection in tools
+- Support for sync/async functions with automatic schema generation
+
+**Output Handling**
+- `TextOutput`: Plain text responses
+- `ToolOutput`: Structured data via tool calls
+- `NativeOutput`: Provider-specific structured output
+- `PromptedOutput`: Prompt-based structured extraction
+
+### Key Design Patterns
+
+**Dependency Injection**
+```python
+@dataclass
+class MyDeps:
+    database: DatabaseConn
+
+agent = Agent('openai:gpt-4o', deps_type=MyDeps)
+
+@agent.tool
+async def get_data(ctx: RunContext[MyDeps]) -> str:
+    return await ctx.deps.database.fetch_data()
+```
+
+**Type-Safe Agents**
+```python
+class OutputModel(BaseModel):
+    result: str
+    confidence: float
+
+agent: Agent[MyDeps, OutputModel] = Agent(
+    'openai:gpt-4o',
+    deps_type=MyDeps,
+    output_type=OutputModel
+)
+```
+
+## Workspace Structure
+
+This is a uv workspace with multiple packages:
+- **`pydantic_ai_slim/`**: Core framework (minimal dependencies)
+- **`pydantic_evals/`**: Evaluation system
+- **`pydantic_graph/`**: Graph execution engine
+- **`examples/`**: Example applications
+- **`clai/`**: CLI tool
+- **`mcp-run-python/`**: MCP server implementation (Deno/TypeScript)
+
+## Testing Strategy
+
+- **Unit tests**: `tests/` directory with comprehensive model and component coverage
+- **VCR cassettes**: `tests/cassettes/` for recorded LLM API interactions
+- **Test models**: Use `TestModel` for deterministic testing
+- **Examples testing**: `tests/test_examples.py` validates all documentation examples
+- **Multi-version testing**: Python 3.9-3.13 support
+
+## Key Configuration Files
+
+- **`pyproject.toml`**: Main workspace configuration with dependency groups
+- **`pydantic_ai_slim/pyproject.toml`**: Core package with model optional dependencies
+- **`Makefile`**: Development task automation
+- **`uv.lock`**: Locked dependencies for reproducible builds
+
+## Important Implementation Notes
+
+- **Model Provider Integration**: Each provider in `models/` directory implements the `Model` abstract base class
+- **Message System**: Vendor-agnostic message format in `messages.py` with rich content type support
+- **Streaming Architecture**: Real-time response processing with validation during streaming
+- **Error Handling**: Specific exception types with retry mechanisms at multiple levels
+- **OpenTelemetry Integration**: Built-in observability support
+
+## Documentation Development
+
+- **Local docs**: `make docs-serve` (serves at http://localhost:8000)
+- **Docs source**: `docs/` directory (MkDocs with Material theme)
+- **API reference**: Auto-generated from docstrings using mkdocstrings
+
+## Dependencies Management
+
+- **Package manager**: uv (fast Python package manager)
+- **Lock file**: `uv.lock` (commit this file)
+- **Sync command**: `make sync` to update dependencies
+- **Optional extras**: Define groups in `pyproject.toml` optional-dependencies

pydantic_ai_slim/pydantic_ai/_function_schema.py

@@ -96,16 +96,20 @@ def function_schema(  # noqa: C901
     config = ConfigDict(title=function.__name__, use_attribute_docstrings=True)
     config_wrapper = ConfigWrapper(config)
     gen_schema = _generate_schema.GenerateSchema(config_wrapper)
+    errors: list[str] = []
 
-    sig = signature(function)
+    try:
+        sig = signature(function)
+    except ValueError as e:
+        errors.append(str(e))
+        sig = signature(lambda: None)
 
     type_hints = _typing_extra.get_function_type_hints(function)
 
     var_kwargs_schema: core_schema.CoreSchema | None = None
     fields: dict[str, core_schema.TypedDictField] = {}
     positional_fields: list[str] = []
     var_positional_field: str | None = None
-    errors: list[str] = []
     decorators = _decorators.DecoratorInfos()
 
     description, field_descriptions = doc_descriptions(function, sig, docstring_format=docstring_format)
@@ -235,14 +239,19 @@ def _takes_ctx(function: TargetFunc[P, R]) -> TypeIs[WithCtx[P, R]]:
     Returns:
         `True` if the function takes a `RunContext` as first argument, `False` otherwise.
     """
-    sig = signature(function)
+    try:
+        sig = signature(function)
+    except ValueError:  # pragma: no cover
+        return False  # pragma: no cover
     try:
         first_param_name = next(iter(sig.parameters.keys()))
     except StopIteration:
         return False
     else:
         type_hints = _typing_extra.get_function_type_hints(function)
-        annotation = type_hints[first_param_name]
+        annotation = type_hints.get(first_param_name)
+        if annotation is None:
+            return False  # pragma: no cover
         return True is not sig.empty and _is_call_ctx(annotation)
 
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -36,7 +36,7 @@
 from .models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
 from .output import OutputDataT, OutputSpec
 from .profiles import ModelProfile
-from .result import FinalResult, StreamedRunResult
+from .result import AgentStream, FinalResult, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDepsT,
@@ -1127,29 +1127,15 @@ async def main():
             while True:
                 if self.is_model_request_node(node):
                     graph_ctx = agent_run.ctx
-                    async with node._stream(graph_ctx) as streamed_response:  # pyright: ignore[reportPrivateUsage]
-
-                        async def stream_to_final(
-                            s: models.StreamedResponse,
-                        ) -> FinalResult[models.StreamedResponse] | None:
-                            output_schema = graph_ctx.deps.output_schema
-                            async for maybe_part_event in streamed_response:
-                                if isinstance(maybe_part_event, _messages.PartStartEvent):
-                                    new_part = maybe_part_event.part
-                                    if isinstance(new_part, _messages.TextPart) and isinstance(
-                                        output_schema, _output.TextOutputSchema
-                                    ):
-                                        return FinalResult(s, None, None)
-                                    elif isinstance(new_part, _messages.ToolCallPart) and (
-                                        tool_def := graph_ctx.deps.tool_manager.get_tool_def(new_part.tool_name)
-                                    ):
-                                        if tool_def.kind == 'output':
-                                            return FinalResult(s, new_part.tool_name, new_part.tool_call_id)
-                                        elif tool_def.kind == 'deferred':
-                                            return FinalResult(s, None, None)
+                    async with node.stream(graph_ctx) as stream:
+
+                        async def stream_to_final(s: AgentStream) -> FinalResult[AgentStream] | None:
+                            async for event in stream:
+                                if isinstance(event, _messages.FinalResultEvent):
+                                    return FinalResult(s, event.tool_name, event.tool_call_id)
                             return None
 
-                        final_result = await stream_to_final(streamed_response)
+                        final_result = await stream_to_final(stream)
                         if final_result is not None:
                             if yielded:
                                 raise exceptions.AgentRunError('Agent run produced final results')  # pragma: no cover
@@ -1184,14 +1170,8 @@ async def on_complete() -> None:
                             yield StreamedRunResult(
                                 messages,
                                 graph_ctx.deps.new_message_index,
-                                graph_ctx.deps.usage_limits,
-                                streamed_response,
-                                graph_ctx.deps.output_schema,
-                                _agent_graph.build_run_context(graph_ctx),
-                                graph_ctx.deps.output_validators,
-                                final_result.tool_name,
+                                stream,
                                 on_complete,
-                                graph_ctx.deps.tool_manager,
                             )
                             break
                 next_node = await agent_run.next(node)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -758,7 +758,7 @@ async def download_item(
 
     data_type = media_type
     if type_format == 'extension':
-        data_type = data_type.split('/')[1]
+        data_type = item.format
 
     data = response.content
     if data_format in ('base64', 'base64_uri'):

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -16,9 +16,7 @@
 from .. import _utils, usage
 from .._utils import PeekableAsyncStream
 from ..messages import (
-    AudioUrl,
     BinaryContent,
-    ImageUrl,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -345,18 +343,19 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
 def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
     if not content:
         return 0
+
     if isinstance(content, str):
-        return len(re.split(r'[\s",.:]+', content.strip()))
-    else:
-        tokens = 0
-        for part in content:
-            if isinstance(part, str):
-                tokens += len(re.split(r'[\s",.:]+', part.strip()))
-            # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
-            if isinstance(part, (AudioUrl, ImageUrl)):
-                tokens += 0
-            elif isinstance(part, BinaryContent):
-                tokens += len(part.data)
-            else:
-                tokens += 0
-        return tokens
+        return len(_TOKEN_SPLIT_RE.split(content.strip()))
+
+    tokens = 0
+    for part in content:
+        if isinstance(part, str):
+            tokens += len(_TOKEN_SPLIT_RE.split(part.strip()))
+        elif isinstance(part, BinaryContent):
+            tokens += len(part.data)
+        # TODO(Marcelo): We need to study how we can estimate the tokens for AudioUrl or ImageUrl.
+
+    return tokens
+
+
+_TOKEN_SPLIT_RE = re.compile(r'[\s",.:]+')

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -52,6 +52,7 @@
         CompletionChunk as MistralCompletionChunk,
         Content as MistralContent,
         ContentChunk as MistralContentChunk,
+        DocumentURLChunk as MistralDocumentURLChunk,
         FunctionCall as MistralFunctionCall,
         ImageURL as MistralImageURL,
         ImageURLChunk as MistralImageURLChunk,
@@ -539,10 +540,19 @@ def _map_user_prompt(self, part: UserPromptPart) -> MistralUserMessage:
                     if item.is_image:
                         image_url = MistralImageURL(url=f'data:{item.media_type};base64,{base64_encoded}')
                         content.append(MistralImageURLChunk(image_url=image_url, type='image_url'))
+                    elif item.media_type == 'application/pdf':
+                        content.append(
+                            MistralDocumentURLChunk(
+                                document_url=f'data:application/pdf;base64,{base64_encoded}', type='document_url'
+                            )
+                        )
                     else:
-                        raise RuntimeError('Only image binary content is supported for Mistral.')
+                        raise RuntimeError('BinaryContent other than image or PDF is not supported in Mistral.')
                 elif isinstance(item, DocumentUrl):
-                    raise RuntimeError('DocumentUrl is not supported in Mistral.')  # pragma: no cover
+                    if item.media_type == 'application/pdf':
+                        content.append(MistralDocumentURLChunk(document_url=item.url, type='document_url'))
+                    else:
+                        raise RuntimeError('DocumentUrl other than PDF is not supported in Mistral.')
                 elif isinstance(item, VideoUrl):
                     raise RuntimeError('VideoUrl is not supported in Mistral.')
                 else:  # pragma: no cover