microsoft--agent-framework/python/packages/ag-ui/tests/ag_ui/conftest.py

# Copyright (c) Microsoft. All rights reserved.

"""Shared test fixtures and stubs for AG-UI tests."""

import sys
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableSequence, Sequence
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Generic, Literal, TypedDict, cast, overload  # noqa: F401

import pytest
from agent_framework import (
    AgentResponse,
    AgentResponseUpdate,
    AgentSession,
    BaseChatClient,
    ChatOptions,
    ChatResponse,
    ChatResponseUpdate,
    Content,
    Message,
    ServiceSessionId,
    SupportsAgentRun,
    SupportsChatGetResponse,
)
from agent_framework._clients import OptionsCoT
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._tools import FunctionInvocationLayer
from agent_framework._types import ResponseStream
from agent_framework.observability import ChatTelemetryLayer

if sys.version_info >= (3, 12):
    from typing import override  # type: ignore # pragma: no cover
else:
    from typing_extensions import override  # type: ignore[import] # pragma: no cover

StreamFn = Callable[..., AsyncIterable[ChatResponseUpdate]]
ResponseFn = Callable[..., Awaitable[ChatResponse]]


def pytest_configure() -> None:
    """Ensure this test directory is on sys.path so helper modules can be imported by name."""
    test_dir = str(Path(__file__).resolve().parent)
    if test_dir not in sys.path:
        sys.path.insert(0, test_dir)


