ronantakizawa/codereview-bench huggingface dataset:
Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/_run_context.py (Python)

Code before:
    """Dependencies for the agent."""
    model: Model
    """The model used in this run."""
    usage: RunUsage
    """LLM usage associated with the run."""
    prompt: str | Sequence[_messages.UserContent] | None = None
    """The original user prompt passed to the run."""
    messages: list[_messages.ModelMessage] = field(default_factory=list)
    """Messages exchanged in the conversation so far."""
    tracer: Tracer = field(default_factory=NoOpTracer)
    """The tracer to use for tracing the run."""
    trace_include_content: bool = False
    """Whether to include the content of the messages in the trace."""
    retries: dict[str, int] = field(default_factory=dict)
    """Number of retries for each tool so far."""
    tool_call_id: str | None = None
    """The ID of the tool call."""
    tool_name: str | None = None
    """Name of the tool being called."""
    retry: int = 0
    """Number of retries so far."""
    run_step: int = 0
    """The current step in the run."""
    tool_call_approved: bool = False
    """Whether a tool call that required approval has now been approved."""
    usage_limits: UsageLimits | None = None
    """Usage limits for the run."""

    __repr__ = _utils.dataclasses_no_defaults_repr


Review comment (suggestion):
Rather than adding this to `RunContext`, I'd prefer to pass it into the `ToolManager` directly if that's possible

Code after:
    model: Model
    """The model used in this run."""
    usage: RunUsage
    """LLM usage associated with the run."""
    prompt: str | Sequence[_messages.UserContent] | None = None
    """The original user prompt passed to the run."""
    messages: list[_messages.ModelMessage] = field(default_factory=list)
    """Messages exchanged in the conversation so far."""
    tracer: Tracer = field(default_factory=NoOpTracer)
    """The tracer to use for tracing the run."""
    trace_include_content: bool = False
    """Whether to include the content of the messages in the trace."""
    retries: dict[str, int] = field(default_factory=dict)
    """Number of retries for each tool so far."""
    tool_call_id: str | None = None
    """The ID of the tool call."""
    tool_name: str | None = None
    """Name of the tool being called."""
    retry: int = 0
    """Number of retries so far."""
    run_step: int = 0
    """The current step in the run."""
    tool_call_approved: bool = False
    """Whether a tool call that required approval has now been approved."""

    __repr__ = _utils.dataclasses_no_defaults_repr
Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/mcp.py (Python)

Code before:
    | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
)
"""The result type of an MCP tool call."""

CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[ToolResult]]
"""A function type that represents a tool call."""

ProcessToolCallback = Callable[
    [
        RunContext[Any],
        CallToolFunc,
        str,
        dict[str, Any],
    ],
    Awaitable[ToolResult],
]
"""A process tool callback.

It accepts a run context, the original tool call function, a tool name, and arguments.

Allows wrapping an MCP server tool call to customize it, including adding extra request
metadata.
"""


def mcp_server_discriminator(value: dict[str, Any]) -> str | None:
    if 'url' in value:
        if value['url'].endswith('/sse'):
            return 'sse'
        return 'streamable-http'
    return 'stdio'


class MCPServerConfig(BaseModel):
    mcp_servers: Annotated[
        dict[
            str,
            Annotated[
                Annotated[MCPServerStdio, Tag('stdio')]
                | Annotated[MCPServerStreamableHTTP, Tag('streamable-http')]
                | Annotated[MCPServerSSE, Tag('sse')],
                Discriminator(mcp_server_discriminator),
            ],
        ],
        Field(alias='mcpServers'),
    ]


def load_mcp_servers(config_path: str | Path) -> list[MCPServerStdio | MCPServerStreamableHTTP | MCPServerSSE]:
    config_path = Path(config_path)


Review comment (suggestion):
```suggestion
def _mcp_server_discriminator(value: dict[str, Any]) -> str | None:
```

Maybe make it private

Code after:
                    'headers': core_schema.typed_dict_field(
                        core_schema.dict_schema(core_schema.str_schema(), core_schema.str_schema()), required=False
                    ),
                }
            ),
        )

    @property
    def _transport_client(self):
        return streamablehttp_client  # pragma: no cover

    def __eq__(self, value: object, /) -> bool:
        if not isinstance(value, MCPServerStreamableHTTP):
            return False  # pragma: no cover
        return self.url == value.url


