Skip to content

Commit 64c4137

Browse files
Web search tool support
1 parent 471cd03 commit 64c4137

10 files changed

Lines changed: 927 additions & 56 deletions

File tree

.claude/settings.local.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
]
1414
},
1515
"hooks": {
16-
"SessionEnd": [
16+
"Stop": [
1717
{
1818
"hooks": [
1919
{

pyproject.toml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,14 @@ dependencies = [
1414
"uvicorn>=0.32.0",
1515
]
1616

17+
[tool.pytest.ini_options]
18+
markers = [
19+
"live: tests that hit the real OpenAI API (require valid OPENAI_API_KEY)",
20+
]
21+
1722
[dependency-groups]
1823
dev = [
1924
"pytest-playwright>=0.7.2",
25+
"ruff>=0.15.5",
2026
"ty>=0.0.1a21",
2127
]

routers/chat.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
ResponseCodeInterpreterCallCompletedEvent, ResponseMcpListToolsInProgressEvent,
2323
ResponseMcpListToolsFailedEvent, ResponseMcpListToolsCompletedEvent,
2424
ResponseMcpCallArgumentsDoneEvent, ResponseMcpCallCompletedEvent,
25-
ResponseMcpCallInProgressEvent, ResponseMcpCallArgumentsDeltaEvent
25+
ResponseMcpCallInProgressEvent, ResponseMcpCallArgumentsDeltaEvent,
26+
ResponseWebSearchCallInProgressEvent, ResponseWebSearchCallSearchingEvent,
27+
ResponseWebSearchCallCompletedEvent,
2628
)
2729
from openai.types.responses.response_output_item import McpApprovalRequest
2830
from openai.types.responses import ResponseFunctionToolCall
@@ -135,6 +137,28 @@ async def event_generator() -> AsyncGenerator[str, None]:
135137
tools.extend(tool_defs)
136138
if "mcp" in enabled_tools:
137139
tools.extend(TOOL_CONFIG.mcp_servers)
140+
if "web_search" in enabled_tools:
141+
ws_tool: Dict[str, Any] = {"type": "web_search_preview"}
142+
ctx_size = os.getenv("WEB_SEARCH_CONTEXT_SIZE", "medium").strip()
143+
if ctx_size in {"low", "medium", "high"}:
144+
ws_tool["search_context_size"] = ctx_size
145+
# Build user_location only if at least one field is set
146+
country = os.getenv("WEB_SEARCH_LOCATION_COUNTRY", "").strip()
147+
city = os.getenv("WEB_SEARCH_LOCATION_CITY", "").strip()
148+
region = os.getenv("WEB_SEARCH_LOCATION_REGION", "").strip()
149+
timezone = os.getenv("WEB_SEARCH_LOCATION_TIMEZONE", "").strip()
150+
if any([country, city, region, timezone]):
151+
loc: Dict[str, str] = {"type": "approximate"}
152+
if country:
153+
loc["country"] = country
154+
if city:
155+
loc["city"] = city
156+
if region:
157+
loc["region"] = region
158+
if timezone:
159+
loc["timezone"] = timezone
160+
ws_tool["user_location"] = loc
161+
tools.append(ws_tool)
138162

139163
stream = await client.responses.create(
140164
input="",
@@ -170,7 +194,9 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
170194
ResponseMcpCallArgumentsDoneEvent() | \
171195
ResponseMcpCallCompletedEvent() | \
172196
ResponseMcpCallInProgressEvent() | \
173-
ResponseFunctionCallArgumentsDoneEvent():
197+
ResponseFunctionCallArgumentsDoneEvent() | \
198+
ResponseWebSearchCallInProgressEvent() | \
199+
ResponseWebSearchCallCompletedEvent():
174200
# Don't need to handle "in progress" or intermediate "done" events
175201
# (though long-running code interpreter interpreting might warrant handling)
176202
continue
@@ -179,6 +205,17 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
179205
# TODO: handle this (currently triggers Network/stream error exception handler)
180206
continue
181207

208+
case ResponseWebSearchCallSearchingEvent():
209+
current_item_id = event.item_id
210+
yield sse_format(
211+
"toolCallCreated",
212+
templates.get_template('components/assistant-step.html').render(
213+
step_type='toolCall',
214+
step_id=event.item_id,
215+
content="Calling web_search tool..."
216+
)
217+
)
218+
182219
case ResponseFileSearchCallSearchingEvent() | ResponseCodeInterpreterCallInProgressEvent():
183220
tool = event.type.split(".")[1].split("_call")[0]
184221
current_item_id = event.item_id
@@ -298,6 +335,11 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
298335
file_url_path = files_router.url_path_for("download_container_file", container_id=container_id, file_id=file_id)
299336
replacement_payload = f"sandbox:{container_file_path}|{file_url_path}"
300337
yield sse_format("textReplacement", wrap_for_oob_swap(current_item_id, replacement_payload))
338+
elif event.annotation["type"] == "url_citation":
339+
url = event.annotation["url"]
340+
title = event.annotation.get("title", url)
341+
citation = f'(<a href="{escape(url)}" target="_blank" rel="noopener noreferrer">{escape(title)}</a>)'
342+
yield sse_format("textDelta", wrap_for_oob_swap(current_item_id, citation))
301343
else:
302344
logger.error(f"Unhandled annotation type: {event.annotation['type']}")
303345

routers/setup.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,13 @@ async def read_setup(
8181
show_detail_env = os.getenv("SHOW_TOOL_CALL_DETAIL", "false")
8282
current_show_tool_call_detail = show_detail_env.lower() in {"1", "true", "yes", "on"}
8383

84+
# Web search config
85+
web_search_context_size = os.getenv("WEB_SEARCH_CONTEXT_SIZE", "medium")
86+
web_search_country = os.getenv("WEB_SEARCH_LOCATION_COUNTRY", "")
87+
web_search_city = os.getenv("WEB_SEARCH_LOCATION_CITY", "")
88+
web_search_region = os.getenv("WEB_SEARCH_LOCATION_REGION", "")
89+
web_search_timezone = os.getenv("WEB_SEARCH_LOCATION_TIMEZONE", "")
90+
8491
if not openai_api_key:
8592
setup_message = "OpenAI API key is missing."
8693

@@ -129,6 +136,11 @@ async def read_setup(
129136
"available_models": available_models,
130137
"existing_registry_entries": read_registry_entries(),
131138
"existing_mcp_servers": existing_mcp_servers,
139+
"web_search_context_size": web_search_context_size,
140+
"web_search_country": web_search_country,
141+
"web_search_city": web_search_city,
142+
"web_search_region": web_search_region,
143+
"web_search_timezone": web_search_timezone,
132144
}
133145
)
134146

@@ -217,7 +229,12 @@ async def save_app_config(
217229
mcp_connector_ids: List[str] = Form(default=[]),
218230
mcp_authorizations: List[str] = Form(default=[]),
219231
mcp_headers_jsons: List[str] = Form(default=[]),
220-
mcp_require_approvals: List[str] = Form(default=[])
232+
mcp_require_approvals: List[str] = Form(default=[]),
233+
web_search_context_size: Optional[str] = Form(default=None),
234+
web_search_country: Optional[str] = Form(default=None),
235+
web_search_city: Optional[str] = Form(default=None),
236+
web_search_region: Optional[str] = Form(default=None),
237+
web_search_timezone: Optional[str] = Form(default=None),
221238
) -> RedirectResponse:
222239
status = "success"
223240
message_text = ""
@@ -311,6 +328,18 @@ def get_or_empty(items: List[str], i: int) -> str:
311328
generate_registry_file(entries, mcp_servers=mcp_servers)
312329
status = "success"
313330
message_text = "Tool config saved."
331+
elif action == "save_web_search_config":
332+
# Save web search configuration
333+
ctx_size = (web_search_context_size or "medium").strip()
334+
if ctx_size not in {"low", "medium", "high"}:
335+
ctx_size = "medium"
336+
update_env_file("WEB_SEARCH_CONTEXT_SIZE", ctx_size)
337+
update_env_file("WEB_SEARCH_LOCATION_COUNTRY", (web_search_country or "").strip())
338+
update_env_file("WEB_SEARCH_LOCATION_CITY", (web_search_city or "").strip())
339+
update_env_file("WEB_SEARCH_LOCATION_REGION", (web_search_region or "").strip())
340+
update_env_file("WEB_SEARCH_LOCATION_TIMEZONE", (web_search_timezone or "").strip())
341+
status = "success"
342+
message_text = "Web search configuration saved."
314343
else:
315344
# Standard app config save
316345
if model is None or instructions is None:

templates/setup.html

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,10 @@ <h3>Model &amp; Instructions</h3>
8787
<input type="checkbox" name="tool_types" value="mcp" {% if "mcp" in current_tools %}checked{% endif %}>
8888
MCP Tools
8989
</label>
90+
<label>
91+
<input type="checkbox" name="tool_types" value="web_search" {% if "web_search" in current_tools %}checked{% endif %}>
92+
Web Search
93+
</label>
9094
</fieldset>
9195

9296
{# ── Display ── #}
@@ -111,6 +115,50 @@ <h3>Model &amp; Instructions</h3>
111115
</div>
112116
{% endif %}
113117

118+
{# ── Web Search Configuration ── #}
119+
{% if "web_search" in current_tools %}
120+
<div class="setupSection">
121+
<h3>Web Search</h3>
122+
<p class="setupMessage">Configure web search context size and optional location for localized results.</p>
123+
<form action="{{ url_for('save_app_config') }}" method="POST">
124+
<input type="hidden" name="action" value="save_web_search_config">
125+
126+
<div>
127+
<label for="web-search-context-size">Search Context Size:</label>
128+
<select name="web_search_context_size" id="web-search-context-size" class="input">
129+
<option value="low" {% if web_search_context_size == "low" %}selected{% endif %}>Low (fast, lower cost)</option>
130+
<option value="medium" {% if web_search_context_size != "low" and web_search_context_size != "high" %}selected{% endif %}>Medium (default)</option>
131+
<option value="high" {% if web_search_context_size == "high" %}selected{% endif %}>High (comprehensive, higher cost)</option>
132+
</select>
133+
</div>
134+
135+
<fieldset class="setupFieldset">
136+
<legend>User Location (optional)</legend>
137+
<div>
138+
<label for="web-search-country">Country (ISO 3166-1 code):</label>
139+
<input type="text" name="web_search_country" id="web-search-country" class="input" placeholder="e.g., US" maxlength="2" value="{{ web_search_country or '' }}">
140+
</div>
141+
<div>
142+
<label for="web-search-city">City:</label>
143+
<input type="text" name="web_search_city" id="web-search-city" class="input" placeholder="e.g., New York" value="{{ web_search_city or '' }}">
144+
</div>
145+
<div>
146+
<label for="web-search-region">Region / State:</label>
147+
<input type="text" name="web_search_region" id="web-search-region" class="input" placeholder="e.g., New York" value="{{ web_search_region or '' }}">
148+
</div>
149+
<div>
150+
<label for="web-search-timezone">Timezone (IANA):</label>
151+
<input type="text" name="web_search_timezone" id="web-search-timezone" class="input" placeholder="e.g., America/New_York" value="{{ web_search_timezone or '' }}">
152+
</div>
153+
</fieldset>
154+
155+
<div class="centered-button-container">
156+
<button type="submit" class="button">Save web search config</button>
157+
</div>
158+
</form>
159+
</div>
160+
{% endif %}
161+
114162
{# ── Custom Functions & MCP Servers ── #}
115163
{% if "function" in current_tools or "mcp" in current_tools %}
116164
<form action="{{ url_for('save_app_config') }}" method="POST" id="tool-config-form">

tests/conftest.py

Lines changed: 101 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import os
2+
import re
23
import sys
34
import time
45
import socket
@@ -14,6 +15,34 @@
1415
if str(PROJECT_ROOT) not in sys.path:
1516
sys.path.insert(0, str(PROJECT_ROOT))
1617

18+
# ---------------------------------------------------------------------------
19+
# Capture the real API key from .env BEFORE any fixtures modify it.
20+
# This is used by live integration tests that need to hit the real API.
21+
# ---------------------------------------------------------------------------
22+
def _read_real_api_key() -> str:
23+
"""Read the real API key from os.environ or .env file.
24+
25+
Called once at conftest import time — before any test fixtures modify
26+
the environment. Falls back to .env on disk if os.environ is empty.
27+
"""
28+
# 1. Check os.environ first (may be set by the shell or a prior load_dotenv)
29+
key = os.environ.get("OPENAI_API_KEY", "")
30+
if key and not key.startswith("sk-fake"):
31+
return key
32+
# 2. Fall back to reading .env directly
33+
env_path = PROJECT_ROOT / ".env"
34+
if not env_path.exists():
35+
return ""
36+
for line in env_path.read_text().splitlines():
37+
m = re.match(r"^OPENAI_API_KEY=(.+)$", line)
38+
if m:
39+
val = m.group(1).strip()
40+
if val and not val.startswith("sk-fake"):
41+
return val
42+
return ""
43+
44+
REAL_API_KEY: str = _read_real_api_key()
45+
1746

1847
@pytest.fixture
1948
def anyio_backend() -> str:
@@ -69,6 +98,27 @@ def _isolate_asyncio_running_loop(request: pytest.FixtureRequest) -> Iterator[No
6998
_aio_events._set_running_loop(saved)
7099

71100

101+
@pytest.fixture(scope="session", autouse=True)
102+
def _backup_dotenv() -> Iterator[None]:
103+
"""Back up .env at session start and restore it at session end.
104+
105+
This ensures a crashed or interrupted test run does not permanently
106+
corrupt the user's .env file.
107+
"""
108+
env_path = PROJECT_ROOT / ".env"
109+
try:
110+
original = env_path.read_text()
111+
except FileNotFoundError:
112+
original = None
113+
114+
yield
115+
116+
if original is not None:
117+
env_path.write_text(original)
118+
else:
119+
env_path.unlink(missing_ok=True)
120+
121+
72122
@pytest.fixture(scope="session")
73123
def app_server() -> Generator[int, None, None]:
74124
"""Start the FastAPI app in a background thread on a free port."""
@@ -108,13 +158,15 @@ def base_url(app_server: int) -> str: # type: ignore[override]
108158

109159

110160
@contextmanager
111-
def _dotenv(overrides: dict[str, str]) -> Iterator[None]:
161+
def _dotenv(overrides: dict[str, str], *, set_fake_api_key: bool = True) -> Iterator[None]:
112162
"""
113163
Write a temporary .env for the duration of a single test, then restore.
114164
115-
Also keeps OPENAI_API_KEY set to the fake value in os.environ so that the
116-
FastAPI Depends(lambda: AsyncOpenAI()) in route signatures never raises
117-
before load_dotenv(override=True) has a chance to run inside the route body.
165+
Saves and restores both the .env file and os.environ for every key in
166+
*overrides*. When *set_fake_api_key* is True (the default, used by
167+
Playwright tests), OPENAI_API_KEY is also force-set in os.environ so
168+
the FastAPI ``Depends(lambda: AsyncOpenAI())`` never raises before
169+
``load_dotenv(override=True)`` has a chance to run inside the route body.
118170
"""
119171
env_path = PROJECT_ROOT / ".env"
120172

@@ -129,7 +181,8 @@ def _dotenv(overrides: dict[str, str]) -> Iterator[None]:
129181
original_osenv = {k: os.environ.get(k) for k in all_keys}
130182

131183
# Always keep a fake key in the process env for the Depends constructor.
132-
os.environ["OPENAI_API_KEY"] = _FAKE_API_KEY
184+
if set_fake_api_key:
185+
os.environ["OPENAI_API_KEY"] = _FAKE_API_KEY
133186

134187
# Write the test-specific .env; the route body's load_dotenv will read it.
135188
env_path.write_text(
@@ -153,6 +206,28 @@ def _dotenv(overrides: dict[str, str]) -> Iterator[None]:
153206
del os.environ[k]
154207

155208

209+
def parse_sse_events(raw: str) -> list[dict[str, str]]:
210+
"""Parse raw SSE text into a list of {event, data} dicts."""
211+
events: list[dict[str, str]] = []
212+
current_event: str | None = None
213+
current_data_lines: list[str] = []
214+
215+
for line in raw.split("\n"):
216+
if line.startswith("event: "):
217+
current_event = line[len("event: "):]
218+
elif line.startswith("data: "):
219+
current_data_lines.append(line[len("data: "):])
220+
elif line == "" and current_event is not None:
221+
events.append({
222+
"event": current_event,
223+
"data": "\n".join(current_data_lines),
224+
})
225+
current_event = None
226+
current_data_lines = []
227+
228+
return events
229+
230+
156231
# ---------------------------------------------------------------------------
157232
# Environment fixtures used by test_setup_page_rendering.py
158233
# ---------------------------------------------------------------------------
@@ -206,7 +281,27 @@ def env_function_and_mcp(app_server: int) -> Iterator[None]:
206281
yield
207282

208283

284+
@pytest.fixture
285+
def env_web_search_only(app_server: int) -> Iterator[None]:
286+
with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "web_search"}):
287+
yield
288+
289+
290+
@pytest.fixture
291+
def env_web_search_with_config(app_server: int) -> Iterator[None]:
292+
with _dotenv({
293+
**_BASE_ENV,
294+
"ENABLED_TOOLS": "web_search",
295+
"WEB_SEARCH_CONTEXT_SIZE": "high",
296+
"WEB_SEARCH_LOCATION_COUNTRY": "US",
297+
"WEB_SEARCH_LOCATION_CITY": "New York",
298+
"WEB_SEARCH_LOCATION_REGION": "New York",
299+
"WEB_SEARCH_LOCATION_TIMEZONE": "America/New_York",
300+
}):
301+
yield
302+
303+
209304
@pytest.fixture
210305
def env_all_tools(app_server: int) -> Iterator[None]:
211-
with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "file_search,function,mcp"}):
306+
with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "file_search,function,mcp,web_search"}):
212307
yield

0 commit comments

Comments
 (0)