diff --git a/CLAUDE.md b/CLAUDE.md index d083bcb..be08dc5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,6 +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, 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 @@ -69,6 +70,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 + `/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) @@ -92,18 +96,20 @@ 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, 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 | 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) | | `/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 @@ -144,6 +150,38 @@ 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 + +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. + ## 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..92fa5b5 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** - 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 @@ -90,21 +91,29 @@ 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) -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 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.** `/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 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/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..0b04410 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -22,11 +22,13 @@ 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, 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 @@ -51,6 +53,7 @@ COMPLETIONS_OPG_SESSION_MAX_SPEND, FACILITATOR_URL, OHTTP_OPG_SESSION_MAX_SPEND, + WEB_SEARCH_OPG_SESSION_MAX_SPEND, ) # --------------------------------------------------------------------------- @@ -321,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( @@ -395,6 +421,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 +487,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 +533,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 +548,10 @@ def health(): "tee_enabled": True, "uptime_seconds": int(time.time() - _started_at), "providers": 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(), }, 200 @@ -567,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. @@ -600,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() @@ -608,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/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..09612df 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -27,8 +27,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, @@ -104,13 +102,12 @@ 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). +def _build_tools_list(chat_request: CreateChatCompletionRequest) -> list: + """Normalize the caller's function tools into bind_tools() form. - 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. + 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: @@ -123,11 +120,6 @@ def _build_tools_list(chat_request: CreateChatCompletionRequest, provider: str) else: tools_list.append(tool) - 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 @@ -252,7 +244,7 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): # 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) + tools_list = _build_tools_list(chat_request) model = get_chat_model_cached( model=chat_request.model, @@ -260,7 +252,6 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): 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 ), @@ -373,12 +364,7 @@ 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 - ) + cost = compute_session_cost(chat_request.model, usage) if cost is not None: openai_response["opengradient"] = cost.model_dump(mode="json") @@ -418,7 +404,7 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): # 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) + tools_list = _build_tools_list(chat_request) model = get_chat_model_cached( model=chat_request.model, @@ -426,7 +412,6 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest): 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 ), @@ -503,10 +488,6 @@ def generate(): 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: @@ -580,13 +561,6 @@ def generate(): chunks_iter = model.stream(langchain_messages) # 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 - ) - # --- Text content --- if chunk.content: if isinstance(chunk.content, str): @@ -793,22 +767,13 @@ 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. 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=web_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 diff --git a/tee_gateway/controllers/completions_controller.py b/tee_gateway/controllers/completions_controller.py index 773dd74..976ed1c 100644 --- a/tee_gateway/controllers/completions_controller.py +++ b/tee_gateway/controllers/completions_controller.py @@ -13,9 +13,6 @@ 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.pricing import compute_session_cost @@ -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 = { @@ -55,21 +55,13 @@ def create_completion(body): 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]) - messages = [HumanMessage(content=body.prompt)] response = 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 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) @@ -106,10 +98,7 @@ 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 - ) + cost = compute_session_cost(body.model, usage) if cost is not None: completion_response["opengradient"] = cost.model_dump(mode="json") return completion_response 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..a18b014 100644 --- a/tee_gateway/definitions.py +++ b/tee_gateway/definitions.py @@ -78,6 +78,14 @@ # 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: 10000000000000000000 = 10 OPG). +# Each search settles at the flat WEB_SEARCH_PRICE_USD (see model_registry.py), +# ~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. # 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 b358957..2d9dbbb 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 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 979cd90..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 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. - 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,20 +79,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 call to the /v1/web_search endpoint. +# +# 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 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 # 4.5, Seedream 5.0 Lite) return the URL response format and require these extra @@ -718,21 +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's native web search. - - 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. - """ - 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")) - - -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 diff --git a/tee_gateway/openapi/openapi.yaml b/tee_gateway/openapi/openapi.yaml index 80cb352..7a954ee 100644 --- a/tee_gateway/openapi/openapi.yaml +++ b/tee_gateway/openapi/openapi.yaml @@ -2886,12 +2886,14 @@ components: $ref: "#/components/schemas/CreateChatCompletionRequest_model" web_search: default: false + deprecated: true 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. + 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 type: boolean @@ -3550,10 +3552,10 @@ components: type: boolean web_search: default: false + deprecated: true 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. + 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 1bd952a..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,15 +96,6 @@ def compute_session_cost( Decimal(out_tok) * cfg.output_price_usd ) - # Native web search is billed per search unit on top of token cost. - 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)) @@ -134,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), @@ -162,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/test/test_web_search.py b/tee_gateway/test/test_web_search.py index 277dcd3..fffef6f 100644 --- a/tee_gateway/test/test_web_search.py +++ b/tee_gateway/test/test_web_search.py @@ -1,198 +1,336 @@ """ -Unit tests for native web search support across providers. +Unit tests for the dedicated in-enclave web search endpoint (Exa). 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 - - pricing.compute_session_cost: per-search surcharge added to token cost - - chat_controller: web_search flag binds the tool and bills the searches + - 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 patch, Mock +from unittest.mock import Mock, patch -from langchain_core.messages import AIMessage +from flask import Flask -from tee_gateway.model_registry import ( - get_web_search_price_usd, - provider_supports_web_search, -) -from tee_gateway.llm_backend import ( - get_web_search_tool, - extract_web_search_count, -) -from tee_gateway.pricing import SessionCost, compute_session_cost -from tee_gateway.controllers.chat_controller import create_chat_completion +from tee_gateway import web_search as ws +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 # --------------------------------------------------------------------------- -# 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 _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 _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 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 __init__(self, responses): + self.responses = list(responses) + self.payloads: list[dict] = [] - 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 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 test_unknown_model_raises(self): - with self.assertRaises(ValueError): - get_web_search_price_usd("not-a-real-model") + +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 +# Availability # --------------------------------------------------------------------------- -class TestGetWebSearchTool(unittest.TestCase): - def test_openai_tool(self): - self.assertEqual(get_web_search_tool("openai"), {"type": "web_search"}) +class TestAvailability(unittest.TestCase): + def test_availability_tracks_the_injected_key(self): + ws.configure_exa_client("test-key") + try: + self.assertTrue(ws.web_search_available()) + finally: + ws.configure_exa_client(None) + self.assertFalse(ws.web_search_available()) - 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_google_tool(self): - self.assertEqual(get_web_search_tool("google"), {"google_search": {}}) +# --------------------------------------------------------------------------- +# Argument coercion +# --------------------------------------------------------------------------- - def test_xai_tool(self): - self.assertEqual(get_web_search_tool("x-ai"), {"type": "web_search"}) - def test_bytedance_has_no_bound_tool(self): - self.assertIsNone(get_web_search_tool("bytedance")) +class TestArgumentCoercion(unittest.TestCase): + def test_missing_or_blank_query_is_an_unbilled_error(self): + 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_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]) # --------------------------------------------------------------------------- -# 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_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.run_web_search("anything") + payload = client.payloads[0] + self.assertEqual( + payload["contents"], {"text": {"maxCharacters": ws.MAX_RESULT_CHARS}} + ) + self.assertEqual(payload["type"], ws.EXA_SEARCH_TYPE) + + 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: + outcome = ws.run_web_search("q") + self.assertEqual(outcome.reported_cost_usd, 0.012) + self.assertTrue(outcome.billable) - def test_plain_text_response_has_no_searches(self): - self.assertEqual(extract_web_search_count(AIMessage(content="hi")), 0) - 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"}, +# --------------------------------------------------------------------------- +# Result formatting +# --------------------------------------------------------------------------- + + +class TestResultFormatting(unittest.TestCase): + def _search(self, results): + patcher, _ = _with_exa(_exa_response(200, {"results": results})) + with patcher: + return ws.run_web_search("test query") + + 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.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": []}, - ] + 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": "Alpha", + "url": "https://a.com", + "published_date": "2026-03-04T10:00:00.000Z", + } + ], + ) + + def test_result_without_url_is_skipped(self): + outcome = self._search( + [{"title": "no url", "text": "t"}, _exa_result("https://a.com")] ) - # Only the server_tool_use (the request) is billed, not the result block. - self.assertEqual(extract_web_search_count(msg), 1) + self.assertEqual(len(outcome.citations), 1) - 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) + def test_per_result_text_is_truncated(self): + outcome = self._search( + [_exa_result("https://a.com", text="y" * (ws.MAX_RESULT_CHARS * 2))] + ) + self.assertIn("…", outcome.content) - 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) + 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, []) # --------------------------------------------------------------------------- -# pricing.compute_session_cost with web search +# Failure modes # --------------------------------------------------------------------------- -def _usage(input_tokens: int = 100, output_tokens: int = 50) -> dict: - return {"prompt_tokens": input_tokens, "completion_tokens": output_tokens} +class TestSearchFailureModes(unittest.TestCase): + def test_no_key_configured_is_an_unbilled_error(self): + with patch.object(ws, "_exa_http_client", None): + outcome = ws.run_web_search("q") + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + + def test_transport_error_is_an_unbilled_error(self): + import httpx + + client = Mock() + client.post.side_effect = httpx.ConnectError("boom") + with patch.object(ws, "_exa_http_client", client): + outcome = ws.run_web_search("q") + self.assertTrue(outcome.is_error) + self.assertFalse(outcome.billable) + self.assertIn("could not reach", outcome.content) + + def test_http_error_surfaces_exa_detail(self): + patcher, _ = _with_exa(_exa_response(401, {"error": "invalid api key"})) + with patcher: + 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_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) -def _call(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) +# --------------------------------------------------------------------------- +# Pricing +# --------------------------------------------------------------------------- + +class _FakePriceFeed: + def __init__(self, price): + self._price = price -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) - self.assertIsInstance(base, SessionCost) - self.assertIsInstance(searched, SessionCost) - self.assertGreater(searched.cost_opg, base.cost_opg) + def get_price(self): + if isinstance(self._price, Exception): + raise self._price + return self._price - 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 - 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) - 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") - self.assertEqual(a.cost_opg, b.cost_opg) +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( + cost.cost_usd, + Decimal(cost.cost_opg) / Decimal(10) ** 18 * Decimal("0.10"), + ) - 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) + 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(cost.cost_opg, expected_opg) # --------------------------------------------------------------------------- -# chat_controller integration +# /v1/web_search controller # --------------------------------------------------------------------------- -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 _mock_tee_keys(): tee = Mock() tee.sign_data.return_value = "bW9ja3NpZ25hdHVyZQ==" @@ -200,78 +338,203 @@ 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 = { - "model": "claude-sonnet-4-5", - "messages": [{"role": "user", "content": "latest news?"}], - "web_search": True, - "stream": False, - } - - # 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 TestWebSearchController(unittest.TestCase): + def setUp(self): + app = Flask(__name__) + app.add_url_rule( + "/v1/web_search", "web-search", create_web_search, methods=["POST"] ) - model = Mock() - model.invoke.return_value = response - 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 - - 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.client = app.test_client() + + self.tee = patch( + "tee_gateway.controllers.web_search_controller.get_tee_keys", + return_value=_mock_tee_keys(), + ) + self.tee.start() + self.addCleanup(self.tee.stop) + + self.feed = patch( + "tee_gateway.price_feed.get_price_feed", + return_value=_FakePriceFeed(Decimal("0.10")), + ) + self.feed.start() + self.addCleanup(self.feed.stop) + + 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: + 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: + body = self._post(request_body).get_json() + + 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"] + ) + self.assertEqual(body["tee_request_hash"], input_hash_hex) + self.assertEqual(body["tee_output_hash"], output_hash_hex) + + def test_zero_results_still_bills_and_tells_the_model(self): + patcher, _ = _with_exa(_exa_response(200, {"results": []})) + with patcher: + 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) + + 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 + + patcher, _ = _with_exa( + _exa_response(200, {"results": [_exa_result("https://a.com")]}) ) - # Billing must receive the detected search count (1 server_tool_use). - self.assertEqual(mock_cost.call_args.kwargs["web_search_count"], 1) - 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, - } - 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 - - 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) + with patcher: + response = self._post({"query": "q"}) + self.assertEqual(response.status_code, 200) + self.assertNotIn("opengradient", response.get_json()) + + +# --------------------------------------------------------------------------- +# OHTTP inner-endpoint dispatch +# --------------------------------------------------------------------------- + + +class _FakeDecap: + plaintext = b"" # set per-test + response_key = b"k" * 32 + response_key_chunked = b"c" * 32 + enc = b"e" * 32 + + +class TestOhttpEndpointDispatch(unittest.TestCase): + """The sealed payload's `endpoint` field picks the inner path.""" + + def setUp(self): + app = Flask(__name__) + app.add_url_rule( + "/v1/ohttp", + "anonymous-chat", + ohttp_controller.create_anonymous_chat_completion, + methods=["POST"], + ) + self.client = app.test_client() + + 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.decap, self.subrequest, _ = [p.start() for p in self.patchers] + self.addCleanup(lambda: [p.stop() for p in self.patchers]) + + 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" + ) + + 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 new file mode 100644 index 0000000..01b2eea --- /dev/null +++ b/tee_gateway/web_search.py @@ -0,0 +1,361 @@ +"""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. + +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 (``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. +""" + +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" + +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 + + +# --------------------------------------------------------------------------- +# Execution +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class WebSearchOutcome: + """The result of one web search. + + ``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 + 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 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 (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(): + 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]