From 2f34dbe372633a1efcf73b8cb40f95a76e05c2f2 Mon Sep 17 00:00:00 2001 From: jonah-berman Date: Sat, 27 Jun 2026 06:41:09 -0700 Subject: [PATCH] Add native Exa-backed web_search support (#29342) Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> --- .../autoregressive/OpenAI/GPT-OSS.mdx | 47 +-- .../advanced_features/server_arguments.mdx | 2 +- python/sglang/srt/entrypoints/http_server.py | 6 + .../entrypoints/openai/serving_responses.py | 16 + .../srt/entrypoints/openai/tool_server.py | 22 +- .../sglang/srt/entrypoints/search/__init__.py | 1 + .../srt/entrypoints/search/exa_client.py | 138 ++++++++ python/sglang/srt/entrypoints/tool.py | 245 ++++++++++++-- python/sglang/srt/environ.py | 8 + .../entrypoints/openai/test_exa_search.py | 303 ++++++++++++++++++ 10 files changed, 737 insertions(+), 51 deletions(-) create mode 100644 python/sglang/srt/entrypoints/search/__init__.py create mode 100644 python/sglang/srt/entrypoints/search/exa_client.py create mode 100644 test/registered/unit/entrypoints/openai/test_exa_search.py diff --git a/docs_new/cookbook/autoregressive/OpenAI/GPT-OSS.mdx b/docs_new/cookbook/autoregressive/OpenAI/GPT-OSS.mdx index 646299464..68b79b21a 100644 --- a/docs_new/cookbook/autoregressive/OpenAI/GPT-OSS.mdx +++ b/docs_new/cookbook/autoregressive/OpenAI/GPT-OSS.mdx @@ -41,12 +41,12 @@ import { GPTOSSDeployment } from "/src/snippets/autoregressive/gpt-oss-deploymen ### 3.2 Configuration Tips -- **Demo tool server:** Launch with `--tool-server demo` to enable the built-in web-search (Exa) and Python interpreter tools. -- **Web search tool:** Requires an Exa API key — set `EXA_API_KEY` in your environment. -- **Python tool:** Runs in a Docker sandbox by default. Set `PYTHON_EXECUTION_BACKEND=UV` to run on the host (model-generated code executes locally — use with care). +- **Native web search:** Set `EXA_API_KEY` in the SGLang server environment to enable built-in web search (Exa). No `--tool-server` is required, and requests are tagged with `x-exa-integration: sglang`. +- **Web search defaults:** `numResults=10`, search `type="auto"`, and `contents.highlights=true`. Override with `SGLANG_EXA_NUM_RESULTS`, `SGLANG_EXA_SEARCH_TYPE`, and `SGLANG_EXA_INCLUDE_HIGHLIGHTS`. +- **Python tool:** Add `--tool-server demo` to enable the Python interpreter. Runs in a Docker sandbox by default; set `PYTHON_EXECUTION_BACKEND=UV` to run on the host (model-generated code executes locally — use with care). - **MCP tool servers:** For production, point SGLang at external MCP SSE servers with `--tool-server ip-1:port-1,ip-2:port-2`. -- **Responses API:** GPT-OSS supports OpenAI's Responses API (`client.responses.create`) in addition to the standard Chat Completions API (see section 4.2.3). -- **Use Python 3.12** when running the demo tools. +- **Responses API:** GPT-OSS supports OpenAI's Responses API (`client.responses.create`) in addition to the standard Chat Completions API (see section 4.2.4). +- **Use Python 3.12** when running the demo Python tool. ## 4.Model Invocation @@ -433,16 +433,19 @@ The spec-v2 overlap scheduler is enabled by default. It improves performance by #### 4.2.4 Responses API and Built-in Tools -GPT-OSS supports the OpenAI Responses API with built-in tool use (web search and Python interpreter). Set up your environment and launch with `--tool-server demo`: +GPT-OSS supports the OpenAI Responses API with built-in tool use (web search and Python interpreter). Set `EXA_API_KEY` to enable native web search; add `--tool-server demo` only when you also want the Python tool: ```shell Command export EXA_API_KEY=YOUR_EXA_KEY +# Optional: server-side Exa tuning (defaults shown) +export SGLANG_EXA_NUM_RESULTS=10 +export SGLANG_EXA_SEARCH_TYPE=auto +export SGLANG_EXA_INCLUDE_HIGHLIGHTS=true # Optional: run Python tool on host instead of Docker (model code executes locally) export PYTHON_EXECUTION_BACKEND=UV python3 -m sglang.launch_server \ --model-path openai/gpt-oss-120b \ - --tool-server demo \ --tp 2 ``` @@ -464,10 +467,8 @@ from openai import OpenAI client = OpenAI(base_url="http://localhost:30000/v1", api_key="sk-123456") -tools = [ - {"type": "code_interpreter"}, - {"type": "web_search_preview"}, -] +search_tools = [{"type": "web_search"}] +python_tools = [{"type": "code_interpreter"}] # Configurable reasoning effort: "high", "medium", or "low" response = client.responses.create( @@ -478,24 +479,24 @@ response = client.responses.create( ) print(response.output_text) -# Python tool usage +# Web search (requires EXA_API_KEY on the SGLang server) response = client.responses.create( model="openai/gpt-oss-120b", - instructions="You are a helpful assistant.", + instructions="You are a helpful assistant, you can search the web when needed.", + input="Search the web for the latest news about Nvidia stock price", + tools=search_tools, +) +print(response.output_text) + +# Python tool (requires launching SGLang with --tool-server demo) +response = client.responses.create( + model="openai/gpt-oss-120b", + instructions="You are a helpful assistant, you could use python tool to execute code.", input="Use python tool to calculate the sum of 29138749187 and 29138749187", - tools=tools, + tools=python_tools, ) print(response.output_text) # Output: The sum is 58,277,498,374. - -# Web search usage -response = client.responses.create( - model="openai/gpt-oss-120b", - instructions="You are a helpful assistant.", - input="Search the web for the latest news about Nvidia stock price", - tools=tools, -) -print(response.output_text) ``` ## 5.Benchmark diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index dd96747aa..8667cafdf 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1096,7 +1096,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--tool-server` - Either 'demo' or a comma-separated list of tool server urls to use for the model. If not specified, no tool server will be used. + Either 'demo' or a comma-separated list of tool server urls to use for the model. If not specified, no external tool server will be used. Native GPT-OSS `web_search` can still be enabled with `EXA_API_KEY`. `None` Type: str diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 9676c9ef4..d543e3a64 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -337,6 +337,10 @@ async def lifespan(fast_api_app: FastAPI): tool_server = MCPToolServer() await tool_server.add_tool_server(server_args.tool_server) + elif envs.EXA_API_KEY.get(): + from sglang.srt.entrypoints.openai.tool_server import NativeToolServer + + tool_server = NativeToolServer() try: from sglang.srt.entrypoints.openai.serving_responses import ( @@ -381,6 +385,8 @@ async def lifespan(fast_api_app: FastAPI): try: yield finally: + if tool_server is not None and hasattr(tool_server, "aclose"): + await tool_server.aclose() warmup_thread.join() diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index 86ed083fc..760ca5c4b 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -137,6 +137,10 @@ class OpenAIServingResponses(OpenAIServingChat): self.background_tasks: dict[str, asyncio.Task] = {} + @staticmethod + def _has_response_tool(request: ResponsesRequest, *tool_types: str) -> bool: + return any(tool.type in tool_types for tool in (request.tools or [])) + # error helpers dedicated for v1/responses def create_error_response( self, @@ -194,6 +198,18 @@ class OpenAIServingResponses(OpenAIServingChat): 'type="function"; other built-in tool types cannot be forced.' ) + if ( + self.use_harmony + and self._has_response_tool(request, "web_search", "web_search_preview") + and not self.supports_browsing + ): + return self.create_error_response( + "web_search requires a browser backend. Set EXA_API_KEY on the " + "SGLang server to enable native Exa-backed web search, or " + "configure a browser MCP tool server. Create an Exa API key at " + "https://dashboard.exa.ai/api-keys." + ) + # Handle the previous response ID prev_response_id = request.previous_response_id if prev_response_id is not None: diff --git a/python/sglang/srt/entrypoints/openai/tool_server.py b/python/sglang/srt/entrypoints/openai/tool_server.py index 71f20c341..e4bebf647 100644 --- a/python/sglang/srt/entrypoints/openai/tool_server.py +++ b/python/sglang/srt/entrypoints/openai/tool_server.py @@ -144,7 +144,7 @@ class MCPToolServer(ToolServer): class DemoToolServer(ToolServer): - def __init__(self): + def __init__(self, *, enable_python: bool = True): from sglang.srt.entrypoints.tool import ( HarmonyBrowserTool, HarmonyPythonTool, @@ -155,9 +155,10 @@ class DemoToolServer(ToolServer): browser_tool = HarmonyBrowserTool() if browser_tool.enabled: self.tools["browser"] = browser_tool - python_tool = HarmonyPythonTool() - if python_tool.enabled: - self.tools["python"] = python_tool + if enable_python: + python_tool = HarmonyPythonTool() + if python_tool.enabled: + self.tools["python"] = python_tool def has_tool(self, tool_name: str): return tool_name in self.tools @@ -175,3 +176,16 @@ class DemoToolServer(ToolServer): @asynccontextmanager async def get_tool_session(self, tool_name: str): yield self.tools[tool_name] + + async def aclose(self): + browser = self.tools.get("browser") + exa_client = getattr(browser, "exa_client", None) if browser else None + if exa_client is not None: + await exa_client.close() + + +class NativeToolServer(DemoToolServer): + """Built-in SGLang hosted tools that do not require an external MCP server.""" + + def __init__(self): + super().__init__(enable_python=False) diff --git a/python/sglang/srt/entrypoints/search/__init__.py b/python/sglang/srt/entrypoints/search/__init__.py new file mode 100644 index 000000000..b67587901 --- /dev/null +++ b/python/sglang/srt/entrypoints/search/__init__.py @@ -0,0 +1 @@ +"""Built-in search backends for SGLang entrypoints.""" diff --git a/python/sglang/srt/entrypoints/search/exa_client.py b/python/sglang/srt/entrypoints/search/exa_client.py new file mode 100644 index 000000000..10c0ea8cc --- /dev/null +++ b/python/sglang/srt/entrypoints/search/exa_client.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +import asyncio +from typing import Any, Optional + +import aiohttp +import msgspec + +from sglang.srt.environ import envs +from sglang.srt.utils import print_warning_once + +EXA_API_BASE_URL = "https://api.exa.ai" +EXA_INTEGRATION_HEADER = "x-exa-integration" +EXA_INTEGRATION_NAME = "sglang" +EXA_DEFAULT_NUM_RESULTS = 10 +EXA_DEFAULT_SEARCH_TYPE = "auto" +EXA_DEFAULT_HIGHLIGHTS = True + +_SEARCH_TYPES = {"instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"} + + +class ExaClientError(RuntimeError): + """Raised when an Exa API request fails.""" + + +class ExaSearchConfig(msgspec.Struct, frozen=True, kw_only=True): + num_results: int = EXA_DEFAULT_NUM_RESULTS + search_type: str = EXA_DEFAULT_SEARCH_TYPE + include_highlights: bool = EXA_DEFAULT_HIGHLIGHTS + + @classmethod + def from_env(cls) -> "ExaSearchConfig": + # EnvField handles type parsing (int/bool) and falls back to its own + # default on a parse error; here we only enforce the semantic bounds + # that the generic descriptors cannot express. + num_results = envs.SGLANG_EXA_NUM_RESULTS.get() + if not 1 <= num_results <= 100: + print_warning_once( + f"Ignoring invalid SGLANG_EXA_NUM_RESULTS={num_results!r}; " + f"expected a value from 1 to 100." + ) + num_results = EXA_DEFAULT_NUM_RESULTS + + search_type = envs.SGLANG_EXA_SEARCH_TYPE.get() + if search_type not in _SEARCH_TYPES: + print_warning_once( + f"Ignoring invalid SGLANG_EXA_SEARCH_TYPE={search_type!r}; " + f"expected one of {', '.join(sorted(_SEARCH_TYPES))}." + ) + search_type = EXA_DEFAULT_SEARCH_TYPE + + return cls( + num_results=num_results, + search_type=search_type, + include_highlights=envs.SGLANG_EXA_INCLUDE_HIGHLIGHTS.get(), + ) + + +class ExaClient: + def __init__( + self, + api_key: str, + *, + config: Optional[ExaSearchConfig] = None, + base_url: str = EXA_API_BASE_URL, + timeout: float = 30.0, + ): + self.api_key = api_key + self.config = config or ExaSearchConfig() + self.base_url = base_url.rstrip("/") + self.timeout = timeout + self._session: Optional[aiohttp.ClientSession] = None + self._session_lock = asyncio.Lock() + + def _headers(self) -> dict[str, str]: + return { + "x-api-key": self.api_key, + "Content-Type": "application/json", + EXA_INTEGRATION_HEADER: EXA_INTEGRATION_NAME, + } + + def _search_payload(self, query: str) -> dict[str, Any]: + payload: dict[str, Any] = { + "query": query, + "numResults": self.config.num_results, + "type": self.config.search_type, + } + if self.config.include_highlights: + payload["contents"] = {"highlights": True} + return payload + + def _contents_payload(self, urls: list[str]) -> dict[str, Any]: + payload: dict[str, Any] = {"urls": urls, "text": True} + if self.config.include_highlights: + payload["highlights"] = True + return payload + + async def search(self, query: str) -> dict[str, Any]: + return await self._post("/search", self._search_payload(query)) + + async def contents(self, urls: list[str]) -> dict[str, Any]: + return await self._post("/contents", self._contents_payload(urls)) + + async def _get_session(self) -> aiohttp.ClientSession: + # Reuse a single session across requests so the connection pool and + # DNS cache survive; aiohttp.ClientSession is intended to be long-lived. + if self._session is None: + async with self._session_lock: + if self._session is None: + self._session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=self.timeout) + ) + return self._session + + async def close(self): + if self._session is not None: + await self._session.close() + self._session = None + + async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]: + session = await self._get_session() + url = f"{self.base_url}{path}" + async with session.post( + url, + json=payload, + headers=self._headers(), + ) as response: + response_text = await response.text() + if response.status >= 400: + raise ExaClientError( + f"Exa API request failed with status {response.status}: " + f"{response_text}" + ) + try: + return await response.json() + except Exception as exc: + raise ExaClientError( + f"Failed to decode Exa API response: {response_text}" + ) from exc diff --git a/python/sglang/srt/entrypoints/tool.py b/python/sglang/srt/entrypoints/tool.py index 45b87ac3a..00b54cfd2 100644 --- a/python/sglang/srt/entrypoints/tool.py +++ b/python/sglang/srt/entrypoints/tool.py @@ -1,9 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 import logging -import os from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any +import orjson + +from sglang.srt.entrypoints.search.exa_client import ( + ExaClient, + ExaSearchConfig, +) +from sglang.srt.environ import envs from sglang.srt.utils import print_info_once, print_warning_once if TYPE_CHECKING: @@ -22,39 +28,232 @@ class Tool(ABC): class HarmonyBrowserTool(Tool): - def __init__(self): + def __init__(self, client: ExaClient | None = None): self.enabled = True - exa_api_key = os.getenv("EXA_API_KEY") - if not exa_api_key: + if client is not None: + self.exa_client = client + print_info_once("Browser tool initialized") + return + + api_key = envs.EXA_API_KEY.get() + if not api_key: self.enabled = False print_warning_once("EXA_API_KEY is not set, browsing is disabled") return - try: - from gpt_oss.tools.simple_browser import SimpleBrowserTool - from gpt_oss.tools.simple_browser.backend import ExaBackend - except ImportError: - self.enabled = False - print_warning_once("gpt_oss is not installed, browsing is disabled") - return - - browser_backend = ExaBackend(source="web", api_key=exa_api_key) - self.browser_tool = SimpleBrowserTool(backend=browser_backend) + self.exa_client = ExaClient(api_key, config=ExaSearchConfig.from_env()) print_info_once("Browser tool initialized") async def get_result(self, context: "ConversationContext") -> Any: - from sglang.srt.entrypoints.context import HarmonyContext + from openai_harmony import Author, Message, Role, TextContent - assert isinstance(context, HarmonyContext) last_msg = context.messages[-1] - tool_output_msgs = [] - async for msg in self.browser_tool.process(last_msg): - tool_output_msgs.append(msg) - return tool_output_msgs + recipient = last_msg.recipient + if recipient is None or not recipient.startswith("browser."): + raise ValueError("No browser tool call found") - @property - def tool_config(self) -> Any: - return self.browser_tool.tool_config + try: + args = orjson.loads(last_msg.content[0].text) + result_text = await self._dispatch_browser_call(context, recipient, args) + except Exception as exc: + logger.exception("Browser tool call failed") + result_text = f"Browser tool call failed: {exc}" + + content = TextContent(text=result_text) + author = Author(role=Role.TOOL, name=recipient) + return [Message(author=author, content=[content], recipient=Role.ASSISTANT)] + + async def _dispatch_browser_call( + self, context: "ConversationContext", recipient: str, args: dict[str, Any] + ) -> str: + if recipient == "browser.search": + query = args.get("query") + if not query: + raise ValueError("browser.search requires a query") + data = await self.exa_client.search(query) + return self._format_search_results(context, query, data) + + if recipient == "browser.open": + url = self._resolve_url(context, args) + data = await self.exa_client.contents([url]) + return self._format_page_contents(context, url, data) + + if recipient == "browser.find": + pattern = args.get("pattern") + if not pattern: + raise ValueError("browser.find requires a pattern") + return await self._find_pattern(context, args, pattern) + + raise ValueError(f"Unknown browser action: {recipient}") + + def _browser_state(self, context: "ConversationContext") -> dict[str, Any]: + state = getattr(context, "_sglang_exa_browser_state", None) + if state is None: + state = {"pages": {}, "page_text": {}} + setattr(context, "_sglang_exa_browser_state", state) + return state + + def _format_search_results( + self, context: "ConversationContext", query: str, data: dict[str, Any] + ) -> str: + state = self._browser_state(context) + state["pages"] = {} + state["page_text"] = {} + + results = data.get("results") or [] + request_id = data.get("requestId") + if request_id: + logger.debug("Exa search request id: %s", request_id) + lines = [f"Search results for: {query}"] + lines.append("Use browser.open with the cursor number to inspect a result.") + + for index, result in enumerate(results, start=1): + cursor = str(index) + state["pages"][cursor] = result + title = result.get("title") or "Untitled" + url = result.get("url") or result.get("id") or "" + snippet = self._best_snippet(result) + lines.append("") + lines.append(f"[{cursor}] {title}") + if url: + lines.append(f"URL: {url}") + if snippet: + lines.append(f"Snippet: {snippet}") + + if not results: + lines.append("No results found.") + return "\n".join(lines) + + def _format_page_contents( + self, context: "ConversationContext", url: str, data: dict[str, Any] + ) -> str: + state = self._browser_state(context) + results = data.get("results") or [] + if not results: + return f"No page contents returned for {url}." + + page = results[0] + cursor = self._cursor_for_url(state, url) + if cursor: + state["pages"][cursor] = page + title = page.get("title") or "Untitled" + page_url = page.get("url") or url + text = page.get("text") or self._best_snippet(page) or page.get("summary") or "" + if cursor: + state["page_text"][cursor] = text + + return "\n".join( + [ + f"Opened page: {title}", + f"URL: {page_url}", + "", + self._truncate(text, 12000) if text else "No page text available.", + ] + ) + + async def _find_pattern( + self, context: "ConversationContext", args: dict[str, Any], pattern: str + ) -> str: + state = self._browser_state(context) + cursor = self._normalize_cursor(args.get("cursor")) + + if cursor: + text = state["page_text"].get(cursor) + if text is None: + url = self._resolve_url(context, args) + data = await self.exa_client.contents([url]) + self._format_page_contents(context, url, data) + text = state["page_text"].get(cursor, "") + return self._format_matches(pattern, text) + + if args.get("url"): + url = self._resolve_url(context, args) + data = await self.exa_client.contents([url]) + results = data.get("results") or [] + if not results: + return f"No page contents returned for {url}." + page = results[0] + text = ( + page.get("text") + or self._best_snippet(page) + or page.get("summary") + or "" + ) + return self._format_matches(pattern, text) + + searchable_text = "\n\n".join( + self._best_snippet(page) for page in state["pages"].values() + ) + return self._format_matches(pattern, searchable_text) + + def _resolve_url(self, context: "ConversationContext", args: dict[str, Any]) -> str: + if args.get("url"): + return str(args["url"]) + + state = self._browser_state(context) + cursor = self._normalize_cursor(args.get("cursor")) + if not cursor: + raise ValueError("browser.open requires a cursor or url") + + page = state["pages"].get(cursor) + if page is None: + raise ValueError(f"Unknown browser cursor: {cursor}") + + url = page.get("url") or page.get("id") + if not url: + raise ValueError(f"No URL recorded for browser cursor: {cursor}") + return str(url) + + def _normalize_cursor(self, cursor: Any) -> str | None: + if cursor is None: + return None + cursor_str = str(cursor) + # GPT-OSS emits 0 as a 1-based cursor; map it to the first result. + if cursor_str == "0": + return "1" + return cursor_str + + def _cursor_for_url(self, state: dict[str, Any], url: str) -> str | None: + for cursor, page in state["pages"].items(): + if page.get("url") == url or page.get("id") == url: + return cursor + return None + + def _best_snippet(self, result: dict[str, Any]) -> str: + highlights = result.get("highlights") or [] + if highlights: + return self._truncate(str(highlights[0]), 1000) + summary = result.get("summary") + if summary: + return self._truncate(str(summary), 1000) + text = result.get("text") + if text: + return self._truncate(str(text), 1000) + return "" + + def _format_matches(self, pattern: str, text: str) -> str: + if not text: + return f"No text available to search for {pattern!r}." + + pattern_lower = pattern.lower() + matches = [] + for line in text.splitlines(): + if pattern_lower in line.lower(): + matches.append(self._truncate(line.strip(), 1000)) + if len(matches) >= 10: + break + + if not matches: + return f"No matches found for {pattern!r}." + + lines = [f"Matches for {pattern!r}:"] + lines.extend(f"- {match}" for match in matches) + return "\n".join(lines) + + def _truncate(self, text: str, max_chars: int) -> str: + if len(text) <= max_chars: + return text + return text[: max_chars - 3] + "..." class HarmonyPythonTool(Tool): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3f78790c3..c9476a38d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -403,6 +403,14 @@ class Envs: # Tool Calling SGLANG_FORWARD_UNKNOWN_TOOLS = EnvBool(False) + # Native web search (Exa). EXA_API_KEY is the vendor BYOK credential + # (kept as-is, not renamed to SGLANG_*); the SGLANG_EXA_* knobs tune the + # request defaults for the built-in GPT-OSS web_search tool. + EXA_API_KEY = EnvStr(None) + SGLANG_EXA_NUM_RESULTS = EnvInt(10) + SGLANG_EXA_SEARCH_TYPE = EnvStr("auto") + SGLANG_EXA_INCLUDE_HIGHLIGHTS = EnvBool(True) + # Hi-Cache SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None) diff --git a/test/registered/unit/entrypoints/openai/test_exa_search.py b/test/registered/unit/entrypoints/openai/test_exa_search.py new file mode 100644 index 000000000..3578e7e3a --- /dev/null +++ b/test/registered/unit/entrypoints/openai/test_exa_search.py @@ -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()