From 39b155596e507f6690f368b3695dee9574d7a3e7 Mon Sep 17 00:00:00 2001 From: Adam Date: Wed, 29 Jul 2026 11:29:21 +0000 Subject: [PATCH 1/3] Replace provider-native web search with in-enclave Exa search MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `web_search` flag used to switch on whichever native search tool the target model's provider happened to ship: OpenAI's Responses-API `web_search`, Anthropic's `web_search_20250305`, Gemini's `google_search` grounding, or xAI's Responses-API `web_search`. That meant four result shapes, four billable units priced from $0.01 to $0.035 (xAI billed per *citation*), and no search at all on ByteDance, Nous or Z.ai. Now the gateway advertises one ordinary `web_search` function tool and executes it itself, inside the enclave, against Exa: - Works on every text model that can call a function, with identical results, excerpt sizes and citations whichever model is picked. - One flat per-search price on every model, so a client can verify its surcharge as `searches x rate`. - The query leaves the TEE only to Exa; the LLM provider sees nothing beyond the result text in the prompt. The client protocol is untouched. Only `web_search` calls are intercepted — a caller's own tools still come back as `tool_calls` for it to run — so one request still yields one answer. Streaming buffers tool calls unconditionally while search is on (a fragment can't be forwarded before we know whose tool it is) and wraps the rounds in a generator, so the chunk handler sees one flat stream. Progress frames carry a top-level `web_search` object and sources ride out-of-band as `citations`, on the same unsigned terms as generated images. Billing accounts for the loop honestly: the flat surcharge per search that actually reached Exa, plus the tokens of every round, since each round re-sends the conversation and the accumulated results. Failed searches are free. Rounds are capped, and the last one runs with the search tool unbound so a model that would keep searching has to answer. EXA_API_KEY is injected via POST /v1/keys like every provider key; without it the flag is a no-op and /health reports web_search_enabled: false. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01V9SCMV8LhnpZ21pzBLjzxQ --- CLAUDE.md | 42 +- README.md | 32 +- scripts/run-enclave.sh | 8 +- tee_gateway/__main__.py | 10 + tee_gateway/config.py | 7 +- tee_gateway/controllers/chat_controller.py | 399 +++++-- .../controllers/completions_controller.py | 54 +- tee_gateway/llm_backend.py | 117 +- tee_gateway/model_registry.py | 60 +- tee_gateway/openapi/openapi.yaml | 36 +- tee_gateway/pricing.py | 6 +- tee_gateway/search_loop.py | 223 ++++ tee_gateway/test/test_web_search.py | 1046 ++++++++++++++--- tee_gateway/web_search.py | 426 +++++++ 14 files changed, 2065 insertions(+), 401 deletions(-) create mode 100644 tee_gateway/search_loop.py create mode 100644 tee_gateway/web_search.py diff --git a/CLAUDE.md b/CLAUDE.md index d083bcb..bfd946b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,6 +12,8 @@ The repo must provide a stable AWS Nitro PCR when the code doesn't change in ord │ ├── llm_backend.py # LLM provider routing via LangChain, HTTP client management │ ├── image_generation.py # Endpoint-based image gen (/images/generations): request shaping, URL→inline-bytes, signed responses │ ├── tee_manager.py # TEE key generation, nitriding registration, response signing +│ ├── web_search.py # In-enclave web search: Exa client, `web_search` tool spec, result formatting +│ ├── search_loop.py # Server-side tool loop that answers web_search calls in-enclave │ ├── model_registry.py # Model config and per-token pricing │ ├── definitions.py # On-chain addresses, network IDs, payment amounts │ ├── facilitator_api.py # x402 facilitator API client @@ -69,6 +71,9 @@ API keys (injected at runtime via POST /v1/keys — do NOT bake into the image): - `ARK_API_KEY` (BytePlus / ByteDance ModelArk; injected as `bytedance_api_key`) - `NOUS_API_KEY` (Nous Research / Nous Portal; injected as `nous_api_key`) - `ZAI_API_KEY` (Z.ai Model API; injected as `zai_api_key`) +- `EXA_API_KEY` (Exa search; injected as `exa_api_key`) — backs the in-enclave + `web_search` tool, not an LLM provider. Without it the `web_search` flag is a + no-op and `/health` reports `web_search_enabled: false`. Server configuration: - `API_SERVER_PORT` (default: 8000) @@ -92,13 +97,15 @@ Server configuration: - **`llm_backend.py`**: LangChain model instantiation, HTTP client management, provider routing from model name - **`model_registry.py`**: Maps model names to providers and per-token USD pricing (used by dynamic cost calculator) - **`definitions.py`**: On-chain constants (addresses, network IDs, payment amounts) — configure here for your deployment +- **`web_search.py`**: Exa HTTP client, the single provider-agnostic `web_search` function-tool spec, and result formatting/citation extraction +- **`search_loop.py`**: the in-enclave tool loop — intercepts `web_search` calls, runs them, feeds results back, bounds the rounds, and sums usage across them - **`util.py`**: `dynamic_session_cost_calculator` converts actual token usage to x402 payment amounts ### API Endpoints | Endpoint | Purpose | |----------|---------| -| `/health` | Health check (status, version, tee_enabled) | +| `/health` | Health check (status, version, tee_enabled, web_search_enabled) | | `/signing-key` | TEE public key (PEM) and tee_id | | `/enclave/attestation` | Nitro attestation document (served by nitriding) | | `/v1/keys` | One-time API key injection (POST, loopback-only) | @@ -144,6 +151,39 @@ skipped rather than dereferenced in the enclave). Per-provider request quirks `model_registry.py`. These models are billed a flat per-image price (see `per_image_price_usd`), not per token. +### Web Search + +The `web_search` request flag does NOT use any provider's native web search +(OpenAI/Anthropic/Google/xAI all have one; those were removed). Instead the +gateway advertises a single provider-agnostic `web_search` function tool +(`web_search.py`) and executes it itself, inside the enclave, against Exa — +`search_loop.py` intercepts the call, runs the search, feeds the results back as +a `ToolMessage`, and lets the model continue. Consequences to keep in mind when +touching this code: + +- **Every text model can search**, including ByteDance, Nous and Z.ai, which had + no native option. `model_supports_web_search` excludes only image models. + Anthropic's structured-output path also skips search, since + `with_structured_output` occupies the tool slot with a forced schema tool. +- **The client protocol is unchanged.** Only `web_search` calls are intercepted; + a caller's own `tools` still come back as `tool_calls` for it to run. A turn + that mixes both is terminal and the caller's tools win — ours are dropped and + the model re-issues them next turn. +- **Streaming buffers tool calls unconditionally** when search is on (a fragment + can't be forwarded until the round ends and we know whose tool it was), and the + rounds are wrapped in a generator so the chunk handler sees one flat stream. + Progress frames carry a top-level `web_search` object; `citations` ride + out-of-band on the final frame (unsigned, like `images`). +- **Billing has two parts**: a flat per-search surcharge + (`WEB_SEARCH_PRICE_USD`, identical on every model, so a client can verify it as + `searches × rate`), plus the token cost of every round — each round re-sends + the conversation plus the accumulated results, and `SearchLoopState` sums usage + across all of them. Failed searches are not counted. Rounds are capped + (`MAX_SEARCH_ROUNDS`), with the final round run against a model that has no + search tool bound so it must answer. +- Exa's self-reported `costDollars` is logged for margin reconciliation only; + settlement never depends on it. + ## Verification Examples - `examples/verify_attestation.py` — Validates AWS Nitro attestation documents against the root CA diff --git a/README.md b/README.md index edd8728..1c6d9d0 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,10 @@ The gateway solves this by running inside a hardware-isolated Nitro Enclave wher - **Request integrity** - SHA256 hash of original request included in signed response - **Streaming support** - SSE streaming for chat completions - **Tool/function calling** - Full support for LLM tool use -- **Native web search** - Opt-in `web_search` flag enables each provider's built-in - web search (OpenAI, Anthropic, Google, xAI); searches are billed per search on top - of token usage +- **In-enclave web search** - Opt-in `web_search` flag lets the model search the + live web via a tool the gateway executes inside the enclave (backed by Exa). + Works identically on every text model regardless of provider; searches are + billed at one flat per-search rate on top of token usage ## Supported Models @@ -90,7 +91,7 @@ curl -X POST http://127.0.0.1:8000/v1/completions \ "prompt": "Explain quantum computing in one sentence" }' -# Native web search (set "web_search": true on any supported model) +# Web search (set "web_search": true on any text model, any provider) curl -X POST http://127.0.0.1:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ @@ -100,11 +101,24 @@ curl -X POST http://127.0.0.1:8000/v1/chat/completions \ }' ``` -> **Web search & billing.** When `web_search` is `true`, the gateway enables the -> provider's built-in web search (OpenAI, Anthropic, Google, xAI). Each search is -> billed per search on top of token usage — at the provider's list price — and the -> charge is reflected in the dynamically-settled x402 amount. Providers without -> native web search (e.g. ByteDance/ModelArk) ignore the flag and are not charged. +> **Web search & billing.** When `web_search` is `true`, the gateway offers the +> model a `web_search` function tool and runs it itself, inside the enclave, +> against Exa — it does not use any provider's built-in search. So it works on +> every text model (ByteDance, Nous and Z.ai models included), returns the same +> results and the same `citations` whichever model you pick, and the query never +> leaves the TEE except to the search backend. +> +> The loop is invisible from the outside: one request still returns one answer, +> and any `tool_calls` you get back are only for tools you supplied yourself. +> Streaming responses also emit progress frames carrying a top-level +> `web_search` object you can render or ignore. +> +> Each search adds a flat per-search surcharge — the same on every model, so you +> can verify it as `searches x rate` — reflected in the dynamically-settled x402 +> amount. Because every round re-sends the conversation plus the accumulated +> results, the reported token usage covers all rounds. Failed searches are not +> charged. Requires `EXA_API_KEY` to be injected; check `web_search_enabled` on +> `/health`. ## Deployment to Nitro Enclave diff --git a/scripts/run-enclave.sh b/scripts/run-enclave.sh index 3928511..e9d6e55 100755 --- a/scripts/run-enclave.sh +++ b/scripts/run-enclave.sh @@ -92,6 +92,8 @@ if [ -f "$ENV_FILE" ]; then ARK_API_KEY="$(grep -E '^ARK_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" NOUS_API_KEY="$(grep -E '^NOUS_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" ZAI_API_KEY="$(grep -E '^ZAI_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" + # Backs the in-enclave `web_search` tool (Exa), not an LLM provider. + EXA_API_KEY="$(grep -E '^EXA_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" # FACILITATOR_URL is used for both x402 payment verification and the heartbeat relay. # HEARTBEAT_CONTRACT_ADDRESS and TEE_HEARTBEAT_INTERVAL are optional heartbeat parameters. @@ -110,6 +112,7 @@ if [ -f "$ENV_FILE" ]; then --arg bytedance "$ARK_API_KEY" \ --arg nous "$NOUS_API_KEY" \ --arg zai "$ZAI_API_KEY" \ + --arg exa "$EXA_API_KEY" \ --arg hb_contract "$HEARTBEAT_CONTRACT_ADDRESS" \ --arg facilitator "$FACILITATOR_URL" \ --arg hb_interval "$TEE_HEARTBEAT_INTERVAL" \ @@ -120,7 +123,8 @@ if [ -f "$ENV_FILE" ]; then xai_api_key: $xai, bytedance_api_key: $bytedance, nous_api_key: $nous, - zai_api_key: $zai + zai_api_key: $zai, + exa_api_key: $exa } + if $hb_contract != "" then {heartbeat_contract_address: $hb_contract} else {} end + if $facilitator != "" then {facilitator_url: $facilitator} else {} end @@ -151,7 +155,7 @@ if [ -f "$ENV_FILE" ]; then # Clear key variables from this shell immediately after use unset OPENAI_API_KEY GOOGLE_API_KEY ANTHROPIC_API_KEY XAI_API_KEY ARK_API_KEY - unset NOUS_API_KEY ZAI_API_KEY + unset NOUS_API_KEY ZAI_API_KEY EXA_API_KEY unset HEARTBEAT_CONTRACT_ADDRESS FACILITATOR_URL TEE_HEARTBEAT_INTERVAL fi else diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 3b317db..4e14ed0 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -22,6 +22,7 @@ DEFAULT_HEARTBEAT_INTERVAL, ) from tee_gateway.llm_backend import get_provider_config, set_provider_config +from tee_gateway.web_search import web_search_available from tee_gateway.heartbeat import create_heartbeat_service from tee_gateway.controllers.ohttp_controller import ( create_anonymous_chat_completion, @@ -395,6 +396,7 @@ def set_provider_keys(): bytedance_api_key=body.get("bytedance_api_key") or None, nous_api_key=body.get("nous_api_key") or None, zai_api_key=body.get("zai_api_key") or None, + exa_api_key=body.get("exa_api_key") or None, ) set_provider_config(provider_config) @@ -460,6 +462,9 @@ def _set(val: str | None) -> str: logger.info( " zai_api_key : %s", _set(provider_config.zai_api_key) ) + logger.info( + " exa_api_key (web search) : %s", _set(provider_config.exa_api_key) + ) logger.info(" facilitator_url : %s", facilitator_url) logger.info( " heartbeat_contract_address : %s", @@ -503,6 +508,7 @@ def _set(val: str | None) -> str: "status": "ok", "providers_initialized": providers_set, "heartbeat_enabled": heartbeat_config is not None, + "web_search_enabled": bool(provider_config.exa_api_key), } ), 200 @@ -517,6 +523,10 @@ def health(): "tee_enabled": True, "uptime_seconds": int(time.time() - _started_at), "providers": providers, + # Whether the `web_search` request flag will actually search. Not a + # provider capability — the gateway searches in-enclave for every model — + # so it is reported separately from `providers`. + "web_search_enabled": web_search_available(), "facilitator_url": _active_facilitator_url, "price_feed": _price_feed.get_status(), }, 200 diff --git a/tee_gateway/config.py b/tee_gateway/config.py index 18f977e..aa017f5 100644 --- a/tee_gateway/config.py +++ b/tee_gateway/config.py @@ -21,7 +21,7 @@ @dataclass(frozen=True) class ProviderConfig: - """API keys for each supported LLM provider.""" + """API keys for each supported LLM provider, plus the web-search backend.""" openai_api_key: Optional[str] = None anthropic_api_key: Optional[str] = None @@ -30,6 +30,11 @@ class ProviderConfig: bytedance_api_key: Optional[str] = None nous_api_key: Optional[str] = None zai_api_key: Optional[str] = None + # Exa, which backs the `web_search` tool the gateway executes inside the + # enclave for every model (see web_search.py). Not an LLM provider, so it is + # deliberately absent from initialized_providers() — /health reports it + # separately as `web_search_enabled`. + exa_api_key: Optional[str] = None def initialized_providers(self) -> list[str]: """Return provider names whose API key is set (non-empty).""" diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index 5d042ee..60f1926 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -5,6 +5,7 @@ import connexion from flask import Response +from dataclasses import dataclass, field from typing import Any from tee_gateway.models.create_chat_completion_request import ( @@ -27,8 +28,6 @@ from tee_gateway.llm_backend import ( get_provider_from_model, get_chat_model_cached, - get_web_search_tool, - extract_web_search_count, convert_messages, extract_usage, validate_attachments, @@ -39,8 +38,21 @@ create_image_generation_response, create_image_generation_streaming_response, ) -from tee_gateway.model_registry import get_model_config +from tee_gateway.model_registry import get_model_config, model_supports_web_search from tee_gateway.pricing import compute_session_cost +from tee_gateway.web_search import ( + WEB_SEARCH_TOOL_NAME, + get_web_search_tool, + web_search_available, +) +from tee_gateway.search_loop import ( + MAX_SEARCH_ROUNDS, + SearchLoopState, + execute_search_calls, + run_search_loop, + split_tool_calls, + strip_search_tool_calls, +) logger = logging.getLogger(__name__) @@ -80,6 +92,97 @@ def _split_text_and_images(content: Any) -> tuple[str, list[str]]: return ("".join(text_parts), images) +@dataclass +class _SearchEvent: + """Streaming-only marker yielded between rounds of the search loop. + + The chunk handler in ``_create_streaming_response`` buffers tool-call + fragments without knowing whose tool they belong to; this tells it. Carried + in-band on the chunk iterator so that handler needs no other knowledge of the + loop. + """ + + # Queries to announce to the client before the (blocking) searches run. + queries: list[str] = field(default_factory=list) + # The buffered calls were all ours: forget them and keep streaming. + clear_buffer: bool = False + # Terminal turn mixing our tool with the caller's: strip ours, forward theirs. + drop_search_calls: bool = False + + +def _search_status_frame(model: str, query: str) -> dict[str, Any]: + """An SSE frame telling the client which query is being searched. + + Shaped as an ordinary empty-delta chunk so a client that doesn't know about + `web_search` ignores it harmlessly, with the status hung off a top-level key + (the same convention `images` and `citations` use on the final frame). + """ + return { + "choices": [{"delta": {}, "index": 0, "finish_reason": None}], + "model": model, + "web_search": {"status": "searching", "query": query}, + } + + +def _accumulate_round( + chunk: Any, round_text: list[str], round_calls: dict[int, dict[str, Any]] +) -> None: + """Collect one round's text and tool-call fragments inside the search loop. + + Separate from the outer chunk handler's buffering because the two answer + different questions: this one decides whether to search again, that one + decides what to forward to the client. + """ + content = getattr(chunk, "content", None) + if isinstance(content, str): + round_text.append(content) + elif isinstance(content, list): + round_text.extend( + item.get("text", "") for item in content if isinstance(item, dict) + ) + + for fragment in getattr(chunk, "tool_call_chunks", None) or []: + index = fragment.get("index", 0) + entry = round_calls.setdefault(index, {"id": "", "name": "", "args": ""}) + if fragment.get("id"): + entry["id"] = fragment["id"] + if fragment.get("name"): + entry["name"] = fragment["name"] + args = fragment.get("args") + if args: + entry["args"] += args if isinstance(args, str) else json.dumps(args) + + +def _round_tool_calls(round_calls: dict[int, dict[str, Any]]) -> list[dict[str, Any]]: + """Turn buffered fragments into LangChain-shaped tool calls. + + Arguments arrive as a concatenated JSON string; a call whose arguments never + parse is still returned with empty args so it is classified (and, if it is a + web_search, answered with a "query is required" error the model can recover + from) rather than silently vanishing. + """ + calls: list[dict[str, Any]] = [] + for index in sorted(round_calls): + entry = round_calls[index] + if not entry["name"]: + continue + try: + args = json.loads(entry["args"]) if entry["args"].strip() else {} + except ValueError: + logger.warning( + "Could not parse streamed arguments for tool %r", entry["name"] + ) + args = {} + calls.append( + { + "id": entry["id"], + "name": entry["name"], + "args": args if isinstance(args, dict) else {}, + } + ) + return calls + + def create_chat_completion(body): """Create a chat completion (streaming or non-streaming).""" if not connexion.request.is_json: @@ -104,14 +207,8 @@ def create_chat_completion(body): return _create_non_streaming_response(chat_request) -def _build_tools_list(chat_request: CreateChatCompletionRequest, provider: str) -> list: - """Build the combined tools list (user tools + native web search tool). - - User-supplied function tools and the provider's built-in web search tool are - bound together in a single bind_tools() call. xAI configures search at - construction time (no tool), and providers without native web search return - no tool — in both cases only user tools are included. - """ +def _build_user_tools_list(chat_request: CreateChatCompletionRequest) -> list: + """Normalize the caller's own function tools into bind_tools() form.""" tools_list: list = [] if chat_request.tools: for tool in chat_request.tools: @@ -122,13 +219,39 @@ def _build_tools_list(chat_request: CreateChatCompletionRequest, provider: str) ) else: tools_list.append(tool) + return tools_list - if getattr(chat_request, "web_search", False): - ws_tool = get_web_search_tool(provider) - if ws_tool is not None: - tools_list.append(ws_tool) - return tools_list +def _search_enabled( + chat_request: CreateChatCompletionRequest, + provider: str, + anthropic_structured: bool, +) -> bool: + """Whether to bind and run the gateway's web_search tool for this request. + + Off when the caller didn't ask, when no Exa key was injected (rather than + advertising a tool that always fails), for image models, and for Anthropic's + structured-output path — ``with_structured_output`` occupies the tool slot + with a forced schema tool, so a search tool bound beside it would never be + callable. Every other model gets search, regardless of provider. + """ + if not getattr(chat_request, "web_search", False): + return False + if anthropic_structured and provider == "anthropic": + logger.info( + "web_search requested with Anthropic structured output; skipping " + "search (structured output uses a forced tool)" + ) + return False + if not model_supports_web_search(chat_request.model): + return False + if not web_search_available(): + logger.warning( + "web_search requested but no Exa API key is configured; answering " + "without search" + ) + return False + return True def _needs_responses_api_for_tools(provider: str, cfg, tools_list: list) -> bool: @@ -250,37 +373,51 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): if cfg.image_generation: return create_image_generation_response(chat_request, request_bytes) + # response_format is resolved before the model is built: whether Anthropic + # takes the structured-output path decides whether web search can run at + # all (see _search_enabled). + rf_dict: dict | None = None + if chat_request.response_format: + rf = _normalize_response_format(chat_request.response_format) + if rf.get("type", "text") != "text": + rf_dict = rf + # Build the tools list first: some OpenAI models (gpt-5.6 family) must be # constructed against the Responses API when function tools are bound. - tools_list = _build_tools_list(chat_request, provider) + user_tools = _build_user_tools_list(chat_request) + search_enabled = _search_enabled(chat_request, provider, rf_dict is not None) + tools_list = ( + user_tools + [get_web_search_tool()] if search_enabled else user_tools + ) - model = get_chat_model_cached( + base_model = get_chat_model_cached( model=chat_request.model, temperature=float(chat_request.temperature) if chat_request.temperature is not None else 0.0, max_tokens=chat_request.max_tokens or 4096, - web_search=bool(chat_request.web_search), force_responses_api=_needs_responses_api_for_tools( provider, cfg, tools_list ), ) - # Bind user tools and/or the native web search tool if requested. - if tools_list: - model = model.bind_tools(tools_list) + model = base_model.bind_tools(tools_list) if tools_list else base_model + # The search loop's final round runs without the search tool bound, which + # is what turns the round cap into "answer now" instead of "hand back an + # unanswerable search request". + model_without_search = ( + (base_model.bind_tools(user_tools) if user_tools else base_model) + if search_enabled + else model + ) # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no # Anthropic native equivalent and raises a clear error). - rf_dict: dict | None = None - if chat_request.response_format: - rf = _normalize_response_format(chat_request.response_format) - if rf.get("type", "text") != "text": - rf_dict = rf - if provider != "anthropic": - model = model.bind(response_format=rf_dict) + if rf_dict is not None and provider != "anthropic": + model = model.bind(response_format=rf_dict) + model_without_search = model_without_search.bind(response_format=rf_dict) langchain_messages = convert_messages(chat_request.messages) @@ -293,7 +430,18 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): SystemMessage(content="Respond in JSON format.") ] + langchain_messages - if rf_dict and provider == "anthropic": + search_state = SearchLoopState() + if search_enabled: + # Run searches to completion inside the enclave, then answer. The + # caller sent one request and gets one answer; the rounds in between + # are invisible except in the token usage they add. + response = run_search_loop( + model, + model_without_search, + langchain_messages, + search_state, + ) + elif rf_dict and provider == "anthropic": response = _invoke_anthropic_structured(model, rf_dict, langchain_messages) else: response = model.invoke(langchain_messages) @@ -310,9 +458,22 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): # of the signed output hash (see _split_text_and_images). if generated_images: message_dict["images"] = generated_images + # Sources the model searched, surfaced out-of-band alongside images and + # on the same terms: the answer text is signed, this metadata about what + # informed it rides inside the OHTTP envelope unsigned. + if search_state.citations: + message_dict["citations"] = search_state.citations finish_reason = "stop" - if hasattr(response, "tool_calls") and response.tool_calls: + # Any web_search calls left on a terminal turn belong to a turn that also + # called one of the caller's tools; they are dropped rather than handed to + # a client that has no way to run them. + client_tool_calls = ( + strip_search_tool_calls(response) + if search_enabled + else (getattr(response, "tool_calls", None) or []) + ) + if client_tool_calls: finish_reason = "tool_calls" message_dict["tool_calls"] = [ { @@ -323,7 +484,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): "arguments": json.dumps(tc.get("args", {})), }, } - for tc in response.tool_calls + for tc in client_tool_calls ] # For tool-call responses, hash the serialized tool calls so the @@ -364,7 +525,10 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): ) # TODO: If no usage is returned, we should compute it here. - usage = extract_usage(response) + # With search on, `usage` is the sum over every round of the loop — each + # round re-sent the conversation plus the accumulated search results, and + # the caller is charged for all of those tokens, not just the last round's. + usage = search_state.usage if search_enabled else extract_usage(response) if usage: # Surface the standard OpenAI usage triple on the response; the # reasoning split rides along to the cost calculator via `usage`. @@ -373,11 +537,10 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): "completion_tokens": usage["completion_tokens"], "total_tokens": usage["total_tokens"], } - web_search_count = ( - extract_web_search_count(response) if chat_request.web_search else 0 - ) cost = compute_session_cost( - chat_request.model, usage, web_search_count=web_search_count + chat_request.model, + usage, + web_search_count=search_state.search_count, ) if cost is not None: openai_response["opengradient"] = cost.model_dump(mode="json") @@ -399,9 +562,6 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): try: provider = get_provider_from_model(chat_request.model) cfg = get_model_config(chat_request.model) - # OpenAI and Anthropic stream tool calls as fragments that must be - # buffered and flushed once complete. Gemini emits complete tool calls. - buffer_tool_calls = provider in ["openai", "anthropic"] # Gemini inline-image models return a single image rather than a token # stream — invoke once and emit the result inside the SSE envelope. image_output_model = cfg.image_output @@ -416,38 +576,60 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): chat_request, request_bytes ) + # response_format is resolved before the model is built: whether Anthropic + # takes the structured-output path decides whether web search can run at + # all (see _search_enabled). + rf_dict: dict | None = None + if chat_request.response_format: + rf = _normalize_response_format(chat_request.response_format) + if rf.get("type", "text") != "text": + rf_dict = rf + anthropic_structured_rf: dict | None = ( + rf_dict if rf_dict is not None and provider == "anthropic" else None + ) + # Build the tools list first: some OpenAI models (gpt-5.6 family) must be # constructed against the Responses API when function tools are bound. - tools_list = _build_tools_list(chat_request, provider) + user_tools = _build_user_tools_list(chat_request) + search_enabled = _search_enabled(chat_request, provider, rf_dict is not None) + tools_list = ( + user_tools + [get_web_search_tool()] if search_enabled else user_tools + ) + + # OpenAI and Anthropic stream tool calls as fragments that must be + # buffered and flushed once complete. Gemini emits complete tool calls. + # With search on, ALWAYS buffer: a fragment can't be forwarded to the + # client until the round ends and we know whether the call was a + # web_search this gateway will answer itself. + buffer_tool_calls = provider in ["openai", "anthropic"] or search_enabled - model = get_chat_model_cached( + base_model = get_chat_model_cached( model=chat_request.model, temperature=float(chat_request.temperature) if chat_request.temperature is not None else 0.0, max_tokens=chat_request.max_tokens or 4096, - web_search=bool(chat_request.web_search), force_responses_api=_needs_responses_api_for_tools( provider, cfg, tools_list ), ) - # Bind user tools and/or the native web search tool if requested. - if tools_list: - model = model.bind_tools(tools_list) + model = base_model.bind_tools(tools_list) if tools_list else base_model + # The search loop's final round runs without the search tool bound, so a + # model that would keep searching has to answer with what it has. + model_without_search = ( + (base_model.bind_tools(user_tools) if user_tools else base_model) + if search_enabled + else model + ) # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no # Anthropic native equivalent and raises a clear error). - anthropic_structured_rf: dict | None = None - if chat_request.response_format: - rf = _normalize_response_format(chat_request.response_format) - if rf.get("type", "text") != "text": - if provider == "anthropic": - anthropic_structured_rf = rf - else: - model = model.bind(response_format=rf) + if rf_dict is not None and anthropic_structured_rf is None: + model = model.bind(response_format=rf_dict) + model_without_search = model_without_search.bind(response_format=rf_dict) langchain_messages = convert_messages(chat_request.messages) @@ -497,16 +679,73 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): anthropic_structured_content = None anthropic_structured_usage = None + search_state = SearchLoopState() + + def _streamed_search_rounds(): + """Yield chunks across every round of the in-enclave search loop. + + Written as a generator wrapping ``model.stream`` so the chunk handling + below is identical whether or not search is on — one round or five, it + sees one flat stream of chunks. Between rounds it yields a + ``_SearchEvent`` telling that handler what to do with the tool-call + fragments it just buffered, since only this generator knows whether + they were ``web_search`` calls it answered itself. + """ + loop_messages = list(langchain_messages) + rounds = MAX_SEARCH_ROUNDS + 1 if search_enabled else 1 + + for round_index in range(rounds): + last_round = search_enabled and round_index == MAX_SEARCH_ROUNDS + active_model = model_without_search if last_round else model + + round_text: list[str] = [] + round_calls: dict[int, dict[str, Any]] = {} + + for chunk in active_model.stream(loop_messages): + if search_enabled: + _accumulate_round(chunk, round_text, round_calls) + yield chunk + + if not search_enabled: + return + + ours, theirs = split_tool_calls(_round_tool_calls(round_calls)) + if theirs or not ours: + # Terminal turn. A turn that asked for both our search and one + # of the caller's tools can't be served by either side alone, + # so the caller's tools win and ours are dropped downstream. + if ours: + logger.info( + "Dropping %d streamed web_search call(s) from a turn " + "that also called %d client tool(s)", + len(ours), + len(theirs), + ) + yield _SearchEvent(drop_search_calls=True) + return + + # Announce the queries before running them: the Exa call blocks + # for a second or two and the client should say why it is waiting. + yield _SearchEvent( + queries=[ + q + for q in ((c.get("args") or {}).get("query") for c in ours) + if isinstance(q, str) and q.strip() + ], + clear_buffer=True, + ) + + loop_messages.append( + AIMessage(content="".join(round_text), tool_calls=ours) + ) + loop_messages.extend(execute_search_calls(ours, search_state)) + def generate(): full_content = "" final_usage = None buffered_tool_calls = {} finish_reason = "stop" generated_images: list[str] = [] - # Accumulate streamed chunks so native web-search activity (content - # blocks, citations, grounding metadata) can be counted for billing - # once the stream completes. - merged_chunk = None try: if image_output_model: @@ -577,15 +816,30 @@ def generate(): yield f"data: {json.dumps(data)}\n\n" chunks_iter = [] else: - chunks_iter = model.stream(langchain_messages) # type: ignore[assignment] + chunks_iter = _streamed_search_rounds() # type: ignore[assignment] for chunk in chunks_iter: - # Accumulate for post-stream web-search billing (cheap: merges - # content/metadata deltas into a single AIMessageChunk). - if chat_request.web_search: - merged_chunk = ( - chunk if merged_chunk is None else merged_chunk + chunk - ) + # --- Search-round boundary (in-enclave web search) --- + # Not a model chunk: an instruction about the tool-call + # fragments buffered so far, plus the queries to tell the + # client about. + if isinstance(chunk, _SearchEvent): + if chunk.clear_buffer: + # Those calls were web_search calls this gateway is + # answering itself — the client must never see them. + buffered_tool_calls = {} + finish_reason = "stop" + if chunk.drop_search_calls: + buffered_tool_calls = { + index: tc + for index, tc in buffered_tool_calls.items() + if tc["function"]["name"] != WEB_SEARCH_TOOL_NAME + } + if not buffered_tool_calls: + finish_reason = "stop" + for query in chunk.queries: + yield f"data: {json.dumps(_search_status_frame(chat_request.model, query))}\n\n" + continue # --- Text content --- if chunk.content: @@ -781,6 +1035,10 @@ def generate(): # are not part of the signed output hash. if generated_images: final_data["images"] = generated_images + # Likewise the sources the model searched: the answer text is + # signed, this metadata about what informed it is not. + if search_state.citations: + final_data["citations"] = search_state.citations logger.debug( f"Response Final\n\tTEE Signature: {tee_signature}\n\tTEE request hash: {input_hash_hex}\n\tTEE output hash: {output_hash_hex}\n\tTEE timestamp: {timestamp}\n\tTEE ID: 0x{tee_keys.get_tee_id()}" @@ -793,13 +1051,11 @@ def generate(): "completion_tokens": final_usage.get("output_tokens", 0), "total_tokens": final_usage.get("total_tokens", 0), } - web_search_count = ( - extract_web_search_count(merged_chunk) - if chat_request.web_search - else 0 - ) # Pass thinking tokens to the cost calculator (for the image # dual-rate split) without polluting the OpenAI usage triple. + # `final_usage` already sums every round of the search loop — + # each round re-sent the conversation plus the accumulated + # results, and the caller pays for all of those input tokens. cost_usage = dict( final_data["usage"], reasoning_tokens=final_usage.get("reasoning", 0), @@ -807,7 +1063,7 @@ def generate(): cost = compute_session_cost( chat_request.model, cost_usage, - web_search_count=web_search_count, + web_search_count=search_state.search_count, ) if cost is not None: # final_data is hand-serialized to SSE via json.dumps below, @@ -817,6 +1073,7 @@ def generate(): logger.info( f"Stream completed — usage: {final_data['usage']}, " f"finish: {finish_reason}, " + f"searches: {search_state.search_count}, " f"inputHash: {input_hash_hex[:16]}..., outputHash: {output_hash_hex[:16]}..." ) diff --git a/tee_gateway/controllers/completions_controller.py b/tee_gateway/controllers/completions_controller.py index 773dd74..472c270 100644 --- a/tee_gateway/controllers/completions_controller.py +++ b/tee_gateway/controllers/completions_controller.py @@ -13,12 +13,12 @@ from tee_gateway.tee_manager import get_tee_keys, compute_tee_msg_hash from tee_gateway.llm_backend import ( get_chat_model_cached, - get_provider_from_model, - get_web_search_tool, - extract_web_search_count, extract_usage, ) +from tee_gateway.model_registry import model_supports_web_search from tee_gateway.pricing import compute_session_cost +from tee_gateway.search_loop import SearchLoopState, run_search_loop +from tee_gateway.web_search import get_web_search_tool, web_search_available logger = logging.getLogger(__name__) @@ -49,27 +49,43 @@ def create_completion(body): request_bytes = json.dumps(request_dict, sort_keys=True).encode("utf-8") - model = get_chat_model_cached( + base_model = get_chat_model_cached( model=body.model, temperature=float(body.temperature) if body.temperature is not None else 0.0, max_tokens=body.max_tokens or 4096, - web_search=web_search, ) - # Enable the provider's native web search tool where binding is required - # (xAI configures it at construction; unsupported providers return None). - if web_search: - ws_tool = get_web_search_tool(get_provider_from_model(body.model)) - if ws_tool is not None: - model = model.bind_tools([ws_tool]) + # Web search is the gateway's own function tool, executed in-enclave for + # any model that can call a function — see web_search.py. Skipped when no + # Exa key was injected rather than advertising a tool that always fails. + search_enabled = ( + web_search + and model_supports_web_search(body.model) + and web_search_available() + ) + if web_search and not search_enabled: + logger.warning( + "web_search requested for %s but search is unavailable; " + "completing without it", + body.model, + ) - messages = [HumanMessage(content=body.prompt)] - response = model.invoke(messages) + messages: list[Any] = [HumanMessage(content=body.prompt)] + search_state = SearchLoopState() + if search_enabled: + response = run_search_loop( + base_model.bind_tools([get_web_search_tool()]), + base_model, + messages, + search_state, + ) + else: + response = base_model.invoke(messages) - # Web search (OpenAI Responses API / Gemini) can return content as a list - # of blocks; flatten to the text the caller expects. + # Some providers return content as a list of blocks; flatten to the text + # the caller expects. if isinstance(response.content, list): response_content = "".join( item.get("text", "") if isinstance(item, dict) else str(item) @@ -77,7 +93,8 @@ def create_completion(body): ) else: response_content = response.content or "" - usage = extract_usage(response) + # With search on, usage is the sum over every round of the loop. + usage = search_state.usage if search_enabled else extract_usage(response) timestamp = int(time.time()) msg_hash, input_hash_hex, output_hash_hex = compute_tee_msg_hash( @@ -106,12 +123,13 @@ def create_completion(body): "tee_id": f"0x{tee_keys.get_tee_id()}", } if usage: - web_search_count = extract_web_search_count(response) if web_search else 0 cost = compute_session_cost( - body.model, usage, web_search_count=web_search_count + body.model, usage, web_search_count=search_state.search_count ) if cost is not None: completion_response["opengradient"] = cost.model_dump(mode="json") + if search_state.citations: + completion_response["citations"] = search_state.citations return completion_response except Exception as e: diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index b358957..52a02e1 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -27,6 +27,7 @@ from tee_gateway.config import ProviderConfig from tee_gateway.model_registry import get_model_config +from tee_gateway.web_search import configure_exa_client logger = logging.getLogger(__name__) @@ -130,6 +131,10 @@ def set_provider_config(config: ProviderConfig) -> None: follow_redirects=False, ) + # Web search runs inside the enclave against Exa rather than through any + # provider's native tool, so its client is built here alongside them. + configure_exa_client(config.exa_api_key) + get_chat_model_cached.cache_clear() _provider_config = config @@ -160,25 +165,22 @@ def get_chat_model_cached( model: str, temperature: float, max_tokens: int, - web_search: bool = False, force_responses_api: bool = False, ): """Get cached chat model instance using the injected ProviderConfig. - Models are cached by (model, temperature, max_tokens, web_search, - force_responses_api) tuple. Cache is cleared by set_provider_config() after - key injection. + Models are cached by (model, temperature, max_tokens, force_responses_api) + tuple. Cache is cleared by set_provider_config() after key injection. - When ``web_search`` is True, provider-specific native web search is enabled. - Some providers (OpenAI, xAI) require search configuration at construction - time; others (Anthropic, Google) enable it by binding a tool — see - ``get_web_search_tool``. Providers without native web search ignore the flag. + Web search needs nothing here: it is a plain function tool the gateway binds + and executes itself (see web_search.py), not a provider feature that has to + be switched on at construction time. When ``force_responses_api`` is True, OpenAI models are constructed against - the Responses API (like the web-search path). The gpt-5.6 family rejects - function tools combined with its default ``reasoning_effort`` on Chat - Completions, so the chat controller sets this flag for those models when - tools are bound. Ignored by non-OpenAI providers. + the Responses API. The gpt-5.6 family rejects function tools combined with + its default ``reasoning_effort`` on Chat Completions, so the chat controller + sets this flag for those models when tools are bound. Ignored by non-OpenAI + providers. """ config = _provider_config if config is None: @@ -225,14 +227,11 @@ def get_chat_model_cached( if openai_http_client is None: raise RuntimeError("OpenAI HTTP client has not been initialized") - # Switch to the Responses API when either (a) web search is requested — - # OpenAI's built-in web_search tool is only available there — or (b) the - # caller forced it because this model rejects function tools + its - # default reasoning_effort on Chat Completions (gpt-5.6 family). The - # responses/v1 output format surfaces web_search_call items in the - # message content, which we count for billing. + # Switch to the Responses API when the caller forced it because this + # model rejects function tools + its default reasoning_effort on Chat + # Completions (gpt-5.6 family). openai_kwargs: dict[str, Any] = {} - if web_search or force_responses_api: + if force_responses_api: openai_kwargs["use_responses_api"] = True openai_kwargs["output_version"] = "responses/v1" @@ -274,12 +273,6 @@ def get_chat_model_cached( if xai_http_client is None: raise RuntimeError("XAI HTTP client has not been initialized") - # xAI deprecated Live Search on Chat Completions. Web search now lives on - # the Responses API and is enabled by binding the built-in web_search tool. - xai_kwargs: dict[str, Any] = {} - if web_search: - xai_kwargs["use_responses_api"] = True - return ChatXAI( model=api_name, api_key=SecretStr(config.xai_api_key), @@ -288,7 +281,6 @@ def get_chat_model_cached( http_client=xai_http_client, streaming=True, stream_usage=True, - **xai_kwargs, ) elif provider == "bytedance": @@ -681,73 +673,6 @@ def extract_usage(response) -> Optional[Dict[str, int]]: return None -# Anthropic's server-side web search tool. The dated type string is the current -# tool version; max_uses caps searches per request to bound cost. -ANTHROPIC_WEB_SEARCH_TOOL: Dict[str, Any] = { - "type": "web_search_20250305", - "name": "web_search", - "max_uses": 5, -} - - -def get_web_search_tool(provider: str) -> Optional[Dict[str, Any]]: - """Return the provider-specific web search tool spec to bind, or None. - - OpenAI/Anthropic/Google/xAI enable web search by binding a built-in tool. - xAI requires this on the Responses API path; its old Live Search - ``search_parameters`` path is deprecated. - """ - if provider == "openai": - return {"type": "web_search"} - if provider == "anthropic": - return dict(ANTHROPIC_WEB_SEARCH_TOOL) - if provider == "google": - return {"google_search": {}} - if provider == "x-ai": - return {"type": "web_search"} - # bytedance has no native web search. - return None - - -def extract_web_search_count(message) -> int: - """Best-effort count of billable web-search units the model performed. - - Each provider reports search activity differently and bills a different - unit, so we count the unit that matches that provider's list price: - - - OpenAI (Responses API): ``web_search_call`` content blocks (per call) - - Anthropic: ``server_tool_use`` web_search content blocks (per request) - - xAI/OpenAI Responses API: ``web_search_call`` content blocks (per call) - or legacy xAI ``citations`` in additional_kwargs (per source) - - Google: 1 per grounded response (per grounded request) - - Works on a completed AIMessage or an accumulated AIMessageChunk. - """ - if message is None: - return 0 - - count = 0 - - content = getattr(message, "content", None) - if isinstance(content, list): - for block in content: - if not isinstance(block, dict): - continue - btype = block.get("type") - if btype == "web_search_call": - count += 1 - elif btype == "server_tool_use" and block.get("name") == "web_search": - count += 1 - - # xAI returns the list of sources it grounded on as `citations`. - additional_kwargs = getattr(message, "additional_kwargs", None) or {} - citations = additional_kwargs.get("citations") - if citations: - count += len(citations) - - # Google reports grounding via response_metadata; billed per grounded request. - response_metadata = getattr(message, "response_metadata", None) or {} - if response_metadata.get("grounding_metadata"): - count += 1 - - return count +# Web search is not a provider feature here — the gateway executes it itself +# against Exa and bills one flat rate on every model. See web_search.py for the +# tool spec and search_loop.py for the loop that runs it. diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 979cd90..f8804a9 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -68,9 +68,9 @@ class ModelConfig: # ``output_tokens`` count and only breaks out thinking (``reasoning``), so the # billing splits reasoning at ``output_price_usd`` and the remainder here. image_output_price_usd: Optional[Decimal] = None - # Per-search USD surcharge billed when native web search is used. ``None`` - # means "use the provider default" (see WEB_SEARCH_PRICE_USD_BY_PROVIDER); - # set an explicit value here to override a single model's web-search price. + # Per-search USD surcharge override. Web search is one flat rate on every + # model (``WEB_SEARCH_PRICE_USD``) because the gateway runs the search + # itself; set this only to price a single model's searches differently. web_search_price_usd: Optional[Decimal] = None # OpenAI's newest reasoning models (the gpt-5.6 family) apply a default # ``reasoning_effort`` that the Chat Completions endpoint rejects when @@ -83,20 +83,20 @@ class ModelConfig: responses_api_for_tools: bool = False -# Default per-search USD price charged when a model uses native web search. -# The billable "unit" differs per provider (see extract_web_search_count in -# llm_backend.py) and these mirror each provider's public list price: -# - OpenAI: per web_search tool call (~$10 / 1k calls) -# - Anthropic: per web_search request ($10 / 1k searches) -# - xAI: per web_search tool call / source ($25 / 1k units) -# - Google: per grounded request ($35 / 1k requests) -# Providers without native web search are omitted (charged nothing). -WEB_SEARCH_PRICE_USD_BY_PROVIDER: dict[str, Decimal] = { - "openai": Decimal("0.01"), - "anthropic": Decimal("0.01"), - "x-ai": Decimal("0.025"), - "google": Decimal("0.035"), -} +# Flat USD price per web search, identical on every model. +# +# The gateway runs searches itself against Exa (see web_search.py), so there is +# one cost to pass through instead of four provider list prices with four +# different billable units. At our request shape — one Exa search plus page text +# for up to `MAX_NUM_RESULTS` results — Exa charges $7/1k requests and $1/1k +# pages per content type, i.e. ~$0.013 for a 6-result search. This rate covers +# that with a small margin, and is below what three of the four native searches +# used to cost (xAI $0.025/unit, Google $0.035/request). +# +# The billable unit is "one search the model asked for that reached Exa", so a +# client can verify its surcharge as `searches * this rate`. Searches that +# failed or were malformed are not counted (see WebSearchOutcome.billable). +WEB_SEARCH_PRICE_USD: Decimal = Decimal("0.015") # ByteDance ModelArk image *deployment* endpoints (api_name "ep-…", e.g. Seedance # 4.5, Seedream 5.0 Lite) return the URL response format and require these extra @@ -721,18 +721,28 @@ def get_rate_card(model: str) -> dict[str, Decimal]: def get_web_search_price_usd(model: str) -> Decimal: - """Return the per-search USD surcharge for a model's native web search. + """Return the per-search USD surcharge for a model. - Falls back to the provider default when the model does not override it, and - to ``Decimal("0")`` for providers with no native web search support. Raises - ValueError if the model is unknown. + The flat ``WEB_SEARCH_PRICE_USD`` unless the model overrides it. Image models + are free: they never reach the chat path that can search. Raises ValueError + if the model is unknown. """ cfg = get_model_config(model) if cfg.web_search_price_usd is not None: return cfg.web_search_price_usd - return WEB_SEARCH_PRICE_USD_BY_PROVIDER.get(cfg.provider, Decimal("0")) + if cfg.image_generation or cfg.image_output: + return Decimal("0") + return WEB_SEARCH_PRICE_USD + +def model_supports_web_search(model: str) -> bool: + """Whether ``web_search`` can be enabled for a model. -def provider_supports_web_search(provider: str) -> bool: - """Whether the given provider has native web search the gateway can enable.""" - return provider in WEB_SEARCH_PRICE_USD_BY_PROVIDER + True for every text model in the registry — the gateway supplies the search + tool itself, so this is a question about function calling, not about which + provider shipped a search feature. Image models are excluded: generation + models are served off the chat path entirely, and image-*output* models are + invoked in a single non-streaming shot with no tool loop around them. + """ + cfg = get_model_config(model) + return not (cfg.image_generation or cfg.image_output) diff --git a/tee_gateway/openapi/openapi.yaml b/tee_gateway/openapi/openapi.yaml index 80cb352..fc15712 100644 --- a/tee_gateway/openapi/openapi.yaml +++ b/tee_gateway/openapi/openapi.yaml @@ -2887,11 +2887,27 @@ components: web_search: default: false description: | - Enable the model's native web search capability. When `true`, the - gateway turns on the provider's built-in web search (OpenAI, Anthropic, - Google, and xAI). Searches performed are billed per search on top of - token usage. Models whose provider has no native web search ignore - this flag and are not charged for it. + Let the model search the web. When `true`, the gateway offers the + model a `web_search` tool and executes it inside the enclave against + its own search backend, feeding the results back so the model can + answer from them. This works on every text model regardless of + provider, and behaves identically on all of them. + + The search rounds are invisible to the caller: one request still + yields one answer, and the `tool_calls` you get back are only ever + for tools you supplied yourself. Sources used are returned + out-of-band on the response message as `citations`. Streaming + responses additionally emit progress frames carrying a top-level + `web_search` object (`{"status": "searching", "query": "..."}`) that + clients may render or ignore. + + Billing: each search the model runs adds a flat per-search surcharge + (identical on every model) on top of token usage, and because each + round re-sends the conversation plus the accumulated results, the + reported token usage covers every round. Searches that fail are not + charged. Image models, and gateways with no search backend + configured, ignore this flag and are not charged for it — check + `web_search_enabled` on `/health`. nullable: true title: web_search type: boolean @@ -3551,9 +3567,13 @@ components: web_search: default: false description: | - Enable the model's native web search capability. Searches performed - are billed per search on top of token usage. Models whose provider - has no native web search ignore this flag and are not charged for it. + Let the model search the web. The gateway offers the model a + `web_search` tool and executes it inside the enclave, feeding the + results back before the completion is produced. Works on every text + model regardless of provider. Each search adds a flat per-search + surcharge on top of token usage, and the reported token usage covers + every round of the loop. Failed searches are not charged. Sources + are returned on the response as `citations`. nullable: true title: web_search type: boolean diff --git a/tee_gateway/pricing.py b/tee_gateway/pricing.py index 1bd952a..d6d97a0 100644 --- a/tee_gateway/pricing.py +++ b/tee_gateway/pricing.py @@ -96,7 +96,11 @@ def compute_session_cost( Decimal(out_tok) * cfg.output_price_usd ) - # Native web search is billed per search unit on top of token cost. + # Web search is billed per search on top of token cost, at one flat rate + # for every model (the gateway runs the search itself — see web_search.py). + # Note the token cost already reflects the search: a search round re-sends + # the conversation plus the results, and search_loop.SearchLoopState sums + # the usage of every round into `usage`. searches = max(0, int(web_search_count)) web_search_usd = ( Decimal(searches) * get_web_search_price_usd(model) diff --git a/tee_gateway/search_loop.py b/tee_gateway/search_loop.py new file mode 100644 index 0000000..4fa25bb --- /dev/null +++ b/tee_gateway/search_loop.py @@ -0,0 +1,223 @@ +"""The in-enclave web-search tool loop. + +The gateway advertises ``web_search`` (see ``web_search.get_web_search_tool``) as +an ordinary function tool, which means the model asks for a search the same way +it asks for any other tool — and something has to answer it. That something is +this module: it runs the search inside the enclave, feeds the results back, and +lets the model continue, all within one client request. + +Two properties this preserves, both of which the previous provider-native search +gave up: + + * The client's tool protocol is untouched. A caller that passes its own + ``tools`` still gets tool calls handed back to execute; only ``web_search`` + calls are intercepted. A caller that passes no tools never learns a loop + happened — it sends one request and gets one answer. + * Billing stays honest about a loop's real cost. Each round re-sends the whole + conversation *plus* every prior search result, so the input tokens are + genuinely spent several times over. ``SearchLoopState`` accumulates usage + across every round so the caller is charged for all of it rather than for + the last round alone. + +The loop is bounded (``MAX_SEARCH_ROUNDS``) and the final round is run with the +search tool unbound, so a model that would otherwise keep searching is forced to +answer with what it has instead of spending the caller's money indefinitely. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +from langchain_core.messages import AIMessage, ToolMessage + +from tee_gateway.web_search import ( + WEB_SEARCH_TOOL_NAME, + execute_web_search_call, +) + +logger = logging.getLogger(__name__) + +# How many times the model may search before it must answer. Each round costs a +# full prompt re-send, so this is a cost ceiling as much as a latency one: four +# rounds is enough for "search, refine, cross-check" without letting a model +# that has decided to keep googling run up an unbounded bill. +MAX_SEARCH_ROUNDS = 4 + + +@dataclass +class SearchLoopState: + """Accumulator threaded through every round of one request's loop. + + Kept separate from the loop functions because the streaming controller runs + its rounds itself (it has to forward SSE frames as they arrive) while the + non-streaming controller delegates the whole loop — both share this state. + """ + + search_count: int = 0 + citations: list[dict[str, str]] = field(default_factory=list) + # Running token totals across every round, in the shape extract_usage + # returns. None until some round actually reports usage, matching the + # "provider reported nothing, so do not charge" convention elsewhere. + usage: Optional[dict[str, int]] = None + rounds: int = 0 + + def add_usage(self, round_usage: Optional[dict[str, int]]) -> None: + """Fold one round's token usage into the running totals.""" + if not round_usage: + return + if self.usage is None: + self.usage = {} + for key, value in round_usage.items(): + if isinstance(value, (int, float)): + self.usage[key] = self.usage.get(key, 0) + int(value) + + def add_citations(self, citations: list[dict[str, str]]) -> None: + """Append citations from one search, de-duplicated by URL. + + A refining second search very often re-surfaces the best hit from the + first, and showing the same source twice reads as a bug. + """ + seen = {c.get("url") for c in self.citations} + for citation in citations: + url = citation.get("url") + if url and url not in seen: + seen.add(url) + self.citations.append(citation) + + +def split_tool_calls( + tool_calls: Optional[list[dict[str, Any]]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Partition a turn's tool calls into (ours, the client's). + + "Ours" are ``web_search`` calls this gateway executes; the rest belong to + tools the caller supplied and must be handed back for the caller to run. + """ + ours: list[dict[str, Any]] = [] + theirs: list[dict[str, Any]] = [] + for call in tool_calls or []: + name = call.get("name") if isinstance(call, dict) else None + if name == WEB_SEARCH_TOOL_NAME: + ours.append(call) + else: + theirs.append(call) + return ours, theirs + + +def execute_search_calls( + calls: list[dict[str, Any]], + state: SearchLoopState, + on_search: Optional[Callable[[str], None]] = None, +) -> list[ToolMessage]: + """Run each ``web_search`` call and build the ToolMessages to feed back. + + ``on_search`` is invoked with each query before it runs, so the streaming + controller can tell the client what is being searched for while it waits. + Only searches that actually reached Exa are counted as billable. + """ + messages: list[ToolMessage] = [] + for call in calls: + args = call.get("args") + if not isinstance(args, dict): + args = {} + query = args.get("query") + if on_search is not None and isinstance(query, str) and query.strip(): + try: + on_search(query.strip()) + except Exception: + # A status callback is cosmetic; never let it kill the search. + logger.debug("web_search status callback failed", exc_info=True) + + outcome = execute_web_search_call(args) + if outcome.billable: + state.search_count += 1 + state.add_citations(outcome.citations) + + messages.append( + ToolMessage( + content=outcome.content, + tool_call_id=call.get("id") or "", + name=WEB_SEARCH_TOOL_NAME, + status="error" if outcome.is_error else "success", + ) + ) + return messages + + +def strip_search_tool_calls(message: AIMessage) -> list[dict[str, Any]]: + """Client-facing tool calls for a turn that mixes our tool with theirs. + + A turn asking for both ``web_search`` and one of the caller's tools cannot be + completed by either side alone: we cannot run their tool, and they cannot run + ours. The caller's tools win — their loop is the outer one and will come back + to us — so our calls are dropped here and the model re-issues them on the + next turn if it still wants to search. Rare in practice; logged when it + happens so it does not stay invisible. + """ + ours, theirs = split_tool_calls(getattr(message, "tool_calls", None)) + if ours: + logger.info( + "Dropping %d web_search call(s) from a turn that also called %d " + "client tool(s); the model can re-issue them next turn", + len(ours), + len(theirs), + ) + return theirs + + +def run_search_loop( + model: Any, + model_without_search: Any, + messages: list[Any], + state: SearchLoopState, + invoke: Optional[Callable[[Any, list[Any]], AIMessage]] = None, + on_search: Optional[Callable[[str], None]] = None, + max_rounds: int = MAX_SEARCH_ROUNDS, +) -> AIMessage: + """Drive the loop to a terminal turn and return it (non-streaming callers). + + Terminal means: a plain answer, or a turn calling one of the *caller's* + tools. ``messages`` is extended in place with each round's assistant turn and + search results, so the caller can inspect the full trajectory afterwards. + + ``model`` has the search tool bound; ``model_without_search`` does not and is + used for the final round, which is what converts the round cap into "answer + now" rather than "return an unanswerable search request". ``invoke`` lets a + caller substitute its own invocation (the Anthropic structured-output path + does not use plain ``.invoke``). + """ + call_model = invoke if invoke is not None else (lambda m, msgs: m.invoke(msgs)) + + for round_index in range(max_rounds + 1): + last_round = round_index == max_rounds + active_model = model_without_search if last_round else model + + response = call_model(active_model, messages) + state.rounds = round_index + 1 + state.add_usage(_message_usage(response)) + + ours, theirs = split_tool_calls(getattr(response, "tool_calls", None)) + if theirs or not ours: + # Terminal: either a plain answer or the caller's tools to run. + return response + + messages.append(response) + messages.extend(execute_search_calls(ours, state, on_search)) + + # Unreachable: the last iteration binds no search tool, so `ours` is empty + # and the loop returns above. + raise RuntimeError("search loop exited without a terminal response") + + +def _message_usage(message: Any) -> Optional[dict[str, int]]: + """Token usage for one round, in the shape the cost calculator expects.""" + metadata = getattr(message, "usage_metadata", None) + if not metadata: + return None + details = metadata.get("output_token_details") or {} + return { + "prompt_tokens": metadata.get("input_tokens", 0), + "completion_tokens": metadata.get("output_tokens", 0), + "total_tokens": metadata.get("total_tokens", 0), + "reasoning_tokens": details.get("reasoning", 0), + } diff --git a/tee_gateway/test/test_web_search.py b/tee_gateway/test/test_web_search.py index 277dcd3..30ba47a 100644 --- a/tee_gateway/test/test_web_search.py +++ b/tee_gateway/test/test_web_search.py @@ -1,145 +1,388 @@ """ -Unit tests for native web search support across providers. +Unit tests for in-enclave web search (Exa) across providers. Covers: - - model_registry: per-search pricing lookup and provider support predicate - - llm_backend.get_web_search_tool: provider-specific tool specs - - llm_backend.extract_web_search_count: counting billable search units from - each provider's response shape + - web_search: tool spec, argument coercion, Exa request shaping, result + formatting, and every failure mode of the Exa call + - model_registry: one flat per-search price and which models can search + - search_loop: tool-call partitioning, round accumulation, the round cap, and + terminal conditions - pricing.compute_session_cost: per-search surcharge added to token cost - - chat_controller: web_search flag binds the tool and bills the searches + - chat_controller: the web_search flag binds the tool, searches are executed + in-enclave rather than handed to the client, and every round is billed """ +import json import unittest from decimal import Decimal from types import SimpleNamespace -from unittest.mock import patch, Mock +from unittest.mock import Mock, patch -from langchain_core.messages import AIMessage +from langchain_core.messages import AIMessage, AIMessageChunk, ToolMessage +from tee_gateway import web_search as ws from tee_gateway.model_registry import ( + WEB_SEARCH_PRICE_USD, get_web_search_price_usd, - provider_supports_web_search, -) -from tee_gateway.llm_backend import ( - get_web_search_tool, - extract_web_search_count, + model_supports_web_search, ) from tee_gateway.pricing import SessionCost, compute_session_cost +from tee_gateway.search_loop import ( + MAX_SEARCH_ROUNDS, + SearchLoopState, + execute_search_calls, + run_search_loop, + split_tool_calls, + strip_search_tool_calls, +) from tee_gateway.controllers.chat_controller import create_chat_completion # --------------------------------------------------------------------------- -# model_registry pricing +# Exa test doubles # --------------------------------------------------------------------------- -class TestWebSearchPricing(unittest.TestCase): - def test_provider_support_predicate(self): - for provider in ("openai", "anthropic", "google", "x-ai"): - self.assertTrue(provider_supports_web_search(provider)) - self.assertFalse(provider_supports_web_search("bytedance")) - - def test_price_uses_provider_default(self): - # gpt-4.1 -> openai default ($0.01/search) - self.assertEqual(get_web_search_price_usd("gpt-4.1"), Decimal("0.01")) - # grok-4 -> xAI default ($0.025/search unit) - self.assertEqual(get_web_search_price_usd("grok-4"), Decimal("0.025")) - # gemini -> google default ($0.035/grounded request) - self.assertEqual(get_web_search_price_usd("gemini-2.5-flash"), Decimal("0.035")) - - def test_unsupported_provider_is_free(self): - # ByteDance has no native web search -> no charge - self.assertEqual(get_web_search_price_usd("seed-1.6"), Decimal("0")) +def _exa_result(url: str, title: str = "A title", text: str = "Some body text"): + return { + "title": title, + "url": url, + "publishedDate": "2026-03-04T10:00:00.000Z", + "author": "An author", + "id": url, + "text": text, + } - def test_unknown_model_raises(self): - with self.assertRaises(ValueError): - get_web_search_price_usd("not-a-real-model") + +def _exa_response(status: int = 200, body: dict | None = None): + """A stand-in for httpx.Response carrying just what the code reads.""" + response = Mock() + response.status_code = status + response.json.return_value = body if body is not None else {} + response.text = json.dumps(body) if body is not None else "" + response.reason_phrase = "OK" if status == 200 else "Error" + return response + + +class _ExaClient: + """Records the payloads posted to /search and replays queued responses.""" + + def __init__(self, responses): + self.responses = list(responses) + self.payloads: list[dict] = [] + + def post(self, path, json=None): # noqa: A002 - matches httpx.Client.post + self.payloads.append(json) + return self.responses.pop(0) if self.responses else _exa_response(200, {}) + + +def _with_exa(*responses): + """Patch in a fake Exa client, returning it so payloads can be asserted.""" + client = _ExaClient(responses) + return patch.object(ws, "_exa_http_client", client), client # --------------------------------------------------------------------------- -# llm_backend.get_web_search_tool +# Tool specification # --------------------------------------------------------------------------- -class TestGetWebSearchTool(unittest.TestCase): - def test_openai_tool(self): - self.assertEqual(get_web_search_tool("openai"), {"type": "web_search"}) +class TestWebSearchToolSpec(unittest.TestCase): + def test_single_provider_agnostic_function_tool(self): + """One spec for every provider — no per-provider variants any more.""" + tool = ws.get_web_search_tool() + self.assertEqual(tool["type"], "function") + self.assertEqual(tool["function"]["name"], "web_search") + + def test_schema_is_flat_and_only_requires_a_query(self): + """Nested/exotic schemas are where provider support diverges.""" + params = ws.get_web_search_tool()["function"]["parameters"] + self.assertEqual(params["required"], ["query"]) + self.assertEqual( + set(params["properties"]), {"query", "num_results", "recency_days"} + ) + for prop in params["properties"].values(): + self.assertIn(prop["type"], {"string", "integer"}) - def test_anthropic_tool(self): - tool = get_web_search_tool("anthropic") - self.assertEqual(tool["type"], "web_search_20250305") - self.assertEqual(tool["name"], "web_search") + def test_availability_tracks_the_injected_key(self): + ws.configure_exa_client("test-key") + self.assertTrue(ws.web_search_available()) + ws.configure_exa_client(None) + self.assertFalse(ws.web_search_available()) + + +# --------------------------------------------------------------------------- +# Argument coercion +# --------------------------------------------------------------------------- - def test_google_tool(self): - self.assertEqual(get_web_search_tool("google"), {"google_search": {}}) - def test_xai_tool(self): - self.assertEqual(get_web_search_tool("x-ai"), {"type": "web_search"}) +class TestArgumentCoercion(unittest.TestCase): + def test_missing_or_blank_query_is_an_unbilled_error(self): + for args in ({}, {"query": " "}, {"query": 42}): + outcome = ws.execute_web_search_call(args) + self.assertTrue(outcome.is_error, args) + self.assertFalse(outcome.billable, args) + self.assertIn("query", outcome.content) - def test_bytedance_has_no_bound_tool(self): - self.assertIsNone(get_web_search_tool("bytedance")) + def test_num_results_is_clamped_not_rejected(self): + """Models pass these as strings and floats; be forgiving.""" + self.assertEqual(ws._clamp_int("3", 6, 1, 10), 3) + self.assertEqual(ws._clamp_int(4.7, 6, 1, 10), 4) + self.assertEqual(ws._clamp_int(99, 6, 1, 10), 10) + self.assertEqual(ws._clamp_int(0, 6, 1, 10), 1) + self.assertEqual(ws._clamp_int("nonsense", 6, 1, 10), 6) + self.assertEqual(ws._clamp_int(None, 6, 1, 10), 6) + self.assertEqual(ws._clamp_int(True, 6, 1, 10), 6) # --------------------------------------------------------------------------- -# llm_backend.extract_web_search_count +# Exa request shaping # --------------------------------------------------------------------------- -class TestExtractWebSearchCount(unittest.TestCase): - def test_none_message(self): - self.assertEqual(extract_web_search_count(None), 0) +class TestExaRequestShaping(unittest.TestCase): + def test_default_request_asks_for_text_only(self): + """highlights/summary are each billed per page; text is what's used.""" + patcher, client = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + with patcher: + ws.execute_web_search_call({"query": "who won"}) + + payload = client.payloads[0] + self.assertEqual(payload["query"], "who won") + self.assertEqual(payload["type"], ws.EXA_SEARCH_TYPE) + self.assertEqual(payload["numResults"], ws.DEFAULT_NUM_RESULTS) + self.assertEqual( + payload["contents"], {"text": {"maxCharacters": ws.MAX_RESULT_CHARS}} + ) + self.assertNotIn("startPublishedDate", payload) - def test_plain_text_response_has_no_searches(self): - self.assertEqual(extract_web_search_count(AIMessage(content="hi")), 0) + def test_num_results_is_forwarded_and_capped(self): + patcher, client = _with_exa( + _exa_response(200, {"results": []}), + _exa_response(200, {"results": []}), + ) + with patcher: + ws.execute_web_search_call({"query": "q", "num_results": 3}) + ws.execute_web_search_call({"query": "q", "num_results": 500}) - def test_openai_web_search_call_blocks(self): - msg = AIMessage( - content=[ - {"type": "web_search_call", "id": "ws_1"}, - {"type": "text", "text": "answer"}, - {"type": "web_search_call", "id": "ws_2"}, - ] + self.assertEqual(client.payloads[0]["numResults"], 3) + self.assertEqual(client.payloads[1]["numResults"], ws.MAX_NUM_RESULTS) + + def test_recency_days_becomes_a_published_date_floor(self): + patcher, client = _with_exa(_exa_response(200, {"results": []})) + with patcher: + ws.execute_web_search_call({"query": "q", "recency_days": 7}) + + cutoff = client.payloads[0]["startPublishedDate"] + self.assertRegex(cutoff, r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.000Z$") + + +# --------------------------------------------------------------------------- +# Result formatting +# --------------------------------------------------------------------------- + + +class TestResultFormatting(unittest.TestCase): + def test_results_are_numbered_with_urls_and_citations(self): + patcher, _ = _with_exa( + _exa_response( + 200, + { + "results": [ + _exa_result("https://a.com", "First", "Body A"), + _exa_result("https://b.com", "Second", "Body B"), + ] + }, + ) + ) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertTrue(outcome.billable) + self.assertFalse(outcome.is_error) + self.assertIn("[1] First", outcome.content) + self.assertIn("https://a.com", outcome.content) + self.assertIn("[2] Second", outcome.content) + self.assertIn("Body B", outcome.content) + self.assertEqual( + outcome.citations, + [ + { + "title": "First", + "url": "https://a.com", + "published_date": "2026-03-04T10:00:00.000Z", + }, + { + "title": "Second", + "url": "https://b.com", + "published_date": "2026-03-04T10:00:00.000Z", + }, + ], ) - self.assertEqual(extract_web_search_count(msg), 2) - def test_anthropic_server_tool_use_blocks(self): - msg = AIMessage( - content=[ - {"type": "text", "text": "let me search"}, - {"type": "server_tool_use", "name": "web_search", "id": "srv_1"}, - {"type": "web_search_tool_result", "content": []}, - ] + def test_results_without_a_url_are_skipped(self): + patcher, _ = _with_exa( + _exa_response( + 200, + {"results": [{"title": "No URL"}, _exa_result("https://ok.com")]}, + ) + ) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertEqual([c["url"] for c in outcome.citations], ["https://ok.com"]) + + def test_total_size_is_capped_and_citations_match_what_was_shown(self): + """One verbose page must not crowd out the rest, or balloon the bill.""" + big = "x" * ws.MAX_RESULT_CHARS + results = [_exa_result(f"https://a{i}.com", f"T{i}", big) for i in range(20)] + patcher, _ = _with_exa(_exa_response(200, {"results": results})) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertLessEqual(len(outcome.content), ws.MAX_TOTAL_CHARS + 500) + self.assertLess(len(outcome.citations), 20) + # Every citation corresponds to a block actually put in front of the model. + for citation in outcome.citations: + self.assertIn(citation["url"], outcome.content) + + def test_long_excerpts_are_truncated(self): + patcher, _ = _with_exa( + _exa_response( + 200, + {"results": [_exa_result("https://a.com", "T", "y" * 9_000)]}, + ) ) - # Only the server_tool_use (the request) is billed, not the result block. - self.assertEqual(extract_web_search_count(msg), 1) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) - def test_xai_citations_counted_as_sources(self): - msg = AIMessage(content="answer") - msg.additional_kwargs = { - "citations": ["https://a.com", "https://b.com", "https://c.com"] - } - self.assertEqual(extract_web_search_count(msg), 3) + self.assertNotIn("y" * (ws.MAX_RESULT_CHARS + 1), outcome.content) + + +# --------------------------------------------------------------------------- +# Failure modes +# --------------------------------------------------------------------------- - def test_google_grounding_counts_as_one_request(self): - msg = AIMessage(content="answer") - msg.response_metadata = { - "grounding_metadata": {"web_search_queries": ["q1", "q2"]} - } - # Google bills per grounded request, not per query. - self.assertEqual(extract_web_search_count(msg), 1) + +class TestSearchFailureModes(unittest.TestCase): + def test_no_key_injected_is_a_recoverable_error(self): + with patch.object(ws, "_exa_http_client", None): + outcome = ws.execute_web_search_call({"query": "q"}) + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + self.assertIn("not configured", outcome.content) + + def test_http_error_surfaces_the_provider_detail_and_is_not_billed(self): + patcher, _ = _with_exa(_exa_response(401, {"error": "invalid api key"})) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + self.assertIn("401", outcome.content) + self.assertIn("invalid api key", outcome.content) + + def test_transport_error_is_not_billed(self): + import httpx + + client = Mock() + client.post.side_effect = httpx.ConnectError("no route") + with patch.object(ws, "_exa_http_client", client): + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + + def test_malformed_json_is_not_billed(self): + response = _exa_response(200) + response.json.side_effect = ValueError("nope") + patcher, _ = _with_exa(response) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + + def test_zero_results_is_billable_but_tells_the_model(self): + """The Exa request was consumed, and the model must not fake an answer.""" + patcher, _ = _with_exa(_exa_response(200, {"results": []})) + with patcher: + outcome = ws.execute_web_search_call({"query": "obscure thing"}) + + self.assertTrue(outcome.billable) + self.assertFalse(outcome.is_error) + self.assertIn("No web results", outcome.content) + self.assertEqual(outcome.citations, []) + + def test_reported_cost_is_captured_for_reconciliation_only(self): + patcher, _ = _with_exa( + _exa_response( + 200, + { + "results": [_exa_result("https://a.com")], + "costDollars": {"total": 0.008}, + }, + ) + ) + with patcher: + outcome = ws.execute_web_search_call({"query": "q"}) + + self.assertEqual(outcome.reported_cost_usd, 0.008) # --------------------------------------------------------------------------- -# pricing.compute_session_cost with web search +# Pricing # --------------------------------------------------------------------------- +class TestWebSearchPricing(unittest.TestCase): + def test_every_text_model_supports_search(self): + for model in ( + "gpt-4.1", + "claude-sonnet-4-5", + "gemini-2.5-flash", + "grok-4", + "seed-1.6", + "hermes-4-405b", + "glm-5.2", + ): + self.assertTrue(model_supports_web_search(model), model) + + def test_image_models_do_not(self): + for model in ("grok-2-image", "gemini-2.5-flash-image"): + self.assertFalse(model_supports_web_search(model), model) + + def test_one_flat_price_across_providers(self): + """The whole point: a search costs the same wherever it runs.""" + prices = { + model: get_web_search_price_usd(model) + for model in ( + "gpt-4.1", + "claude-sonnet-4-5", + "gemini-2.5-flash", + "grok-4", + "seed-1.6", + "hermes-4-405b", + "glm-5.2", + ) + } + self.assertEqual(set(prices.values()), {WEB_SEARCH_PRICE_USD}) + + def test_image_models_are_free(self): + self.assertEqual(get_web_search_price_usd("grok-2-image"), Decimal("0")) + + def test_unknown_model_raises(self): + with self.assertRaises(ValueError): + get_web_search_price_usd("not-a-real-model") + + def _usage(input_tokens: int = 100, output_tokens: int = 50) -> dict: return {"prompt_tokens": input_tokens, "completion_tokens": output_tokens} -def _call(usage, model, web_search_count=0, price=Decimal("0.10")): +def _cost(usage, model, web_search_count=0, price=Decimal("0.10")): feed = SimpleNamespace(get_price=lambda: price) with patch("tee_gateway.price_feed.get_price_feed", return_value=feed): return compute_session_cost(model, usage, web_search_count=web_search_count) @@ -147,50 +390,251 @@ def _call(usage, model, web_search_count=0, price=Decimal("0.10")): class TestSessionCostWithWebSearch(unittest.TestCase): def test_web_search_increases_cost(self): - base = _call(_usage(), "gpt-4.1") - searched = _call(_usage(), "gpt-4.1", web_search_count=2) + base = _cost(_usage(), "gpt-4.1") + searched = _cost(_usage(), "gpt-4.1", web_search_count=2) self.assertIsInstance(base, SessionCost) self.assertIsInstance(searched, SessionCost) self.assertGreater(searched.cost_opg, base.cost_opg) - def test_web_search_surcharge_amount(self): - """Two openai searches at $0.01 each add $0.02 of USD cost.""" - base = _call(_usage(), "gpt-4.1") - searched = _call(_usage(), "gpt-4.1", web_search_count=2) - # cost_usd is reconciled from rounded OPG; compare via the underlying math. - # At price $0.10/OPG, $0.02 surcharge ≈ 0.2 OPG = 2e17 smallest units. - delta_opg = searched.cost_opg - base.cost_opg + def test_surcharge_is_exactly_searches_times_the_flat_rate(self): + """The client-verifiable property: surcharge == searches * rate.""" + base = _cost(_usage(), "gpt-4.1") + searched = _cost(_usage(), "gpt-4.1", web_search_count=3) scale = Decimal(10) ** 18 - delta_usd = (Decimal(delta_opg) / scale) * Decimal("0.10") - # Allow a tiny rounding tolerance from ceiling rounding on each call. - self.assertAlmostEqual(delta_usd, Decimal("0.02"), places=6) + delta_usd = (Decimal(searched.cost_opg - base.cost_opg) / scale) * Decimal( + "0.10" + ) + self.assertAlmostEqual(delta_usd, 3 * WEB_SEARCH_PRICE_USD, places=6) + + def test_same_surcharge_on_a_different_provider(self): + deltas = [] + for model in ("gpt-4.1", "seed-1.6"): + base = _cost(_usage(), model) + searched = _cost(_usage(), model, web_search_count=2) + deltas.append(searched.cost_opg - base.cost_opg) + self.assertEqual(deltas[0], deltas[1]) def test_zero_searches_matches_no_web_search(self): - a = _call(_usage(), "gpt-4.1", web_search_count=0) - b = _call(_usage(), "gpt-4.1") + a = _cost(_usage(), "gpt-4.1", web_search_count=0) + b = _cost(_usage(), "gpt-4.1") self.assertEqual(a.cost_opg, b.cost_opg) - def test_unsupported_provider_not_charged_for_search(self): - # Even if a count slips through, bytedance price is 0 -> no surcharge. - base = _call(_usage(), "seed-1.6") - searched = _call(_usage(), "seed-1.6", web_search_count=5) - self.assertEqual(searched.cost_opg, base.cost_opg) - # --------------------------------------------------------------------------- -# chat_controller integration +# search_loop # --------------------------------------------------------------------------- -class _MockResponse: - def __init__(self, content="", tool_calls=None, usage=None): - self.content = content - self.tool_calls = tool_calls or [] - self.usage_metadata = usage or { - "input_tokens": 10, - "output_tokens": 5, - "total_tokens": 15, - } +def _search_call(query="q", call_id="call_1"): + return {"name": "web_search", "args": {"query": query}, "id": call_id} + + +def _client_call(name="get_weather", call_id="call_2"): + return {"name": name, "args": {}, "id": call_id} + + +def _ai(content="", tool_calls=None, tokens=(10, 5)): + message = AIMessage(content=content, tool_calls=tool_calls or []) + message.usage_metadata = { + "input_tokens": tokens[0], + "output_tokens": tokens[1], + "total_tokens": sum(tokens), + } + return message + + +class _ScriptedModel: + """Returns queued AIMessages, recording the messages it was invoked with.""" + + def __init__(self, *responses): + self.responses = list(responses) + self.calls: list[list] = [] + + def invoke(self, messages): + self.calls.append(list(messages)) + return self.responses.pop(0) + + +class TestSplitToolCalls(unittest.TestCase): + def test_partitions_ours_from_the_callers(self): + ours, theirs = split_tool_calls([_search_call(), _client_call()]) + self.assertEqual([c["name"] for c in ours], ["web_search"]) + self.assertEqual([c["name"] for c in theirs], ["get_weather"]) + + def test_handles_none_and_empty(self): + self.assertEqual(split_tool_calls(None), ([], [])) + self.assertEqual(split_tool_calls([]), ([], [])) + + +class TestSearchLoopState(unittest.TestCase): + def test_usage_sums_across_rounds(self): + state = SearchLoopState() + state.add_usage({"prompt_tokens": 100, "completion_tokens": 10}) + state.add_usage({"prompt_tokens": 400, "completion_tokens": 20}) + assert state.usage is not None + self.assertEqual(state.usage["prompt_tokens"], 500) + self.assertEqual(state.usage["completion_tokens"], 30) + + def test_usage_stays_none_when_nothing_reported(self): + state = SearchLoopState() + state.add_usage(None) + self.assertIsNone(state.usage) + + def test_citations_are_deduped_by_url(self): + state = SearchLoopState() + state.add_citations([{"title": "A", "url": "https://a.com"}]) + state.add_citations( + [ + {"title": "A again", "url": "https://a.com"}, + {"title": "B", "url": "https://b.com"}, + ] + ) + self.assertEqual( + [c["url"] for c in state.citations], ["https://a.com", "https://b.com"] + ) + + +class TestExecuteSearchCalls(unittest.TestCase): + def test_builds_tool_messages_and_counts_billable_searches(self): + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + state = SearchLoopState() + with patcher: + messages = execute_search_calls([_search_call("news", "abc")], state) + + self.assertEqual(len(messages), 1) + self.assertIsInstance(messages[0], ToolMessage) + self.assertEqual(messages[0].tool_call_id, "abc") + self.assertEqual(messages[0].name, "web_search") + self.assertEqual(messages[0].status, "success") + self.assertEqual(state.search_count, 1) + self.assertEqual(len(state.citations), 1) + + def test_failed_search_yields_an_error_tool_message_and_no_charge(self): + patcher, _ = _with_exa(_exa_response(500, {"error": "boom"})) + state = SearchLoopState() + with patcher: + messages = execute_search_calls([_search_call()], state) + + self.assertEqual(messages[0].status, "error") + self.assertEqual(state.search_count, 0) + + def test_status_callback_receives_each_query(self): + patcher, _ = _with_exa( + _exa_response(200, {"results": []}), _exa_response(200, {"results": []}) + ) + seen: list[str] = [] + with patcher: + execute_search_calls( + [_search_call("first", "1"), _search_call("second", "2")], + SearchLoopState(), + on_search=seen.append, + ) + self.assertEqual(seen, ["first", "second"]) + + def test_a_throwing_status_callback_does_not_break_the_search(self): + patcher, _ = _with_exa(_exa_response(200, {"results": []})) + state = SearchLoopState() + with patcher: + messages = execute_search_calls( + [_search_call()], + state, + on_search=Mock(side_effect=RuntimeError("ui gone")), + ) + self.assertEqual(len(messages), 1) + self.assertEqual(state.search_count, 1) + + +class TestRunSearchLoop(unittest.TestCase): + def test_plain_answer_returns_immediately(self): + model = _ScriptedModel(_ai("just an answer")) + state = SearchLoopState() + result = run_search_loop(model, model, [], state) + self.assertEqual(result.content, "just an answer") + self.assertEqual(state.search_count, 0) + self.assertEqual(state.rounds, 1) + + def test_searches_then_answers_and_feeds_results_back(self): + model = _ScriptedModel( + _ai("", [_search_call("og price")]), + _ai("The answer, with sources."), + ) + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + state = SearchLoopState() + messages: list = [] + with patcher: + result = run_search_loop(model, model, messages, state) + + self.assertEqual(result.content, "The answer, with sources.") + self.assertEqual(state.search_count, 1) + self.assertEqual(state.rounds, 2) + # The second invocation saw the assistant turn plus the search results. + second_round = model.calls[1] + self.assertIsInstance(second_round[-1], ToolMessage) + self.assertIn("https://a.com", second_round[-1].content) + + def test_every_round_is_billed_not_just_the_last(self): + """Each round re-sends the conversation; the caller pays for all of it.""" + model = _ScriptedModel( + _ai("", [_search_call("a")], tokens=(100, 10)), + _ai("", [_search_call("b")], tokens=(600, 12)), + _ai("done", tokens=(1200, 40)), + ) + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}), + _exa_response(200, {"results": [_exa_result("https://b.com")]}), + ) + state = SearchLoopState() + with patcher: + run_search_loop(model, model, [], state) + + assert state.usage is not None + self.assertEqual(state.usage["prompt_tokens"], 1900) + self.assertEqual(state.usage["completion_tokens"], 62) + self.assertEqual(state.search_count, 2) + + def test_client_tool_call_is_terminal(self): + model = _ScriptedModel(_ai("", [_client_call()])) + state = SearchLoopState() + result = run_search_loop(model, model, [], state) + self.assertEqual([c["name"] for c in result.tool_calls], ["get_weather"]) + self.assertEqual(state.rounds, 1) + + def test_round_cap_forces_an_answer_with_the_search_tool_unbound(self): + """A model that keeps searching must still terminate.""" + searching = _ScriptedModel( + *[_ai("", [_search_call(f"q{i}")]) for i in range(MAX_SEARCH_ROUNDS)] + ) + answering = _ScriptedModel(_ai("forced answer")) + patcher, _ = _with_exa( + *[ + _exa_response(200, {"results": [_exa_result(f"https://a{i}.com")]}) + for i in range(MAX_SEARCH_ROUNDS) + ] + ) + state = SearchLoopState() + with patcher: + result = run_search_loop(searching, answering, [], state) + + self.assertEqual(result.content, "forced answer") + self.assertEqual(state.search_count, MAX_SEARCH_ROUNDS) + self.assertEqual(state.rounds, MAX_SEARCH_ROUNDS + 1) + # The final round went to the model without the search tool bound. + self.assertEqual(len(answering.calls), 1) + + def test_mixed_turn_drops_our_calls_and_keeps_the_callers(self): + message = _ai("", [_search_call(), _client_call()]) + self.assertEqual( + [c["name"] for c in strip_search_tool_calls(message)], ["get_weather"] + ) + + +# --------------------------------------------------------------------------- +# chat_controller integration +# --------------------------------------------------------------------------- def _mock_tee_keys(): @@ -200,78 +644,342 @@ def _mock_tee_keys(): return tee -class TestChatControllerWebSearch(unittest.TestCase): - @patch("tee_gateway.controllers.chat_controller.compute_session_cost") - @patch("tee_gateway.controllers.chat_controller.get_tee_keys") - @patch("tee_gateway.controllers.chat_controller.get_chat_model_cached") - @patch("tee_gateway.controllers.chat_controller.connexion") - def test_web_search_flag_binds_tool_and_bills( - self, mock_connexion, mock_get_model, mock_get_tee_keys, mock_cost - ): - mock_connexion.request.is_json = True - mock_connexion.request.get_json.return_value = { +class _ControllerHarness(unittest.TestCase): + """Shared patching for the chat_controller tests.""" + + def setUp(self): + self.patchers = [ + patch("tee_gateway.controllers.chat_controller.compute_session_cost"), + patch("tee_gateway.controllers.chat_controller.get_tee_keys"), + patch("tee_gateway.controllers.chat_controller.get_chat_model_cached"), + patch("tee_gateway.controllers.chat_controller.connexion"), + patch( + "tee_gateway.controllers.chat_controller.web_search_available", + return_value=True, + ), + ] + ( + self.cost, + self.tee, + self.get_model, + self.connexion, + self.available, + ) = [p.start() for p in self.patchers] + self.addCleanup(lambda: [p.stop() for p in self.patchers]) + self.tee.return_value = _mock_tee_keys() + self.cost.return_value = None + + def request(self, **overrides): + body = { "model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "latest news?"}], - "web_search": True, "stream": False, } + body.update(overrides) + self.connexion.request.is_json = True + self.connexion.request.get_json.return_value = body - # Anthropic response with a server_tool_use web_search block. - response = _MockResponse( - content=[ - {"type": "text", "text": "Here is the news."}, - {"type": "server_tool_use", "name": "web_search", "id": "srv_1"}, - ] - ) + +class TestChatControllerNonStreaming(_ControllerHarness): + def test_flag_binds_our_function_tool_on_an_anthropic_model(self): + self.request(web_search=True) model = Mock() - model.invoke.return_value = response + model.invoke.return_value = _ai("Here is the news.") model.bind_tools.return_value = model - mock_get_model.return_value = model - mock_get_tee_keys.return_value = _mock_tee_keys() - mock_cost.return_value = None + self.get_model.return_value = model - result = create_chat_completion(None) + patcher, _ = _with_exa() + with patcher: + result = create_chat_completion(None) - # Model must be constructed with web_search=True. - self.assertTrue(mock_get_model.call_args.kwargs["web_search"]) - # The anthropic web search tool must be bound. bound = model.bind_tools.call_args[0][0] - self.assertTrue( - any( - isinstance(t, dict) and t.get("type") == "web_search_20250305" - for t in bound - ) + self.assertEqual( + [t["function"]["name"] for t in bound if isinstance(t, dict)], + ["web_search"], ) - # Billing must receive the detected search count (1 server_tool_use). - self.assertEqual(mock_cost.call_args.kwargs["web_search_count"], 1) + # No provider-native tool types any more. + for tool in bound: + self.assertEqual(tool.get("type"), "function") self.assertIn("choices", result) - @patch("tee_gateway.controllers.chat_controller.compute_session_cost") - @patch("tee_gateway.controllers.chat_controller.get_tee_keys") - @patch("tee_gateway.controllers.chat_controller.get_chat_model_cached") - @patch("tee_gateway.controllers.chat_controller.connexion") - def test_no_web_search_does_not_bind_or_bill_search( - self, mock_connexion, mock_get_model, mock_get_tee_keys, mock_cost - ): - mock_connexion.request.is_json = True - mock_connexion.request.get_json.return_value = { - "model": "gpt-4.1", - "messages": [{"role": "user", "content": "hello"}], - "stream": False, - } + def test_search_runs_in_enclave_and_is_never_handed_to_the_client(self): + self.request(web_search=True) + model = Mock() + model.bind_tools.return_value = model + model.invoke.side_effect = [ + _ai("", [_search_call("og token price")]), + _ai("It trades at $X."), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://coin.com")]}) + ) + with patcher: + result = create_chat_completion(None) + + choice = result["choices"][0] + self.assertEqual(choice["finish_reason"], "stop") + self.assertNotIn("tool_calls", choice["message"]) + self.assertEqual(choice["message"]["content"], "It trades at $X.") + self.assertEqual( + [c["url"] for c in choice["message"]["citations"]], ["https://coin.com"] + ) + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) + + def test_bytedance_can_search_too(self): + """The case the old native-search implementation could not serve at all.""" + self.request(model="seed-1.6", web_search=True) + model = Mock() + model.bind_tools.return_value = model + model.invoke.side_effect = [ + _ai("", [_search_call("q")]), + _ai("answer"), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + with patcher: + result = create_chat_completion(None) + + self.assertEqual(result["choices"][0]["message"]["content"], "answer") + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) + + def test_billed_usage_is_the_sum_over_rounds(self): + self.request(web_search=True) + model = Mock() + model.bind_tools.return_value = model + model.invoke.side_effect = [ + _ai("", [_search_call("q")], tokens=(100, 10)), + _ai("answer", tokens=(700, 30)), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + with patcher: + result = create_chat_completion(None) + + self.assertEqual(result["usage"]["prompt_tokens"], 800) + self.assertEqual(result["usage"]["completion_tokens"], 40) + + def test_client_tools_still_come_back_for_the_client_to_run(self): + self.request( + web_search=True, + tools=[ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {}}, + } + ], + ) model = Mock() - model.invoke.return_value = _MockResponse(content="hi") model.bind_tools.return_value = model - mock_get_model.return_value = model - mock_get_tee_keys.return_value = _mock_tee_keys() - mock_cost.return_value = None + model.invoke.return_value = _ai("", [_client_call()]) + self.get_model.return_value = model + + patcher, _ = _with_exa() + with patcher: + result = create_chat_completion(None) + + choice = result["choices"][0] + self.assertEqual(choice["finish_reason"], "tool_calls") + self.assertEqual( + [tc["function"]["name"] for tc in choice["message"]["tool_calls"]], + ["get_weather"], + ) + + def test_no_flag_binds_nothing_and_bills_no_search(self): + self.request(model="gpt-4.1", messages=[{"role": "user", "content": "hi"}]) + model = Mock() + model.invoke.return_value = _ai("hi") + model.bind_tools.return_value = model + self.get_model.return_value = model create_chat_completion(None) - self.assertFalse(mock_get_model.call_args.kwargs["web_search"]) - # No tools and no web search -> bind_tools must not be called. model.bind_tools.assert_not_called() - self.assertEqual(mock_cost.call_args.kwargs["web_search_count"], 0) + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) + + def test_missing_exa_key_answers_without_searching(self): + """Better a plain answer than a tool that always fails.""" + self.available.return_value = False + self.request(web_search=True) + model = Mock() + model.invoke.return_value = _ai("answer from memory") + model.bind_tools.return_value = model + self.get_model.return_value = model + + result = create_chat_completion(None) + + model.bind_tools.assert_not_called() + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) + self.assertEqual( + result["choices"][0]["message"]["content"], "answer from memory" + ) + + +def _chunk(content="", tool_call_chunks=None, usage=None): + message = AIMessageChunk(content=content, tool_call_chunks=tool_call_chunks or []) + if usage: + message.usage_metadata = { + "input_tokens": usage[0], + "output_tokens": usage[1], + "total_tokens": sum(usage), + } + return message + + +def _search_chunk(query="q", call_id="call_1"): + return _chunk( + tool_call_chunks=[ + { + "name": "web_search", + "args": json.dumps({"query": query}), + "id": call_id, + "index": 0, + } + ] + ) + + +def _sse_frames(response): + """Parse a Flask SSE response into the list of JSON data frames.""" + raw = "".join( + part.decode("utf-8") if isinstance(part, bytes) else part + for part in response.response + ) + frames = [] + for line in raw.split("\n\n"): + line = line.strip() + if line.startswith("data: ") and line != "data: [DONE]": + frames.append(json.loads(line[len("data: ") :])) + return frames + + +class TestChatControllerStreaming(_ControllerHarness): + def test_search_rounds_are_invisible_to_the_client(self): + """The client sees status, then the answer — never our tool calls.""" + self.request(web_search=True, stream=True) + model = Mock() + model.bind_tools.return_value = model + model.stream.side_effect = [ + iter([_search_chunk("og price"), _chunk(usage=(100, 10))]), + iter([_chunk("It "), _chunk("trades."), _chunk(usage=(700, 20))]), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://coin.com")]}) + ) + with patcher: + frames = _sse_frames(create_chat_completion(None)) + + # A status frame naming the query reached the client. + statuses = [f["web_search"] for f in frames if "web_search" in f] + self.assertEqual(statuses, [{"status": "searching", "query": "og price"}]) + + # No tool_calls delta was ever forwarded. + for frame in frames: + delta = frame.get("choices", [{}])[0].get("delta", {}) + self.assertNotIn("tool_calls", delta) + + final = frames[-1] + self.assertEqual(final["choices"][0]["finish_reason"], "stop") + self.assertEqual([c["url"] for c in final["citations"]], ["https://coin.com"]) + # Both rounds are billed. + self.assertEqual(final["usage"]["prompt_tokens"], 800) + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) + + def test_streamed_text_is_forwarded_and_signed(self): + self.request(web_search=True, stream=True) + model = Mock() + model.bind_tools.return_value = model + model.stream.side_effect = [ + iter([_search_chunk(), _chunk(usage=(10, 1))]), + iter([_chunk("Hello "), _chunk("world"), _chunk(usage=(20, 2))]), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + with patcher: + frames = _sse_frames(create_chat_completion(None)) + + text = "".join( + f["choices"][0]["delta"].get("content", "") + for f in frames + if f.get("choices") + ) + self.assertEqual(text, "Hello world") + self.assertIn("tee_signature", frames[-1]) + + def test_client_tool_calls_still_stream_through(self): + self.request( + web_search=True, + stream=True, + tools=[ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {}}, + } + ], + ) + model = Mock() + model.bind_tools.return_value = model + model.stream.side_effect = [ + iter( + [ + _chunk( + tool_call_chunks=[ + { + "name": "get_weather", + "args": "{}", + "id": "c1", + "index": 0, + } + ] + ), + _chunk(usage=(10, 1)), + ] + ) + ] + self.get_model.return_value = model + + patcher, _ = _with_exa() + with patcher: + frames = _sse_frames(create_chat_completion(None)) + + names = [ + tc["function"]["name"] + for f in frames + for tc in f.get("choices", [{}])[0].get("delta", {}).get("tool_calls", []) + if tc.get("function", {}).get("name") + ] + self.assertEqual(names, ["get_weather"]) + self.assertEqual(frames[-1]["choices"][0]["finish_reason"], "tool_calls") + + def test_failed_search_still_produces_an_answer(self): + self.request(web_search=True, stream=True) + model = Mock() + model.bind_tools.return_value = model + model.stream.side_effect = [ + iter([_search_chunk(), _chunk(usage=(10, 1))]), + iter([_chunk("Couldn't verify that."), _chunk(usage=(20, 2))]), + ] + self.get_model.return_value = model + + patcher, _ = _with_exa(_exa_response(503, {"error": "unavailable"})) + with patcher: + frames = _sse_frames(create_chat_completion(None)) + + self.assertEqual(frames[-1]["choices"][0]["finish_reason"], "stop") + self.assertNotIn("citations", frames[-1]) + self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) if __name__ == "__main__": diff --git a/tee_gateway/web_search.py b/tee_gateway/web_search.py new file mode 100644 index 0000000..f5572b6 --- /dev/null +++ b/tee_gateway/web_search.py @@ -0,0 +1,426 @@ +"""In-enclave web search, backed by Exa. + +This module replaces the four provider-native web search tools that used to back +the ``web_search`` request flag (OpenAI's Responses-API ``web_search``, +Anthropic's ``web_search_20250305``, Gemini's ``google_search`` grounding, and +xAI's Responses-API ``web_search``). Those differed in every dimension that +matters: which models could use them, what the results looked like, what the +response reported back, and what a "search" cost ($0.01–$0.035, with xAI billing +per *citation*). Models on providers without one (ByteDance, Nous, Z.ai) simply +could not search at all. + +Instead the gateway advertises ONE ordinary function tool — see +``get_web_search_tool`` — and executes it itself, inside the enclave, against +Exa. Consequences worth stating plainly: + + * It works on every model that can call a function, which is every non-image + model in the registry. There is no per-provider capability matrix. + * Results, excerpt sizes, and citations are identical across models, so answer + quality stops depending on whose search backend the model happened to ship. + * A search is one flat price on every model (``WEB_SEARCH_PRICE_USD``), so + clients can verify the surcharge as ``searches * price`` rather than + reverse-engineering a provider's billable unit. + * The query never leaves the TEE except to Exa. Nothing about the search is + visible to the LLM provider beyond the result text the model is shown, and + nothing is visible to the gateway operator at all. + +The Exa API key is injected at runtime via ``POST /v1/keys`` like every provider +key; it is never baked into the image. +""" + +import logging +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import Any, Optional + +import httpx + +logger = logging.getLogger(__name__) + +# Exa's search API. Single fixed host — this client is never pointed at a +# caller-supplied URL, so it needs none of the SSRF guarding that the image +# fetcher in image_generation.py does. +EXA_BASE_URL = "https://api.exa.ai" + +# Exa search mode. "auto" lets Exa pick between its neural and keyword indices +# per query, which behaves best on the open-ended questions chat users ask. +# The "deep*" modes cost ~2x and add seconds of latency; not worth it inside an +# interactive chat turn. +EXA_SEARCH_TYPE = "auto" + +# The tool name the model sees. Also the name the controllers match on to decide +# a tool call is ours to execute rather than the client's. +WEB_SEARCH_TOOL_NAME = "web_search" + +DEFAULT_NUM_RESULTS = 6 +MAX_NUM_RESULTS = 10 + +# Per-result excerpt cap, and a cap on the whole formatted block. Search results +# are the single largest thing we inject into a prompt, and the caller pays for +# every one of those input tokens on every subsequent round of the tool loop — +# so this bounds both context pressure and cost. ~12k chars ≈ 3k tokens. +MAX_RESULT_CHARS = 1_500 +MAX_TOTAL_CHARS = 12_000 + +# Exa's ceiling on the published-date filter we map `recency_days` onto. +MAX_RECENCY_DAYS = 3_650 + +_EXA_TIMEOUT = httpx.Timeout(timeout=30.0, connect=10.0, read=25.0, write=10.0) +_EXA_LIMITS = httpx.Limits(max_keepalive_connections=5, max_connections=20) + +_exa_http_client: Optional[httpx.Client] = None + + +def configure_exa_client(api_key: Optional[str]) -> None: + """Build (or tear down) the shared Exa HTTP client after key injection. + + Called from ``llm_backend.set_provider_config`` alongside the provider + clients. Passing an empty key leaves web search unavailable, which the + controllers surface as a clear error rather than silently answering without + searching. + """ + global _exa_http_client + + old = _exa_http_client + if api_key: + _exa_http_client = httpx.Client( + base_url=EXA_BASE_URL, + headers={ + "x-api-key": api_key, + "content-type": "application/json", + }, + timeout=_EXA_TIMEOUT, + limits=_EXA_LIMITS, + follow_redirects=False, + ) + else: + _exa_http_client = None + + if old is not None: + old.close() + + +def web_search_available() -> bool: + """Whether an Exa key was injected, i.e. whether searches can run.""" + return _exa_http_client is not None + + +# --------------------------------------------------------------------------- +# Tool specification +# --------------------------------------------------------------------------- + + +def get_web_search_tool() -> dict[str, Any]: + """The ``web_search`` function tool, in OpenAI function-calling format. + + One spec for every provider: langchain converts this to each provider's own + tool format on ``bind_tools``. The schema is deliberately flat — three + scalar parameters, one required — because nested objects, unions, and + ``additionalProperties`` are exactly where the providers' function-calling + schema subsets diverge (Gemini's is the narrowest). Flat and boring is what + makes this work identically on all of them. + """ + return { + "type": "function", + "function": { + "name": WEB_SEARCH_TOOL_NAME, + "description": ( + "Search the live web and get back ranked results with an " + "excerpt of each page. Use this whenever the answer depends on " + "information you may not have: current events, anything after " + "your training cutoff, prices, releases, documentation, or any " + "specific fact you are not confident about. Prefer searching " + "over guessing. Write a focused natural-language query rather " + "than keywords, and call this again with a refined query if the " + "first set of results does not answer the question." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "What to search for, as a focused natural-language query." + ), + }, + "num_results": { + "type": "integer", + "description": ( + f"How many results to return, 1-{MAX_NUM_RESULTS}. " + f"Defaults to {DEFAULT_NUM_RESULTS}. Ask for more " + "only when the question needs broad coverage." + ), + }, + "recency_days": { + "type": "integer", + "description": ( + "Only return pages published within this many days. " + "Omit unless the question is genuinely " + "time-sensitive — it discards older pages that are " + "often the best sources." + ), + }, + }, + "required": ["query"], + }, + }, + } + + +# --------------------------------------------------------------------------- +# Execution +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class WebSearchOutcome: + """The result of one ``web_search`` tool call. + + ``content`` is the text handed back to the model as the tool result. + ``citations`` are surfaced to the client out-of-band (like generated + images) so a UI can show its sources. ``billable`` is False for failures + and for empty/invalid calls, so a caller is never charged for a search that + produced nothing. + """ + + content: str + citations: list[dict[str, str]] = field(default_factory=list) + billable: bool = False + is_error: bool = False + # Exa's own reported price for this call. Logged for margin reconciliation + # only — the client is billed the flat published rate, never this number, + # so settlement never depends on a third party's self-reported figure. + reported_cost_usd: Optional[float] = None + + +def execute_web_search_call(args: dict[str, Any]) -> WebSearchOutcome: + """Run one ``web_search`` tool call from its (already-parsed) arguments. + + Never raises: every failure mode comes back as an error outcome whose + ``content`` reads as an instruction to the model, so a flaky search degrades + into "the model was told the search failed" rather than a dead request. + """ + query = args.get("query") + if not isinstance(query, str) or not query.strip(): + return WebSearchOutcome( + content=( + "Web search error: a non-empty `query` string is required. " + "Call the tool again with a query." + ), + is_error=True, + ) + query = query.strip() + + num_results = _clamp_int( + args.get("num_results"), DEFAULT_NUM_RESULTS, 1, MAX_NUM_RESULTS + ) + recency_days = _clamp_int(args.get("recency_days"), 0, 0, MAX_RECENCY_DAYS) + + return run_web_search(query, num_results, recency_days or None) + + +def run_web_search( + query: str, + num_results: int = DEFAULT_NUM_RESULTS, + recency_days: Optional[int] = None, +) -> WebSearchOutcome: + """Query Exa and format the results for model consumption.""" + client = _exa_http_client + if client is None: + logger.error("web_search requested but no Exa API key was injected") + return WebSearchOutcome( + content=( + "Web search error: search is not configured on this gateway. " + "Answer from your own knowledge and say that you could not " + "verify it against the web." + ), + is_error=True, + ) + + payload: dict[str, Any] = { + "query": query, + "type": EXA_SEARCH_TYPE, + "numResults": num_results, + # Ask for page text only. `highlights` and `summary` are each billed as + # another content type per page, and the excerpt is what the model + # actually reasons over. + "contents": {"text": {"maxCharacters": MAX_RESULT_CHARS}}, + } + if recency_days: + cutoff = datetime.now(timezone.utc) - timedelta(days=recency_days) + payload["startPublishedDate"] = cutoff.strftime("%Y-%m-%dT%H:%M:%S.000Z") + + try: + response = client.post("/search", json=payload) + except httpx.HTTPError as exc: + logger.warning("Exa search transport error for %r: %s", query, exc) + return WebSearchOutcome( + content=f"Web search error: could not reach the search service ({exc}).", + is_error=True, + ) + + if response.status_code != 200: + # Surface the provider's own message: Exa's 4xx bodies name the actual + # problem (bad key, quota, malformed filter), which is what makes this + # debuggable from enclave logs where we cannot attach a debugger. + detail = _error_detail(response) + logger.warning( + "Exa search failed for %r: HTTP %s %s", query, response.status_code, detail + ) + return WebSearchOutcome( + content=( + f"Web search error: the search service returned " + f"HTTP {response.status_code} ({detail})." + ), + is_error=True, + ) + + try: + body = response.json() + except ValueError as exc: + logger.warning("Exa returned non-JSON for %r: %s", query, exc) + return WebSearchOutcome( + content="Web search error: the search service returned a malformed response.", + is_error=True, + ) + + results = body.get("results") + if not isinstance(results, list): + results = [] + + cost = body.get("costDollars") + reported_cost = ( + float(cost["total"]) + if isinstance(cost, dict) and isinstance(cost.get("total"), (int, float)) + else None + ) + + if not results: + # A search that ran but matched nothing still consumed an Exa request, so + # it is billable — and the model needs to be told, or it will silently + # answer as if it had searched successfully. + logger.info("Exa search for %r returned no results", query) + return WebSearchOutcome( + content=( + f'No web results were found for "{query}". Try a broader or ' + "differently-worded query, or tell the user you could not find " + "anything." + ), + billable=True, + reported_cost_usd=reported_cost, + ) + + content, citations = _format_results(query, results) + logger.info( + "Exa search ok — query=%r results=%d chars=%d reported_cost_usd=%s", + query, + len(citations), + len(content), + reported_cost, + ) + return WebSearchOutcome( + content=content, + citations=citations, + billable=True, + reported_cost_usd=reported_cost, + ) + + +# --------------------------------------------------------------------------- +# Formatting helpers +# --------------------------------------------------------------------------- + + +def _format_results(query: str, results: list[Any]) -> tuple[str, list[dict[str, str]]]: + """Render Exa results as numbered, citable blocks for the model. + + Stops adding results once MAX_TOTAL_CHARS is reached so one verbose page + cannot crowd out the rest (and cannot balloon the input-token bill on every + later round of the tool loop). Citations are collected only for results that + actually made it into the text, so what the UI shows as a source is exactly + what the model was shown. + """ + header = f'Web search results for "{query}":\n' + blocks: list[str] = [] + citations: list[dict[str, str]] = [] + used = len(header) + + for item in results: + if not isinstance(item, dict): + continue + url = item.get("url") + if not isinstance(url, str) or not url: + continue + + title = _clean_str(item.get("title")) or url + published = _clean_str(item.get("publishedDate")) + author = _clean_str(item.get("author")) + text = _clean_str(item.get("text")) + if len(text) > MAX_RESULT_CHARS: + text = text[:MAX_RESULT_CHARS].rstrip() + "…" + + index = len(citations) + 1 + lines = [f"[{index}] {title}", f"URL: {url}"] + if published: + lines.append(f"Published: {published[:10]}") + if author: + lines.append(f"Author: {author}") + if text: + lines.append(text) + block = "\n".join(lines) + + if citations and used + len(block) + 2 > MAX_TOTAL_CHARS: + break + + blocks.append(block) + used += len(block) + 2 + citation: dict[str, str] = {"title": title, "url": url} + if published: + citation["published_date"] = published + citations.append(citation) + + if not blocks: + return ( + f'No usable web results were found for "{query}".', + [], + ) + + footer = ( + "\nCite the sources you relied on by their URL, and say so plainly if " + "the results do not answer the question." + ) + return (header + "\n" + "\n\n".join(blocks) + "\n" + footer, citations) + + +def _clean_str(value: Any) -> str: + """Coerce an Exa field to a stripped string (fields are nullable).""" + return value.strip() if isinstance(value, str) else "" + + +def _clamp_int(value: Any, default: int, low: int, high: int) -> int: + """Coerce a model-supplied number to an int in [low, high]. + + Models pass these as strings and floats often enough that being strict here + would turn a usable call into a retry loop. + """ + if isinstance(value, bool) or value is None: + return default + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return max(low, min(high, parsed)) + + +def _error_detail(response: httpx.Response) -> str: + """Best-effort human-readable detail from an Exa error response.""" + try: + body = response.json() + except ValueError: + return response.text[:200].strip() or response.reason_phrase + if isinstance(body, dict): + for key in ("error", "message", "detail"): + value = body.get(key) + if isinstance(value, str) and value: + return value[:200] + return str(body)[:200] From 83177dbe5cf49eb55463944c30c127f5f0a6799e Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 12:30:49 +0000 Subject: [PATCH 2/3] Serve web search as a dedicated /v1/web_search endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the in-enclave search tool loop with a plain endpoint the client's own tool loop calls. The gateway no longer binds a web_search tool, runs search rounds, or accumulates cross-round usage — chat and completions are back to stateless request/response, and the web_search request flag is a deprecated no-op (still accepted and still hashed, so old clients' request signatures verify unchanged). - POST /v1/web_search: runs one Exa search inside the enclave, returns model-ready content plus structured citations, signed with the standard tee_* fields (request hash over the canonical JSON body, output hash over content), billed via x402 at the flat WEB_SEARCH_PRICE_USD per search. Failures (400/502/503) carry no cost block and are never settled. - OHTTP: the sealed payload's `endpoint` field routes the inner request ("web_search" -> /v1/web_search); absent still means chat, so existing OHTTP clients are unaffected and relay billing reuses the existing cost-header/billing-frame channel. - Removed search_loop.py and the streaming buffering/search-frame machinery from the chat controller; compute_session_cost no longer takes a search count (search billing left the chat path entirely). - web_search.py slims to the Exa client, execution, and formatting; the tool spec moves to the client. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01V9SCMV8LhnpZ21pzBLjzxQ --- CLAUDE.md | 66 +- README.md | 47 +- tee_gateway/__main__.py | 48 +- tee_gateway/controllers/chat_controller.py | 380 +----- .../controllers/completions_controller.py | 49 +- tee_gateway/controllers/ohttp_controller.py | 29 +- .../controllers/web_search_controller.py | 98 ++ tee_gateway/definitions.py | 6 + tee_gateway/llm_backend.py | 6 +- tee_gateway/model_registry.py | 54 +- tee_gateway/openapi/openapi.yaml | 36 +- tee_gateway/pricing.py | 70 +- tee_gateway/search_loop.py | 223 ---- tee_gateway/test/test_web_search.py | 1121 +++++------------ tee_gateway/web_search.py | 109 +- 15 files changed, 712 insertions(+), 1630 deletions(-) create mode 100644 tee_gateway/controllers/web_search_controller.py delete mode 100644 tee_gateway/search_loop.py diff --git a/CLAUDE.md b/CLAUDE.md index bfd946b..be08dc5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,8 +12,7 @@ The repo must provide a stable AWS Nitro PCR when the code doesn't change in ord │ ├── llm_backend.py # LLM provider routing via LangChain, HTTP client management │ ├── image_generation.py # Endpoint-based image gen (/images/generations): request shaping, URL→inline-bytes, signed responses │ ├── tee_manager.py # TEE key generation, nitriding registration, response signing -│ ├── web_search.py # In-enclave web search: Exa client, `web_search` tool spec, result formatting -│ ├── search_loop.py # Server-side tool loop that answers web_search calls in-enclave +│ ├── web_search.py # In-enclave web search: Exa client, execution, result formatting │ ├── model_registry.py # Model config and per-token pricing │ ├── definitions.py # On-chain addresses, network IDs, payment amounts │ ├── facilitator_api.py # x402 facilitator API client @@ -72,8 +71,8 @@ API keys (injected at runtime via POST /v1/keys — do NOT bake into the image): - `NOUS_API_KEY` (Nous Research / Nous Portal; injected as `nous_api_key`) - `ZAI_API_KEY` (Z.ai Model API; injected as `zai_api_key`) - `EXA_API_KEY` (Exa search; injected as `exa_api_key`) — backs the in-enclave - `web_search` tool, not an LLM provider. Without it the `web_search` flag is a - no-op and `/health` reports `web_search_enabled: false`. + `/v1/web_search` endpoint, not an LLM provider. Without it the endpoint + returns 503 and `/health` reports `web_search_enabled: false`. Server configuration: - `API_SERVER_PORT` (default: 8000) @@ -97,8 +96,7 @@ Server configuration: - **`llm_backend.py`**: LangChain model instantiation, HTTP client management, provider routing from model name - **`model_registry.py`**: Maps model names to providers and per-token USD pricing (used by dynamic cost calculator) - **`definitions.py`**: On-chain constants (addresses, network IDs, payment amounts) — configure here for your deployment -- **`web_search.py`**: Exa HTTP client, the single provider-agnostic `web_search` function-tool spec, and result formatting/citation extraction -- **`search_loop.py`**: the in-enclave tool loop — intercepts `web_search` calls, runs them, feeds results back, bounds the rounds, and sums usage across them +- **`web_search.py`**: Exa HTTP client, search execution, and result formatting/citation extraction (serves `/v1/web_search`) - **`util.py`**: `dynamic_session_cost_calculator` converts actual token usage to x402 payment amounts ### API Endpoints @@ -111,6 +109,7 @@ Server configuration: | `/v1/keys` | One-time API key injection (POST, loopback-only) | | `/v1/completions` | Text completion (signed) | | `/v1/chat/completions` | Chat completion with tool support (signed) | +| `/v1/web_search` | In-enclave Exa web search (signed, flat per-search price) | ### TEE Integration @@ -153,34 +152,33 @@ price (see `per_image_price_usd`), not per token. ### Web Search -The `web_search` request flag does NOT use any provider's native web search -(OpenAI/Anthropic/Google/xAI all have one; those were removed). Instead the -gateway advertises a single provider-agnostic `web_search` function tool -(`web_search.py`) and executes it itself, inside the enclave, against Exa — -`search_loop.py` intercepts the call, runs the search, feeds the results back as -a `ToolMessage`, and lets the model continue. Consequences to keep in mind when -touching this code: - -- **Every text model can search**, including ByteDance, Nous and Z.ai, which had - no native option. `model_supports_web_search` excludes only image models. - Anthropic's structured-output path also skips search, since - `with_structured_output` occupies the tool slot with a forced schema tool. -- **The client protocol is unchanged.** Only `web_search` calls are intercepted; - a caller's own `tools` still come back as `tool_calls` for it to run. A turn - that mixes both is terminal and the caller's tools win — ours are dropped and - the model re-issues them next turn. -- **Streaming buffers tool calls unconditionally** when search is on (a fragment - can't be forwarded until the round ends and we know whose tool it was), and the - rounds are wrapped in a generator so the chunk handler sees one flat stream. - Progress frames carry a top-level `web_search` object; `citations` ride - out-of-band on the final frame (unsigned, like `images`). -- **Billing has two parts**: a flat per-search surcharge - (`WEB_SEARCH_PRICE_USD`, identical on every model, so a client can verify it as - `searches × rate`), plus the token cost of every round — each round re-sends - the conversation plus the accumulated results, and `SearchLoopState` sums usage - across all of them. Failed searches are not counted. Rounds are capped - (`MAX_SEARCH_ROUNDS`), with the final round run against a model that has no - search tool bound so it must answer. +Web search is a dedicated endpoint — `POST /v1/web_search` — not a chat feature. +It does NOT use any provider's native web search (OpenAI/Anthropic/Google/xAI +all have one; those were removed), and the gateway runs no tool loop of its own: +the client advertises a `web_search` function tool to its model, calls this +endpoint when the model invokes it, and feeds the returned `content` back as the +tool result. The search runs inside the enclave against Exa (`web_search.py`), +so a query rides the same encrypted OHTTP channel as chat and is never visible +to the relay or the gateway operator. Points to keep in mind: + +- **The chat/completions `web_search` request flag is a deprecated no-op.** It + is still accepted (and still part of the signed request hash when sent) so + old clients' requests parse and verify, but it binds nothing and bills + nothing. +- **Every text model can search** — the tool lives in the client, so this is + purely a question of function calling, not of provider search support. +- **Request/response**: `{"query", "num_results"?, "recency_days"?}` in; + `content` (model-ready numbered results), `citations` (structured sources), + and the standard `tee_*` signing fields out. The request hash covers the + canonical (sorted-keys) JSON body; the output hash covers `content`. +- **Reachable through OHTTP**: the inner payload's `endpoint` field + (`"web_search"`) routes the sealed request; absent means chat, so existing + OHTTP clients are unaffected. Billing flows through the same outer + cost-header / billing-frame channel the relay already consumes. +- **Billing is one flat rate** (`WEB_SEARCH_PRICE_USD`) per search that reached + Exa, settled from the response's `opengradient` block like every paid + endpoint. Validation failures (400/503) and Exa failures (502) return no cost + block and are never settled. A search that ran but matched nothing IS billed. - Exa's self-reported `costDollars` is logged for margin reconciliation only; settlement never depends on it. diff --git a/README.md b/README.md index 1c6d9d0..92fa5b5 100644 --- a/README.md +++ b/README.md @@ -24,10 +24,10 @@ The gateway solves this by running inside a hardware-isolated Nitro Enclave wher - **Request integrity** - SHA256 hash of original request included in signed response - **Streaming support** - SSE streaming for chat completions - **Tool/function calling** - Full support for LLM tool use -- **In-enclave web search** - Opt-in `web_search` flag lets the model search the - live web via a tool the gateway executes inside the enclave (backed by Exa). - Works identically on every text model regardless of provider; searches are - billed at one flat per-search rate on top of token usage +- **In-enclave web search** - Dedicated `/v1/web_search` endpoint (backed by + Exa) for clients running their own tool loop. The query never leaves the + enclave except to the search backend; every search is billed at one flat + per-search rate ## Supported Models @@ -91,34 +91,29 @@ curl -X POST http://127.0.0.1:8000/v1/completions \ "prompt": "Explain quantum computing in one sentence" }' -# Web search (set "web_search": true on any text model, any provider) -curl -X POST http://127.0.0.1:8000/v1/chat/completions \ +# Web search (dedicated endpoint; advertise a `web_search` tool to your model +# and call this when the model invokes it) +curl -X POST http://127.0.0.1:8000/v1/web_search \ -H "Content-Type: application/json" \ -d '{ - "model": "claude-sonnet-4-5", - "messages": [{"role": "user", "content": "What happened in the news today?"}], - "web_search": true + "query": "What happened in the news today?", + "num_results": 6 }' ``` -> **Web search & billing.** When `web_search` is `true`, the gateway offers the -> model a `web_search` function tool and runs it itself, inside the enclave, -> against Exa — it does not use any provider's built-in search. So it works on -> every text model (ByteDance, Nous and Z.ai models included), returns the same -> results and the same `citations` whichever model you pick, and the query never -> leaves the TEE except to the search backend. -> -> The loop is invisible from the outside: one request still returns one answer, -> and any `tool_calls` you get back are only for tools you supplied yourself. -> Streaming responses also emit progress frames carrying a top-level -> `web_search` object you can render or ignore. +> **Web search & billing.** `/v1/web_search` runs the search inside the enclave +> against Exa — it does not use any provider's built-in search, and the gateway +> runs no tool loop: feed the returned `content` back to your model as the tool +> result, and show `citations` to your user. The query never leaves the TEE +> except to the search backend, and via OHTTP (inner `"endpoint": +> "web_search"`) it is also invisible to the relay. > -> Each search adds a flat per-search surcharge — the same on every model, so you -> can verify it as `searches x rate` — reflected in the dynamically-settled x402 -> amount. Because every round re-sends the conversation plus the accumulated -> results, the reported token usage covers all rounds. Failed searches are not -> charged. Requires `EXA_API_KEY` to be injected; check `web_search_enabled` on -> `/health`. +> Each search is one flat price, settled via x402 from the response's +> `opengradient` block like any other paid endpoint; failed searches return no +> cost block and are not charged. The response is signed with the same `tee_*` +> fields as chat (request hash over the canonical JSON body, output hash over +> `content`). Requires `EXA_API_KEY` to be injected; check `web_search_enabled` +> on `/health`. The old chat-request `web_search` flag is a deprecated no-op. ## Deployment to Nitro Enclave diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 4e14ed0..0b04410 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -28,6 +28,7 @@ create_anonymous_chat_completion, get_hpke_config, ) +from tee_gateway.controllers.web_search_controller import create_web_search from x402.http import FacilitatorConfig, HTTPFacilitatorClientSync, PaymentOption from x402.http.middleware.flask import payment_middleware @@ -52,6 +53,7 @@ COMPLETIONS_OPG_SESSION_MAX_SPEND, FACILITATOR_URL, OHTTP_OPG_SESSION_MAX_SPEND, + WEB_SEARCH_OPG_SESSION_MAX_SPEND, ) # --------------------------------------------------------------------------- @@ -322,6 +324,29 @@ def _init_payment_middleware(facilitator_url: str) -> None: mime_type="application/json", description="Completion", ), + "POST /v1/web_search": RouteConfig( + accepts=[ + PaymentOption( + scheme="upto", + pay_to=EVM_PAYMENT_ADDRESS, + price=AssetAmount( + amount=WEB_SEARCH_OPG_SESSION_MAX_SPEND, + asset=BASE_MAINNET_OPG_ADDRESS, + extra={ + "name": "OpenGradient", + "version": "1", + "assetTransferMethod": "permit2", + }, + ), + network=BASE_MAINNET_NETWORK, + ), + ], + extensions={ + **declare_erc20_approval_gas_sponsoring_extension(), + }, + mime_type="application/json", + description="Web search", + ), "POST /v1/ohttp": RouteConfig( accepts=[ PaymentOption( @@ -523,9 +548,9 @@ def health(): "tee_enabled": True, "uptime_seconds": int(time.time() - _started_at), "providers": providers, - # Whether the `web_search` request flag will actually search. Not a - # provider capability — the gateway searches in-enclave for every model — - # so it is reported separately from `providers`. + # Whether /v1/web_search can serve searches (an Exa key was injected). + # Not a provider capability — the search endpoint has no model — so it + # is reported separately from `providers`. "web_search_enabled": web_search_available(), "facilitator_url": _active_facilitator_url, "price_feed": _price_feed.get_status(), @@ -577,6 +602,13 @@ def create_app(): "/v1/ohttp/config", "ohttp-config", get_hpke_config, methods=["GET"] ) + # Dedicated in-enclave web search (Exa). Mounted outside the OpenAI spec + # like the other gateway-specific endpoints; the client's tool loop calls + # it when its model asks to search (see web_search_controller.py). + app.app.add_url_rule( + "/v1/web_search", "web-search", create_web_search, methods=["POST"] + ) + # Initialize TEE here so it runs under both Gunicorn and direct execution. # This is the single TEEKeyManager instance — the same key both registers # with nitriding and signs all LLM responses. @@ -610,7 +642,12 @@ def create_app(): @application.before_request def _check_pricing_ready(): - if request.path not in ("/v1/chat/completions", "/v1/completions", "/v1/ohttp"): + if request.path not in ( + "/v1/chat/completions", + "/v1/completions", + "/v1/ohttp", + "/v1/web_search", + ): return try: _price_feed.get_price() @@ -618,7 +655,8 @@ def _check_pricing_ready(): logger.warning("Rejecting inference request — price feed unavailable: %s", exc) return jsonify({"error": f"Pricing unavailable: {exc}"}), 503 - if request.path == "/v1/ohttp": + # OHTTP carries its model inside the ciphertext; web search has no model. + if request.path in ("/v1/ohttp", "/v1/web_search"): return body = request.get_json(silent=True, cache=True) or {} diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index 60f1926..09612df 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -5,7 +5,6 @@ import connexion from flask import Response -from dataclasses import dataclass, field from typing import Any from tee_gateway.models.create_chat_completion_request import ( @@ -38,21 +37,8 @@ create_image_generation_response, create_image_generation_streaming_response, ) -from tee_gateway.model_registry import get_model_config, model_supports_web_search +from tee_gateway.model_registry import get_model_config from tee_gateway.pricing import compute_session_cost -from tee_gateway.web_search import ( - WEB_SEARCH_TOOL_NAME, - get_web_search_tool, - web_search_available, -) -from tee_gateway.search_loop import ( - MAX_SEARCH_ROUNDS, - SearchLoopState, - execute_search_calls, - run_search_loop, - split_tool_calls, - strip_search_tool_calls, -) logger = logging.getLogger(__name__) @@ -92,97 +78,6 @@ def _split_text_and_images(content: Any) -> tuple[str, list[str]]: return ("".join(text_parts), images) -@dataclass -class _SearchEvent: - """Streaming-only marker yielded between rounds of the search loop. - - The chunk handler in ``_create_streaming_response`` buffers tool-call - fragments without knowing whose tool they belong to; this tells it. Carried - in-band on the chunk iterator so that handler needs no other knowledge of the - loop. - """ - - # Queries to announce to the client before the (blocking) searches run. - queries: list[str] = field(default_factory=list) - # The buffered calls were all ours: forget them and keep streaming. - clear_buffer: bool = False - # Terminal turn mixing our tool with the caller's: strip ours, forward theirs. - drop_search_calls: bool = False - - -def _search_status_frame(model: str, query: str) -> dict[str, Any]: - """An SSE frame telling the client which query is being searched. - - Shaped as an ordinary empty-delta chunk so a client that doesn't know about - `web_search` ignores it harmlessly, with the status hung off a top-level key - (the same convention `images` and `citations` use on the final frame). - """ - return { - "choices": [{"delta": {}, "index": 0, "finish_reason": None}], - "model": model, - "web_search": {"status": "searching", "query": query}, - } - - -def _accumulate_round( - chunk: Any, round_text: list[str], round_calls: dict[int, dict[str, Any]] -) -> None: - """Collect one round's text and tool-call fragments inside the search loop. - - Separate from the outer chunk handler's buffering because the two answer - different questions: this one decides whether to search again, that one - decides what to forward to the client. - """ - content = getattr(chunk, "content", None) - if isinstance(content, str): - round_text.append(content) - elif isinstance(content, list): - round_text.extend( - item.get("text", "") for item in content if isinstance(item, dict) - ) - - for fragment in getattr(chunk, "tool_call_chunks", None) or []: - index = fragment.get("index", 0) - entry = round_calls.setdefault(index, {"id": "", "name": "", "args": ""}) - if fragment.get("id"): - entry["id"] = fragment["id"] - if fragment.get("name"): - entry["name"] = fragment["name"] - args = fragment.get("args") - if args: - entry["args"] += args if isinstance(args, str) else json.dumps(args) - - -def _round_tool_calls(round_calls: dict[int, dict[str, Any]]) -> list[dict[str, Any]]: - """Turn buffered fragments into LangChain-shaped tool calls. - - Arguments arrive as a concatenated JSON string; a call whose arguments never - parse is still returned with empty args so it is classified (and, if it is a - web_search, answered with a "query is required" error the model can recover - from) rather than silently vanishing. - """ - calls: list[dict[str, Any]] = [] - for index in sorted(round_calls): - entry = round_calls[index] - if not entry["name"]: - continue - try: - args = json.loads(entry["args"]) if entry["args"].strip() else {} - except ValueError: - logger.warning( - "Could not parse streamed arguments for tool %r", entry["name"] - ) - args = {} - calls.append( - { - "id": entry["id"], - "name": entry["name"], - "args": args if isinstance(args, dict) else {}, - } - ) - return calls - - def create_chat_completion(body): """Create a chat completion (streaming or non-streaming).""" if not connexion.request.is_json: @@ -207,8 +102,13 @@ def create_chat_completion(body): return _create_non_streaming_response(chat_request) -def _build_user_tools_list(chat_request: CreateChatCompletionRequest) -> list: - """Normalize the caller's own function tools into bind_tools() form.""" +def _build_tools_list(chat_request: CreateChatCompletionRequest) -> list: + """Normalize the caller's function tools into bind_tools() form. + + Web search is deliberately NOT handled here: the ``web_search`` request flag + is a deprecated no-op (search moved to the dedicated ``/v1/web_search`` + endpoint, driven by the client's own tool loop). + """ tools_list: list = [] if chat_request.tools: for tool in chat_request.tools: @@ -219,39 +119,8 @@ def _build_user_tools_list(chat_request: CreateChatCompletionRequest) -> list: ) else: tools_list.append(tool) - return tools_list - - -def _search_enabled( - chat_request: CreateChatCompletionRequest, - provider: str, - anthropic_structured: bool, -) -> bool: - """Whether to bind and run the gateway's web_search tool for this request. - Off when the caller didn't ask, when no Exa key was injected (rather than - advertising a tool that always fails), for image models, and for Anthropic's - structured-output path — ``with_structured_output`` occupies the tool slot - with a forced schema tool, so a search tool bound beside it would never be - callable. Every other model gets search, regardless of provider. - """ - if not getattr(chat_request, "web_search", False): - return False - if anthropic_structured and provider == "anthropic": - logger.info( - "web_search requested with Anthropic structured output; skipping " - "search (structured output uses a forced tool)" - ) - return False - if not model_supports_web_search(chat_request.model): - return False - if not web_search_available(): - logger.warning( - "web_search requested but no Exa API key is configured; answering " - "without search" - ) - return False - return True + return tools_list def _needs_responses_api_for_tools(provider: str, cfg, tools_list: list) -> bool: @@ -373,24 +242,11 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): if cfg.image_generation: return create_image_generation_response(chat_request, request_bytes) - # response_format is resolved before the model is built: whether Anthropic - # takes the structured-output path decides whether web search can run at - # all (see _search_enabled). - rf_dict: dict | None = None - if chat_request.response_format: - rf = _normalize_response_format(chat_request.response_format) - if rf.get("type", "text") != "text": - rf_dict = rf - # Build the tools list first: some OpenAI models (gpt-5.6 family) must be # constructed against the Responses API when function tools are bound. - user_tools = _build_user_tools_list(chat_request) - search_enabled = _search_enabled(chat_request, provider, rf_dict is not None) - tools_list = ( - user_tools + [get_web_search_tool()] if search_enabled else user_tools - ) + tools_list = _build_tools_list(chat_request) - base_model = get_chat_model_cached( + model = get_chat_model_cached( model=chat_request.model, temperature=float(chat_request.temperature) if chat_request.temperature is not None @@ -401,23 +257,21 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): ), ) - model = base_model.bind_tools(tools_list) if tools_list else base_model - # The search loop's final round runs without the search tool bound, which - # is what turns the round cap into "answer now" instead of "hand back an - # unanswerable search request". - model_without_search = ( - (base_model.bind_tools(user_tools) if user_tools else base_model) - if search_enabled - else model - ) + # Bind user tools and/or the native web search tool if requested. + if tools_list: + model = model.bind_tools(tools_list) # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no # Anthropic native equivalent and raises a clear error). - if rf_dict is not None and provider != "anthropic": - model = model.bind(response_format=rf_dict) - model_without_search = model_without_search.bind(response_format=rf_dict) + rf_dict: dict | None = None + if chat_request.response_format: + rf = _normalize_response_format(chat_request.response_format) + if rf.get("type", "text") != "text": + rf_dict = rf + if provider != "anthropic": + model = model.bind(response_format=rf_dict) langchain_messages = convert_messages(chat_request.messages) @@ -430,18 +284,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): SystemMessage(content="Respond in JSON format.") ] + langchain_messages - search_state = SearchLoopState() - if search_enabled: - # Run searches to completion inside the enclave, then answer. The - # caller sent one request and gets one answer; the rounds in between - # are invisible except in the token usage they add. - response = run_search_loop( - model, - model_without_search, - langchain_messages, - search_state, - ) - elif rf_dict and provider == "anthropic": + if rf_dict and provider == "anthropic": response = _invoke_anthropic_structured(model, rf_dict, langchain_messages) else: response = model.invoke(langchain_messages) @@ -458,22 +301,9 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): # of the signed output hash (see _split_text_and_images). if generated_images: message_dict["images"] = generated_images - # Sources the model searched, surfaced out-of-band alongside images and - # on the same terms: the answer text is signed, this metadata about what - # informed it rides inside the OHTTP envelope unsigned. - if search_state.citations: - message_dict["citations"] = search_state.citations finish_reason = "stop" - # Any web_search calls left on a terminal turn belong to a turn that also - # called one of the caller's tools; they are dropped rather than handed to - # a client that has no way to run them. - client_tool_calls = ( - strip_search_tool_calls(response) - if search_enabled - else (getattr(response, "tool_calls", None) or []) - ) - if client_tool_calls: + if hasattr(response, "tool_calls") and response.tool_calls: finish_reason = "tool_calls" message_dict["tool_calls"] = [ { @@ -484,7 +314,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): "arguments": json.dumps(tc.get("args", {})), }, } - for tc in client_tool_calls + for tc in response.tool_calls ] # For tool-call responses, hash the serialized tool calls so the @@ -525,10 +355,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): ) # TODO: If no usage is returned, we should compute it here. - # With search on, `usage` is the sum over every round of the loop — each - # round re-sent the conversation plus the accumulated search results, and - # the caller is charged for all of those tokens, not just the last round's. - usage = search_state.usage if search_enabled else extract_usage(response) + usage = extract_usage(response) if usage: # Surface the standard OpenAI usage triple on the response; the # reasoning split rides along to the cost calculator via `usage`. @@ -537,11 +364,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): "completion_tokens": usage["completion_tokens"], "total_tokens": usage["total_tokens"], } - cost = compute_session_cost( - chat_request.model, - usage, - web_search_count=search_state.search_count, - ) + cost = compute_session_cost(chat_request.model, usage) if cost is not None: openai_response["opengradient"] = cost.model_dump(mode="json") @@ -562,6 +385,9 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): try: provider = get_provider_from_model(chat_request.model) cfg = get_model_config(chat_request.model) + # OpenAI and Anthropic stream tool calls as fragments that must be + # buffered and flushed once complete. Gemini emits complete tool calls. + buffer_tool_calls = provider in ["openai", "anthropic"] # Gemini inline-image models return a single image rather than a token # stream — invoke once and emit the result inside the SSE envelope. image_output_model = cfg.image_output @@ -576,34 +402,11 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): chat_request, request_bytes ) - # response_format is resolved before the model is built: whether Anthropic - # takes the structured-output path decides whether web search can run at - # all (see _search_enabled). - rf_dict: dict | None = None - if chat_request.response_format: - rf = _normalize_response_format(chat_request.response_format) - if rf.get("type", "text") != "text": - rf_dict = rf - anthropic_structured_rf: dict | None = ( - rf_dict if rf_dict is not None and provider == "anthropic" else None - ) - # Build the tools list first: some OpenAI models (gpt-5.6 family) must be # constructed against the Responses API when function tools are bound. - user_tools = _build_user_tools_list(chat_request) - search_enabled = _search_enabled(chat_request, provider, rf_dict is not None) - tools_list = ( - user_tools + [get_web_search_tool()] if search_enabled else user_tools - ) - - # OpenAI and Anthropic stream tool calls as fragments that must be - # buffered and flushed once complete. Gemini emits complete tool calls. - # With search on, ALWAYS buffer: a fragment can't be forwarded to the - # client until the round ends and we know whether the call was a - # web_search this gateway will answer itself. - buffer_tool_calls = provider in ["openai", "anthropic"] or search_enabled + tools_list = _build_tools_list(chat_request) - base_model = get_chat_model_cached( + model = get_chat_model_cached( model=chat_request.model, temperature=float(chat_request.temperature) if chat_request.temperature is not None @@ -614,22 +417,22 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): ), ) - model = base_model.bind_tools(tools_list) if tools_list else base_model - # The search loop's final round runs without the search tool bound, so a - # model that would keep searching has to answer with what it has. - model_without_search = ( - (base_model.bind_tools(user_tools) if user_tools else base_model) - if search_enabled - else model - ) + # Bind user tools and/or the native web search tool if requested. + if tools_list: + model = model.bind_tools(tools_list) # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no # Anthropic native equivalent and raises a clear error). - if rf_dict is not None and anthropic_structured_rf is None: - model = model.bind(response_format=rf_dict) - model_without_search = model_without_search.bind(response_format=rf_dict) + anthropic_structured_rf: dict | None = None + if chat_request.response_format: + rf = _normalize_response_format(chat_request.response_format) + if rf.get("type", "text") != "text": + if provider == "anthropic": + anthropic_structured_rf = rf + else: + model = model.bind(response_format=rf) langchain_messages = convert_messages(chat_request.messages) @@ -679,67 +482,6 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): anthropic_structured_content = None anthropic_structured_usage = None - search_state = SearchLoopState() - - def _streamed_search_rounds(): - """Yield chunks across every round of the in-enclave search loop. - - Written as a generator wrapping ``model.stream`` so the chunk handling - below is identical whether or not search is on — one round or five, it - sees one flat stream of chunks. Between rounds it yields a - ``_SearchEvent`` telling that handler what to do with the tool-call - fragments it just buffered, since only this generator knows whether - they were ``web_search`` calls it answered itself. - """ - loop_messages = list(langchain_messages) - rounds = MAX_SEARCH_ROUNDS + 1 if search_enabled else 1 - - for round_index in range(rounds): - last_round = search_enabled and round_index == MAX_SEARCH_ROUNDS - active_model = model_without_search if last_round else model - - round_text: list[str] = [] - round_calls: dict[int, dict[str, Any]] = {} - - for chunk in active_model.stream(loop_messages): - if search_enabled: - _accumulate_round(chunk, round_text, round_calls) - yield chunk - - if not search_enabled: - return - - ours, theirs = split_tool_calls(_round_tool_calls(round_calls)) - if theirs or not ours: - # Terminal turn. A turn that asked for both our search and one - # of the caller's tools can't be served by either side alone, - # so the caller's tools win and ours are dropped downstream. - if ours: - logger.info( - "Dropping %d streamed web_search call(s) from a turn " - "that also called %d client tool(s)", - len(ours), - len(theirs), - ) - yield _SearchEvent(drop_search_calls=True) - return - - # Announce the queries before running them: the Exa call blocks - # for a second or two and the client should say why it is waiting. - yield _SearchEvent( - queries=[ - q - for q in ((c.get("args") or {}).get("query") for c in ours) - if isinstance(q, str) and q.strip() - ], - clear_buffer=True, - ) - - loop_messages.append( - AIMessage(content="".join(round_text), tool_calls=ours) - ) - loop_messages.extend(execute_search_calls(ours, search_state)) - def generate(): full_content = "" final_usage = None @@ -816,31 +558,9 @@ def generate(): yield f"data: {json.dumps(data)}\n\n" chunks_iter = [] else: - chunks_iter = _streamed_search_rounds() # type: ignore[assignment] + chunks_iter = model.stream(langchain_messages) # type: ignore[assignment] for chunk in chunks_iter: - # --- Search-round boundary (in-enclave web search) --- - # Not a model chunk: an instruction about the tool-call - # fragments buffered so far, plus the queries to tell the - # client about. - if isinstance(chunk, _SearchEvent): - if chunk.clear_buffer: - # Those calls were web_search calls this gateway is - # answering itself — the client must never see them. - buffered_tool_calls = {} - finish_reason = "stop" - if chunk.drop_search_calls: - buffered_tool_calls = { - index: tc - for index, tc in buffered_tool_calls.items() - if tc["function"]["name"] != WEB_SEARCH_TOOL_NAME - } - if not buffered_tool_calls: - finish_reason = "stop" - for query in chunk.queries: - yield f"data: {json.dumps(_search_status_frame(chat_request.model, query))}\n\n" - continue - # --- Text content --- if chunk.content: if isinstance(chunk.content, str): @@ -1035,10 +755,6 @@ def generate(): # are not part of the signed output hash. if generated_images: final_data["images"] = generated_images - # Likewise the sources the model searched: the answer text is - # signed, this metadata about what informed it is not. - if search_state.citations: - final_data["citations"] = search_state.citations logger.debug( f"Response Final\n\tTEE Signature: {tee_signature}\n\tTEE request hash: {input_hash_hex}\n\tTEE output hash: {output_hash_hex}\n\tTEE timestamp: {timestamp}\n\tTEE ID: 0x{tee_keys.get_tee_id()}" @@ -1053,18 +769,11 @@ def generate(): } # Pass thinking tokens to the cost calculator (for the image # dual-rate split) without polluting the OpenAI usage triple. - # `final_usage` already sums every round of the search loop — - # each round re-sent the conversation plus the accumulated - # results, and the caller pays for all of those input tokens. cost_usage = dict( final_data["usage"], reasoning_tokens=final_usage.get("reasoning", 0), ) - cost = compute_session_cost( - chat_request.model, - cost_usage, - web_search_count=search_state.search_count, - ) + cost = compute_session_cost(chat_request.model, cost_usage) if cost is not None: # final_data is hand-serialized to SSE via json.dumps below, # which doesn't go through Flask's JSONEncoder — so do the @@ -1073,7 +782,6 @@ def generate(): logger.info( f"Stream completed — usage: {final_data['usage']}, " f"finish: {finish_reason}, " - f"searches: {search_state.search_count}, " f"inputHash: {input_hash_hex[:16]}..., outputHash: {output_hash_hex[:16]}..." ) diff --git a/tee_gateway/controllers/completions_controller.py b/tee_gateway/controllers/completions_controller.py index 472c270..976ed1c 100644 --- a/tee_gateway/controllers/completions_controller.py +++ b/tee_gateway/controllers/completions_controller.py @@ -15,10 +15,7 @@ get_chat_model_cached, extract_usage, ) -from tee_gateway.model_registry import model_supports_web_search from tee_gateway.pricing import compute_session_cost -from tee_gateway.search_loop import SearchLoopState, run_search_loop -from tee_gateway.web_search import get_web_search_tool, web_search_available logger = logging.getLogger(__name__) @@ -31,6 +28,9 @@ def create_completion(body): return {"error": "Request must be application/json"}, 415 try: + # The web_search flag is a deprecated no-op: search moved to the + # dedicated /v1/web_search endpoint. It stays in the hashed request + # dict so signatures from clients that still send it verify unchanged. web_search = bool(getattr(body, "web_search", False)) request_dict = { @@ -49,7 +49,7 @@ def create_completion(body): request_bytes = json.dumps(request_dict, sort_keys=True).encode("utf-8") - base_model = get_chat_model_cached( + model = get_chat_model_cached( model=body.model, temperature=float(body.temperature) if body.temperature is not None @@ -57,35 +57,11 @@ def create_completion(body): max_tokens=body.max_tokens or 4096, ) - # Web search is the gateway's own function tool, executed in-enclave for - # any model that can call a function — see web_search.py. Skipped when no - # Exa key was injected rather than advertising a tool that always fails. - search_enabled = ( - web_search - and model_supports_web_search(body.model) - and web_search_available() - ) - if web_search and not search_enabled: - logger.warning( - "web_search requested for %s but search is unavailable; " - "completing without it", - body.model, - ) - - messages: list[Any] = [HumanMessage(content=body.prompt)] - search_state = SearchLoopState() - if search_enabled: - response = run_search_loop( - base_model.bind_tools([get_web_search_tool()]), - base_model, - messages, - search_state, - ) - else: - response = base_model.invoke(messages) + messages = [HumanMessage(content=body.prompt)] + response = model.invoke(messages) - # Some providers return content as a list of blocks; flatten to the text - # the caller expects. + # Some providers can return content as a list of blocks; flatten to the + # text the caller expects. if isinstance(response.content, list): response_content = "".join( item.get("text", "") if isinstance(item, dict) else str(item) @@ -93,8 +69,7 @@ def create_completion(body): ) else: response_content = response.content or "" - # With search on, usage is the sum over every round of the loop. - usage = search_state.usage if search_enabled else extract_usage(response) + usage = extract_usage(response) timestamp = int(time.time()) msg_hash, input_hash_hex, output_hash_hex = compute_tee_msg_hash( @@ -123,13 +98,9 @@ def create_completion(body): "tee_id": f"0x{tee_keys.get_tee_id()}", } if usage: - cost = compute_session_cost( - body.model, usage, web_search_count=search_state.search_count - ) + cost = compute_session_cost(body.model, usage) if cost is not None: completion_response["opengradient"] = cost.model_dump(mode="json") - if search_state.citations: - completion_response["citations"] = search_state.citations return completion_response except Exception as e: diff --git a/tee_gateway/controllers/ohttp_controller.py b/tee_gateway/controllers/ohttp_controller.py index adb6d51..e5e758f 100644 --- a/tee_gateway/controllers/ohttp_controller.py +++ b/tee_gateway/controllers/ohttp_controller.py @@ -2,10 +2,12 @@ Oblivious HTTP endpoint for anonymous inference (relay-pays model). This handler is a thin shell: it HPKE-decapsulates the inner request, re-issues -it as an in-process WSGI sub-request against the enclave's own -``/v1/chat/completions``, then encapsulates the response. All x402 payment, -LangChain routing, cost settlement and TEE response signing reuse the public -chat code paths — there is no duplicated routing or pricing logic here. +it as an in-process WSGI sub-request against one of the enclave's own paid +endpoints — ``/v1/chat/completions`` by default, or the endpoint named by the +inner payload's ``endpoint`` field (currently also ``web_search`` → +``/v1/web_search``) — then encapsulates the response. All x402 payment, routing, +cost settlement and TEE response signing reuse the public code paths — there is +no duplicated routing or pricing logic here. Two response modes are supported, dispatched by the inner ``stream`` flag: * stream=false → single-shot OHTTP response (RFC 9458 §4.5), @@ -100,6 +102,16 @@ # the upstream LLM provider. _IDENTIFYING_FIELDS = ("user", "metadata", "x-request-id", "request_id") +# Inner endpoints reachable through the OHTTP envelope, keyed by the payload's +# `endpoint` discriminator. Absent means chat (the original OHTTP contract, so +# existing clients keep working unchanged). The field is popped before the +# sub-dispatch: it is routing metadata, not part of the endpoint's request body +# (or its signed request hash). +_INNER_ENDPOINT_PATHS = { + "chat.completions": "/v1/chat/completions", + "web_search": "/v1/web_search", +} + # Response headers we propagate from the inner /v1/chat/completions response # back through the relay to the client. _FORWARDED_HEADER_PREFIXES = ("x-payment", "x-upto", "x-settlement", "x-tee") @@ -154,12 +166,19 @@ def create_anonymous_chat_completion(): if not isinstance(chat_body, dict): return _error(400, "inner payload must be a JSON object") + endpoint = chat_body.pop("endpoint", "chat.completions") + inner_path = _INNER_ENDPOINT_PATHS.get(endpoint) + if inner_path is None: + return _sealed_error( + flask_request, decap, 400, f"unknown inner endpoint {endpoint!r}" + ) + chat_body = _scrub(chat_body) _set_inner_cost_context(flask_request, request_json=chat_body) body_bytes = json.dumps(chat_body, separators=(",", ":")).encode("utf-8") sub_status, sub_headers, sub_iter = _wsgi_subrequest( - path="/v1/chat/completions", + path=inner_path, body_bytes=body_bytes, ) diff --git a/tee_gateway/controllers/web_search_controller.py b/tee_gateway/controllers/web_search_controller.py new file mode 100644 index 0000000..14c21a9 --- /dev/null +++ b/tee_gateway/controllers/web_search_controller.py @@ -0,0 +1,98 @@ +"""Dedicated web-search endpoint (POST /v1/web_search). + +The gateway runs no tool loop of its own: a client whose model asks to search +calls this endpoint, feeds ``content`` back to its model as the tool result, +and shows ``citations`` to its user. The search itself runs inside the enclave +against Exa (see web_search.py), so a query rides the same encrypted channel as +a chat request and is never visible to the relay or the gateway operator. + +Request body:: + + {"query": "...", "num_results": 6, "recency_days": 30} + +``query`` is required; the other two are optional and clamped to the module's +bounds. Any other fields are ignored (but still part of the signed request +hash, which covers the body exactly as received). + +The response is signed like every other paid endpoint — RSA-PSS over +``keccak256(requestHash || outputHash || timestamp)`` with the request hash +computed over the canonical (sorted-keys) JSON body and the output hash over +``content`` — and carries an ``opengradient`` cost block settled by x402 at the +flat ``WEB_SEARCH_PRICE_USD`` rate. Failures (Exa unreachable, provider error) +return 502 without a cost block, so they are never settled; a missing Exa key +returns 503. +""" + +import json +import logging +import time +import uuid + +from flask import request + +from tee_gateway.pricing import compute_web_search_cost +from tee_gateway.tee_manager import get_tee_keys, compute_tee_msg_hash +from tee_gateway.web_search import ( + execute_web_search_call, + web_search_available, +) + +logger = logging.getLogger(__name__) + + +def create_web_search(): + """POST /v1/web_search — run one Exa search and return signed results.""" + body = request.get_json(silent=True) + if not isinstance(body, dict): + return {"error": "Request must be a JSON object"}, 415 + + if not web_search_available(): + return {"error": "Web search is not configured on this gateway"}, 503 + + query = body.get("query") + if not isinstance(query, str) or not query.strip(): + return {"error": "A non-empty `query` string is required"}, 400 + + # Hash the body exactly as the client sent it (canonicalized), so the + # client can recompute the request hash from what it built. The `endpoint` + # discriminator the OHTTP dispatcher pops never reaches this handler. + request_bytes = json.dumps(body, sort_keys=True).encode("utf-8") + + outcome = execute_web_search_call(body) + + if outcome.is_error: + # Not billable, so no cost block: x402 skips settlement on this + # request. `content` is still model-readable ("the search failed"), so + # a client that wants its model to recover can relay it. + return {"error": outcome.content}, 502 + + timestamp = int(time.time()) + msg_hash, input_hash_hex, output_hash_hex = compute_tee_msg_hash( + request_bytes, outcome.content, timestamp + ) + tee_keys = get_tee_keys() + signature = tee_keys.sign_data(msg_hash) + + response = { + "id": f"websearch-{uuid.uuid4()}", + "object": "web_search.result", + "created": timestamp, + "query": query.strip(), + "content": outcome.content, + "citations": outcome.citations, + "tee_signature": signature, + "tee_request_hash": input_hash_hex, + "tee_output_hash": output_hash_hex, + "tee_timestamp": timestamp, + "tee_id": f"0x{tee_keys.get_tee_id()}", + } + + # A search that ran but matched nothing still consumed an Exa request and + # is billable (outcome.billable is True for both). Cost-calculation + # failure (price feed down) logs CRITICAL and simply omits the block — + # the request goes unsettled, matching the chat path's fail-open contract. + cost = compute_web_search_cost() + if cost is not None: + response["opengradient"] = cost.model_dump(mode="json") + + return response diff --git a/tee_gateway/definitions.py b/tee_gateway/definitions.py index d91ae8a..0ec64c1 100644 --- a/tee_gateway/definitions.py +++ b/tee_gateway/definitions.py @@ -78,6 +78,12 @@ # so the relay-paid encrypted endpoint needs a higher per-session cap. OHTTP_OPG_SESSION_MAX_SPEND: str = "5000000000000000000" +# /v1/web_search — maximum OPG spend per session (18 decimals: 1000000000000000000 = 1 OPG). +# Each search settles at the flat WEB_SEARCH_PRICE_USD (see model_registry.py), +# so at the fallback OPG price this cap covers several searches per session +# with headroom for OPG price swings. +WEB_SEARCH_OPG_SESSION_MAX_SPEND: str = "1000000000000000000" + # /v1/completions — maximum OPG spend per session (18 decimals: 100000000000000000 = 0.1 OPG). # This is the upper-bound amount presented to the client during the x402 pre-check handshake. # The x402 "upto" scheme allows the actual charge to be any value up to this cap; diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index 52a02e1..2d9dbbb 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -673,6 +673,6 @@ def extract_usage(response) -> Optional[Dict[str, int]]: return None -# Web search is not a provider feature here — the gateway executes it itself -# against Exa and bills one flat rate on every model. See web_search.py for the -# tool spec and search_loop.py for the loop that runs it. +# Web search is not a provider feature here — the gateway serves it as its own +# Exa-backed endpoint (/v1/web_search, see web_search.py) that the client's tool +# loop calls, billed at one flat per-search rate. diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index f8804a9..5e74883 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -68,10 +68,6 @@ class ModelConfig: # ``output_tokens`` count and only breaks out thinking (``reasoning``), so the # billing splits reasoning at ``output_price_usd`` and the remainder here. image_output_price_usd: Optional[Decimal] = None - # Per-search USD surcharge override. Web search is one flat rate on every - # model (``WEB_SEARCH_PRICE_USD``) because the gateway runs the search - # itself; set this only to price a single model's searches differently. - web_search_price_usd: Optional[Decimal] = None # OpenAI's newest reasoning models (the gpt-5.6 family) apply a default # ``reasoning_effort`` that the Chat Completions endpoint rejects when # function tools are also present ("Function tools with reasoning_effort are @@ -83,19 +79,19 @@ class ModelConfig: responses_api_for_tools: bool = False -# Flat USD price per web search, identical on every model. +# Flat USD price per call to the /v1/web_search endpoint. # -# The gateway runs searches itself against Exa (see web_search.py), so there is -# one cost to pass through instead of four provider list prices with four -# different billable units. At our request shape — one Exa search plus page text -# for up to `MAX_NUM_RESULTS` results — Exa charges $7/1k requests and $1/1k -# pages per content type, i.e. ~$0.013 for a 6-result search. This rate covers -# that with a small margin, and is below what three of the four native searches -# used to cost (xAI $0.025/unit, Google $0.035/request). +# The gateway runs searches against Exa (see web_search.py), so there is one +# cost to pass through, independent of any model. At our request shape — one +# Exa search plus page text for up to `MAX_NUM_RESULTS` results — Exa charges +# $7/1k requests and $1/1k pages per content type, i.e. ~$0.013 for a 6-result +# search. This rate covers that with a small margin, and is below what three of +# the four native provider searches used to cost (xAI $0.025/unit, Google +# $0.035/request). # -# The billable unit is "one search the model asked for that reached Exa", so a -# client can verify its surcharge as `searches * this rate`. Searches that -# failed or were malformed are not counted (see WebSearchOutcome.billable). +# The billable unit is "one search that reached Exa": a request that fails +# validation or errors out at Exa returns without a cost block and is not +# settled (see WebSearchOutcome.billable). WEB_SEARCH_PRICE_USD: Decimal = Decimal("0.015") # ByteDance ModelArk image *deployment* endpoints (api_name "ep-…", e.g. Seedance @@ -718,31 +714,3 @@ def get_rate_card(model: str) -> dict[str, Decimal]: """Return {"input": ..., "output": ...} pricing for a model. Raises on unknown.""" cfg = get_model_config(model) return {"input": cfg.input_price_usd, "output": cfg.output_price_usd} - - -def get_web_search_price_usd(model: str) -> Decimal: - """Return the per-search USD surcharge for a model. - - The flat ``WEB_SEARCH_PRICE_USD`` unless the model overrides it. Image models - are free: they never reach the chat path that can search. Raises ValueError - if the model is unknown. - """ - cfg = get_model_config(model) - if cfg.web_search_price_usd is not None: - return cfg.web_search_price_usd - if cfg.image_generation or cfg.image_output: - return Decimal("0") - return WEB_SEARCH_PRICE_USD - - -def model_supports_web_search(model: str) -> bool: - """Whether ``web_search`` can be enabled for a model. - - True for every text model in the registry — the gateway supplies the search - tool itself, so this is a question about function calling, not about which - provider shipped a search feature. Image models are excluded: generation - models are served off the chat path entirely, and image-*output* models are - invoked in a single non-streaming shot with no tool loop around them. - """ - cfg = get_model_config(model) - return not (cfg.image_generation or cfg.image_output) diff --git a/tee_gateway/openapi/openapi.yaml b/tee_gateway/openapi/openapi.yaml index fc15712..7a954ee 100644 --- a/tee_gateway/openapi/openapi.yaml +++ b/tee_gateway/openapi/openapi.yaml @@ -2886,27 +2886,13 @@ components: $ref: "#/components/schemas/CreateChatCompletionRequest_model" web_search: default: false + deprecated: true description: | - Let the model search the web. When `true`, the gateway offers the - model a `web_search` tool and executes it inside the enclave against - its own search backend, feeding the results back so the model can - answer from them. This works on every text model regardless of - provider, and behaves identically on all of them. - - The search rounds are invisible to the caller: one request still - yields one answer, and the `tool_calls` you get back are only ever - for tools you supplied yourself. Sources used are returned - out-of-band on the response message as `citations`. Streaming - responses additionally emit progress frames carrying a top-level - `web_search` object (`{"status": "searching", "query": "..."}`) that - clients may render or ignore. - - Billing: each search the model runs adds a flat per-search surcharge - (identical on every model) on top of token usage, and because each - round re-sends the conversation plus the accumulated results, the - reported token usage covers every round. Searches that fail are not - charged. Image models, and gateways with no search backend - configured, ignore this flag and are not charged for it — check + Deprecated no-op, accepted for wire compatibility. Web search is a + dedicated endpoint (`POST /v1/web_search`, in-enclave, Exa-backed): + advertise a `web_search` function tool to the model yourself, call + the endpoint when the model invokes it, and feed the returned + `content` back as the tool result. Availability is reported as `web_search_enabled` on `/health`. nullable: true title: web_search @@ -3566,14 +3552,10 @@ components: type: boolean web_search: default: false + deprecated: true description: | - Let the model search the web. The gateway offers the model a - `web_search` tool and executes it inside the enclave, feeding the - results back before the completion is produced. Works on every text - model regardless of provider. Each search adds a flat per-search - surcharge on top of token usage, and the reported token usage covers - every round of the loop. Failed searches are not charged. Sources - are returned on the response as `citations`. + Deprecated no-op, accepted for wire compatibility. Web search is a + dedicated endpoint — see `web_search` on the chat request. nullable: true title: web_search type: boolean diff --git a/tee_gateway/pricing.py b/tee_gateway/pricing.py index d6d97a0..c2f8a88 100644 --- a/tee_gateway/pricing.py +++ b/tee_gateway/pricing.py @@ -17,7 +17,7 @@ ASSET_DECIMALS_BY_ADDRESS, BASE_MAINNET_OPG_ADDRESS, ) -from tee_gateway.model_registry import get_model_config, get_web_search_price_usd +from tee_gateway.model_registry import WEB_SEARCH_PRICE_USD, get_model_config logger = logging.getLogger("llm_server.dynamic_pricing") @@ -50,7 +50,7 @@ def _serialize_decimal(self, value: Decimal) -> str: def compute_session_cost( - model: str, usage: dict, web_search_count: int = 0, image_count: int = 0 + model: str, usage: dict, image_count: int = 0 ) -> SessionCost | None: """Compute the settled cost for a completed inference request. @@ -96,19 +96,6 @@ def compute_session_cost( Decimal(out_tok) * cfg.output_price_usd ) - # Web search is billed per search on top of token cost, at one flat rate - # for every model (the gateway runs the search itself — see web_search.py). - # Note the token cost already reflects the search: a search round re-sends - # the conversation plus the results, and search_loop.SearchLoopState sums - # the usage of every round into `usage`. - searches = max(0, int(web_search_count)) - web_search_usd = ( - Decimal(searches) * get_web_search_price_usd(model) - if searches - else Decimal(0) - ) - raw_usd += web_search_usd - # Image-generation models (xAI Grok, ByteDance Seedream) are billed a flat # price per generated image rather than per token; token prices are 0. images = max(0, int(image_count)) @@ -138,13 +125,11 @@ def compute_session_cost( logger.info( "DYNAMIC_SESSION_COST model=%s input_tokens=%d output_tokens=%d " - "web_searches=%d web_search_usd=%s images=%d image_usd=%s raw_usd=%s " - "settled_usd=%s token_price_usd=%s decimals=%d cost=%d", + "images=%d image_usd=%s raw_usd=%s settled_usd=%s " + "token_price_usd=%s decimals=%d cost=%d", model, in_tok, out_tok, - searches, - str(web_search_usd), images, str(image_usd), str(raw_usd), @@ -166,3 +151,50 @@ def compute_session_cost( exc_info=True, ) return None + + +def compute_web_search_cost() -> SessionCost | None: + """Compute the flat cost of one /v1/web_search call. + + Same OPG conversion and ceiling-rounding as :func:`compute_session_cost`, + with ``WEB_SEARCH_PRICE_USD`` as the raw USD amount — the endpoint has no + token or model dimension. Returns ``None`` when the price feed is down + (matching the token path: the request goes unsettled and the client is not + charged, logged CRITICAL for reconciliation). + """ + from tee_gateway.price_feed import get_price_feed + + try: + token_price_usd = get_price_feed().get_price() + if token_price_usd <= 0: + raise ValueError(f"Token price is non-positive: {token_price_usd}") + + scale = Decimal(10) ** _OPG_DECIMALS + cost_smallest_units = max( + 0, + int( + ((WEB_SEARCH_PRICE_USD / token_price_usd) * scale).to_integral_value( + rounding=ROUND_CEILING + ) + ), + ) + settled_usd = (Decimal(cost_smallest_units) / scale) * token_price_usd + logger.info( + "WEB_SEARCH_COST raw_usd=%s settled_usd=%s token_price_usd=%s cost=%d", + str(WEB_SEARCH_PRICE_USD), + str(settled_usd), + str(token_price_usd), + cost_smallest_units, + ) + return SessionCost( + cost_opg=cost_smallest_units, + cost_usd=settled_usd, + opg_price_usd=token_price_usd, + ) + except Exception as exc: + logger.critical( + "Web search cost calculation failed — client will NOT be charged: %s", + exc, + exc_info=True, + ) + return None diff --git a/tee_gateway/search_loop.py b/tee_gateway/search_loop.py deleted file mode 100644 index 4fa25bb..0000000 --- a/tee_gateway/search_loop.py +++ /dev/null @@ -1,223 +0,0 @@ -"""The in-enclave web-search tool loop. - -The gateway advertises ``web_search`` (see ``web_search.get_web_search_tool``) as -an ordinary function tool, which means the model asks for a search the same way -it asks for any other tool — and something has to answer it. That something is -this module: it runs the search inside the enclave, feeds the results back, and -lets the model continue, all within one client request. - -Two properties this preserves, both of which the previous provider-native search -gave up: - - * The client's tool protocol is untouched. A caller that passes its own - ``tools`` still gets tool calls handed back to execute; only ``web_search`` - calls are intercepted. A caller that passes no tools never learns a loop - happened — it sends one request and gets one answer. - * Billing stays honest about a loop's real cost. Each round re-sends the whole - conversation *plus* every prior search result, so the input tokens are - genuinely spent several times over. ``SearchLoopState`` accumulates usage - across every round so the caller is charged for all of it rather than for - the last round alone. - -The loop is bounded (``MAX_SEARCH_ROUNDS``) and the final round is run with the -search tool unbound, so a model that would otherwise keep searching is forced to -answer with what it has instead of spending the caller's money indefinitely. -""" - -import logging -from dataclasses import dataclass, field -from typing import Any, Callable, Optional - -from langchain_core.messages import AIMessage, ToolMessage - -from tee_gateway.web_search import ( - WEB_SEARCH_TOOL_NAME, - execute_web_search_call, -) - -logger = logging.getLogger(__name__) - -# How many times the model may search before it must answer. Each round costs a -# full prompt re-send, so this is a cost ceiling as much as a latency one: four -# rounds is enough for "search, refine, cross-check" without letting a model -# that has decided to keep googling run up an unbounded bill. -MAX_SEARCH_ROUNDS = 4 - - -@dataclass -class SearchLoopState: - """Accumulator threaded through every round of one request's loop. - - Kept separate from the loop functions because the streaming controller runs - its rounds itself (it has to forward SSE frames as they arrive) while the - non-streaming controller delegates the whole loop — both share this state. - """ - - search_count: int = 0 - citations: list[dict[str, str]] = field(default_factory=list) - # Running token totals across every round, in the shape extract_usage - # returns. None until some round actually reports usage, matching the - # "provider reported nothing, so do not charge" convention elsewhere. - usage: Optional[dict[str, int]] = None - rounds: int = 0 - - def add_usage(self, round_usage: Optional[dict[str, int]]) -> None: - """Fold one round's token usage into the running totals.""" - if not round_usage: - return - if self.usage is None: - self.usage = {} - for key, value in round_usage.items(): - if isinstance(value, (int, float)): - self.usage[key] = self.usage.get(key, 0) + int(value) - - def add_citations(self, citations: list[dict[str, str]]) -> None: - """Append citations from one search, de-duplicated by URL. - - A refining second search very often re-surfaces the best hit from the - first, and showing the same source twice reads as a bug. - """ - seen = {c.get("url") for c in self.citations} - for citation in citations: - url = citation.get("url") - if url and url not in seen: - seen.add(url) - self.citations.append(citation) - - -def split_tool_calls( - tool_calls: Optional[list[dict[str, Any]]], -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Partition a turn's tool calls into (ours, the client's). - - "Ours" are ``web_search`` calls this gateway executes; the rest belong to - tools the caller supplied and must be handed back for the caller to run. - """ - ours: list[dict[str, Any]] = [] - theirs: list[dict[str, Any]] = [] - for call in tool_calls or []: - name = call.get("name") if isinstance(call, dict) else None - if name == WEB_SEARCH_TOOL_NAME: - ours.append(call) - else: - theirs.append(call) - return ours, theirs - - -def execute_search_calls( - calls: list[dict[str, Any]], - state: SearchLoopState, - on_search: Optional[Callable[[str], None]] = None, -) -> list[ToolMessage]: - """Run each ``web_search`` call and build the ToolMessages to feed back. - - ``on_search`` is invoked with each query before it runs, so the streaming - controller can tell the client what is being searched for while it waits. - Only searches that actually reached Exa are counted as billable. - """ - messages: list[ToolMessage] = [] - for call in calls: - args = call.get("args") - if not isinstance(args, dict): - args = {} - query = args.get("query") - if on_search is not None and isinstance(query, str) and query.strip(): - try: - on_search(query.strip()) - except Exception: - # A status callback is cosmetic; never let it kill the search. - logger.debug("web_search status callback failed", exc_info=True) - - outcome = execute_web_search_call(args) - if outcome.billable: - state.search_count += 1 - state.add_citations(outcome.citations) - - messages.append( - ToolMessage( - content=outcome.content, - tool_call_id=call.get("id") or "", - name=WEB_SEARCH_TOOL_NAME, - status="error" if outcome.is_error else "success", - ) - ) - return messages - - -def strip_search_tool_calls(message: AIMessage) -> list[dict[str, Any]]: - """Client-facing tool calls for a turn that mixes our tool with theirs. - - A turn asking for both ``web_search`` and one of the caller's tools cannot be - completed by either side alone: we cannot run their tool, and they cannot run - ours. The caller's tools win — their loop is the outer one and will come back - to us — so our calls are dropped here and the model re-issues them on the - next turn if it still wants to search. Rare in practice; logged when it - happens so it does not stay invisible. - """ - ours, theirs = split_tool_calls(getattr(message, "tool_calls", None)) - if ours: - logger.info( - "Dropping %d web_search call(s) from a turn that also called %d " - "client tool(s); the model can re-issue them next turn", - len(ours), - len(theirs), - ) - return theirs - - -def run_search_loop( - model: Any, - model_without_search: Any, - messages: list[Any], - state: SearchLoopState, - invoke: Optional[Callable[[Any, list[Any]], AIMessage]] = None, - on_search: Optional[Callable[[str], None]] = None, - max_rounds: int = MAX_SEARCH_ROUNDS, -) -> AIMessage: - """Drive the loop to a terminal turn and return it (non-streaming callers). - - Terminal means: a plain answer, or a turn calling one of the *caller's* - tools. ``messages`` is extended in place with each round's assistant turn and - search results, so the caller can inspect the full trajectory afterwards. - - ``model`` has the search tool bound; ``model_without_search`` does not and is - used for the final round, which is what converts the round cap into "answer - now" rather than "return an unanswerable search request". ``invoke`` lets a - caller substitute its own invocation (the Anthropic structured-output path - does not use plain ``.invoke``). - """ - call_model = invoke if invoke is not None else (lambda m, msgs: m.invoke(msgs)) - - for round_index in range(max_rounds + 1): - last_round = round_index == max_rounds - active_model = model_without_search if last_round else model - - response = call_model(active_model, messages) - state.rounds = round_index + 1 - state.add_usage(_message_usage(response)) - - ours, theirs = split_tool_calls(getattr(response, "tool_calls", None)) - if theirs or not ours: - # Terminal: either a plain answer or the caller's tools to run. - return response - - messages.append(response) - messages.extend(execute_search_calls(ours, state, on_search)) - - # Unreachable: the last iteration binds no search tool, so `ours` is empty - # and the loop returns above. - raise RuntimeError("search loop exited without a terminal response") - - -def _message_usage(message: Any) -> Optional[dict[str, int]]: - """Token usage for one round, in the shape the cost calculator expects.""" - metadata = getattr(message, "usage_metadata", None) - if not metadata: - return None - details = metadata.get("output_token_details") or {} - return { - "prompt_tokens": metadata.get("input_tokens", 0), - "completion_tokens": metadata.get("output_tokens", 0), - "total_tokens": metadata.get("total_tokens", 0), - "reasoning_tokens": details.get("reasoning", 0), - } diff --git a/tee_gateway/test/test_web_search.py b/tee_gateway/test/test_web_search.py index 30ba47a..fffef6f 100644 --- a/tee_gateway/test/test_web_search.py +++ b/tee_gateway/test/test_web_search.py @@ -1,41 +1,29 @@ """ -Unit tests for in-enclave web search (Exa) across providers. +Unit tests for the dedicated in-enclave web search endpoint (Exa). Covers: - - web_search: tool spec, argument coercion, Exa request shaping, result - formatting, and every failure mode of the Exa call - - model_registry: one flat per-search price and which models can search - - search_loop: tool-call partitioning, round accumulation, the round cap, and - terminal conditions - - pricing.compute_session_cost: per-search surcharge added to token cost - - chat_controller: the web_search flag binds the tool, searches are executed - in-enclave rather than handed to the client, and every round is billed + - web_search: argument coercion, Exa request shaping, result formatting, and + every failure mode of the Exa call + - pricing: the flat per-search cost (compute_web_search_cost) and that chat + token pricing no longer carries a search surcharge + - web_search_controller: request validation, signed response shape, the + opengradient cost block, and the unbilled failure paths + - ohttp_controller: the inner `endpoint` discriminator routes a sealed + request to /v1/web_search (and defaults to chat for existing clients) """ import json import unittest from decimal import Decimal -from types import SimpleNamespace from unittest.mock import Mock, patch -from langchain_core.messages import AIMessage, AIMessageChunk, ToolMessage +from flask import Flask from tee_gateway import web_search as ws -from tee_gateway.model_registry import ( - WEB_SEARCH_PRICE_USD, - get_web_search_price_usd, - model_supports_web_search, -) -from tee_gateway.pricing import SessionCost, compute_session_cost -from tee_gateway.search_loop import ( - MAX_SEARCH_ROUNDS, - SearchLoopState, - execute_search_calls, - run_search_loop, - split_tool_calls, - strip_search_tool_calls, -) -from tee_gateway.controllers.chat_controller import create_chat_completion +from tee_gateway.controllers import ohttp_controller +from tee_gateway.controllers.web_search_controller import create_web_search +from tee_gateway.model_registry import WEB_SEARCH_PRICE_USD +from tee_gateway.pricing import compute_session_cost, compute_web_search_cost # --------------------------------------------------------------------------- @@ -83,31 +71,17 @@ def _with_exa(*responses): # --------------------------------------------------------------------------- -# Tool specification +# Availability # --------------------------------------------------------------------------- -class TestWebSearchToolSpec(unittest.TestCase): - def test_single_provider_agnostic_function_tool(self): - """One spec for every provider — no per-provider variants any more.""" - tool = ws.get_web_search_tool() - self.assertEqual(tool["type"], "function") - self.assertEqual(tool["function"]["name"], "web_search") - - def test_schema_is_flat_and_only_requires_a_query(self): - """Nested/exotic schemas are where provider support diverges.""" - params = ws.get_web_search_tool()["function"]["parameters"] - self.assertEqual(params["required"], ["query"]) - self.assertEqual( - set(params["properties"]), {"query", "num_results", "recency_days"} - ) - for prop in params["properties"].values(): - self.assertIn(prop["type"], {"string", "integer"}) - +class TestAvailability(unittest.TestCase): def test_availability_tracks_the_injected_key(self): ws.configure_exa_client("test-key") - self.assertTrue(ws.web_search_available()) - ws.configure_exa_client(None) + try: + self.assertTrue(ws.web_search_available()) + finally: + ws.configure_exa_client(None) self.assertFalse(ws.web_search_available()) @@ -118,21 +92,38 @@ def test_availability_tracks_the_injected_key(self): class TestArgumentCoercion(unittest.TestCase): def test_missing_or_blank_query_is_an_unbilled_error(self): - for args in ({}, {"query": " "}, {"query": 42}): + for args in ({}, {"query": ""}, {"query": " "}, {"query": 42}): outcome = ws.execute_web_search_call(args) self.assertTrue(outcome.is_error, args) self.assertFalse(outcome.billable, args) self.assertIn("query", outcome.content) - def test_num_results_is_clamped_not_rejected(self): - """Models pass these as strings and floats; be forgiving.""" - self.assertEqual(ws._clamp_int("3", 6, 1, 10), 3) - self.assertEqual(ws._clamp_int(4.7, 6, 1, 10), 4) - self.assertEqual(ws._clamp_int(99, 6, 1, 10), 10) - self.assertEqual(ws._clamp_int(0, 6, 1, 10), 1) - self.assertEqual(ws._clamp_int("nonsense", 6, 1, 10), 6) - self.assertEqual(ws._clamp_int(None, 6, 1, 10), 6) - self.assertEqual(ws._clamp_int(True, 6, 1, 10), 6) + def test_num_results_is_clamped_and_coerced(self): + patcher, client = _with_exa( + _exa_response(200, {"results": []}), + _exa_response(200, {"results": []}), + _exa_response(200, {"results": []}), + ) + with patcher: + ws.execute_web_search_call({"query": "q", "num_results": 99}) + ws.execute_web_search_call({"query": "q", "num_results": "3"}) + ws.execute_web_search_call({"query": "q", "num_results": "junk"}) + self.assertEqual( + [p["numResults"] for p in client.payloads], + [ws.MAX_NUM_RESULTS, 3, ws.DEFAULT_NUM_RESULTS], + ) + + def test_recency_days_maps_to_a_published_date_floor(self): + patcher, client = _with_exa(_exa_response(200, {"results": []})) + with patcher: + ws.execute_web_search_call({"query": "q", "recency_days": 7}) + self.assertIn("startPublishedDate", client.payloads[0]) + + def test_recency_days_omitted_by_default(self): + patcher, client = _with_exa(_exa_response(200, {"results": []})) + with patcher: + ws.execute_web_search_call({"query": "q"}) + self.assertNotIn("startPublishedDate", client.payloads[0]) # --------------------------------------------------------------------------- @@ -141,42 +132,31 @@ def test_num_results_is_clamped_not_rejected(self): class TestExaRequestShaping(unittest.TestCase): - def test_default_request_asks_for_text_only(self): - """highlights/summary are each billed per page; text is what's used.""" - patcher, client = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}) - ) + def test_request_asks_for_text_only_with_char_cap(self): + """`highlights`/`summary` are each billed as another content type.""" + patcher, client = _with_exa(_exa_response(200, {"results": []})) with patcher: - ws.execute_web_search_call({"query": "who won"}) - + ws.run_web_search("anything") payload = client.payloads[0] - self.assertEqual(payload["query"], "who won") - self.assertEqual(payload["type"], ws.EXA_SEARCH_TYPE) - self.assertEqual(payload["numResults"], ws.DEFAULT_NUM_RESULTS) self.assertEqual( payload["contents"], {"text": {"maxCharacters": ws.MAX_RESULT_CHARS}} ) - self.assertNotIn("startPublishedDate", payload) + self.assertEqual(payload["type"], ws.EXA_SEARCH_TYPE) - def test_num_results_is_forwarded_and_capped(self): - patcher, client = _with_exa( - _exa_response(200, {"results": []}), - _exa_response(200, {"results": []}), + def test_reported_cost_is_captured_but_not_authoritative(self): + patcher, _ = _with_exa( + _exa_response( + 200, + { + "results": [_exa_result("https://a.com")], + "costDollars": {"total": 0.012}, + }, + ) ) with patcher: - ws.execute_web_search_call({"query": "q", "num_results": 3}) - ws.execute_web_search_call({"query": "q", "num_results": 500}) - - self.assertEqual(client.payloads[0]["numResults"], 3) - self.assertEqual(client.payloads[1]["numResults"], ws.MAX_NUM_RESULTS) - - def test_recency_days_becomes_a_published_date_floor(self): - patcher, client = _with_exa(_exa_response(200, {"results": []})) - with patcher: - ws.execute_web_search_call({"query": "q", "recency_days": 7}) - - cutoff = client.payloads[0]["startPublishedDate"] - self.assertRegex(cutoff, r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.000Z$") + outcome = ws.run_web_search("q") + self.assertEqual(outcome.reported_cost_usd, 0.012) + self.assertTrue(outcome.billable) # --------------------------------------------------------------------------- @@ -185,80 +165,67 @@ def test_recency_days_becomes_a_published_date_floor(self): class TestResultFormatting(unittest.TestCase): - def test_results_are_numbered_with_urls_and_citations(self): - patcher, _ = _with_exa( - _exa_response( - 200, - { - "results": [ - _exa_result("https://a.com", "First", "Body A"), - _exa_result("https://b.com", "Second", "Body B"), - ] - }, - ) - ) + def _search(self, results): + patcher, _ = _with_exa(_exa_response(200, {"results": results})) with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) + return ws.run_web_search("test query") - self.assertTrue(outcome.billable) - self.assertFalse(outcome.is_error) - self.assertIn("[1] First", outcome.content) - self.assertIn("https://a.com", outcome.content) - self.assertIn("[2] Second", outcome.content) - self.assertIn("Body B", outcome.content) + def test_results_are_numbered_and_carry_url_and_date(self): + outcome = self._search( + [ + _exa_result("https://a.com", title="Alpha"), + _exa_result("https://b.com", title="Beta"), + ] + ) + self.assertIn("[1] Alpha", outcome.content) + self.assertIn("[2] Beta", outcome.content) + self.assertIn("URL: https://a.com", outcome.content) + self.assertIn("Published: 2026-03-04", outcome.content) + + def test_citations_match_only_results_shown_to_the_model(self): + long_text = "x" * ws.MAX_RESULT_CHARS + results = [ + _exa_result(f"https://site{i}.com", text=long_text) for i in range(30) + ] + outcome = self._search(results) + self.assertLess(len(outcome.citations), 30) + for citation in outcome.citations: + self.assertIn(citation["url"], outcome.content) + self.assertLessEqual( + len(outcome.content), ws.MAX_TOTAL_CHARS + ws.MAX_RESULT_CHARS + ) + + def test_citation_shape(self): + outcome = self._search([_exa_result("https://a.com", title="Alpha")]) self.assertEqual( outcome.citations, [ { - "title": "First", + "title": "Alpha", "url": "https://a.com", "published_date": "2026-03-04T10:00:00.000Z", - }, - { - "title": "Second", - "url": "https://b.com", - "published_date": "2026-03-04T10:00:00.000Z", - }, + } ], ) - def test_results_without_a_url_are_skipped(self): - patcher, _ = _with_exa( - _exa_response( - 200, - {"results": [{"title": "No URL"}, _exa_result("https://ok.com")]}, - ) + def test_result_without_url_is_skipped(self): + outcome = self._search( + [{"title": "no url", "text": "t"}, _exa_result("https://a.com")] ) - with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) - - self.assertEqual([c["url"] for c in outcome.citations], ["https://ok.com"]) - - def test_total_size_is_capped_and_citations_match_what_was_shown(self): - """One verbose page must not crowd out the rest, or balloon the bill.""" - big = "x" * ws.MAX_RESULT_CHARS - results = [_exa_result(f"https://a{i}.com", f"T{i}", big) for i in range(20)] - patcher, _ = _with_exa(_exa_response(200, {"results": results})) - with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) - - self.assertLessEqual(len(outcome.content), ws.MAX_TOTAL_CHARS + 500) - self.assertLess(len(outcome.citations), 20) - # Every citation corresponds to a block actually put in front of the model. - for citation in outcome.citations: - self.assertIn(citation["url"], outcome.content) + self.assertEqual(len(outcome.citations), 1) - def test_long_excerpts_are_truncated(self): - patcher, _ = _with_exa( - _exa_response( - 200, - {"results": [_exa_result("https://a.com", "T", "y" * 9_000)]}, - ) + def test_per_result_text_is_truncated(self): + outcome = self._search( + [_exa_result("https://a.com", text="y" * (ws.MAX_RESULT_CHARS * 2))] ) - with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) + self.assertIn("…", outcome.content) - self.assertNotIn("y" * (ws.MAX_RESULT_CHARS + 1), outcome.content) + def test_zero_results_is_billable_but_says_so(self): + outcome = self._search([]) + self.assertTrue(outcome.billable) + self.assertFalse(outcome.is_error) + self.assertIn("No web results", outcome.content) + self.assertEqual(outcome.citations, []) # --------------------------------------------------------------------------- @@ -267,69 +234,41 @@ def test_long_excerpts_are_truncated(self): class TestSearchFailureModes(unittest.TestCase): - def test_no_key_injected_is_a_recoverable_error(self): + def test_no_key_configured_is_an_unbilled_error(self): with patch.object(ws, "_exa_http_client", None): - outcome = ws.execute_web_search_call({"query": "q"}) + outcome = ws.run_web_search("q") self.assertTrue(outcome.is_error) self.assertFalse(outcome.billable) - self.assertIn("not configured", outcome.content) - - def test_http_error_surfaces_the_provider_detail_and_is_not_billed(self): - patcher, _ = _with_exa(_exa_response(401, {"error": "invalid api key"})) - with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) - self.assertTrue(outcome.is_error) - self.assertFalse(outcome.billable) - self.assertIn("401", outcome.content) - self.assertIn("invalid api key", outcome.content) - - def test_transport_error_is_not_billed(self): + def test_transport_error_is_an_unbilled_error(self): import httpx client = Mock() - client.post.side_effect = httpx.ConnectError("no route") + client.post.side_effect = httpx.ConnectError("boom") with patch.object(ws, "_exa_http_client", client): - outcome = ws.execute_web_search_call({"query": "q"}) - + outcome = ws.run_web_search("q") self.assertTrue(outcome.is_error) self.assertFalse(outcome.billable) + self.assertIn("could not reach", outcome.content) - def test_malformed_json_is_not_billed(self): - response = _exa_response(200) - response.json.side_effect = ValueError("nope") - patcher, _ = _with_exa(response) + def test_http_error_surfaces_exa_detail(self): + patcher, _ = _with_exa(_exa_response(401, {"error": "invalid api key"})) with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) - + outcome = ws.run_web_search("q") self.assertTrue(outcome.is_error) self.assertFalse(outcome.billable) + self.assertIn("401", outcome.content) + self.assertIn("invalid api key", outcome.content) - def test_zero_results_is_billable_but_tells_the_model(self): - """The Exa request was consumed, and the model must not fake an answer.""" - patcher, _ = _with_exa(_exa_response(200, {"results": []})) - with patcher: - outcome = ws.execute_web_search_call({"query": "obscure thing"}) - - self.assertTrue(outcome.billable) - self.assertFalse(outcome.is_error) - self.assertIn("No web results", outcome.content) - self.assertEqual(outcome.citations, []) - - def test_reported_cost_is_captured_for_reconciliation_only(self): - patcher, _ = _with_exa( - _exa_response( - 200, - { - "results": [_exa_result("https://a.com")], - "costDollars": {"total": 0.008}, - }, - ) - ) - with patcher: - outcome = ws.execute_web_search_call({"query": "q"}) - - self.assertEqual(outcome.reported_cost_usd, 0.008) + def test_malformed_json_is_an_unbilled_error(self): + response = Mock() + response.status_code = 200 + response.json.side_effect = ValueError("bad json") + client = _ExaClient([response]) + with patch.object(ws, "_exa_http_client", client): + outcome = ws.run_web_search("q") + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) # --------------------------------------------------------------------------- @@ -337,303 +276,58 @@ def test_reported_cost_is_captured_for_reconciliation_only(self): # --------------------------------------------------------------------------- -class TestWebSearchPricing(unittest.TestCase): - def test_every_text_model_supports_search(self): - for model in ( - "gpt-4.1", - "claude-sonnet-4-5", - "gemini-2.5-flash", - "grok-4", - "seed-1.6", - "hermes-4-405b", - "glm-5.2", - ): - self.assertTrue(model_supports_web_search(model), model) - - def test_image_models_do_not(self): - for model in ("grok-2-image", "gemini-2.5-flash-image"): - self.assertFalse(model_supports_web_search(model), model) - - def test_one_flat_price_across_providers(self): - """The whole point: a search costs the same wherever it runs.""" - prices = { - model: get_web_search_price_usd(model) - for model in ( - "gpt-4.1", - "claude-sonnet-4-5", - "gemini-2.5-flash", - "grok-4", - "seed-1.6", - "hermes-4-405b", - "glm-5.2", - ) - } - self.assertEqual(set(prices.values()), {WEB_SEARCH_PRICE_USD}) - - def test_image_models_are_free(self): - self.assertEqual(get_web_search_price_usd("grok-2-image"), Decimal("0")) - - def test_unknown_model_raises(self): - with self.assertRaises(ValueError): - get_web_search_price_usd("not-a-real-model") - - -def _usage(input_tokens: int = 100, output_tokens: int = 50) -> dict: - return {"prompt_tokens": input_tokens, "completion_tokens": output_tokens} - - -def _cost(usage, model, web_search_count=0, price=Decimal("0.10")): - feed = SimpleNamespace(get_price=lambda: price) - with patch("tee_gateway.price_feed.get_price_feed", return_value=feed): - return compute_session_cost(model, usage, web_search_count=web_search_count) - - -class TestSessionCostWithWebSearch(unittest.TestCase): - def test_web_search_increases_cost(self): - base = _cost(_usage(), "gpt-4.1") - searched = _cost(_usage(), "gpt-4.1", web_search_count=2) - self.assertIsInstance(base, SessionCost) - self.assertIsInstance(searched, SessionCost) - self.assertGreater(searched.cost_opg, base.cost_opg) - - def test_surcharge_is_exactly_searches_times_the_flat_rate(self): - """The client-verifiable property: surcharge == searches * rate.""" - base = _cost(_usage(), "gpt-4.1") - searched = _cost(_usage(), "gpt-4.1", web_search_count=3) - scale = Decimal(10) ** 18 - delta_usd = (Decimal(searched.cost_opg - base.cost_opg) / scale) * Decimal( - "0.10" - ) - self.assertAlmostEqual(delta_usd, 3 * WEB_SEARCH_PRICE_USD, places=6) - - def test_same_surcharge_on_a_different_provider(self): - deltas = [] - for model in ("gpt-4.1", "seed-1.6"): - base = _cost(_usage(), model) - searched = _cost(_usage(), model, web_search_count=2) - deltas.append(searched.cost_opg - base.cost_opg) - self.assertEqual(deltas[0], deltas[1]) - - def test_zero_searches_matches_no_web_search(self): - a = _cost(_usage(), "gpt-4.1", web_search_count=0) - b = _cost(_usage(), "gpt-4.1") - self.assertEqual(a.cost_opg, b.cost_opg) - - -# --------------------------------------------------------------------------- -# search_loop -# --------------------------------------------------------------------------- - - -def _search_call(query="q", call_id="call_1"): - return {"name": "web_search", "args": {"query": query}, "id": call_id} - - -def _client_call(name="get_weather", call_id="call_2"): - return {"name": name, "args": {}, "id": call_id} - - -def _ai(content="", tool_calls=None, tokens=(10, 5)): - message = AIMessage(content=content, tool_calls=tool_calls or []) - message.usage_metadata = { - "input_tokens": tokens[0], - "output_tokens": tokens[1], - "total_tokens": sum(tokens), - } - return message +class _FakePriceFeed: + def __init__(self, price): + self._price = price + def get_price(self): + if isinstance(self._price, Exception): + raise self._price + return self._price -class _ScriptedModel: - """Returns queued AIMessages, recording the messages it was invoked with.""" - def __init__(self, *responses): - self.responses = list(responses) - self.calls: list[list] = [] - - def invoke(self, messages): - self.calls.append(list(messages)) - return self.responses.pop(0) - - -class TestSplitToolCalls(unittest.TestCase): - def test_partitions_ours_from_the_callers(self): - ours, theirs = split_tool_calls([_search_call(), _client_call()]) - self.assertEqual([c["name"] for c in ours], ["web_search"]) - self.assertEqual([c["name"] for c in theirs], ["get_weather"]) - - def test_handles_none_and_empty(self): - self.assertEqual(split_tool_calls(None), ([], [])) - self.assertEqual(split_tool_calls([]), ([], [])) - - -class TestSearchLoopState(unittest.TestCase): - def test_usage_sums_across_rounds(self): - state = SearchLoopState() - state.add_usage({"prompt_tokens": 100, "completion_tokens": 10}) - state.add_usage({"prompt_tokens": 400, "completion_tokens": 20}) - assert state.usage is not None - self.assertEqual(state.usage["prompt_tokens"], 500) - self.assertEqual(state.usage["completion_tokens"], 30) - - def test_usage_stays_none_when_nothing_reported(self): - state = SearchLoopState() - state.add_usage(None) - self.assertIsNone(state.usage) - - def test_citations_are_deduped_by_url(self): - state = SearchLoopState() - state.add_citations([{"title": "A", "url": "https://a.com"}]) - state.add_citations( - [ - {"title": "A again", "url": "https://a.com"}, - {"title": "B", "url": "https://b.com"}, - ] - ) +class TestWebSearchPricing(unittest.TestCase): + def test_flat_cost_converts_usd_to_opg(self): + feed = _FakePriceFeed(Decimal("0.10")) + with patch("tee_gateway.price_feed.get_price_feed", return_value=feed): + cost = compute_web_search_cost() + self.assertIsNotNone(cost) + # $0.015 at $0.10/OPG => 0.15 OPG = 15e16 smallest units. + expected_opg = int((WEB_SEARCH_PRICE_USD / Decimal("0.10")) * Decimal(10) ** 18) + self.assertEqual(cost.cost_opg, expected_opg) + # The USD figure reconciles from the rounded OPG value. self.assertEqual( - [c["url"] for c in state.citations], ["https://a.com", "https://b.com"] - ) - - -class TestExecuteSearchCalls(unittest.TestCase): - def test_builds_tool_messages_and_counts_billable_searches(self): - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}) - ) - state = SearchLoopState() - with patcher: - messages = execute_search_calls([_search_call("news", "abc")], state) - - self.assertEqual(len(messages), 1) - self.assertIsInstance(messages[0], ToolMessage) - self.assertEqual(messages[0].tool_call_id, "abc") - self.assertEqual(messages[0].name, "web_search") - self.assertEqual(messages[0].status, "success") - self.assertEqual(state.search_count, 1) - self.assertEqual(len(state.citations), 1) - - def test_failed_search_yields_an_error_tool_message_and_no_charge(self): - patcher, _ = _with_exa(_exa_response(500, {"error": "boom"})) - state = SearchLoopState() - with patcher: - messages = execute_search_calls([_search_call()], state) - - self.assertEqual(messages[0].status, "error") - self.assertEqual(state.search_count, 0) - - def test_status_callback_receives_each_query(self): - patcher, _ = _with_exa( - _exa_response(200, {"results": []}), _exa_response(200, {"results": []}) - ) - seen: list[str] = [] - with patcher: - execute_search_calls( - [_search_call("first", "1"), _search_call("second", "2")], - SearchLoopState(), - on_search=seen.append, + cost.cost_usd, + Decimal(cost.cost_opg) / Decimal(10) ** 18 * Decimal("0.10"), + ) + + def test_price_feed_failure_returns_none(self): + feed = _FakePriceFeed(ValueError("feed down")) + with patch("tee_gateway.price_feed.get_price_feed", return_value=feed): + self.assertIsNone(compute_web_search_cost()) + + def test_chat_token_cost_carries_no_search_surcharge(self): + """Search billing left the chat path entirely with the loop.""" + feed = _FakePriceFeed(Decimal("0.10")) + usage = {"prompt_tokens": 1000, "completion_tokens": 100} + with patch("tee_gateway.price_feed.get_price_feed", return_value=feed): + cost = compute_session_cost("gpt-4.1", usage) + self.assertIsNotNone(cost) + # Recompute from the rate card alone: tokens only, nothing else. + from tee_gateway.model_registry import get_model_config + + cfg = get_model_config("gpt-4.1") + raw_usd = 1000 * cfg.input_price_usd + 100 * cfg.output_price_usd + expected_opg = int( + ((raw_usd / Decimal("0.10")) * Decimal(10) ** 18).to_integral_value( + rounding="ROUND_CEILING" ) - self.assertEqual(seen, ["first", "second"]) - - def test_a_throwing_status_callback_does_not_break_the_search(self): - patcher, _ = _with_exa(_exa_response(200, {"results": []})) - state = SearchLoopState() - with patcher: - messages = execute_search_calls( - [_search_call()], - state, - on_search=Mock(side_effect=RuntimeError("ui gone")), - ) - self.assertEqual(len(messages), 1) - self.assertEqual(state.search_count, 1) - - -class TestRunSearchLoop(unittest.TestCase): - def test_plain_answer_returns_immediately(self): - model = _ScriptedModel(_ai("just an answer")) - state = SearchLoopState() - result = run_search_loop(model, model, [], state) - self.assertEqual(result.content, "just an answer") - self.assertEqual(state.search_count, 0) - self.assertEqual(state.rounds, 1) - - def test_searches_then_answers_and_feeds_results_back(self): - model = _ScriptedModel( - _ai("", [_search_call("og price")]), - _ai("The answer, with sources."), - ) - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}) - ) - state = SearchLoopState() - messages: list = [] - with patcher: - result = run_search_loop(model, model, messages, state) - - self.assertEqual(result.content, "The answer, with sources.") - self.assertEqual(state.search_count, 1) - self.assertEqual(state.rounds, 2) - # The second invocation saw the assistant turn plus the search results. - second_round = model.calls[1] - self.assertIsInstance(second_round[-1], ToolMessage) - self.assertIn("https://a.com", second_round[-1].content) - - def test_every_round_is_billed_not_just_the_last(self): - """Each round re-sends the conversation; the caller pays for all of it.""" - model = _ScriptedModel( - _ai("", [_search_call("a")], tokens=(100, 10)), - _ai("", [_search_call("b")], tokens=(600, 12)), - _ai("done", tokens=(1200, 40)), - ) - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}), - _exa_response(200, {"results": [_exa_result("https://b.com")]}), - ) - state = SearchLoopState() - with patcher: - run_search_loop(model, model, [], state) - - assert state.usage is not None - self.assertEqual(state.usage["prompt_tokens"], 1900) - self.assertEqual(state.usage["completion_tokens"], 62) - self.assertEqual(state.search_count, 2) - - def test_client_tool_call_is_terminal(self): - model = _ScriptedModel(_ai("", [_client_call()])) - state = SearchLoopState() - result = run_search_loop(model, model, [], state) - self.assertEqual([c["name"] for c in result.tool_calls], ["get_weather"]) - self.assertEqual(state.rounds, 1) - - def test_round_cap_forces_an_answer_with_the_search_tool_unbound(self): - """A model that keeps searching must still terminate.""" - searching = _ScriptedModel( - *[_ai("", [_search_call(f"q{i}")]) for i in range(MAX_SEARCH_ROUNDS)] - ) - answering = _ScriptedModel(_ai("forced answer")) - patcher, _ = _with_exa( - *[ - _exa_response(200, {"results": [_exa_result(f"https://a{i}.com")]}) - for i in range(MAX_SEARCH_ROUNDS) - ] - ) - state = SearchLoopState() - with patcher: - result = run_search_loop(searching, answering, [], state) - - self.assertEqual(result.content, "forced answer") - self.assertEqual(state.search_count, MAX_SEARCH_ROUNDS) - self.assertEqual(state.rounds, MAX_SEARCH_ROUNDS + 1) - # The final round went to the model without the search tool bound. - self.assertEqual(len(answering.calls), 1) - - def test_mixed_turn_drops_our_calls_and_keeps_the_callers(self): - message = _ai("", [_search_call(), _client_call()]) - self.assertEqual( - [c["name"] for c in strip_search_tool_calls(message)], ["get_weather"] ) + self.assertEqual(cost.cost_opg, expected_opg) # --------------------------------------------------------------------------- -# chat_controller integration +# /v1/web_search controller # --------------------------------------------------------------------------- @@ -644,342 +338,203 @@ def _mock_tee_keys(): return tee -class _ControllerHarness(unittest.TestCase): - """Shared patching for the chat_controller tests.""" - +class TestWebSearchController(unittest.TestCase): def setUp(self): - self.patchers = [ - patch("tee_gateway.controllers.chat_controller.compute_session_cost"), - patch("tee_gateway.controllers.chat_controller.get_tee_keys"), - patch("tee_gateway.controllers.chat_controller.get_chat_model_cached"), - patch("tee_gateway.controllers.chat_controller.connexion"), - patch( - "tee_gateway.controllers.chat_controller.web_search_available", - return_value=True, - ), - ] - ( - self.cost, - self.tee, - self.get_model, - self.connexion, - self.available, - ) = [p.start() for p in self.patchers] - self.addCleanup(lambda: [p.stop() for p in self.patchers]) - self.tee.return_value = _mock_tee_keys() - self.cost.return_value = None - - def request(self, **overrides): - body = { - "model": "claude-sonnet-4-5", - "messages": [{"role": "user", "content": "latest news?"}], - "stream": False, - } - body.update(overrides) - self.connexion.request.is_json = True - self.connexion.request.get_json.return_value = body - - -class TestChatControllerNonStreaming(_ControllerHarness): - def test_flag_binds_our_function_tool_on_an_anthropic_model(self): - self.request(web_search=True) - model = Mock() - model.invoke.return_value = _ai("Here is the news.") - model.bind_tools.return_value = model - self.get_model.return_value = model - - patcher, _ = _with_exa() - with patcher: - result = create_chat_completion(None) - - bound = model.bind_tools.call_args[0][0] - self.assertEqual( - [t["function"]["name"] for t in bound if isinstance(t, dict)], - ["web_search"], + app = Flask(__name__) + app.add_url_rule( + "/v1/web_search", "web-search", create_web_search, methods=["POST"] ) - # No provider-native tool types any more. - for tool in bound: - self.assertEqual(tool.get("type"), "function") - self.assertIn("choices", result) - - def test_search_runs_in_enclave_and_is_never_handed_to_the_client(self): - self.request(web_search=True) - model = Mock() - model.bind_tools.return_value = model - model.invoke.side_effect = [ - _ai("", [_search_call("og token price")]), - _ai("It trades at $X."), - ] - self.get_model.return_value = model + self.client = app.test_client() - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://coin.com")]}) + self.tee = patch( + "tee_gateway.controllers.web_search_controller.get_tee_keys", + return_value=_mock_tee_keys(), ) - with patcher: - result = create_chat_completion(None) + self.tee.start() + self.addCleanup(self.tee.stop) - choice = result["choices"][0] - self.assertEqual(choice["finish_reason"], "stop") - self.assertNotIn("tool_calls", choice["message"]) - self.assertEqual(choice["message"]["content"], "It trades at $X.") - self.assertEqual( - [c["url"] for c in choice["message"]["citations"]], ["https://coin.com"] + self.feed = patch( + "tee_gateway.price_feed.get_price_feed", + return_value=_FakePriceFeed(Decimal("0.10")), ) - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) - - def test_bytedance_can_search_too(self): - """The case the old native-search implementation could not serve at all.""" - self.request(model="seed-1.6", web_search=True) - model = Mock() - model.bind_tools.return_value = model - model.invoke.side_effect = [ - _ai("", [_search_call("q")]), - _ai("answer"), - ] - self.get_model.return_value = model + self.feed.start() + self.addCleanup(self.feed.stop) - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}) + def _post(self, body): + return self.client.post("/v1/web_search", json=body) + + def test_successful_search_returns_signed_result_with_cost(self): + patcher, exa = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com", "Alpha")]}) ) with patcher: - result = create_chat_completion(None) - - self.assertEqual(result["choices"][0]["message"]["content"], "answer") - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) - - def test_billed_usage_is_the_sum_over_rounds(self): - self.request(web_search=True) - model = Mock() - model.bind_tools.return_value = model - model.invoke.side_effect = [ - _ai("", [_search_call("q")], tokens=(100, 10)), - _ai("answer", tokens=(700, 30)), - ] - self.get_model.return_value = model + response = self._post({"query": "latest news"}) + + self.assertEqual(response.status_code, 200) + body = response.get_json() + self.assertEqual(body["object"], "web_search.result") + self.assertEqual(body["query"], "latest news") + self.assertIn("[1] Alpha", body["content"]) + self.assertEqual(body["citations"][0]["url"], "https://a.com") + # Signed like every other paid endpoint. + for field in ( + "tee_signature", + "tee_request_hash", + "tee_output_hash", + "tee_timestamp", + "tee_id", + ): + self.assertIn(field, body) + # Billed at the flat per-search rate. + expected_opg = int((WEB_SEARCH_PRICE_USD / Decimal("0.10")) * Decimal(10) ** 18) + self.assertEqual(body["opengradient"]["cost_opg"], str(expected_opg)) + self.assertEqual(exa.payloads[0]["query"], "latest news") + def test_request_hash_covers_the_canonical_body(self): + """The client can recompute the hash from exactly what it sent.""" + from tee_gateway.tee_manager import compute_tee_msg_hash + + request_body = {"query": "q", "num_results": 2} patcher, _ = _with_exa( _exa_response(200, {"results": [_exa_result("https://a.com")]}) ) with patcher: - result = create_chat_completion(None) - - self.assertEqual(result["usage"]["prompt_tokens"], 800) - self.assertEqual(result["usage"]["completion_tokens"], 40) + body = self._post(request_body).get_json() - def test_client_tools_still_come_back_for_the_client_to_run(self): - self.request( - web_search=True, - tools=[ - { - "type": "function", - "function": {"name": "get_weather", "parameters": {}}, - } - ], + request_bytes = json.dumps(request_body, sort_keys=True).encode("utf-8") + _, input_hash_hex, output_hash_hex = compute_tee_msg_hash( + request_bytes, body["content"], body["tee_timestamp"] ) - model = Mock() - model.bind_tools.return_value = model - model.invoke.return_value = _ai("", [_client_call()]) - self.get_model.return_value = model + self.assertEqual(body["tee_request_hash"], input_hash_hex) + self.assertEqual(body["tee_output_hash"], output_hash_hex) - patcher, _ = _with_exa() + def test_zero_results_still_bills_and_tells_the_model(self): + patcher, _ = _with_exa(_exa_response(200, {"results": []})) with patcher: - result = create_chat_completion(None) - - choice = result["choices"][0] - self.assertEqual(choice["finish_reason"], "tool_calls") - self.assertEqual( - [tc["function"]["name"] for tc in choice["message"]["tool_calls"]], - ["get_weather"], - ) - - def test_no_flag_binds_nothing_and_bills_no_search(self): - self.request(model="gpt-4.1", messages=[{"role": "user", "content": "hi"}]) - model = Mock() - model.invoke.return_value = _ai("hi") - model.bind_tools.return_value = model - self.get_model.return_value = model - - create_chat_completion(None) + response = self._post({"query": "obscure thing"}) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertIn("No web results", body["content"]) + self.assertEqual(body["citations"], []) + self.assertIn("opengradient", body) + + def test_missing_query_is_400(self): + with patch.object(ws, "_exa_http_client", Mock()): + for payload in ({}, {"query": ""}, {"query": 7}): + response = self._post(payload) + self.assertEqual(response.status_code, 400, payload) + + def test_no_exa_key_is_503(self): + with patch.object(ws, "_exa_http_client", None): + response = self._post({"query": "q"}) + self.assertEqual(response.status_code, 503) - model.bind_tools.assert_not_called() - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) + def test_exa_failure_is_502_without_cost_block(self): + patcher, _ = _with_exa(_exa_response(500, {"error": "upstream broke"})) + with patcher: + response = self._post({"query": "q"}) + self.assertEqual(response.status_code, 502) + body = response.get_json() + self.assertNotIn("opengradient", body) + self.assertIn("upstream broke", body["error"]) + + def test_price_feed_outage_returns_result_without_cost_block(self): + """Fail-open like chat: the client gets its answer, unsettled.""" + self.feed.stop() + feed = patch( + "tee_gateway.price_feed.get_price_feed", + return_value=_FakePriceFeed(ValueError("down")), + ) + feed.start() + self.addCleanup(feed.stop) + # Re-arm the harness patcher reference so cleanup doesn't double-stop. + self.feed = feed - def test_missing_exa_key_answers_without_searching(self): - """Better a plain answer than a tool that always fails.""" - self.available.return_value = False - self.request(web_search=True) - model = Mock() - model.invoke.return_value = _ai("answer from memory") - model.bind_tools.return_value = model - self.get_model.return_value = model + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) + ) + with patcher: + response = self._post({"query": "q"}) + self.assertEqual(response.status_code, 200) + self.assertNotIn("opengradient", response.get_json()) - result = create_chat_completion(None) - model.bind_tools.assert_not_called() - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) - self.assertEqual( - result["choices"][0]["message"]["content"], "answer from memory" - ) +# --------------------------------------------------------------------------- +# OHTTP inner-endpoint dispatch +# --------------------------------------------------------------------------- -def _chunk(content="", tool_call_chunks=None, usage=None): - message = AIMessageChunk(content=content, tool_call_chunks=tool_call_chunks or []) - if usage: - message.usage_metadata = { - "input_tokens": usage[0], - "output_tokens": usage[1], - "total_tokens": sum(usage), - } - return message - - -def _search_chunk(query="q", call_id="call_1"): - return _chunk( - tool_call_chunks=[ - { - "name": "web_search", - "args": json.dumps({"query": query}), - "id": call_id, - "index": 0, - } - ] - ) - - -def _sse_frames(response): - """Parse a Flask SSE response into the list of JSON data frames.""" - raw = "".join( - part.decode("utf-8") if isinstance(part, bytes) else part - for part in response.response - ) - frames = [] - for line in raw.split("\n\n"): - line = line.strip() - if line.startswith("data: ") and line != "data: [DONE]": - frames.append(json.loads(line[len("data: ") :])) - return frames - - -class TestChatControllerStreaming(_ControllerHarness): - def test_search_rounds_are_invisible_to_the_client(self): - """The client sees status, then the answer — never our tool calls.""" - self.request(web_search=True, stream=True) - model = Mock() - model.bind_tools.return_value = model - model.stream.side_effect = [ - iter([_search_chunk("og price"), _chunk(usage=(100, 10))]), - iter([_chunk("It "), _chunk("trades."), _chunk(usage=(700, 20))]), - ] - self.get_model.return_value = model +class _FakeDecap: + plaintext = b"" # set per-test + response_key = b"k" * 32 + response_key_chunked = b"c" * 32 + enc = b"e" * 32 - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://coin.com")]}) - ) - with patcher: - frames = _sse_frames(create_chat_completion(None)) - - # A status frame naming the query reached the client. - statuses = [f["web_search"] for f in frames if "web_search" in f] - self.assertEqual(statuses, [{"status": "searching", "query": "og price"}]) - - # No tool_calls delta was ever forwarded. - for frame in frames: - delta = frame.get("choices", [{}])[0].get("delta", {}) - self.assertNotIn("tool_calls", delta) - - final = frames[-1] - self.assertEqual(final["choices"][0]["finish_reason"], "stop") - self.assertEqual([c["url"] for c in final["citations"]], ["https://coin.com"]) - # Both rounds are billed. - self.assertEqual(final["usage"]["prompt_tokens"], 800) - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 1) - - def test_streamed_text_is_forwarded_and_signed(self): - self.request(web_search=True, stream=True) - model = Mock() - model.bind_tools.return_value = model - model.stream.side_effect = [ - iter([_search_chunk(), _chunk(usage=(10, 1))]), - iter([_chunk("Hello "), _chunk("world"), _chunk(usage=(20, 2))]), - ] - self.get_model.return_value = model - patcher, _ = _with_exa( - _exa_response(200, {"results": [_exa_result("https://a.com")]}) - ) - with patcher: - frames = _sse_frames(create_chat_completion(None)) +class TestOhttpEndpointDispatch(unittest.TestCase): + """The sealed payload's `endpoint` field picks the inner path.""" - text = "".join( - f["choices"][0]["delta"].get("content", "") - for f in frames - if f.get("choices") - ) - self.assertEqual(text, "Hello world") - self.assertIn("tee_signature", frames[-1]) - - def test_client_tool_calls_still_stream_through(self): - self.request( - web_search=True, - stream=True, - tools=[ - { - "type": "function", - "function": {"name": "get_weather", "parameters": {}}, - } - ], + def setUp(self): + app = Flask(__name__) + app.add_url_rule( + "/v1/ohttp", + "anonymous-chat", + ohttp_controller.create_anonymous_chat_completion, + methods=["POST"], ) - model = Mock() - model.bind_tools.return_value = model - model.stream.side_effect = [ - iter( - [ - _chunk( - tool_call_chunks=[ - { - "name": "get_weather", - "args": "{}", - "id": "c1", - "index": 0, - } - ] - ), - _chunk(usage=(10, 1)), - ] - ) - ] - self.get_model.return_value = model - - patcher, _ = _with_exa() - with patcher: - frames = _sse_frames(create_chat_completion(None)) + self.client = app.test_client() - names = [ - tc["function"]["name"] - for f in frames - for tc in f.get("choices", [{}])[0].get("delta", {}).get("tool_calls", []) - if tc.get("function", {}).get("name") - ] - self.assertEqual(names, ["get_weather"]) - self.assertEqual(frames[-1]["choices"][0]["finish_reason"], "tool_calls") - - def test_failed_search_still_produces_an_answer(self): - self.request(web_search=True, stream=True) - model = Mock() - model.bind_tools.return_value = model - model.stream.side_effect = [ - iter([_search_chunk(), _chunk(usage=(10, 1))]), - iter([_chunk("Couldn't verify that."), _chunk(usage=(20, 2))]), + tee = Mock() + tee.hpke_private_key = object() + self.patchers = [ + patch.object(ohttp_controller, "get_tee_keys", return_value=tee), + patch.object(ohttp_controller.ohttp, "decapsulate_request"), + patch.object(ohttp_controller, "_wsgi_subrequest"), + patch.object( + ohttp_controller.ohttp, "encapsulate_response", return_value=b"sealed" + ), ] - self.get_model.return_value = model + _, self.decap, self.subrequest, _ = [p.start() for p in self.patchers] + self.addCleanup(lambda: [p.stop() for p in self.patchers]) - patcher, _ = _with_exa(_exa_response(503, {"error": "unavailable"})) - with patcher: - frames = _sse_frames(create_chat_completion(None)) + self.subrequest.return_value = ( + 200, + [("Content-Type", "application/json")], + iter([b'{"ok": true}']), + ) + + def _post_inner(self, inner: dict): + decap = _FakeDecap() + decap.plaintext = json.dumps(inner).encode("utf-8") + self.decap.return_value = decap + return self.client.post( + "/v1/ohttp", + data=b"ciphertext", + content_type="message/ohttp-req", + ) + + def test_web_search_endpoint_routes_to_the_search_path(self): + response = self._post_inner({"endpoint": "web_search", "query": "q"}) + self.assertEqual(response.status_code, 200) + self.assertEqual(self.subrequest.call_args.kwargs["path"], "/v1/web_search") + # The discriminator is routing metadata, not part of the inner body. + forwarded = json.loads(self.subrequest.call_args.kwargs["body_bytes"]) + self.assertEqual(forwarded, {"query": "q"}) + + def test_absent_endpoint_still_means_chat(self): + """The original OHTTP contract: existing clients keep working.""" + response = self._post_inner({"model": "gpt-4.1", "messages": []}) + self.assertEqual(response.status_code, 200) + self.assertEqual( + self.subrequest.call_args.kwargs["path"], "/v1/chat/completions" + ) - self.assertEqual(frames[-1]["choices"][0]["finish_reason"], "stop") - self.assertNotIn("citations", frames[-1]) - self.assertEqual(self.cost.call_args.kwargs["web_search_count"], 0) + def test_unknown_endpoint_is_rejected_without_dispatch(self): + response = self._post_inner({"endpoint": "nope", "query": "q"}) + # Sealed error: outer 200 carrying an encapsulated {status: 400, ...}. + self.assertEqual(response.status_code, 200) + self.assertEqual(response.mimetype, "message/ohttp-res") + self.subrequest.assert_not_called() if __name__ == "__main__": diff --git a/tee_gateway/web_search.py b/tee_gateway/web_search.py index f5572b6..01b2eea 100644 --- a/tee_gateway/web_search.py +++ b/tee_gateway/web_search.py @@ -9,20 +9,21 @@ per *citation*). Models on providers without one (ByteDance, Nous, Z.ai) simply could not search at all. -Instead the gateway advertises ONE ordinary function tool — see -``get_web_search_tool`` — and executes it itself, inside the enclave, against -Exa. Consequences worth stating plainly: - - * It works on every model that can call a function, which is every non-image - model in the registry. There is no per-provider capability matrix. +Search is now a dedicated, model-free endpoint — ``POST /v1/web_search``, served +by ``controllers/web_search_controller.py`` on top of this module — that the +client's own tool loop calls when its model asks to search. The gateway runs no +tool loop of its own. Consequences worth stating plainly: + + * It works with every model that can call a function, which is every non-image + model in the registry: the client advertises the tool, the model calls it, + and this endpoint answers it. There is no per-provider capability matrix. * Results, excerpt sizes, and citations are identical across models, so answer quality stops depending on whose search backend the model happened to ship. - * A search is one flat price on every model (``WEB_SEARCH_PRICE_USD``), so - clients can verify the surcharge as ``searches * price`` rather than - reverse-engineering a provider's billable unit. - * The query never leaves the TEE except to Exa. Nothing about the search is - visible to the LLM provider beyond the result text the model is shown, and - nothing is visible to the gateway operator at all. + * A search is one flat price (``WEB_SEARCH_PRICE_USD``) with no model or token + dimension, settled like any other paid endpoint. + * The query never leaves the TEE except to Exa: the endpoint rides the same + encrypted OHTTP channel as chat, so the relay sees ciphertext and the + gateway operator sees nothing. The Exa API key is injected at runtime via ``POST /v1/keys`` like every provider key; it is never baked into the image. @@ -48,10 +49,6 @@ # interactive chat turn. EXA_SEARCH_TYPE = "auto" -# The tool name the model sees. Also the name the controllers match on to decide -# a tool call is ours to execute rather than the client's. -WEB_SEARCH_TOOL_NAME = "web_search" - DEFAULT_NUM_RESULTS = 6 MAX_NUM_RESULTS = 10 @@ -105,68 +102,6 @@ def web_search_available() -> bool: return _exa_http_client is not None -# --------------------------------------------------------------------------- -# Tool specification -# --------------------------------------------------------------------------- - - -def get_web_search_tool() -> dict[str, Any]: - """The ``web_search`` function tool, in OpenAI function-calling format. - - One spec for every provider: langchain converts this to each provider's own - tool format on ``bind_tools``. The schema is deliberately flat — three - scalar parameters, one required — because nested objects, unions, and - ``additionalProperties`` are exactly where the providers' function-calling - schema subsets diverge (Gemini's is the narrowest). Flat and boring is what - makes this work identically on all of them. - """ - return { - "type": "function", - "function": { - "name": WEB_SEARCH_TOOL_NAME, - "description": ( - "Search the live web and get back ranked results with an " - "excerpt of each page. Use this whenever the answer depends on " - "information you may not have: current events, anything after " - "your training cutoff, prices, releases, documentation, or any " - "specific fact you are not confident about. Prefer searching " - "over guessing. Write a focused natural-language query rather " - "than keywords, and call this again with a refined query if the " - "first set of results does not answer the question." - ), - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": ( - "What to search for, as a focused natural-language query." - ), - }, - "num_results": { - "type": "integer", - "description": ( - f"How many results to return, 1-{MAX_NUM_RESULTS}. " - f"Defaults to {DEFAULT_NUM_RESULTS}. Ask for more " - "only when the question needs broad coverage." - ), - }, - "recency_days": { - "type": "integer", - "description": ( - "Only return pages published within this many days. " - "Omit unless the question is genuinely " - "time-sensitive — it discards older pages that are " - "often the best sources." - ), - }, - }, - "required": ["query"], - }, - }, - } - - # --------------------------------------------------------------------------- # Execution # --------------------------------------------------------------------------- @@ -174,13 +109,12 @@ def get_web_search_tool() -> dict[str, Any]: @dataclass(frozen=True) class WebSearchOutcome: - """The result of one ``web_search`` tool call. + """The result of one web search. - ``content`` is the text handed back to the model as the tool result. - ``citations`` are surfaced to the client out-of-band (like generated - images) so a UI can show its sources. ``billable`` is False for failures - and for empty/invalid calls, so a caller is never charged for a search that - produced nothing. + ``content`` is model-ready text: the client feeds it back to its model as + the ``web_search`` tool result. ``citations`` are structured so a UI can + show its sources. ``billable`` is False for failures and for empty/invalid + calls, so a caller is never charged for a search that produced nothing. """ content: str @@ -194,11 +128,12 @@ class WebSearchOutcome: def execute_web_search_call(args: dict[str, Any]) -> WebSearchOutcome: - """Run one ``web_search`` tool call from its (already-parsed) arguments. + """Run one web search from loosely-typed request arguments. Never raises: every failure mode comes back as an error outcome whose - ``content`` reads as an instruction to the model, so a flaky search degrades - into "the model was told the search failed" rather than a dead request. + ``content`` reads as an instruction to the model (the client relays it as + the tool result), so a flaky search degrades into "the model was told the + search failed" rather than a dead request. """ query = args.get("query") if not isinstance(query, str) or not query.strip(): From 9ec8779b0909b6b43b50d0ab29fb88d5e7611459 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 12:46:53 +0000 Subject: [PATCH 3/3] Raise the /v1/web_search x402 session cap to 10 OPG Costs past the cap go unbilled, so the pre-authorization ceiling should sit well above any realistic session's search count. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01V9SCMV8LhnpZ21pzBLjzxQ --- tee_gateway/definitions.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tee_gateway/definitions.py b/tee_gateway/definitions.py index 0ec64c1..a18b014 100644 --- a/tee_gateway/definitions.py +++ b/tee_gateway/definitions.py @@ -78,11 +78,13 @@ # so the relay-paid encrypted endpoint needs a higher per-session cap. OHTTP_OPG_SESSION_MAX_SPEND: str = "5000000000000000000" -# /v1/web_search — maximum OPG spend per session (18 decimals: 1000000000000000000 = 1 OPG). +# /v1/web_search — maximum OPG spend per session (18 decimals: 10000000000000000000 = 10 OPG). # Each search settles at the flat WEB_SEARCH_PRICE_USD (see model_registry.py), -# so at the fallback OPG price this cap covers several searches per session -# with headroom for OPG price swings. -WEB_SEARCH_OPG_SESSION_MAX_SPEND: str = "1000000000000000000" +# ~0.15 OPG at the fallback price, so this cap covers ~60 searches per session +# with headroom for OPG price swings. Costs past the cap go UNBILLED (x402 +# add_cost rejects over-cap), so err high — the cap is a pre-authorization +# ceiling, never an amount actually charged. +WEB_SEARCH_OPG_SESSION_MAX_SPEND: str = "10000000000000000000" # /v1/completions — maximum OPG spend per session (18 decimals: 100000000000000000 = 0.1 OPG). # This is the upper-bound amount presented to the client during the x402 pre-check handshake.