mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test(integration): streamed web search turn capped by max_agentic_loops ends the turn with snippets and ordered blocks (Pylon #7230)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
Devin AI
parent
bda8149000
commit
b22a69abea
1 changed files with 171 additions and 0 deletions
171
tests/integration/providers/test_websearch_interception_wire.py
Normal file
171
tests/integration/providers/test_websearch_interception_wire.py
Normal file
|
|
@ -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]
|
||||
Loading…
Add table
Reference in a new issue