litellm/tests/unit/interactions/test_gemini_interactions_transformation.py

"""
Tests for Gemini Interactions API transformation.

Covers:
- validate_environment: x-goog-api-key header, Api-Revision header
- get_complete_url: API key excluded from URL
- get/delete/cancel interaction request URLs
- transform_request: response_mime_type coalescing, image_config migration
"""

from unittest.mock import MagicMock, patch

import pytest


from litellm.interactions.litellm_responses_transformation.streaming_iterator import (
    LiteLLMResponsesInteractionsStreamingIterator,
)
from litellm.llms.gemini.interactions.transformation import (
    GoogleAIStudioInteractionsConfig,
)
from litellm.types.llms.openai import (
    OutputTextDeltaEvent,
    ResponseCompletedEvent,
    ResponseCreatedEvent,
)
from litellm.types.router import GenericLiteLLMParams

_PATCH_GET_API_KEY = "litellm.llms.gemini.common_utils.GeminiModelInfo.get_api_key"


@pytest.fixture
def config():
    return GoogleAIStudioInteractionsConfig()


class TestValidateEnvironment:
    def test_sets_x_goog_api_key_header(self, config):
        litellm_params = GenericLiteLLMParams(api_key="test-api-key-123")

        headers = config.validate_environment(
            headers={},
            model="gemini-2.5-flash",
            litellm_params=litellm_params,
        )

        assert headers["x-goog-api-key"] == "test-api-key-123"
        assert headers["Content-Type"] == "application/json"

    def test_no_api_key_skips_header(self, config):
        litellm_params = GenericLiteLLMParams(api_key=None)

        with patch(_PATCH_GET_API_KEY, return_value=None):
            headers = config.validate_environment(
                headers={},
                model="gemini-2.5-flash",
                litellm_params=litellm_params,
            )

        assert "x-goog-api-key" not in headers
        assert headers["Content-Type"] == "application/json"

    def test_no_litellm_params_skips_header(self, config):
        headers = config.validate_environment(
            headers={},
            model="gemini-2.5-flash",
            litellm_params=None,
        )

        assert "x-goog-api-key" not in headers
        assert headers["Content-Type"] == "application/json"

    def test_preserves_existing_headers(self, config):
        litellm_params = GenericLiteLLMParams(api_key="test-key")

        headers = config.validate_environment(
            headers={"X-Custom": "value"},
            model="gemini-2.5-flash",
            litellm_params=litellm_params,
        )

        assert headers["X-Custom"] == "value"
        assert headers["x-goog-api-key"] == "test-key"

    def test_sets_api_revision_header(self, config):
        headers = config.validate_environment(headers={}, model="gemini-2.5-flash", litellm_params=None)
        assert headers["Api-Revision"] == "2026-05-20"


class TestGetCompleteUrl:
    def test_url_excludes_api_key(self, config):
        with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
            url = config.get_complete_url(
                api_base=None,
                model="gemini-2.5-flash",
                litellm_params={"api_key": "secret-key"},
            )

        assert "key=" not in url
        assert "secret-key" not in url
        assert url.endswith("/interactions")

    def test_stream_url_has_alt_sse_only(self, config):
        with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
            url = config.get_complete_url(
                api_base=None,
                model="gemini-2.5-flash",
                litellm_params={"api_key": "secret-key"},
                stream=True,
            )

        assert "key=" not in url
        assert "secret-key" not in url
        assert "alt=sse" in url

    def test_raises_without_api_key(self, config):
        with patch(_PATCH_GET_API_KEY, return_value=None):
            with pytest.raises(ValueError, match="Google API key is required"):
                config.get_complete_url(
                    api_base=None,
                    model="gemini-2.5-flash",
                    litellm_params={"api_key": None},
                )


