diff --git a/tests/integration/providers/test_websearch_interception_wire.py b/tests/integration/providers/test_websearch_interception_wire.py new file mode 100644 index 00000000000..77ec219be78 --- /dev/null +++ b/tests/integration/providers/test_websearch_interception_wire.py @@ -0,0 +1,171 @@ +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]