litellm/tests/integration/providers/test_websearch_interception_wire.py

import json
from pathlib import Path
from typing import Final

import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server

BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
INVOKE_TARGET: Final = f"/model/{BEDROCK_MODEL}/invoke"
SEARCH_TARGET: Final = "/tavily/search"
SEARCH_RESULT: Final = {
    "title": "Synthetic result",
    "url": "https://example.test/result",
    "content": "the snippet text",
}


def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]:
    frames: Final = tuple(frame for frame in text.split("\n\n") if frame.strip())
    return tuple(
        (
            next(line.removeprefix("event: ") for line in frame.splitlines() if line.startswith("event: ")),
            json.loads(next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))),
        )
        for frame in frames
    )


@pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results")
def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks(
    gateway: Gateway, tmp_path: Path
) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST", request.target
        body: Final = json.loads(request.body)
        if request.target == SEARCH_TARGET:
            assert request.headers["authorization"] == "Bearer synthetic-tavily-key"
            assert body["query"] == "query-0", body
            return Reply(body=json.dumps({"query": "query-0", "results": [SEARCH_RESULT]}).encode())
        assert request.target == INVOKE_TARGET
        assert request.headers["authorization"] == "Bearer synthetic-bedrock-token"
        assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"]
        assert "stream" not in body, body
        depth: Final = sum(
            1
            for message in body["messages"]
            if isinstance(message["content"], list)
            for block in message["content"]
            if block["type"] == "tool_result"
        )
        if depth == 1:
            assert body["messages"][2]["content"] == [
                {
                    "type": "tool_result",
                    "tool_use_id": "toolu_0",
                    "content": "Title: Synthetic result\nURL: https://example.test/result\nSnippet: the snippet text",
                }
            ], body["messages"]
        return Reply(
            body=json.dumps(
                {
                    "id": f"msg_{depth}",
                    "type": "message",
                    "role": "assistant",
                    "model": BEDROCK_MODEL,
                    "content": [
                        {"type": "text", "text": f"turn-{depth}"},
                        {
                            "type": "tool_use",
                            "id": f"toolu_{depth}",
                            "name": "litellm_web_search",
                            "input": {"query": f"query-{depth}"},
                        },
                    ],
                    "stop_reason": "tool_use",
                    "stop_sequence": None,
                    "usage": {"input_tokens": 10, "output_tokens": 4},
                }
            ).encode()
        )

    with wire_server(respond) as wire:
        config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
        config["search_tools"] = [
            {
                "search_tool_name": "integration-search",
                "litellm_params": {
                    "search_provider": "tavily",
                    "api_key": "synthetic-tavily-key",
                    "api_base": wire.url + "/tavily",
                },
            }
        ]
        config["litellm_settings"].update(
            {
                "callbacks": ["websearch_interception"],
                "websearch_interception_params": {
                    "enabled_providers": ["bedrock"],
                    "search_tool_name": "integration-search",
                    "max_agentic_loops": 1,
                },
            }
        )
        path: Final = tmp_path / "websearch.yaml"
        path.write_text(yaml.safe_dump(config))
        with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
            model: Final = scenario.model(
                model=f"bedrock/{BEDROCK_MODEL}",
                api_key="synthetic-bedrock-token",
                api_base=wire.url,
                aws_region_name="us-east-1",
                aws_bedrock_runtime_endpoint=wire.url,
            )
            response: Final = candidate.request(
                "POST",
                "/v1/messages",
                {
                    "model": model,
                    "max_tokens": 64,
                    "stream": True,
                    "messages": [{"role": "user", "content": "search control"}],
                    "tools": [{"type": "web_search_20250305", "name": "web_search"}],
                },
            )
            assert response.status_code == 200, response.text
            events: Final = sse_events(response.text)
            assert [name for name, _ in events][:1] == ["message_start"], response.text
            assert [name for name, _ in events][-2:] == ["message_delta", "message_stop"], response.text
            for position, (name, event) in enumerate(events):
                if name == "content_block_stop":
                    assert event["index"] in {
                        earlier_event["index"]
                        for earlier, earlier_event in events[:position]
                        if earlier == "content_block_start"
                    }, response.text
            started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start")
            search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use")
            assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text
            assert started[-1] == {"type": "text", "text": ""}, response.text
            assert set(json.dumps(block, sort_keys=True) for block in started[:-1]) == {
                json.dumps(block, sort_keys=True)
                for search_id in search_ids
                for block in (
                    {"type": "server_tool_use", "id": search_id, "name": "web_search", "input": {"query": "query-0"}},
                    {
                        "type": "web_search_tool_result",
                        "tool_use_id": search_id,
                        "content": [
                            {
                                "type": "web_search_result",
                                "url": "https://example.test/result",
                                "title": "Synthetic result",
                                "page_age": None,
                                "encrypted_content": "",
                                "snippet": "the snippet text",
                            }
                        ],
                    },
                )
            }, response.text
            assert (
                "".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1"
            ), response.text
            assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [
                "end_turn"
            ], response.text
            assert "litellm_web_search" not in response.text, response.text
            assert [request.target for request in wire.drain()] == [INVOKE_TARGET, SEARCH_TARGET, INVOKE_TARGET]