class TestTransformRequest:
    def test_passes_environment_to_request_body(self, config):
        request_body = config.transform_request(
            model=None,
            agent="my-custom-slides-agent",
            input=[
                {
                    "type": "text",
                    "text": "Create a 5-slide presentation about AI trends.",
                }
            ],
            optional_params={
                "environment": "remote",
                "stream": False,
            },
            litellm_params=GenericLiteLLMParams(api_key="test-api-key"),
            headers={},
        )

        assert request_body["agent"] == "my-custom-slides-agent"
        assert request_body["environment"] == "remote"
        assert request_body["stream"] is False
        assert request_body["input"] == [{"type": "text", "text": "Create a 5-slide presentation about AI trends."}]

    def test_passes_environment_object_to_request_body(self, config):
        environment_config = {
            "type": "remote",
            "sources": [{"type": "gcs", "uri": "gs://bucket/skills.zip"}],
            "network": {"egress": "allow_all"},
        }
        request_body = config.transform_request(
            model=None,
            agent="waverunner",
            input="What is 2 + 2?",
            optional_params={"environment": environment_config},
            litellm_params=GenericLiteLLMParams(api_key="test-api-key"),
            headers={},
        )

        assert request_body["environment"] == environment_config

    def test_passes_existing_environment_id_to_request_body(self, config):
        env_id = "env-abc123"
        request_body = config.transform_request(
            model=None,
            agent="my-custom-slides-agent",
            input="Continue the presentation.",
            optional_params={"environment": env_id},
            litellm_params=GenericLiteLLMParams(api_key="test-api-key"),
            headers={},
        )

        assert request_body["environment"] == env_id

    def test_stream_param_included_in_request_body(self, config):
        """When stream=True is in optional_params, the request body must include it
        so the proxy forwards the SSE streaming flag to Google's backend."""
        body = config.transform_request(
            model="gemini-2.5-flash",
            agent=None,
            input="Hello",
            optional_params={"stream": True},
            litellm_params=GenericLiteLLMParams(api_key="test-key"),
            headers={},
        )

        assert body.get("stream") is True
        assert body.get("input") == "Hello"

    def test_stream_false_not_included_when_absent(self, config):
        body = config.transform_request(
            model="gemini-2.5-flash",
            agent=None,
            input="Hello",
            optional_params={},
            litellm_params=GenericLiteLLMParams(api_key="test-key"),
            headers={},
        )

        assert "stream" not in body


