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
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Built-in search backends for SGLang entrypoints."""
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user