pydantic--pydantic-ai/tests/test_direct.py

import re
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager, contextmanager
from datetime import timezone
from typing import cast
from unittest.mock import AsyncMock

import pytest

from pydantic_ai import Agent
from pydantic_ai.direct import (
    StreamedResponseSync,
    _prepare_model,  # pyright: ignore[reportPrivateUsage]
    model_request,
    model_request_stream,
    model_request_stream_sync,
    model_request_sync,
)
from pydantic_ai.messages import (
    FinalResultEvent,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    ModelResponseStreamEvent,
    PartDeltaEvent,
    PartEndEvent,
    PartStartEvent,
    TextPart,
    TextPartDelta,
    ToolCallPart,
)
from pydantic_ai.models import ModelRequestParameters, StreamedResponse
from pydantic_ai.models.instrumented import InstrumentedModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.usage import RequestUsage

from ._inline_snapshot import snapshot
from .conftest import IsDatetime, IsNow, IsStr

pytestmark = pytest.mark.anyio


async def test_model_request():
    model_response = await model_request('test', [ModelRequest.user_text_prompt('x')])
    assert model_response == snapshot(
        ModelResponse(
            parts=[TextPart(content='success (no tool calls)')],
            model_name='test',
            timestamp=IsNow(tz=timezone.utc),
            usage=RequestUsage(input_tokens=51, output_tokens=4),
            provider_name='test',
        )
    )


async def test_model_request_tool_call():
    model_response = await model_request(
        'test',
        [ModelRequest.user_text_prompt('x')],
        model_request_parameters=ModelRequestParameters(
            function_tools=[ToolDefinition(name='tool_name', parameters_json_schema={'type': 'object'})],
            allow_text_output=False,
        ),
    )
    assert model_response == snapshot(
        ModelResponse(
            parts=[ToolCallPart(tool_name='tool_name', args={}, tool_call_id=IsStr(regex='pyd_ai_.*'))],
            model_name='test',
            timestamp=IsNow(tz=timezone.utc),
            usage=RequestUsage(input_tokens=51, output_tokens=2),
            provider_name='test',
        )
    )


def test_model_request_sync():
    model_response = model_request_sync('test', [ModelRequest.user_text_prompt('x')])
    assert model_response == snapshot(
        ModelResponse(
            parts=[TextPart(content='success (no tool calls)')],
            model_name='test',
            timestamp=IsNow(tz=timezone.utc),
            usage=RequestUsage(input_tokens=51, output_tokens=4),
            provider_name='test',
        )
    )


def test_model_request_stream_sync():
    with model_request_stream_sync('test', [ModelRequest.user_text_prompt('x')]) as stream:
        chunks = list(stream)
        assert chunks == snapshot(
            [
                PartStartEvent(index=0, part=TextPart(content='')),
                FinalResultEvent(tool_name=None, tool_call_id=None),
                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='success ')),
                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='(no ')),
                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='tool ')),
                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='calls)')),
                PartEndEvent(index=0, part=TextPart(content='success (no tool calls)')),
            ]
        )
        assert stream.response == snapshot(
            ModelResponse(
                parts=[TextPart(content='success (no tool calls)')],
                usage=RequestUsage(input_tokens=51, output_tokens=4),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
            )
        )

        repr_str = repr(stream)
        assert 'TestStreamedResponse' in repr_str
        assert 'test' in repr_str


async def test_model_request_stream():
    async with model_request_stream('test', [ModelRequest.user_text_prompt('x')]) as stream:
        chunks = [chunk async for chunk in stream]
    assert chunks == snapshot(
        [
            PartStartEvent(index=0, part=TextPart(content='')),
            FinalResultEvent(tool_name=None, tool_call_id=None),
            PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='success ')),
            PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='(no ')),
            PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='tool ')),
            PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='calls)')),
            PartEndEvent(index=0, part=TextPart(content='success (no tool calls)')),
        ]
    )


def test_model_request_stream_sync_without_context_manager():
    """Test that accessing properties or iterating without context manager raises RuntimeError."""
    messages: list[ModelMessage] = [ModelRequest.user_text_prompt('x')]

    expected_error_msg = re.escape(
        'StreamedResponseSync must be used as a context manager. Use: `with model_request_stream_sync(...) as stream:`'
    )

    stream_cm = model_request_stream_sync('test', messages)

    stream_repr = repr(stream_cm)
    assert 'StreamedResponseSync' in stream_repr
    assert 'context_entered=False' in stream_repr

    with pytest.raises(RuntimeError, match=expected_error_msg):
        _ = stream_cm.model_name

    with pytest.raises(RuntimeError, match=expected_error_msg):
        _ = stream_cm.timestamp

    with pytest.raises(RuntimeError, match=expected_error_msg):
        _ = stream_cm.response

    with pytest.raises(RuntimeError, match=expected_error_msg):
        _ = stream_cm.usage

    with pytest.raises(RuntimeError, match=expected_error_msg):
        list(stream_cm)

    with pytest.raises(RuntimeError, match=expected_error_msg):
        for _ in stream_cm:
            break  # pragma: no cover