class TestStreamingIterator:
    def _make_iterator(self) -> LiteLLMResponsesInteractionsStreamingIterator:
        return LiteLLMResponsesInteractionsStreamingIterator(
            model="gpt-5.4",
            litellm_custom_stream_wrapper=MagicMock(),
            request_input="hi",
            optional_params={},
        )

    def _make_text_delta(self, text: str, item_id: str = "item_1") -> OutputTextDeltaEvent:
        event = MagicMock(spec=OutputTextDeltaEvent)
        event.delta = text
        event.item_id = item_id
        return event

    def _make_response_created(self) -> ResponseCreatedEvent:
        event = MagicMock(spec=ResponseCreatedEvent)
        event.response = MagicMock(id="resp_123")
        return event

    def test_step_delta_includes_type_field(self):
        """step.delta events must carry delta.type='text' so the UI can display them."""
        it = self._make_iterator()
        it.sent_interaction_start = True
        it.sent_content_start = True

        chunk = it._transform_responses_chunk_to_interactions_chunk(self._make_text_delta("Hello"))

        assert chunk is not None
        assert chunk.event_type == "step.delta"
        assert chunk.delta == {"type": "text", "text": "Hello"}

    def test_response_created_emits_interaction_created(self):
        it = self._make_iterator()

        chunk = it._transform_responses_chunk_to_interactions_chunk(self._make_response_created())

        assert chunk is not None
        assert chunk.event_type == "interaction.created"
        assert chunk.id == "resp_123"
        assert it.sent_interaction_start is True

    def test_text_delta_sequence(self):
        """First chunk yields created + step.start + step.delta; later chunks yield step.delta."""
        it = self._make_iterator()

        first_events = it._events_for_chunk(self._make_text_delta("Hello"))
        assert [e.event_type for e in first_events] == [
            "interaction.created",
            "step.start",
            "step.delta",
        ]
        assert first_events[-1].delta == {"type": "text", "text": "Hello"}
        assert it.sent_interaction_start is True
        assert it.sent_content_start is True

        second_events = it._events_for_chunk(self._make_text_delta(" World"))
        assert [e.event_type for e in second_events] == ["step.delta"]
        assert second_events[0].delta == {"type": "text", "text": " World"}

        third_events = it._events_for_chunk(self._make_text_delta("!"))
        assert [e.event_type for e in third_events] == ["step.delta"]
        assert third_events[0].delta == {"type": "text", "text": "!"}

    def test_first_text_delta_without_item_id_uses_fallback_id(self):
        it = self._make_iterator()
        event = self._make_text_delta("Hi")
        event.item_id = None

        events = it._events_for_chunk(event)

        assert events[0].event_type == "interaction.created"
        assert events[0].id == f"interaction_{id(it)}"

    def test_first_text_delta_emits_text_via_compat_shim(self):
        """The legacy single-chunk shim must surface the synthetic events AND the delta."""
        it = self._make_iterator()

        first = it._transform_responses_chunk_to_interactions_chunk(self._make_text_delta("Hello"))
        assert first is not None
        assert first.event_type == "interaction.created"

        second = it.__next__() if it._pending_events else None
        assert second is not None
        assert second.event_type == "step.start"

        third = it.__next__() if it._pending_events else None
        assert third is not None
        assert third.event_type == "step.delta"
        assert third.delta == {"type": "text", "text": "Hello"}

    def test_response_created_then_text_delta_emits_step_start_and_delta(self):
        """Realistic flow: response.created arrives first, then text delta."""
        it = self._make_iterator()

        first = it._events_for_chunk(self._make_response_created())
        assert [e.event_type for e in first] == ["interaction.created"]

        second = it._events_for_chunk(self._make_text_delta("Hello"))
        assert [e.event_type for e in second] == ["step.start", "step.delta"]
        assert second[-1].delta == {"type": "text", "text": "Hello"}

    def test_no_text_token_is_dropped_during_streaming(self):
        """Concatenated step.delta payloads must equal the upstream text."""
        it = self._make_iterator()

        chunks = ["Hello", " ", "world", "!"]
        emitted_text = ""
        for c in chunks:
            for ev in it._events_for_chunk(self._make_text_delta(c)):
                if ev.event_type == "step.delta":
                    assert ev.delta is not None
                    emitted_text += ev.delta["text"]

        assert emitted_text == "Hello world!"

    def test_stop_iteration_fallback_emits_completion_event(self):
        """If upstream ends without ResponseCompletedEvent, terminal events still flow."""
        from unittest.mock import MagicMock

        text_event = self._make_text_delta("hi")
        sync_iter = MagicMock()
        sync_iter.__iter__ = lambda self: self
        sync_iter.__next__ = MagicMock(side_effect=[text_event, StopIteration])

        it = LiteLLMResponsesInteractionsStreamingIterator(
            model="gpt-5.4",
            litellm_custom_stream_wrapper=sync_iter,
            request_input="hi",
            optional_params={},
        )

        emitted: list = []
        try:
            while True:
                emitted.append(next(it))
        except StopIteration:
            pass

        event_types = [e.event_type for e in emitted]
        assert event_types == [
            "interaction.created",
            "step.start",
            "step.delta",
            "step.stop",
            "interaction.completed",
        ]
        terminal = emitted[-1]
        assert terminal.steps == [
            {
                "type": "model_output",
                "content": [{"type": "text", "text": "hi"}],
            }
        ]
        # EOF-flushed terminal event must carry the same id as interaction.created.
        assert terminal.id == emitted[0].id == "item_1"

    def test_response_completed_emits_stop_then_completion(self):
        """ResponseCompletedEvent expands into step.stop + interaction.completed."""
        from unittest.mock import MagicMock

        text_event = self._make_text_delta("hi")
        completed = MagicMock(spec=ResponseCompletedEvent)
        completed.response = MagicMock(id="resp_999")

        sync_iter = MagicMock()
        sync_iter.__iter__ = lambda self: self
        sync_iter.__next__ = MagicMock(side_effect=[text_event, completed])

        it = LiteLLMResponsesInteractionsStreamingIterator(
            model="gpt-5.4",
            litellm_custom_stream_wrapper=sync_iter,
            request_input="hi",
            optional_params={},
        )

        emitted: list = []
        try:
            while True:
                emitted.append(next(it))
        except StopIteration:
            pass

        event_types = [e.event_type for e in emitted]
        assert event_types == [
            "interaction.created",
            "step.start",
            "step.delta",
            "step.stop",
            "interaction.completed",
        ]
        # StopIteration fallback path must NOT add a duplicate completion event.
        assert event_types.count("interaction.completed") == 1
        # When the stream starts directly with a text delta (no preceding
        # response.created), the terminal events must reuse the id derived from
        # the first chunk's item_id rather than switching to response.id, so
        # consumers can correlate the start and completion events by id.
        assert emitted[0].id == "item_1"
        assert emitted[-1].id == "item_1"