ToolResult = (
    str
    | messages.BinaryContent
    | dict[str, Any]
    | list[Any]
    | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
)
"""The result type of an MCP tool call."""

CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[ToolResult]]
"""A function type that represents a tool call."""

ProcessToolCallback = Callable[
    [
        RunContext[Any],
        CallToolFunc,
        str,
        dict[str, Any],
    ],
    Awaitable[ToolResult],
]
"""A process tool callback.

It accepts a run context, the original tool call function, a tool name, and arguments.

Allows wrapping an MCP server tool call to customize it, including adding extra request
metadata.
"""


def _mcp_server_discriminator(value: dict[str, Any]) -> str | None:
    if 'url' in value:
        if value['url'].endswith('/sse'):
            return 'sse'
Repository: pydantic/pydantic-ai
File: pydantic_evals/pydantic_evals/evaluators/__init__.py (Python)

Code before:
    'EqualsExpected',
    'Contains',
    'IsInstance',
    'MaxDuration',
    'LLMJudge',
    'HasMatchingSpan',
    'OutputConfig',
    # context
    'EvaluatorContext',
    # evaluator
    'Evaluator',
    'EvaluationReason',
    'EvaluatorFailure',
    'EvaluatorOutput',
    'EvaluatorSpec',
    'EvaluationReason',
    'EvaluationResult',
)


def __getattr__(name: str):
    if name == 'Python':
        raise ImportError(  # pragma: no cover
            'The `Python` evaluator has been removed for security reasons. See https://github.com/pydantic/pydantic-ai/pull/2808 for more details and a workaround.'
        )
    raise AttributeError(f'module {__name__} has no attribute {name}')


Review comment (suggestion):
```suggestion
    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
```

Code after:
    'EqualsExpected',
    'Contains',
    'IsInstance',
    'MaxDuration',
    'LLMJudge',
    'HasMatchingSpan',
    'OutputConfig',
    # context
    'EvaluatorContext',
    # evaluator
    'Evaluator',
    'EvaluationReason',
    'EvaluatorFailure',
    'EvaluatorOutput',
    'EvaluatorSpec',
    'EvaluationReason',
    'EvaluationResult',
)


def __getattr__(name: str):
    if name == 'Python':
        raise ImportError(
            'The `Python` evaluator has been removed for security reasons. See https://github.com/pydantic/pydantic-ai/pull/2808 for more details and a workaround.'
        )
    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/providers/ovhcloud.py (Python)

Code before:
from __future__ import annotations as _annotations

import os
from typing import overload

import httpx

from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import cached_async_http_client
from pydantic_ai.providers import Provider

try:
    from openai import AsyncOpenAI
except ImportError as _import_error:  # pragma: no cover
    raise ImportError(
        'Please install the `openai` package to use OVHcloud AI Endpoints provider.'
        'You can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
    ) from _import_error


class OVHcloudAIEndpointsProvider(Provider[AsyncOpenAI]):
    """Provider for OVHcloud AI Endpoints."""

    @property
    def name(self) -> str:
        return 'ovhcloud'

    @property
    def base_url(self) -> str:
        return 'https://oai.endpoints.kepler.ai.cloud.ovh.net/v1'

    @property
    def client(self) -> AsyncOpenAI:
        return self._client

    def model_profile(self, model_name: str) -> ModelProfile | None:
        return None

    @overload
    def __init__(self) -> None: ...

    @overload
    def __init__(self, *, api_key: str) -> None: ...

    @overload
    def __init__(self, *, api_key: str, http_client: httpx.AsyncClient) -> None: ...

Review comment (suggestion):
See above

```suggestion
class OVHcloudProvider(Provider[AsyncOpenAI]):
```

Code after:
from __future__ import annotations as _annotations

import os
from typing import overload

import httpx

from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import cached_async_http_client
from pydantic_ai.profiles.deepseek import deepseek_model_profile
from pydantic_ai.profiles.harmony import harmony_model_profile
from pydantic_ai.profiles.meta import meta_model_profile
from pydantic_ai.profiles.mistral import mistral_model_profile
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
from pydantic_ai.profiles.qwen import qwen_model_profile
from pydantic_ai.providers import Provider

