Add native Exa-backed web_search support (#29342)
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
jonah-berman
co-authored by
Xinyuan Tong
Xinyuan Tong
parent
3306233961
commit
2f34dbe372
@@ -0,0 +1,303 @@
|
||||
import asyncio
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from utils import make_serving
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import ResponsesRequest
|
||||
from sglang.srt.entrypoints.openai.tool_server import NativeToolServer
|
||||
from sglang.srt.entrypoints.search.exa_client import (
|
||||
EXA_INTEGRATION_HEADER,
|
||||
EXA_INTEGRATION_NAME,
|
||||
ExaClient,
|
||||
ExaSearchConfig,
|
||||
)
|
||||
from sglang.srt.entrypoints.tool import HarmonyBrowserTool
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class ExaClientTestCase(unittest.TestCase):
|
||||
def test_headers_include_sglang_integration_tag(self):
|
||||
client = ExaClient("test-key")
|
||||
|
||||
headers = client._headers()
|
||||
|
||||
self.assertEqual(headers["x-api-key"], "test-key")
|
||||
self.assertEqual(headers[EXA_INTEGRATION_HEADER], EXA_INTEGRATION_NAME)
|
||||
self.assertEqual(headers["Content-Type"], "application/json")
|
||||
|
||||
def test_default_search_payload_uses_server_side_defaults(self):
|
||||
client = ExaClient("test-key")
|
||||
|
||||
payload = client._search_payload("SGLang native web search")
|
||||
|
||||
self.assertEqual(payload["numResults"], 10)
|
||||
self.assertEqual(payload["type"], "auto")
|
||||
self.assertEqual(payload["contents"], {"highlights": True})
|
||||
|
||||
def test_contents_payload_requests_text_and_highlights(self):
|
||||
client = ExaClient("test-key")
|
||||
|
||||
payload = client._contents_payload(["https://example.com"])
|
||||
|
||||
self.assertEqual(payload["urls"], ["https://example.com"])
|
||||
self.assertTrue(payload["text"])
|
||||
self.assertTrue(payload["highlights"])
|
||||
|
||||
def test_config_can_be_set_from_server_environment(self):
|
||||
env = {
|
||||
"SGLANG_EXA_NUM_RESULTS": "7",
|
||||
"SGLANG_EXA_SEARCH_TYPE": "fast",
|
||||
"SGLANG_EXA_INCLUDE_HIGHLIGHTS": "false",
|
||||
}
|
||||
with patch.dict(os.environ, env, clear=False):
|
||||
config = ExaSearchConfig.from_env()
|
||||
|
||||
self.assertEqual(config.num_results, 7)
|
||||
self.assertEqual(config.search_type, "fast")
|
||||
self.assertFalse(config.include_highlights)
|
||||
|
||||
def test_post_sends_integration_header_without_network(self):
|
||||
captured = {}
|
||||
|
||||
class FakeResponse:
|
||||
status = 200
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
async def text(self):
|
||||
return '{"ok": true}'
|
||||
|
||||
async def json(self):
|
||||
return {"ok": True}
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, timeout):
|
||||
captured["timeout"] = timeout
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
def post(self, url, json, headers):
|
||||
captured["url"] = url
|
||||
captured["json"] = json
|
||||
captured["headers"] = headers
|
||||
return FakeResponse()
|
||||
|
||||
client = ExaClient("test-key")
|
||||
with patch(
|
||||
"sglang.srt.entrypoints.search.exa_client.aiohttp.ClientSession",
|
||||
FakeSession,
|
||||
):
|
||||
result = asyncio.run(client._post("/search", {"query": "sglang"}))
|
||||
|
||||
self.assertEqual(result, {"ok": True})
|
||||
self.assertEqual(captured["url"], "https://api.exa.ai/search")
|
||||
self.assertEqual(captured["json"], {"query": "sglang"})
|
||||
self.assertEqual(captured["headers"][EXA_INTEGRATION_HEADER], "sglang")
|
||||
|
||||
|
||||
class ResponsesNativeWebSearchTestCase(unittest.TestCase):
|
||||
def test_harmony_web_search_requires_configured_backend(self):
|
||||
serving = make_serving()
|
||||
serving.use_harmony = True
|
||||
serving.supports_browsing = False
|
||||
request = ResponsesRequest(
|
||||
model="x",
|
||||
input="search the web",
|
||||
tools=[{"type": "web_search"}],
|
||||
store=False,
|
||||
)
|
||||
|
||||
result = asyncio.run(serving.create_responses(request, raw_request=None))
|
||||
|
||||
self.assertEqual(getattr(result, "status_code", None), 400)
|
||||
self.assertIn("EXA_API_KEY", result.body.decode())
|
||||
self.assertIn("https://dashboard.exa.ai/api-keys", result.body.decode())
|
||||
|
||||
|
||||
class NativeWebSearchIntegrationTestCase(unittest.TestCase):
|
||||
def test_native_tool_server_hits_exa_client_with_server_api_key(self):
|
||||
captured_calls = []
|
||||
|
||||
class FakeResponse:
|
||||
status = 200
|
||||
|
||||
def __init__(self, url, payload):
|
||||
self.url = url
|
||||
self.payload = payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
async def text(self):
|
||||
return "{}"
|
||||
|
||||
async def json(self):
|
||||
if self.url.endswith("/search"):
|
||||
return {
|
||||
"requestId": "mock_req_search",
|
||||
"results": [
|
||||
{
|
||||
"title": "Mock SGLang Result",
|
||||
"url": "https://example.com/sglang",
|
||||
"highlights": ["SGLang native web search via Exa."],
|
||||
}
|
||||
],
|
||||
}
|
||||
return {
|
||||
"requestId": "mock_req_contents",
|
||||
"results": [
|
||||
{
|
||||
"title": "Mock SGLang Result",
|
||||
"url": self.payload["urls"][0],
|
||||
"text": "Opened content returned through Exa contents.",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, timeout):
|
||||
self.timeout = timeout
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
def post(self, url, json, headers):
|
||||
captured_calls.append({"url": url, "json": json, "headers": headers})
|
||||
return FakeResponse(url, json)
|
||||
|
||||
async def run_tool_flow():
|
||||
native_tool_server = NativeToolServer()
|
||||
self.assertTrue(native_tool_server.has_tool("browser"))
|
||||
async with native_tool_server.get_tool_session("browser") as browser_tool:
|
||||
context = FakeContext()
|
||||
search_result = await browser_tool._dispatch_browser_call(
|
||||
context,
|
||||
"browser.search",
|
||||
{"query": "test native web search"},
|
||||
)
|
||||
open_result = await browser_tool._dispatch_browser_call(
|
||||
context, "browser.open", {"cursor": 1}
|
||||
)
|
||||
return search_result, open_result
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, {"EXA_API_KEY": "mock-sglang-key"}, clear=False),
|
||||
patch(
|
||||
"sglang.srt.entrypoints.search.exa_client.aiohttp.ClientSession",
|
||||
FakeSession,
|
||||
),
|
||||
):
|
||||
search_result, open_result = asyncio.run(run_tool_flow())
|
||||
|
||||
self.assertIn("Mock SGLang Result", search_result)
|
||||
self.assertIn("SGLang native web search via Exa.", search_result)
|
||||
self.assertIn("Opened content returned through Exa contents.", open_result)
|
||||
self.assertEqual(len(captured_calls), 2)
|
||||
|
||||
search_call, contents_call = captured_calls
|
||||
self.assertEqual(search_call["url"], "https://api.exa.ai/search")
|
||||
self.assertEqual(search_call["headers"]["x-api-key"], "mock-sglang-key")
|
||||
self.assertEqual(search_call["headers"][EXA_INTEGRATION_HEADER], "sglang")
|
||||
self.assertEqual(search_call["json"]["numResults"], 10)
|
||||
self.assertEqual(search_call["json"]["type"], "auto")
|
||||
self.assertEqual(search_call["json"]["contents"], {"highlights": True})
|
||||
|
||||
self.assertEqual(contents_call["url"], "https://api.exa.ai/contents")
|
||||
self.assertEqual(contents_call["headers"]["x-api-key"], "mock-sglang-key")
|
||||
self.assertEqual(contents_call["headers"][EXA_INTEGRATION_HEADER], "sglang")
|
||||
self.assertEqual(contents_call["json"]["urls"], ["https://example.com/sglang"])
|
||||
self.assertTrue(contents_call["json"]["text"])
|
||||
self.assertTrue(contents_call["json"]["highlights"])
|
||||
|
||||
|
||||
class FakeExaClient:
|
||||
def __init__(self):
|
||||
self.search_queries = []
|
||||
self.content_urls = []
|
||||
|
||||
async def search(self, query):
|
||||
self.search_queries.append(query)
|
||||
return {
|
||||
"requestId": "req_123",
|
||||
"results": [
|
||||
{
|
||||
"title": "SGLang",
|
||||
"url": "https://example.com/sglang",
|
||||
"highlights": ["Native web search powered by Exa."],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
async def contents(self, urls):
|
||||
self.content_urls.append(urls)
|
||||
return {
|
||||
"results": [
|
||||
{
|
||||
"title": "SGLang",
|
||||
"url": urls[0],
|
||||
"text": "SGLang native web search uses Exa for retrieval.",
|
||||
"highlights": ["Exa for retrieval."],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
class FakeContext:
|
||||
pass
|
||||
|
||||
|
||||
class HarmonyBrowserToolTestCase(unittest.TestCase):
|
||||
def test_search_and_open_use_exa_client_with_request_scoped_state(self):
|
||||
client = FakeExaClient()
|
||||
tool = HarmonyBrowserTool(client=client)
|
||||
context = FakeContext()
|
||||
|
||||
search_result = asyncio.run(
|
||||
tool._dispatch_browser_call(
|
||||
context, "browser.search", {"query": "SGLang web search"}
|
||||
)
|
||||
)
|
||||
open_result = asyncio.run(
|
||||
tool._dispatch_browser_call(context, "browser.open", {"cursor": 1})
|
||||
)
|
||||
find_result = asyncio.run(
|
||||
tool._dispatch_browser_call(
|
||||
context,
|
||||
"browser.find",
|
||||
{"url": "https://example.com/direct", "pattern": "retrieval"},
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(client.search_queries, ["SGLang web search"])
|
||||
self.assertEqual(
|
||||
client.content_urls,
|
||||
[["https://example.com/sglang"], ["https://example.com/direct"]],
|
||||
)
|
||||
self.assertNotIn("req_123", search_result)
|
||||
self.assertIn("[1] SGLang", search_result)
|
||||
self.assertIn("Snippet: Native web search powered by Exa.", search_result)
|
||||
self.assertIn("Opened page: SGLang", open_result)
|
||||
self.assertIn("SGLang native web search uses Exa", open_result)
|
||||
self.assertIn("SGLang native web search uses Exa", find_result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user