class TestInteractionOperationUrls:
    """Test that get/delete/cancel interaction URLs exclude API key."""

    @pytest.mark.parametrize(
        "method_name,interaction_id,expected_suffix",
        [
            ("transform_get_interaction_request", "interaction-123", "interaction-123"),
            (
                "transform_delete_interaction_request",
                "interaction-456",
                "interaction-456",
            ),
            (
                "transform_cancel_interaction_request",
                "interaction-789",
                "interaction-789:cancel",
            ),
        ],
    )
    def test_url_excludes_key(self, config, method_name, interaction_id, expected_suffix):
        with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
            url, params = getattr(config, method_name)(
                interaction_id=interaction_id,
                api_base="https://generativelanguage.googleapis.com",
                litellm_params=GenericLiteLLMParams(api_key="secret-key"),
                headers={},
            )

        assert "key=" not in url
        assert "secret-key" not in url
        assert expected_suffix in url

    def test_interaction_id_is_encoded_as_one_path_segment(self, config):
        with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
            url, params = config.transform_cancel_interaction_request(
                interaction_id="../../interactions/other?x=1#frag",
                api_base="https://generativelanguage.googleapis.com",
                litellm_params=GenericLiteLLMParams(api_key="secret-key"),
                headers={},
            )

        assert (
            url
            == "https://generativelanguage.googleapis.com/v1beta/interactions/..%2F..%2Finteractions%2Fother%3Fx%3D1%23frag:cancel"
        )
        assert params == {}

    def test_get_interaction_raises_without_key(self, config):
        with patch(_PATCH_GET_API_KEY, return_value=None):
            with pytest.raises(ValueError, match="Google API key is required"):
                config.transform_get_interaction_request(
                    interaction_id="interaction-123",
                    api_base="https://generativelanguage.googleapis.com",
                    litellm_params=GenericLiteLLMParams(api_key=None),
                    headers={},
                )


class TestTransformRequestSchemaCoalescing:
    """Test new-schema request coalescing (Api-Revision: 2026-05-20)."""

    def test_response_mime_type_folded_into_response_format(self, config):
        body = config.transform_request(
            model="gemini/gemini-2.5-flash",
            agent=None,
            input="summarise",
            optional_params={
                "response_mime_type": "application/json",
                "response_format": {"type": "object", "properties": {}},
            },
            litellm_params=GenericLiteLLMParams(),
            headers={},
        )

        # response_mime_type must not appear as a top-level body key
        assert "response_mime_type" not in body
        rf = body["response_format"]
        assert rf["type"] == "text"
        assert rf["mime_type"] == "application/json"
        assert "schema" in rf

    def test_image_config_moved_to_response_format(self, config):
        body = config.transform_request(
            model="gemini/gemini-2.5-flash",
            agent=None,
            input="draw a sunset",
            optional_params={
                "generation_config": {
                    "temperature": 0.7,
                    "image_config": {"aspect_ratio": "1:1", "image_size": "1K"},
                }
            },
            litellm_params=GenericLiteLLMParams(),
            headers={},
        )

        # image_config removed from generation_config
        assert "image_config" not in body.get("generation_config", {})
        # moved into response_format with type=image
        rf = body["response_format"]
        assert rf["type"] == "image"
        assert rf["aspect_ratio"] == "1:1"

    def test_response_mime_type_skipped_when_response_format_is_list(self, config):
        """Lists are already polymorphic; do not wrap them into schema."""
        rf_list = [
            {"type": "text", "mime_type": "application/json"},
            {"type": "image", "aspect_ratio": "1:1"},
        ]
        body = config.transform_request(
            model="gemini/gemini-2.5-flash",
            agent=None,
            input="multimodal",
            optional_params={
                "response_format": rf_list,
                "response_mime_type": "application/json",
            },
            litellm_params=GenericLiteLLMParams(),
            headers={},
        )

        assert body["response_format"] == rf_list
        assert "response_mime_type" not in body

    def test_image_config_appended_to_response_format_list_without_mutating_input(
        self,
        config,
    ):
        """When response_format is already a list, image_config must not mutate optional_params."""
        text_rf = {"type": "text", "mime_type": "application/json"}
        optional_params = {
            "response_format": [text_rf],
            "generation_config": {
                "image_config": {"aspect_ratio": "16:9", "image_size": "2K"},
            },
        }
        original_rf = optional_params["response_format"]

        body = config.transform_request(
            model="gemini/gemini-2.5-flash",
            agent=None,
            input="draw and summarise",
            optional_params=optional_params,
            litellm_params=GenericLiteLLMParams(),
            headers={},
        )

        assert optional_params["response_format"] is original_rf
        assert len(optional_params["response_format"]) == 1
        assert body["response_format"] == [
            text_rf,
            {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"},
        ]

        # Retry must not append a second image entry into the caller's list.
        body_retry = config.transform_request(
            model="gemini/gemini-2.5-flash",
            agent=None,
            input="draw and summarise",
            optional_params=optional_params,
            litellm_params=GenericLiteLLMParams(),
            headers={},
        )
        assert len(optional_params["response_format"]) == 1
        assert body_retry["response_format"] == body["response_format"]