try:
    from openai import AsyncOpenAI
except ImportError as _import_error:  # pragma: no cover
    raise ImportError(
        'Please install the `openai` package to use OVHcloud AI Endpoints provider.'
        'You can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
    ) from _import_error


class OVHcloudProvider(Provider[AsyncOpenAI]):
    """Provider for OVHcloud AI Endpoints."""

    @property
    def name(self) -> str:
        return 'ovhcloud'

    @property
    def base_url(self) -> str:
        return 'https://oai.endpoints.kepler.ai.cloud.ovh.net/v1'

    @property
    def client(self) -> AsyncOpenAI:
        return self._client

    def model_profile(self, model_name: str) -> ModelProfile | None:
        model_name = model_name.lower()

        prefix_to_profile = {
            'llama': meta_model_profile,
Repository: pydantic/pydantic-ai
File: tests/test_agent.py (Python)

Code before:
                ),
            ]
        )
        assert run.all_messages_json().startswith(b'[{"parts":[{"content":"Hello",')


def test_message_history_cannot_start_with_model_response():
    """Test that message history starting with ModelResponse raises UserError."""

    def simple_response(_messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart(content='Final response')])  # pragma: no cover

    agent = Agent(FunctionModel(simple_response))

    invalid_history = [
        ModelResponse(parts=[TextPart(content='ai response')]),
    ]

    with pytest.raises(
        UserError,
        match='Message history cannot start with a `ModelResponse`.',
    ):
        agent.run_sync('hello', message_history=invalid_history)


async def test_message_history_starts_with_model_request():
    """Test that valid history starting with ModelRequest works correctly."""

    def llm(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart('ok here is text')])

    agent = Agent(FunctionModel(llm))

    valid_history = [
        ModelRequest(parts=[UserPromptPart(content='Hello')]),
        ModelResponse(parts=[TextPart(content='Hi there!')]),
    ]

    # Should not raise error - valid history starting with ModelRequest
    async with agent.iter('How are you?', message_history=valid_history) as run:
        async for _ in run:
            pass
        # Verify messages are processed correctly
        all_messages = run.all_messages()
        assert len(all_messages) >= 3  # History + new request + response
        assert isinstance(all_messages[0], ModelRequest)  # First message is ModelRequest


async def test_empty_message_history_is_valid():
    """Test that empty message history works fine."""


Review comment (suggestion):
I don't think we need any tests other than the one that verifies the exception is raised, as the non-exception behavior is extensively tested by the existing test suite already.

Code after:
                ),
                ModelResponse(
                    parts=[TextPart(content='ok here is text')],
                    usage=RequestUsage(input_tokens=51, output_tokens=4),
                    model_name='function:llm:',
                    timestamp=IsDatetime(),
                    run_id=IsStr(),
                ),
            ]
        )
        assert run.all_messages_json().startswith(b'[{"parts":[{"content":"Hello",')


def test_message_history_cannot_start_with_model_response():
    """Test that message history starting with ModelResponse raises UserError."""

    agent = Agent('test')

    invalid_history = [
        ModelResponse(parts=[TextPart(content='ai response')]),
    ]

    with pytest.raises(
        UserError,
        match='Message history cannot start with a `ModelResponse`.',
    ):
        agent.run_sync('hello', message_history=invalid_history)
Repository: pydantic/pydantic-ai
File: tests/test_temporal.py (Python)

Code before:
        Args:
            input_value: The input number to process

        Returns:
            List of results from parallel execution
        """
        result = await parallel_test_graph.run(
            state=GraphState(),
            inputs=input_value,
        )
        return result


async def test_beta_graph_parallel_execution_in_workflow(client: Client):
    """Test that beta graph API with parallel execution works in Temporal workflows.

    This test verifies the fix for the bug where parallel task execution in graphs
    wasn't working properly with Temporal workflows due to GraphTask/GraphTaskRequest
    serialization issues.
    """
    async with Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[ParallelGraphWorkflow],
    ):
        output = await client.execute_workflow(  # pyright: ignore[reportUnknownMemberType]
            ParallelGraphWorkflow.run,
            args=[10],
            id=ParallelGraphWorkflow.__name__,
            task_queue=TASK_QUEUE,
        )
        # Results can be in any order due to parallel execution
        # 10 * 2 = 20, 10 * 3 = 30, 10 * 4 = 40
        assert sorted(output) == [20, 30, 40]


Review comment (suggestion):
```suggestion
        output = await client.execute_workflow(
```

Code after:
        Args:
            input_value: The input number to process

        Returns:
            List of results from parallel execution
        """
        result = await parallel_test_graph.run(
            state=GraphState(),
            inputs=input_value,
        )
        return result