class StreamingChatClientStub(
    FunctionInvocationLayer[OptionsCoT],
    ChatMiddlewareLayer[OptionsCoT],
    ChatTelemetryLayer[OptionsCoT],
    BaseChatClient[OptionsCoT],
    Generic[OptionsCoT],
):
    """Typed streaming stub that satisfies SupportsChatGetResponse."""

    def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
        super().__init__(middleware=[])
        self._stream_fn = stream_fn
        self._response_fn = response_fn
        self.last_session: AgentSession | None = None
        self.last_service_session_id: str | ServiceSessionId | None = None

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: ChatOptions[Any],
        **kwargs: Any,
    ) -> Awaitable[ChatResponse[Any]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: OptionsCoT | ChatOptions[None] | None = ...,
        **kwargs: Any,
    ) -> Awaitable[ChatResponse[Any]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[True],
        options: OptionsCoT | ChatOptions[Any] | None = ...,
        **kwargs: Any,
    ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...

    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: bool = False,
        options: OptionsCoT | ChatOptions[Any] | None = None,
        **kwargs: Any,
    ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
        client_kwargs = kwargs.get("client_kwargs")
        if isinstance(client_kwargs, Mapping):
            self.last_session = cast(AgentSession | None, client_kwargs.get("session"))
        else:
            self.last_session = None
        self.last_service_session_id = self.last_session.service_session_id if self.last_session else None
        if stream:
            return super().get_response(
                messages=messages,
                stream=True,
                options=options,
                **kwargs,
            )
        return super().get_response(
            messages=messages,
            stream=False,
            options=options,
            **kwargs,
        )

    @override
    def _inner_get_response(
        self,
        *,
        messages: Sequence[Message],
        stream: bool = False,
        options: Mapping[str, Any],
        **kwargs: Any,
    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
        if stream:

            def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
                return ChatResponse.from_updates(updates)

            return ResponseStream(self._stream_fn(messages, options, **kwargs), finalizer=_finalize)

        return self._get_response_impl(messages, options, **kwargs)

    async def _get_response_impl(
        self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
    ) -> ChatResponse:
        """Non-streaming implementation."""
        if self._response_fn is not None:
            return await self._response_fn(messages, options, **kwargs)

        contents: list[Any] = []
        async for update in self._stream_fn(list(messages), dict(options), **kwargs):
            contents.extend(update.contents)

        return ChatResponse(
            messages=[Message(role="assistant", contents=contents)],
            response_id="stub-response",
        )


def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
    """Create a stream function that yields from a static list of updates."""

    async def _stream(
        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
    ) -> AsyncIterator[ChatResponseUpdate]:
        for update in updates:
            yield update

    return _stream


class StubAgent(SupportsAgentRun):
    """Minimal SupportsAgentRun stub for orchestrator tests."""

    def __init__(
        self,
        updates: list[AgentResponseUpdate] | None = None,
        *,
        agent_id: str = "stub-agent",
        agent_name: str | None = "stub-agent",
        default_options: Any | None = None,
        client: Any | None = None,
    ) -> None:
        self.id = agent_id
        self.name = agent_name
        self.description = "stub agent"
        self.updates = updates or [AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")]
        self.default_options: dict[str, Any] = (
            default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
        )
        self.client = client or SimpleNamespace(function_invocation_configuration=None)
        self.messages_received: list[Any] = []
        self.tools_received: list[Any] | None = None
        self.last_session: AgentSession | None = None

    @overload
    def run(
        self,
        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
        *,
        stream: Literal[False] = ...,
        session: AgentSession | None = None,
        **kwargs: Any,
    ) -> Awaitable[AgentResponse[Any]]: ...

    @overload
    def run(
        self,
        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
        *,
        stream: Literal[True],
        session: AgentSession | None = None,
        **kwargs: Any,
    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...

    def run(
        self,
        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
        *,
        stream: bool = False,
        session: AgentSession | None = None,
        **kwargs: Any,
    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
        if stream:

            async def _stream() -> AsyncIterator[AgentResponseUpdate]:
                if messages is None:
                    self.messages_received = []
                elif isinstance(messages, (str, Content, Message)):
                    self.messages_received = [messages]
                else:
                    self.messages_received = list(messages)
                self.last_session = session
                self.tools_received = kwargs.get("tools")
                for update in self.updates:
                    yield update

            def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
                return AgentResponse.from_updates(updates)

            return ResponseStream(_stream(), finalizer=_finalize)

        async def _get_response() -> AgentResponse[Any]:
            return AgentResponse(messages=[], response_id="stub-response")

        return _get_response()

    def create_session(self, **kwargs: Any) -> AgentSession:
        return AgentSession(session_id=kwargs.get("session_id"))

    def get_session(self, service_session_id: str | ServiceSessionId, *, session_id: str | None = None) -> AgentSession:
        return AgentSession(session_id=session_id, service_session_id=service_session_id)


# Fixtures


@pytest.fixture
def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
    """Return the StreamingChatClientStub class for creating test instances."""
    return StreamingChatClientStub  # type: ignore[return-value]


@pytest.fixture
def stream_from_updates_fixture() -> Callable[[list[ChatResponseUpdate]], StreamFn]:
    """Return the stream_from_updates helper function."""
    return stream_from_updates


@pytest.fixture
def stub_agent() -> type[SupportsAgentRun]:
    """Return the StubAgent class for creating test instances."""
    return StubAgent  # type: ignore[return-value]


# ── Fixtures for golden / integration tests ──


@pytest.fixture
def collect_events() -> Callable[..., Any]:
    """Return an async helper that collects all events from an async generator."""

    async def _collect(async_gen: AsyncIterable[Any]) -> list[Any]:
        return [event async for event in async_gen]

    return _collect


@pytest.fixture
def make_agent_wrapper() -> Callable[..., Any]:
    """Factory that builds an AgentFrameworkAgent from a stream function.

    Usage::

        agent = make_agent_wrapper(
            stream_fn=stream_from_updates(updates),
            state_schema=...,
        )
        events = [e async for e in agent.run(payload)]
    """
    from agent_framework_ag_ui import AgentFrameworkAgent

    def _factory(
        stream_fn: StreamFn,
        *,
        state_schema: Any | None = None,
        predict_state_config: dict[str, dict[str, str]] | None = None,
        require_confirmation: bool = True,
    ) -> Any:
        client = StreamingChatClientStub(stream_fn)
        stub = StubAgent(client=client)
        return AgentFrameworkAgent(
            agent=stub,
            state_schema=state_schema,
            predict_state_config=predict_state_config,
            require_confirmation=require_confirmation,
        )

    return _factory


@pytest.fixture
def make_app() -> Callable[..., Any]:
    """Factory that builds a FastAPI app with an AG-UI endpoint.

    Usage::

        app = make_app(agent_or_wrapper, path="/test")
    """
    from fastapi import FastAPI

    from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint

    def _factory(
        agent: Any,
        *,
        path: str = "/",
        state_schema: Any | None = None,
        predict_state_config: dict[str, dict[str, str]] | None = None,
        default_state: dict[str, Any] | None = None,
    ) -> FastAPI:
        app = FastAPI()
        add_agent_framework_fastapi_endpoint(
            app,
            agent,
            path=path,
            state_schema=state_schema,
            predict_state_config=predict_state_config,
            default_state=default_state,
        )
        return app

    return _factory