def test_model_request_stream_sync_exception_on_enter():
    """A failure entering the async stream surfaces when the context manager is entered."""
    async_stream_mock = AsyncMock()
    async_stream_mock.__aenter__ = AsyncMock(side_effect=ValueError('Stream error'))

    stream_sync = StreamedResponseSync(_async_stream_cm=async_stream_mock)

    with pytest.raises(ValueError, match='Stream error'):
        with stream_sync:
            pass  # pragma: no cover


def test_model_request_stream_sync_exception_during_iteration():
    """An error raised while streaming (after a successful enter) surfaces to the consumer."""

    class FailingStream:
        def __aiter__(self) -> 'FailingStream':
            return self

        async def __anext__(self) -> ModelResponseStreamEvent:
            raise ValueError('Stream error')

    @asynccontextmanager
    async def failing_stream_cm() -> AsyncGenerator[StreamedResponse]:
        yield cast(StreamedResponse, FailingStream())

    with StreamedResponseSync(_async_stream_cm=failing_stream_cm()) as stream_sync:
        with pytest.raises(ValueError, match='Stream error'):
            list(stream_sync)


def test_model_request_stream_sync_intermediate_get():
    """Test getting properties of StreamedResponse before consuming all events."""
    messages: list[ModelMessage] = [ModelRequest.user_text_prompt('x')]

    with model_request_stream_sync('test', messages) as stream:
        response = stream.response
        assert response is not None

        usage = stream.usage
        assert usage is not None

        assert stream.model_name == 'test'
        assert stream.timestamp is not None


@contextmanager
def set_instrument_default(value: bool):
    """Context manager to temporarily set the default instrumentation value."""
    initial_value = Agent._instrument_default  # pyright: ignore[reportPrivateUsage]
    try:
        Agent._instrument_default = value  # pyright: ignore[reportPrivateUsage]
        yield
    finally:
        Agent._instrument_default = initial_value  # pyright: ignore[reportPrivateUsage]


async def test_model_request_with_instructions_on_message():
    """Instructions set on ModelRequest are picked up even without instruction_parts on ModelRequestParameters."""
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    def check_instructions(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.instructions == 'Be concise.'
        return ModelResponse(parts=[TextPart(content='ok')])

    response = await model_request(
        FunctionModel(check_instructions),
        [ModelRequest.user_text_prompt('Hello', instructions='Be concise.')],
    )
    assert response.parts[0].content == 'ok'  # type: ignore[union-attr]


async def test_model_request_with_instruction_parts_on_parameters():
    """When instruction_parts is explicitly set on ModelRequestParameters, it is preserved as-is."""
    from pydantic_ai.messages import InstructionPart
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    def check_instructions(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.instructions == 'From params.'
        return ModelResponse(parts=[TextPart(content='ok')])

    response = await model_request(
        FunctionModel(check_instructions),
        [ModelRequest.user_text_prompt('Hello')],
        model_request_parameters=ModelRequestParameters(
            instruction_parts=[InstructionPart(content='From params.')],
        ),
    )
    assert response.parts[0].content == 'ok'  # type: ignore[union-attr]


async def test_model_request_instructions_fallback_with_tool_return():
    """Instructions from the second-to-last request are used when the last has only tool-return parts."""
    from pydantic_ai.messages import ToolCallPart, ToolReturnPart, UserPromptPart
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    def check_instructions(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.instructions == 'Be helpful.'
        return ModelResponse(parts=[TextPart(content='ok')])

    messages: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='Hello')], instructions='Be helpful.'),
        ModelResponse(parts=[ToolCallPart(tool_name='my_tool', args='{}', tool_call_id='call_1')]),
        ModelRequest(parts=[ToolReturnPart(tool_name='my_tool', content='result', tool_call_id='call_1')]),
    ]

    response = await model_request(FunctionModel(check_instructions), messages)
    assert response.parts[0].content == 'ok'  # type: ignore[union-attr]


async def test_model_request_stream_with_instructions_on_message():
    """Instructions set on ModelRequest are picked up in streaming mode too."""
    response_parts: list[str] = []
    async with model_request_stream(
        'test',
        [ModelRequest.user_text_prompt('Hello', instructions='Be concise.')],
    ) as stream:
        async for event in stream:
            if isinstance(event, PartStartEvent) and isinstance(event.part, TextPart):
                response_parts.append(event.part.content)
    assert len(response_parts) > 0


def test_prepare_model():
    with set_instrument_default(False):
        model = _prepare_model('test', None)
        assert isinstance(model, TestModel)

        model = _prepare_model('test', True)
        assert isinstance(model, InstrumentedModel)

    with set_instrument_default(True):
        model = _prepare_model('test', None)
        assert isinstance(model, InstrumentedModel)

        model = _prepare_model('test', False)
        assert isinstance(model, TestModel)