async def test_beta_graph_parallel_execution_in_workflow(client: Client):
    """Test that beta graph API with parallel execution works in Temporal workflows.

    This test verifies the fix for the bug where parallel task execution in graphs
    wasn't working properly with Temporal workflows due to GraphTask/GraphTaskRequest
    serialization issues.
    """
    async with Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[ParallelGraphWorkflow],
    ):
        output = await client.execute_workflow(
            ParallelGraphWorkflow.run,
            args=[10],
            id=ParallelGraphWorkflow.__name__,
            task_queue=TASK_QUEUE,
        )
        # Results can be in any order due to parallel execution
        # 10 * 2 = 20, 10 * 3 = 30, 10 * 4 = 40
        assert sorted(output) == [20, 30, 40]
Repository: pydantic/pydantic-ai
File: tests/models/test_model_function.py (Python)

Code before:
            )
        )


async def test_pass_neither():
    with pytest.raises(TypeError, match='Either `function` or `stream_function` must be provided'):
        FunctionModel()  # pyright: ignore[reportCallIssue]


async def test_pass_both():
    Agent(FunctionModel(return_last, stream_function=stream_text_function))


async def stream_text_function_empty(_messages: list[ModelMessage], _: AgentInfo) -> AsyncIterator[str]:
    if False:
        yield 'hello '


async def test_return_empty():
    agent = Agent(FunctionModel(stream_function=stream_text_function_empty))
    with pytest.raises(ValueError, match='Stream function must return at least one item'):
        async with agent.run_stream(''):
            pass


async def test_central_content_filter_handling():
    """
    Test that the agent graph correctly raises ContentFilterError
    when a model returns finish_reason='content_filter' AND empty content.
    """

    async def filtered_response(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        # Simulate a model response that was blocked completely
        return ModelResponse(
            parts=[],  # Empty parts triggers the exception
            model_name='test-model',
            finish_reason='content_filter',
        )

    model = FunctionModel(function=filtered_response, model_name='test-model')
    agent = Agent(model)

    with pytest.raises(ContentFilterError, match="Content filter triggered. Finish reason: 'content_filter'"):
        await agent.run('Trigger filter')


async def test_central_content_filter_with_partial_content():
    """
    Test that the agent graph returns partial content (does not raise exception)
    even if finish_reason='content_filter', provided parts are not empty.
    """

Review comment (refactor):
This file is for _testing the function model_, not using the function model to test other things :) Can you move this to `test_agent.py` please?

Code after:
    agent = Agent(FunctionModel(stream_function=stream_structured_function), output_type=Foo)
    async with agent.run_stream('') as result:
        assert await result.get_output() == snapshot(Foo(x=1))
        assert result.usage() == snapshot(
            RunUsage(
                requests=1,
                input_tokens=50,
                output_tokens=4,
            )
        )


async def test_pass_neither():
    with pytest.raises(TypeError, match='Either `function` or `stream_function` must be provided'):
        FunctionModel()  # pyright: ignore[reportCallIssue]


async def test_pass_both():
    Agent(FunctionModel(return_last, stream_function=stream_text_function))


async def stream_text_function_empty(_messages: list[ModelMessage], _: AgentInfo) -> AsyncIterator[str]:
    if False:
        yield 'hello '


async def test_return_empty():
    agent = Agent(FunctionModel(stream_function=stream_text_function_empty))
    with pytest.raises(ValueError, match='Stream function must return at least one item'):
        async with agent.run_stream(''):
            pass
Repository: pydantic/pydantic-ai
File: tests/providers/test_bedrock.py (Python)

Code before:

def test_bedrock_provider_model_profile_us_gov_anthropic(env: TestEnv, mocker: MockerFixture):
    """Test that us-gov. prefixed Anthropic models get the correct profile.

    This specifically tests the us-gov. prefix which was previously broken
    because the provider only handled 2-character prefixes.
    """
    env.set('AWS_DEFAULT_REGION', 'us-east-1')
    provider = BedrockProvider()

    ns = 'pydantic_ai.providers.bedrock'
    anthropic_model_profile_mock = mocker.patch(f'{ns}.anthropic_model_profile', wraps=anthropic_model_profile)

    # Test us-gov. prefix (AWS GovCloud cross-region inference)
    profile = provider.model_profile('us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0')
    anthropic_model_profile_mock.assert_called_with('claude-sonnet-4-5-20250929')
    assert isinstance(profile, BedrockModelProfile)
    assert profile.bedrock_supports_tool_choice is True
    assert profile.bedrock_send_back_thinking_parts is True

    # Test global. prefix
    profile = provider.model_profile('global.anthropic.claude-opus-4-5-20251101-v1:0')
    anthropic_model_profile_mock.assert_called_with('claude-opus-4-5-20251101')
    assert isinstance(profile, BedrockModelProfile)
    assert profile.bedrock_supports_tool_choice is True
    assert profile.bedrock_send_back_thinking_parts is True


Review comment (suggestion):
we're already testing this above, aren't we?

let's add a test to check that we don't have any bedrock shorthand in `KnownModelName` with a geo prefix that we haven't added in the `BEDROCK_GEO_PREFIXES`

Code after:
    assert profile is None, f'model_profile returned {profile} for {model_name}'


def test_latest_bedrock_model_names_geo_prefixes_are_supported():
    """Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES.

    This test prevents adding new model names with geo prefixes that aren't handled
    by the provider's model_profile method.
    """
    model_names = get_args(LatestBedrockModelNames)

    missing_prefixes: set[str] = set()

    for model_name in model_names:
        # Model names with geo prefixes have 3+ dot-separated parts:
        # - No prefix: "anthropic.claude-xxx" (2 parts)
        # - With prefix: "us.anthropic.claude-xxx" (3 parts)
        parts = model_name.split('.')
        if len(parts) >= 3:
            geo_prefix = parts[0]
            if geo_prefix not in BEDROCK_GEO_PREFIXES:  # pragma: no cover
                missing_prefixes.add(geo_prefix)

    if missing_prefixes:  # pragma: no cover
        pytest.fail(
            f'Found geo prefixes in LatestBedrockModelNames that are not in BEDROCK_GEO_PREFIXES: {missing_prefixes}. '
            f'Please add them to BEDROCK_GEO_PREFIXES'
        )
Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/models/bedrock.py (Python)

Code before:
from __future__ import annotations

import functools
import typing
from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime
from itertools import count
from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload

import anyio.to_thread
from botocore.exceptions import ClientError
from mypy_boto3_bedrock_runtime.type_defs import DocumentSourceTypeDef
from typing_extensions import ParamSpec, assert_never

from pydantic_ai import (
    AudioUrl,
    BinaryContent,
    BuiltinToolCallPart,
    BuiltinToolReturnPart,
    CachePoint,
    DocumentUrl,
    FinishReason,
    ImageUrl,
    ModelMessage,
    ModelProfileSpec,
    ModelRequest,
    ModelResponse,
    ModelResponsePart,
    ModelResponseStreamEvent,
    RetryPromptPart,
    SystemPromptPart,
    TextPart,
    ThinkingPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
    VideoUrl,

Review comment (refactor):
Let's move this line into the `if TYPE_CHECKING:` block

Code after:
from __future__ import annotations

import functools
import typing
from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime
from itertools import count
from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
from urllib.parse import parse_qs, urlparse

import anyio.to_thread
from botocore.exceptions import ClientError
from typing_extensions import ParamSpec, assert_never

from pydantic_ai import (
    AudioUrl,
    BinaryContent,
    BuiltinToolCallPart,
    BuiltinToolReturnPart,
    CachePoint,
    DocumentUrl,
    FinishReason,
    ImageUrl,
    ModelMessage,
    ModelProfileSpec,
    ModelRequest,
    ModelResponse,
    ModelResponsePart,
    ModelResponseStreamEvent,
    RetryPromptPart,
    SystemPromptPart,
    TextPart,
    ThinkingPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
    VideoUrl,