diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 673b2f3cc5..2edcae8ab2 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -285,7 +285,15 @@ jobs: # The PR-time CI must validate the code in this PR; PyPI unsloth # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream. - pip install --no-deps unsloth_zoo + # unsloth_zoo from git main mirrors every other CI (Core / MLX / + # install.sh) so PR-time validation sees the same zoo HEAD. + for attempt in 1 2 3; do + if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install --no-deps -e ./unsloth - name: Convert notebooks for AST scan diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 63eb70f7f1..ee5bbe8633 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -144,9 +144,19 @@ jobs: # versions ship a CPU build that imports cleanly on Linux. pip install 'bitsandbytes>=0.45' # unsloth.device_type imports unsloth_zoo.utils.Version at module - # scope, so the conftest preload needs unsloth_zoo even though - # it is an optional dep of unsloth. - pip install 'unsloth_zoo>=2026.5.1' + # scope, so the conftest preload needs unsloth_zoo. Pull from + # git main so this job sees the same zoo HEAD as Core / MLX / + # install.sh do (otherwise a fix on zoo main hides until release). + # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'` + # behaviour so triton etc. still come in for the Repo tests CPU + # collection imports. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install -e . --no-deps - name: Repo tests (CPU, auto-discovered) diff --git a/README.md b/README.md index 3699c50736..ecf0f8a7a3 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,9 @@ Then to launch every time: unsloth studio -p 8888 ``` +#### Advanced launch options +Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. Explicit `OMP_NUM_THREADS` / `MKL_NUM_THREADS` / `OPENBLAS_NUM_THREADS` / `NUMEXPR_NUM_THREADS` still take precedence. + #### Uninstall The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows): diff --git a/install.ps1 b/install.ps1 index 1ed8fafdb4..3bb5fb3f8c 100644 --- a/install.ps1 +++ b/install.ps1 @@ -887,12 +887,19 @@ shell.Run cmd, 0, False } # ── Check winget ── + # winget is only needed to install Python or uv. If both are + # already on PATH (Windows ARM64 GitHub-hosted runners, manual + # python.org + Astral uv installs, corporate locked-down hosts + # without the Store, etc.) the script can proceed without it. + # We defer the hard failure to the Python / uv install branches + # below, where winget is actually invoked. Write-TauriLog "STEP" "Checking system dependencies" - if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { - step "winget" "not available" "Red" - substep "Install it from https://aka.ms/getwinget" "Yellow" - substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow" - return (Exit-InstallFailure "winget is not available") + $script:WingetAvailable = [bool](Get-Command winget -ErrorAction SilentlyContinue) + if ($script:WingetAvailable) { + step "winget" "available" + } else { + step "winget" "not available -- will require Python + uv to be already installed" "Yellow" + substep "Get it from https://aka.ms/getwinget if Python / uv are not already on PATH." "Yellow" } # ── Helper: detect a working Python 3.11-3.13 on the system ── @@ -973,6 +980,12 @@ shell.Run cmd, 0, False step "python" "Python $($DetectedPython.Version) already installed" } if (-not $DetectedPython) { + if (-not $script:WingetAvailable) { + Write-Host "[ERROR] No compatible Python (3.11-3.13) found and winget is unavailable on this host." -ForegroundColor Red + Write-Host " Install Python $PythonVersion from https://www.python.org/downloads/" -ForegroundColor Yellow + Write-Host " and re-run this installer (make sure 'Add Python to PATH' is checked)." -ForegroundColor Yellow + return (Exit-InstallFailure "winget required to install Python on this host") + } substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" # Temporarily lower ErrorActionPreference so that winget stderr @@ -1024,14 +1037,19 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing uv package manager" if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { substep "installing uv package manager..." - $prevEAP = $ErrorActionPreference - $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} - $ErrorActionPreference = $prevEAP - Refresh-SessionPath - # Fallback: if winget didn't put uv on PATH, try the PowerShell installer + if ($script:WingetAvailable) { + $prevEAP = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} + $ErrorActionPreference = $prevEAP + Refresh-SessionPath + } + # Fallback: if winget is unavailable or didn't put uv on PATH, + # use Astral's official PowerShell installer. This is the only + # supported path on hosts without winget (Windows ARM64 runners, + # corporate machines without the Store, etc.). if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - substep "trying alternative uv installer..." "Yellow" + substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") Refresh-SessionPath } diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 8f34bb23fc..8a1edd608b 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -8,10 +8,12 @@ Most registry providers expose OpenAI-compatible /v1/chat/completions endpoints; Anthropic uses native Messages API with translation in this client. """ +import base64 import json as _json +import mimetypes import re import time -from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional +from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional, Union from urllib.parse import urlparse import httpx @@ -506,6 +508,250 @@ def _apply_mistral_reasoning_controls( _http_client = httpx.AsyncClient() +# Cap per-image fetch well below Gemini's ~20 MB total request budget. +_GEMINI_REMOTE_IMAGE_MAX_BYTES = 10 * 1024 * 1024 +_GEMINI_REMOTE_IMAGE_TIMEOUT_S = 15.0 + + +def _safe_fetch_image_for_gemini_sync( + url: str, + fallback_mime: str, + max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES, +) -> Optional[tuple[str, str]]: + """Synchronous IP-pinned HTTPS image fetch with SSRF guards. + + Uses the same pinned-IP + SNI pattern as `tools._fetch_page_text` so + DNS rebinding between validation and the actual connection cannot + redirect us to a private/metadata address. Follows up to 4 hops, + re-validating each redirect target. Returns (mime, base64) or None. + + `max_bytes` is clamped to the per-image cap and additionally lets + the caller pass the remaining per-request budget so an over-budget + URL is rejected via Content-Length (or read short-circuit) instead + of being fully downloaded then discarded after the fact. + """ + import urllib.error + import urllib.request + from urllib.parse import urljoin, urlunparse + + # Refuse upfront if the per-request budget is already spent. + _byte_limit = min(max(0, int(max_bytes)), _GEMINI_REMOTE_IMAGE_MAX_BYTES) + if _byte_limit <= 0: + return None + + # Share tools.py's pinned-IP hardening: validate-once-then-pin. + from .tools import ( + _NoRedirect, + _SNIHTTPSHandler, + _validate_and_resolve_host, + ) + + def _safe_parse_https(raw_url: str) -> Optional[tuple[Any, str, int]]: + """Validate https + hostname + port. Returns (parsed, host, port) or + None. Handles malformed-port and malformed-bracketed-IPv6 URLs that + would otherwise raise ValueError mid-build. + """ + try: + parsed_url = urlparse(raw_url) + host_value = parsed_url.hostname + port_value = parsed_url.port or 443 + except (ValueError, UnicodeError) as _err: + logger.info( + "Gemini image fetch: refusing malformed url err=%s", + type(_err).__name__, + ) + return None + scheme_value = (parsed_url.scheme or "").lower() + if scheme_value != "https": + logger.info( + "Gemini image fetch: refusing non-https scheme=%s", + scheme_value, + ) + return None + if not host_value: + logger.info("Gemini image fetch: refusing url with no hostname") + return None + return parsed_url, host_value, port_value + + parsed_info = _safe_parse_https(url) + if parsed_info is None: + return None + parsed, current_host, current_port = parsed_info + current_url = url + ok, reason, pinned_ip = _validate_and_resolve_host(current_host, current_port) + if not ok: + logger.warning( + "Gemini image fetch: refusing host=%s reason=%s", + current_host, + reason, + ) + return None + + for _hop in range(4): + # Pin to validated IP; SNI + cert still use hostname via _SNIHTTPSHandler. + cp_info = _safe_parse_https(current_url) + if cp_info is None: + return None + cp, _cp_host, _cp_port = cp_info + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str + pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + + opener = urllib.request.build_opener( + _NoRedirect, + _SNIHTTPSHandler(current_host), + ) + req = urllib.request.Request( + pinned_url, + headers = {"Host": current_host}, + method = "GET", + ) + + try: + resp = opener.open(req, timeout = _GEMINI_REMOTE_IMAGE_TIMEOUT_S) + except urllib.error.HTTPError as e: + if e.code not in (301, 302, 303, 307, 308): + logger.info( + "Gemini image fetch: status=%d host=%s", + e.code, + current_host, + ) + return None + location = e.headers.get("Location") + if not location: + return None + try: + current_url = urljoin(current_url, location) + except (ValueError, UnicodeError) as _err: + logger.info( + "Gemini image fetch: refusing malformed redirect err=%s", + type(_err).__name__, + ) + return None + rp_info = _safe_parse_https(current_url) + if rp_info is None: + return None + _rp, current_host, current_port = rp_info + ok2, reason2, pinned_ip = _validate_and_resolve_host( + current_host, current_port + ) + if not ok2: + logger.warning( + "Gemini image fetch: refusing redirect host=%s reason=%s", + current_host, + reason2, + ) + return None + continue + except (urllib.error.URLError, OSError) as _err: + logger.warning( + "Gemini image fetch failed host=%s err=%s", + current_host, + type(_err).__name__, + ) + return None + + with resp: + status = getattr(resp, "status", None) or resp.getcode() + if status != 200: + logger.info( + "Gemini image fetch: status=%s host=%s", status, current_host + ) + return None + _hdr_mime = ( + (resp.headers.get("content-type") or "").split(";")[0].strip().lower() + ) + # Declared non-image MIME is a refusal; missing MIME falls back to caller's. + if _hdr_mime and not _hdr_mime.startswith("image/"): + logger.info( + "Gemini image fetch: non-image content-type=%s host=%s", + _hdr_mime, + current_host, + ) + return None + _final_mime_pre = _hdr_mime if _hdr_mime else fallback_mime + if not isinstance(_final_mime_pre, str) or not _final_mime_pre.startswith( + "image/" + ): + logger.info( + "Gemini image fetch: missing content-type and no image fallback host=%s", + current_host, + ) + return None + _hdr_len = resp.headers.get("content-length") + if _hdr_len and _hdr_len.isdigit() and int(_hdr_len) > _byte_limit: + logger.info( + "Gemini image fetch: declared %s bytes exceeds cap=%s host=%s", + _hdr_len, + _byte_limit, + current_host, + ) + return None + # Read cap+1 to detect oversize without buffering unbounded data. + raw = resp.read(_byte_limit + 1) + if len(raw) > _byte_limit: + logger.info( + "Gemini image fetch: streamed bytes exceed cap=%s host=%s", + _byte_limit, + current_host, + ) + return None + return _final_mime_pre, base64.b64encode(raw).decode("ascii") + + logger.info("Gemini image fetch: too many redirects host=%s", current_host) + return None + + +async def _safe_fetch_image_for_gemini( + url: str, + fallback_mime: str, + max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES, +) -> Optional[tuple[str, str]]: + """Async wrapper running the IP-pinned fetch on a worker thread. + + SSRF guards (https only, pinned IP, per-hop redirect re-check, size + cap, image/* content-type) live in the sync helper. `max_bytes` + carries the remaining per-request budget so over-budget URLs are + rejected up front. + """ + import asyncio + + return await asyncio.to_thread( + _safe_fetch_image_for_gemini_sync, url, fallback_mime, max_bytes + ) + + +# Synthetic-tool names stamped onto outbound _toolEvent.arguments so the +# frontend can distinguish provider-side cards from real user-declared +# tools of the same name. Mirrored on the TS side. +_SERVER_SIDE_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} +) + + +def _stamp_server_tool_marker(payload: dict[str, Any]) -> None: + """Tag synthetic provider-side tool events so the frontend can + distinguish them from real user-declared / local function tools of + the same name. The marker rides on `arguments._server_tool` and is + only added for known server-side builtin names; user-supplied + tool calls echoed back through these helpers (e.g. Kimi + `$web_search`) keep their existing shape because we keep this scoped + to the canonical builtin names. + """ + if not isinstance(payload, dict): + return + if payload.get("type") != "tool_start": + return + name = payload.get("tool_name") + if not isinstance(name, str) or name not in _SERVER_SIDE_BUILTIN_TOOL_NAMES: + return + args = payload.get("arguments") + if not isinstance(args, dict): + args = {} + payload["arguments"] = args + args["_server_tool"] = True + + def _build_kimi_tool_end( synthetic_chunk_fn: Any, tool_call_id: str, @@ -546,16 +792,23 @@ class ExternalProviderClient: ): self.provider_type = provider_type self.base_url = base_url.rstrip("/") + # Strip a legacy `/openai` suffix from Google-hosted bases so + # configs saved before the native switch still route correctly. + # Custom proxy paths ending in `/openai` are left untouched. + if self.provider_type == "gemini": + _parsed_base = urlparse(self.base_url) + if ( + (_parsed_base.hostname or "").lower() + == "generativelanguage.googleapis.com" + and _parsed_base.path.rstrip("/") == "/v1beta/openai" + ): + self.base_url = self.base_url[: -len("/openai")] self.api_key = api_key self._timeout = httpx.Timeout(timeout, connect = 10.0) - # Separate timeout for SSE streams: reasoning-heavy providers - # (Anthropic Opus 4.7 with adaptive thinking, OpenAI gpt-5.x via - # /v1/responses) can pause for tens of seconds between bytes - # while the model is internally thinking. httpx's read timeout is - # the *gap* between successive reads, not a wall clock — so - # disabling it lets long thinks complete without cutting the - # stream prematurely. connect/write/pool keep the 10s / 120s - # bounds so genuine network failures still surface. + # Disable read timeout on SSE streams: reasoning-heavy models + # pause tens of seconds between bytes while thinking, and httpx's + # read timeout is the per-byte gap, not wall clock. connect/write + # bounds still surface real network failures. self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None) def _auth_headers(self) -> dict[str, str]: @@ -566,6 +819,14 @@ class ExternalProviderClient: auth_header = provider_info.get("auth_header", "Authorization") auth_prefix = provider_info.get("auth_prefix", "Bearer ") + # Non-Google Gemini bases (LiteLLM, custom gateways) use OAI-compat + # Bearer auth, not Google's x-goog-api-key. Override the registry default. + if self.provider_type == "gemini": + _host = (urlparse(self.base_url).hostname or "").lower() + if _host != "generativelanguage.googleapis.com": + auth_header = "Authorization" + auth_prefix = "Bearer " + headers = {"Content-Type": "application/json"} # Skip auth header when api_key is empty (optional for local providers); # httpx rejects an empty `Bearer ` value as "Illegal header value". @@ -580,6 +841,12 @@ class ExternalProviderClient: from core.inference.providers import get_provider_info info = get_provider_info(self.provider_type) or {} + # Google-hosted Gemini uses the native translator; non-Google + # bases stay on OAI-compat so LiteLLM / custom proxies still work. + if self.provider_type == "gemini": + _host = (urlparse(self.base_url).hostname or "").lower() + if _host != "generativelanguage.googleapis.com": + return True return info.get("openai_compatible", True) async def stream_chat_completion( @@ -594,11 +861,13 @@ class ExternalProviderClient: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, enabled_tools: Optional[list[str]] = None, - enable_prompt_caching: Optional[bool] = None, + enable_prompt_caching: Optional[Union[bool, str]] = None, openai_code_exec_container_id: Optional[str] = None, anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, fast_mode: Optional[bool] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: @@ -616,7 +885,35 @@ class ExternalProviderClient: ``fast_mode`` only applies to Anthropic Opus 4.6 / 4.7 (silently dropped elsewhere); adds the beta header and ``speed: "fast"``. """ + # tool_choice="none" hard-disables hosted/builtin tools across + # every provider so enabled_tools cannot accidentally bill or leak. + tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + if not self._is_openai_compatible(): + # Gemini speaks its own native REST shape (contents/parts); + # `_stream_gemini` translates request/response into the OpenAI + # Chat Completions chunk format the rest of Studio expects. + # API reference: https://ai.google.dev/gemini-api/docs + if self.provider_type == "gemini": + async for line in self._stream_gemini( + messages, + model, + temperature, + top_p, + max_tokens, + top_k, + presence_penalty, + enabled_tools, + enable_prompt_caching, + enable_thinking, + reasoning_effort, + tools, + tool_choice, + ): + yield line + return async for line in self._stream_anthropic( messages, model, @@ -631,6 +928,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id, prompt_cache_ttl, compaction_threshold, + tool_choice, fast_mode = fast_mode, ): yield line @@ -654,20 +952,25 @@ class ExternalProviderClient: enable_prompt_caching, openai_code_exec_container_id, compaction_threshold, + tools, + tool_choice, ): yield line return - # Kimi's $web_search is a builtin_function that requires a client - # round-trip: the first call returns a tool_calls envelope with - # function.arguments populated; the caller echoes those arguments - # back as a role=tool message; the second call streams the final - # answer with the search incorporated. The doc also mandates - # disabling thinking while $web_search is active. Route to a - # dedicated helper so the default OAI-compat path stays single-pass. - # https://platform.kimi.ai/docs/guide/use-web-search + # Kimi $web_search needs a 2-call round-trip + thinking off; route + # to a helper. Forced-function tool_choice suppresses it. + # https://platform.kimi.ai/docs/guide/use-web-search + _kimi_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) if ( self.provider_type == "kimi" + and not tool_choice_disabled + and not _kimi_tool_choice_forced_function and enabled_tools and "web_search" in enabled_tools ): @@ -694,28 +997,18 @@ class ExternalProviderClient: else: body["max_tokens"] = max_tokens - # Strip body fields a provider's registry entry declares unusable — - # reasoning-class models that lock these to fixed defaults (e.g. - # Kimi k2.5/k2.6 only accept temperature=1, top_p=1) 400 otherwise. - # The frontend capability map already hides the matching sliders; - # this is the matching guard for the pydantic default that the - # route layer would otherwise still fill in. + # Drop fields the registry flags as unusable so reasoning-class + # models with fixed defaults (Kimi k2.6 etc) don't 400 on pydantic + # default values that the route layer still fills in. from core.inference.providers import get_provider_info provider_info = get_provider_info(self.provider_type) or {} for field in provider_info.get("body_omit", ()): body.pop(field, None) - # Kimi (kimi-k2.6, kimi-k2-thinking) accepts a boolean thinking toggle - # via a top-level `thinking` field (the docs show it nested under - # extra_body, but that is an OpenAI Python SDK convention; on the - # wire it merges into the request body). - # - kimi-k2.6 defaults to thinking enabled; clients can pass - # {"type": "disabled"} to suppress it. - # - kimi-k2-thinking is always on; we never send disabled there. - # `keep: all` retains every thinking chunk through the stream, which - # is what we need so our frontend can wrap reasoning_content into - # the chat reasoning panel. + # Kimi thinking is a top-level body field. kimi-k2-thinking is + # always on (ignore the toggle); kimi-k2.6 defaults on, can be + # disabled. `keep: all` preserves every chunk for the UI panel. if self.provider_type == "kimi" and enable_thinking is not None: if model == "kimi-k2-thinking": # Always on; ignore client toggle to avoid an API-level reject. @@ -736,17 +1029,9 @@ class ExternalProviderClient: tpl_kw["enable_thinking"] = bool(enable_thinking) body["chat_template_kwargs"] = tpl_kw - # OpenRouter exposes a unified `reasoning` parameter on every - # chat-completion request — the gateway routes it to whichever - # underlying model actually supports reasoning, and silently - # no-ops for ones that don't. Documented at - # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens - # Shape: `reasoning: {enabled?: bool, effort?: low|medium|high, - # max_tokens?: N, exclude?: bool}` with effort and max_tokens - # mutually exclusive. We forward either an effort level (when - # the user picked one) or a bare {enabled: true}. A small set of - # known routes rejects explicit disable with 400 ("Reasoning is - # mandatory for this endpoint ..."), so only those omit "off". + # OpenRouter's unified `reasoning` field gates per-model thinking. + # Some routes (`*_MANDATORY_REASONING_MODELS`) 400 on explicit off. + # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens if self.provider_type == "openrouter": normalized_or_model = model.strip().lower() if reasoning_effort in ("low", "medium", "high"): @@ -759,17 +1044,22 @@ class ExternalProviderClient: else: body["reasoning"] = {"enabled": False} - # OpenRouter web-search plugin — universal shape that works - # for every model id, including the `openrouter/free` and - # `openrouter/auto` meta-routers. Documented at - # https://openrouter.ai/docs/guides/features/plugins/web-search - # The `:online` model-suffix shortcut is "exactly equivalent - # to" this plugin per the same doc, but only works on - # concrete model ids — meta-routers reject the suffix. - # `plugins: [{id: "web"}]` works everywhere, no model id - # rewrite needed, and idempotent if some future call site - # adds the entry first. - if enabled_tools and "web_search" in enabled_tools: + # OpenRouter web plugin works on every model id including + # meta-routers (unlike the `:online` suffix). Forced-function + # tool_choice suppresses it, matching Gemini/Anthropic. + # https://openrouter.ai/docs/guides/features/plugins/web-search + _or_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + if ( + not tool_choice_disabled + and not _or_tool_choice_forced_function + and enabled_tools + and "web_search" in enabled_tools + ): plugins = list(body.get("plugins") or []) if not any( isinstance(p, dict) and p.get("id") == "web" for p in plugins @@ -781,6 +1071,15 @@ class ExternalProviderClient: body.get("model"), ) + # Forward OpenAI-style function tools / tool_choice on every + # OAI-compat route (incl. custom Gemini OpenAI proxies like + # LiteLLM). Without this, callers that wire user-defined tools + # silently lose function-calling on non-native providers. + if tools: + body["tools"] = tools + if tool_choice is not None: + body["tool_choice"] = tool_choice + url = f"{self.base_url}/chat/completions" logger.info( "Proxying chat completion to %s (provider=%s, model=%s)", @@ -816,30 +1115,22 @@ class ExternalProviderClient: ) return - # NOTE: manual __anext__ loop instead of `async for` is intentional. - # On Python 3.13 + httpcore 1.0.x, `async for` auto-calls aclose() on - # early exit (break/return/GeneratorExit) BEFORE our finally block runs. - # That propagates GeneratorExit into PoolByteStream.__aiter__() while it - # calls `await self.aclose()` inside `with AsyncShieldCancellation()`, - # triggering "RuntimeError: async generator ignored GeneratorExit". - # Fix: call response.aclose() FIRST (sets PoolByteStream._closed=True), - # then lines_gen.aclose() is a no-op and GeneratorExit re-raises cleanly. + # Manual __anext__ (not `async for`) so we can close + # the response BEFORE lines_gen, avoiding the httpcore + # 1.0 GeneratorExit -> RuntimeError path on Python 3.13. lines_gen = response.aiter_lines().__aiter__() - # Best-effort diagnostics for the default OAI-compat path. Without - # this, OpenRouter mid-stream errors (200 OK + error event in the - # SSE body) and OpenRouter-router model selection were invisible - # in the backend logs — the user only saw "Provider returned - # error" in the UI with no trail on the server side. + # Diagnostic counters for the OAI-compat path; surfaces + # OpenRouter mid-stream errors that would otherwise be + # invisible server-side. event_counts: dict[str, int] = {} chosen_model: Optional[str] = None - # Web-search tool-card synthesis for OpenRouter. The gateway - # doesn't emit structured web_search_call events — citations - # come back as `annotations` of type=url_citation on delta / - # message objects. Mirror the OpenAI/Anthropic UX by yielding - # a synthetic tool_start at stream open and tool_end at - # stream close with the collected citation list. + # OpenRouter has no web_search_call events — citations + # arrive as url_citation annotations. Synthesise a + # tool_start/tool_end pair to match the OpenAI/Anthropic UX. web_search_active = ( self.provider_type == "openrouter" + and not tool_choice_disabled + and not _or_tool_choice_forced_function and bool(enabled_tools) and "web_search" in (enabled_tools or []) ) @@ -849,6 +1140,7 @@ class ExternalProviderClient: web_search_tool_ended = False def _emit_synthetic_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": f"chatcmpl-{self.provider_type}-synthetic", "object": "chat.completion.chunk", @@ -1104,6 +1396,7 @@ class ExternalProviderClient: synthetic_id = f"chatcmpl-{self.provider_type}-synthetic" def _synthetic_chunk(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": synthetic_id, "object": "chat.completion.chunk", @@ -1441,6 +1734,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + tool_choice: Optional[Any] = None, *, fast_mode: Optional[bool] = None, ) -> AsyncGenerator[str, None]: @@ -1471,6 +1765,42 @@ class ExternalProviderClient: continue content = msg.get("content") + # OpenAI role="tool" with list content -> Anthropic native + # tool_result block on a user message. Translating in the + # string-content branch only (below) leaves the list-content + # form forwarded as an invalid `role:"tool"` message that + # Anthropic rejects. Handle both upfront. + if msg.get("role") == "tool": + _tr_id = msg.get("tool_call_id") or "" + if isinstance(content, list): + _flat_parts: list[str] = [] + for part in content: + if ( + isinstance(part, dict) + and part.get("type") == "text" + and part.get("text") + ): + _flat_parts.append(str(part["text"])) + _flat_result = "".join(_flat_parts) + elif content is None: + _flat_result = "" + elif isinstance(content, str): + _flat_result = content + else: + _flat_result = _json.dumps(content) + filtered.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": _tr_id, + "content": _flat_result, + } + ], + } + ) + continue if isinstance(content, list): # Translate OpenAI multimodal parts -> Anthropic native shapes. # - `image_url` -> `{type:"image", source:...}` @@ -1583,6 +1913,37 @@ class ExternalProviderClient: if title: doc_block["title"] = title anthropic_parts.append(doc_block) + # Assistant tool_calls -> Anthropic tool_use blocks + # appended to the same message. Anthropic native + # Messages API does not accept OpenAI's top-level + # `tool_calls` field; the call lives inside a content + # block with `{type:"tool_use", id, name, input}`. + if msg.get("role") == "assistant" and isinstance( + msg.get("tool_calls"), list + ): + for _tc in msg["tool_calls"]: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _raw = _fn.get("arguments") or "{}" + try: + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) + except Exception: + _input = {"_raw": _raw} + if not isinstance(_input, dict): + _input = {"value": _input} + anthropic_parts.append( + { + "type": "tool_use", + "id": _tc.get("id") or f"toolu_{time.time_ns()}", + "name": _fn["name"], + "input": _input, + } + ) # Skip whole-message append when nothing usable survived. # An empty content array (e.g. user dropped only an unparseable # `input_document`) would 400 the Anthropic API with @@ -1590,6 +1951,72 @@ class ExternalProviderClient: if anthropic_parts: filtered.append({"role": msg["role"], "content": anthropic_parts}) else: + # role="tool" follow-up -> Anthropic native tool_result + # block on a `user` message. The OpenAI shape + # (role=tool, content=string, tool_call_id) is not a + # valid Anthropic role. + if msg.get("role") == "tool": + _tr_id = msg.get("tool_call_id") or "" + _tr_content = msg.get("content") + if _tr_content is None: + _tr_content = "" + filtered.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": _tr_id, + "content": ( + _tr_content + if isinstance(_tr_content, str) + else _json.dumps(_tr_content) + ), + } + ], + } + ) + continue + # Assistant turn whose content is a plain string but + # also carries OpenAI `tool_calls`: convert into a + # content-array message with a text block + tool_use + # blocks. Without this, the top-level tool_calls leaks + # through unchanged. + if ( + msg.get("role") == "assistant" + and isinstance(msg.get("tool_calls"), list) + and msg["tool_calls"] + ): + _text_content = msg.get("content") + _blocks: list[dict[str, Any]] = [] + if isinstance(_text_content, str) and _text_content: + _blocks.append({"type": "text", "text": _text_content}) + for _tc in msg["tool_calls"]: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _raw = _fn.get("arguments") or "{}" + try: + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) + except Exception: + _input = {"_raw": _raw} + if not isinstance(_input, dict): + _input = {"value": _input} + _blocks.append( + { + "type": "tool_use", + "id": _tc.get("id") or f"toolu_{time.time_ns()}", + "name": _fn["name"], + "input": _input, + } + ) + if _blocks: + filtered.append({"role": "assistant", "content": _blocks}) + continue filtered.append(msg) # Claude 4.7 family removed temperature / top_p / top_k entirely. @@ -1616,34 +2043,16 @@ class ExternalProviderClient: # same as True here (callers that don't set the flag still get # caching). Pass False explicitly to opt out. prompt_caching_enabled = enable_prompt_caching is not False - # Anthropic accepts an optional `ttl` on each cache_control marker - # (default is the 5m ephemeral pool; set "1h" to land in the 1h - # pool instead). Per the prompt-caching docs, 1h cache writes are - # billed at 2x base input vs 1.25x for 5m, but reads are 0.1x for - # both. The 1h pool is the right pick when conversations span - # multiple short bursts more than 5 minutes apart -- the read - # discount makes up for the 1.6x write premium after a single - # additional hit. Anything other than the known TTL strings is - # dropped to avoid sending a malformed marker. - # - # The `extended-cache-ttl-2025-04-11` beta header that originally - # gated 1h TTL has been promoted to GA: as of 2026-05 the live - # API accepts `ttl: "1h"` without any beta opt-in. Verified - # against api.anthropic.com on claude-opus-4-7 (status 200 + - # `ephemeral_1h_input_tokens` populated). The test below pins - # the contract by asserting the header is NOT on the wire so a - # future regression that reintroduces the gate would surface - # before users see a 400. + # Optional 1h cache TTL is GA as of 2026-05 (no beta header). 1h + # writes are 2x vs 5m's 1.25x but reads are 0.1x for both, so 1h + # wins after a single extra hit. Unknown TTL strings drop. cache_marker: dict[str, Any] = {"type": "ephemeral"} if prompt_cache_ttl in ("5m", "1h"): cache_marker["ttl"] = prompt_cache_ttl if system: if prompt_caching_enabled: - # System block is the most stable prefix across turns, so - # it gets its own breakpoint. Skipped when system is - # empty — there's nothing to cache, and an empty marker - # is a no-op. + # System is the most stable cross-turn prefix; own breakpoint. body["system"] = [ { "type": "text", @@ -1749,19 +2158,29 @@ class ExternalProviderClient: if body.get("max_tokens", 0) <= budget_tokens: body["max_tokens"] = budget_tokens + 1024 - # Anthropic server-side web_search — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool - # The tool type is date-pinned per model family. Newer Opus / - # Sonnet 4.6 + 4.7 accept `web_search_20260209` with dynamic - # filtering (Claude writes code to filter results before they - # reach context); everything else uses `web_search_20250305`. - # `_anthropic_web_search_version` picks the right one. Anthropic - # dispatches search calls server-side, returning server_tool_use - # + web_search_tool_result blocks in the SSE stream, plus - # url-citation annotations on text deltas. We translate all of - # that into our local _toolEvent shape so the chat UI renders - # web_search exactly like OpenAI's path. - if enabled_tools and "web_search" in enabled_tools: + # tool_choice="none" or pinned-function suppresses hosted tools + # so a stale UI toggle can't fire server-side search/code-exec. + _anthropic_tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + _anthropic_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _anthropic_hosted_builtins_allowed = ( + not _anthropic_tool_choice_disabled + and not _anthropic_tool_choice_forced_function + ) + + # Anthropic web_search (date-pinned per model family). + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool + if ( + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( { @@ -1772,14 +2191,13 @@ class ExternalProviderClient: ) body["tools"] = anthropic_tools - # Anthropic server-side web_fetch reads a single URL (text/PDF) - # and returns a `web_fetch_tool_result` document block. Opt in - # via `enabled_tools=["web_fetch"]`; no beta header required. - # `_anthropic_web_fetch_version` picks `web_fetch_20260209` - # (dynamic filtering) for Opus 4.6/4.7 + Sonnet 4.6, falling - # back to `web_fetch_20250910` elsewhere; mismatched variants - # return 400 so the per-model picker is required. - web_fetch_enabled = bool(enabled_tools and "web_fetch" in enabled_tools) + # Anthropic web_fetch: only URLs already in conversation. Date-pinned. + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool + web_fetch_enabled = bool( + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_fetch" in enabled_tools + ) if web_fetch_enabled: anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( @@ -1792,24 +2210,13 @@ class ExternalProviderClient: body["tools"] = anthropic_tools # Anthropic server-side code execution — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool - # The tool type is date-pinned per model family. - # `_anthropic_code_execution_version` picks `code_execution_20260120` - # for Opus 4.5+ / Sonnet 4.5+ / Opus 4.7 / Sonnet 4.6 (adds REPL - # state persistence + programmatic tool calling) and falls back - # to `code_execution_20250825` everywhere else. Both versions - # run Python + bash + str_replace file edits inside a 5 GB - # sandboxed container per request, with no internet access, and - # both are unlocked by the same `code-execution-2025-08-25` - # `anthropic-beta` header set further down. On the SSE stream - # Anthropic emits two sub-tool names -- `bash_code_execution` - # and `text_editor_code_execution` -- wrapped in the standard - # server_tool_use / *_tool_result block shape. - # v1 wires the tool only; file uploads (container_upload - # content blocks and generated-file retrieval via the Files - # API) are a deliberate follow-up. + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool + # Date-pinned tool type per model; both unlock via the same + # `code-execution-2025-08-25` beta header set below. code_execution_enabled = bool( - enabled_tools and "code_execution" in enabled_tools + _anthropic_hosted_builtins_allowed + and enabled_tools + and "code_execution" in enabled_tools ) if code_execution_enabled: anthropic_tools = list(body.get("tools") or []) @@ -1820,32 +2227,14 @@ class ExternalProviderClient: } ) body["tools"] = anthropic_tools - # Reuse the prior turn's container so filesystem state - # (files written, packages installed, variables set) - # persists across turns of the same thread. Anthropic - # exposes the container id on the Message object's - # top-level `container.id`; on the SSE stream we latch it - # off `message_start.message.container.id` further down - # and emit a `container_ready` _toolEvent so the chat - # adapter persists it on the thread record. A stale id - # (container expired / not found) surfaces as a 4xx - # below, where we emit `container_invalidated` and let - # the next turn fall back to auto-create. + # Reuse the thread's prior container so filesystem state + # persists. Stale ids 4xx and clear via container_invalidated. if anthropic_code_exec_container_id: body["container"] = anthropic_code_exec_container_id - # Server-side context compaction — see - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Beta as of `compact-2026-01-12`. When `compaction_threshold` is - # provided AND the model accepts compaction (Opus 4.6+ / 4.7, - # Sonnet 4.6, Mythos preview), attach - # `context_management.edits[{type:"compact_20260112", trigger: - # {type:"input_tokens", value:N}}]` to the body. Anthropic runs - # the compaction step server-side once the rendered prompt - # crosses the threshold and replies with a top-level - # `context_management` block plus `usage.iterations[]` so we can - # account per-iteration. Below-min thresholds get clamped up to - # 50K so the request doesn't 400. + # Server-side compaction (beta `compact-2026-01-12`). Clamps + # below-min thresholds to 50K so the request doesn't 400. + # https://platform.claude.com/docs/en/build-with-claude/compaction compaction_active = ( compaction_threshold is not None and compaction_threshold > 0 @@ -1878,10 +2267,8 @@ class ExternalProviderClient: url = f"{self.base_url}/messages" completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}" - # Log the outgoing config keys (not the messages themselves) so we - # can prove which thinking/effort fields actually reached the wire. - # If Anthropic skips reasoning despite a configured effort, this - # tells us whether we sent the field or dropped it on the floor. + # Log outgoing config keys (not messages) to prove which thinking / + # effort fields actually reached the wire. logger.info( "Anthropic request shape (model=%s, has_thinking=%s, thinking=%s, " "output_config=%s, temperature=%s, has_top_p=%s, has_top_k=%s, " @@ -1896,16 +2283,10 @@ class ExternalProviderClient: body.get("max_tokens"), ) - # Translate Anthropic stop reasons onto the OpenAI chat-completions - # `finish_reason` vocabulary. `pause_turn` maps to None so the - # adapter does NOT emit a finish_reason chunk: pause_turn means - # Claude paused a long server-tool turn (web_search / web_fetch) - # and will continue once the user (or our retry) sends back the - # partial assistant message. Forwarding it as "stop" makes the - # OpenAI client think the answer is done and truncates the - # rendered message. `refusal` maps to "content_filter" as the - # nearest semantic match. See - # https://platform.claude.com/docs/en/api/messages#response-stop-reason + # Anthropic stop_reason -> OpenAI finish_reason. `pause_turn` + # maps to None so the UI doesn't treat a paused server-tool turn + # as final. `refusal` -> "content_filter" (closest match). + # https://platform.claude.com/docs/en/api/messages#response-stop-reason _finish_reason_map: dict[str, Optional[str]] = { "end_turn": "stop", "max_tokens": "length", @@ -1918,11 +2299,7 @@ class ExternalProviderClient: logger.info("Proxying Anthropic Messages API to %s (model=%s)", url, model) request_headers = self._auth_headers() - # Anthropic accepts comma-separated beta features in a single - # `anthropic-beta` header. Merge our flags onto whatever the - # registry's extra_headers contributed (currently nothing on - # the beta axis, just anthropic-version) so future betas - # added at the registry level keep working. + # Merge new beta flags onto whatever the registry contributed. existing_beta = request_headers.get("anthropic-beta", "").strip() beta_parts = ( [p.strip() for p in existing_beta.split(",") if p.strip()] @@ -1988,27 +2365,14 @@ class ExternalProviderClient: # "no thinking content" — distinguishes "Anthropic never sent # thinking_delta" from "frontend didn't render the chunks". event_counts: dict[str, int] = {} - # web_search state. Anthropic emits the query inside an - # `input_json_delta` stream on a `server_tool_use` content - # block, then a separate `web_search_tool_result` block - # with the URL list. Unlike OpenAI we get per-call results - # directly, so each tool card carries its own citations. - # `current_server_tool_use`: {id, name, partial_json_buffer} - # `current_result_block`: {tool_use_id, results} - # Both go to None when the matching content_block_stop fires. + # web_search state. Query streams via input_json_delta + # on a server_tool_use block; results land in a separate + # web_search_tool_result block. Per-call citations. current_server_tool_use: Optional[dict[str, Any]] = None current_result_block: Optional[dict[str, Any]] = None web_search_calls: dict[str, dict[str, Any]] = {} - # code_execution state. Anthropic's - # `code_execution_20250825` tool emits the same - # server_tool_use → *_tool_result block shape as - # web_search, but the server_tool_use carries one of - # two sub-tool names (`bash_code_execution` or - # `text_editor_code_execution`) and the result block - # type matches (`bash_code_execution_tool_result` / - # `text_editor_code_execution_tool_result`). Kept - # parallel to web_search state so the two paths don't - # collide when both pills are on in the same turn. + # code_execution state (bash / text_editor sub-tools); + # kept parallel to web_search so concurrent pills don't collide. current_code_exec_use: Optional[dict[str, Any]] = None current_code_exec_result: Optional[dict[str, Any]] = None code_execution_calls: dict[str, dict[str, Any]] = {} @@ -2078,6 +2442,7 @@ class ExternalProviderClient: return f"data: {_json.dumps(chunk)}" def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -2335,18 +2700,8 @@ class ExternalProviderClient: "inner": inner if isinstance(inner, dict) else {}, } elif block_type == "compaction": - # Server-side compaction emits a `compaction` - # content block on the assistant message. - # Anthropic may include the summary text on - # this start event AND/OR stream it via - # text_delta events on the same block. See - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Capture either form; finalize and emit - # on content_block_stop. The chat-adapter - # persists the block onto the assistant - # message so the next turn's request - # carries it back -- Anthropic then skips - # re-compaction from scratch. + # Summary may arrive on start AND/OR via + # text_delta. Capture both; emit on stop. seed = content_block.get("content") or "" current_compaction = { "content": seed if isinstance(seed, str) else "", @@ -2356,12 +2711,7 @@ class ExternalProviderClient: delta = event.get("delta", {}) delta_type = delta.get("type") if delta_type == "thinking_delta": - # Anthropic streams extended-thinking content as - # thinking_delta events on a separate content - # block. Wrap as inline ... so - # the frontend's parseAssistantContent lifts it - # into the reasoning panel — same pattern as - # the OpenAI Responses path. + # Wrap as ... for parseAssistantContent. thinking_text = delta.get("thinking", "") if thinking_text: if not thinking_open: @@ -2631,19 +2981,9 @@ class ExternalProviderClient: delta_usage = event.get("usage") if isinstance(delta_usage, dict): last_usage.update(delta_usage) - # When a fresh compaction has run, Anthropic - # publishes per-iteration token counts in - # `usage.iterations[]`. The top-level - # input_tokens / output_tokens only cover the - # `message` iteration, NOT the compaction - # passes — billing has to sum the whole - # array. See - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Fold the compaction iterations into - # `compaction_input_tokens` / `compaction_output_tokens` - # so the cost surface can add them without - # re-walking the array (and so the closing - # log line names the figures). + # Compaction iterations aren't in top-level + # input/output_tokens; fold them into + # compaction_{input,output}_tokens for billing. iterations = delta_usage.get("iterations") if isinstance(iterations, list): c_in = 0 @@ -2880,6 +3220,1722 @@ class ExternalProviderClient: self.provider_type, ) + async def _stream_gemini( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float, + top_p: float, + max_tokens: Optional[int], + top_k: Optional[int] = None, + presence_penalty: float = 0.0, + enabled_tools: Optional[list[str]] = None, + enable_prompt_caching: Optional[Any] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, + ) -> AsyncGenerator[str, None]: + """ + Call Google's native Gemini API and translate its streaming + ``streamGenerateContent`` response into OpenAI Chat Completions + chunk format. + + Gemini does NOT speak the OpenAI Chat Completions contract on + its primary endpoint. The wire shape is: + + POST /v1beta/models/{model}:streamGenerateContent?alt=sse + { + "contents": [{"role": "user|model", "parts": [{"text": "..."}]}], + "systemInstruction": {"parts": [{"text": "..."}]}, + "generationConfig": {"temperature": 0.7, "topP": 0.95, "topK": 40, + "maxOutputTokens": 1024}, + "tools": [{"googleSearch": {}}, {"codeExecution": {}}], + "cachedContent": "" // optional, see caching docs + } + + Streamed responses are SSE frames carrying partial + ``GenerateContentResponse`` objects: + + {"candidates": [{"content": {"parts": [{"text": "Hello"}]}, + "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 7, "candidatesTokenCount": 3}} + + Image generation uses the same endpoint with model + ``gemini-2.5-flash-image`` (also called Nano Banana); the + response carries an ``inlineData`` part with the base64 PNG + bytes and a ``mimeType``. We surface that through the same + ``tool_start`` / ``tool_end`` ``image_b64`` envelope the OpenAI + image_generation path uses, so the chat UI renders the image + inline with no extra plumbing. + + References: + - https://ai.google.dev/gemini-api/docs/text-generation + - https://ai.google.dev/gemini-api/docs/function-calling + - https://ai.google.dev/gemini-api/docs/grounding + - https://ai.google.dev/gemini-api/docs/caching + - https://ai.google.dev/gemini-api/docs/image-generation + """ + import json as _json + + # Validate the user-controlled model id BEFORE any message + # translation. A model like `../cachedContents/x` is path- + # traversal that lands in `/v1beta/cachedContents/...`; rejecting + # it here also avoids triggering user-controlled outbound fetches + # (remote image_url inlining) on a request we'll error out + # anyway. Documented catalog ids match `[A-Za-z0-9._-]+`. + if not re.fullmatch(r"[A-Za-z0-9._-]+", model): + yield _error_sse_line( + 400, + f"Invalid Gemini model id: {model!r}", + self.provider_type, + ) + return + + # Translate OpenAI messages -> Gemini contents. system role + # promotes to top-level systemInstruction. + system_text_parts: list[str] = [] + contents: list[dict[str, Any]] = [] + # OpenAI may drop `name` from role="tool" follow-ups. Remember + # prior function names so functionResponse isn't sent name-less + # (Gemini 400s on empty names). + tool_call_names: dict[str, str] = {} + # tool_call_ids whose assistant card was dropped (synthetic + # builtin) or already replayed as native parts. Their role="tool" + # follow-up must be skipped to avoid orphan/duplicate responses. + _gemini_skip_tool_result_ids: set[str] = set() + # Per-request image caps. The byte cap counts DECODED bytes; we + # set it to ~14 MB because base64 expansion + prompt overhead + # must fit Gemini's ~20 MB request limit. + _GEMINI_REMOTE_IMAGE_MAX_COUNT = 8 + _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES = 14 * 1024 * 1024 + _remote_image_count = 0 + _remote_image_total_bytes = 0 + for msg in messages: + role = msg.get("role") + content = msg.get("content", "") + if role == "system": + if isinstance(content, str): + if content: + system_text_parts.append(content) + elif isinstance(content, list): + for part in content: + if ( + isinstance(part, dict) + and part.get("type") == "text" + and part.get("text") + ): + system_text_parts.append(part["text"]) + continue + # Map OpenAI roles to Gemini's two-role contract. + gemini_role = "model" if role == "assistant" else "user" + parts: list[dict[str, Any]] = [] + if isinstance(content, str): + if content: + parts.append({"text": content}) + elif isinstance(content, list): + for part in content: + if not isinstance(part, dict): + continue + ptype = part.get("type") + if ptype == "text": + text = part.get("text", "") + if text: + parts.append({"text": text}) + elif ptype == "image_url": + url = part.get("image_url", {}).get("url", "") + if url.startswith("data:"): + header, _, b64data = url.partition(",") + media_type = ( + header.split(";")[0] + .replace("data:", "") + .strip() + .lower() + or "image/jpeg" + ) + # Symmetry with the fetched remote image + # path, which already rejects non-image + # Content-Type. A `data:text/html;base64,...` + # URL otherwise lands as Gemini inlineData + # with mimeType="text/html" and 400s the + # whole request. + if not media_type.startswith("image/"): + logger.info( + "Gemini inlineData: refusing non-image data URL media_type=%s", + media_type, + ) + elif b64data: + # data: URLs share the same caps as fetched + # URLs so inline payloads don't bypass them. + _data_approx_bytes = (len(b64data) * 3) // 4 + if ( + _remote_image_count + >= _GEMINI_REMOTE_IMAGE_MAX_COUNT + ): + logger.info( + "Gemini inlineData: per-request count cap %d reached, dropping image", + _GEMINI_REMOTE_IMAGE_MAX_COUNT, + ) + elif ( + _remote_image_total_bytes + _data_approx_bytes + > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + ): + logger.info( + "Gemini inlineData: per-request byte cap reached, dropping image", + ) + else: + _remote_image_count += 1 + _remote_image_total_bytes += _data_approx_bytes + parts.append( + { + "inlineData": { + "mimeType": media_type, + "data": b64data, + } + } + ) + elif url: + # fileData.fileUri only accepts Files-API URIs + # and YouTube; everything else must be downloaded + # and inlined. Parse fields explicitly so + # attacker URLs like https://evil.com/youtube.com/x + # aren't misclassified as YouTube. + try: + _parsed_image_url = urlparse(url) + except (ValueError, UnicodeError): + _parsed_image_url = None + if _parsed_image_url is None: + _img_scheme = "" + _img_host = "" + _img_path = "" + else: + _img_scheme = (_parsed_image_url.scheme or "").lower() + _img_host = (_parsed_image_url.hostname or "").lower() + _img_path = _parsed_image_url.path or "" + _is_native_uri = ( + _img_scheme == "https" + and _img_host == "generativelanguage.googleapis.com" + and _img_path.startswith("/v1beta/files/") + ) + _is_youtube = _img_scheme == "https" and ( + _img_host == "youtu.be" + or _img_host == "youtube.com" + or _img_host.endswith(".youtube.com") + ) + _guessed, _ = mimetypes.guess_type(_img_path) + _media_type = ( + _guessed + if isinstance(_guessed, str) + and _guessed.startswith("image/") + else "image/jpeg" + ) + if _is_youtube: + # YouTube URIs must use video/mp4; the + # default image/jpeg yields a 400. + parts.append( + { + "fileData": { + "fileUri": url, + "mimeType": "video/mp4", + } + } + ) + elif _is_native_uri: + parts.append( + { + "fileData": { + "fileUri": url, + "mimeType": _media_type, + } + } + ) + elif _remote_image_count >= _GEMINI_REMOTE_IMAGE_MAX_COUNT: + logger.info( + "Gemini image fetch: per-request count cap %d reached, dropping image", + _GEMINI_REMOTE_IMAGE_MAX_COUNT, + ) + else: + # Refuse pre-fetch when the per-request + # byte budget is spent; pass the remainder + # so over-budget URLs reject on Content-Length. + _remaining_bytes = ( + _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + - _remote_image_total_bytes + ) + if _remaining_bytes <= 0: + logger.info( + "Gemini image fetch: per-request byte cap already reached, dropping image", + ) + else: + # Count attempts before awaiting so + # slow URLs don't each burn the timeout. + _remote_image_count += 1 + _fetched = await _safe_fetch_image_for_gemini( + url, + _media_type, + max_bytes = _remaining_bytes, + ) + if _fetched is not None: + _final_mime, _b64 = _fetched + # base64 expands ~4/3 — recover bytes from len(_b64). + _approx_bytes = (len(_b64) * 3) // 4 + if ( + _remote_image_total_bytes + _approx_bytes + > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + ): + logger.info( + "Gemini image fetch: per-request byte cap reached, dropping image", + ) + else: + _remote_image_total_bytes += _approx_bytes + parts.append( + { + "inlineData": { + "mimeType": _final_mime, + "data": _b64, + } + } + ) + # Gemini 3 strict function-calling requires text-part + # thoughtSignatures to be replayed on history; the frontend + # stows the latest one as + # extra_content.google.thought_signature on the assistant + # message and we pin it onto the last text part here. + if role == "assistant" and parts: + _msg_extra = msg.get("extra_content") if isinstance(msg, dict) else None + if isinstance(_msg_extra, dict): + _msg_g = _msg_extra.get("google") or {} + if isinstance(_msg_g, dict): + _msg_sig = _msg_g.get("thought_signature") or _msg_g.get( + "thoughtSignature" + ) + if isinstance(_msg_sig, str) and _msg_sig: + for _idx in range(len(parts) - 1, -1, -1): + if "text" in parts[_idx]: + parts[_idx] = { + **parts[_idx], + "thoughtSignature": _msg_sig, + } + break + # Translate OpenAI tool_calls into Gemini functionCall parts. + # code_execution / image_generation replay their native parts + # (executableCode / codeExecutionResult / inlineData) stowed + # on extra_content.google.native_part. + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if isinstance(tool_calls, list): + for tc in tool_calls: + if not isinstance(tc, dict): + continue + fn = tc.get("function") or {} + if not isinstance(fn, dict): + continue + args_raw = fn.get("arguments") or "{}" + if isinstance(args_raw, str): + try: + args = _json.loads(args_raw) + except Exception: + args = {"_raw": args_raw} + elif isinstance(args_raw, dict): + args = args_raw + else: + args = {} + fn_name = fn.get("name", "") + tc_id = tc.get("id") + if fn_name and isinstance(tc_id, str) and tc_id: + tool_call_names[tc_id] = fn_name + + # Replay native Gemini code_execution / image_generation parts + # from extra_content.google.native_part, with fallback to + # args.google.native_part for OAI-compat round-trips. + _extra = tc.get("extra_content") + _native_part = None + _google_extra: dict[str, Any] = {} + if isinstance(_extra, dict): + _ge = _extra.get("google") or {} + if isinstance(_ge, dict): + _google_extra = _ge + _native_part = _ge.get("native_part") + if _native_part is None and isinstance(args, dict): + _args_google = args.get("google") + if isinstance(_args_google, dict): + _args_np = _args_google.get("native_part") + if isinstance(_args_np, dict): + _native_part = _args_np + if not _google_extra: + _google_extra = _args_google + + # Synthetic builtin cards (web_search/web_fetch) must + # not become fake functionCalls; drop them. Native + # code_execution / image_generation replay below. + _name_lc = fn_name.lower() if isinstance(fn_name, str) else "" + _is_synthetic_server_builtin = ( + _name_lc + in ( + "web_search", + "web_fetch", + "code_execution", + "image_generation", + ) + and isinstance(args, dict) + and ( + args.get("_server_tool") is True + or isinstance( + (args.get("google") or {}).get("native_part"), dict + ) + ) + ) + if _is_synthetic_server_builtin and not ( + _name_lc in ("code_execution", "image_generation") + and isinstance(_native_part, dict) + ): + # No replayable Gemini native part -- skip + # entirely rather than send a fake functionCall. + # Also remember this tool_call_id so a matching + # role="tool" follow-up does not become an + # orphan functionResponse below. + if isinstance(tc_id, str) and tc_id: + _gemini_skip_tool_result_ids.add(tc_id) + tool_call_names.pop(tc_id, None) + continue + if fn_name in ("code_execution", "image_generation") and isinstance( + _native_part, dict + ): + # code_execution/image_generation history is + # replayed as native parts; the matching + # role="tool" must be skipped or Gemini sees a + # functionResponse with no declared function + # name and 400s the turn. + if isinstance(tc_id, str) and tc_id: + _gemini_skip_tool_result_ids.add(tc_id) + # New shape: `native_part.parts` is an ordered list + # of full part wrappers, each carrying its own + # `thoughtSignature`. This preserves Gemini 3's + # strict per-part replay requirement when the + # frontend has merged executableCode + + # codeExecutionResult + inlineData into the same + # tool-call card. + _native_parts_list = _native_part.get("parts") + if isinstance(_native_parts_list, list): + for _entry in _native_parts_list: + if isinstance(_entry, dict): + parts.append(_entry) + continue + # Legacy single-object native_part: fan the shared + # thoughtSignature only when one subpart exists; + # for code+result, prefer executableCode and drop + # the signature elsewhere. + _legacy_sig = _native_part.get( + "thoughtSignature" + ) or _native_part.get("thought_signature") + _legacy_subparts = [ + _k + for _k in ( + "executableCode", + "codeExecutionResult", + "inlineData", + ) + if isinstance(_native_part.get(_k), dict) + ] + for _native_key in ( + "executableCode", + "codeExecutionResult", + "inlineData", + ): + _sub = _native_part.get(_native_key) + if not isinstance(_sub, dict): + continue + _replay_part: dict[str, Any] = {_native_key: _sub} + if isinstance(_legacy_sig, str) and _legacy_sig: + if len(_legacy_subparts) == 1: + _replay_part["thoughtSignature"] = _legacy_sig + elif _native_key == "executableCode": + _replay_part["thoughtSignature"] = _legacy_sig + parts.append(_replay_part) + continue + + # Forward the OpenAI tool_call id into Gemini's + # functionCall.id so a follow-up turn that issues + # multiple calls to the same function (different + # args, same name) can be disambiguated on the + # response side. Gemini accepts the field per + # https://ai.google.dev/gemini-api/docs/function-calling. + function_call_part: dict[str, Any] = { + "name": fn_name, + "args": args, + } + if isinstance(tc_id, str) and tc_id: + function_call_part["id"] = tc_id + # Gemini 3 function-calling requires the prior + # thoughtSignature to be echoed back as a sibling + # of the functionCall part. The translator stows + # it on the assistant tool_call via + # `extra_content.google.thought_signature` (see + # the inbound emit below). + fc_part: dict[str, Any] = {"functionCall": function_call_part} + sig = _google_extra.get("thought_signature") or _google_extra.get( + "thoughtSignature" + ) + if isinstance(sig, str) and sig: + fc_part["thoughtSignature"] = sig + parts.append(fc_part) + if role == "tool": + # If the matching assistant-side tool_call was either + # dropped (synthetic server-tool with no native part) + # or already replayed as Gemini-native parts + # (code_execution/image_generation native_part), drop + # the follow-up too. Emitting it as a functionResponse + # would be orphaned or duplicate the native result. + _tc_id_for_skip = msg.get("tool_call_id") + if ( + isinstance(_tc_id_for_skip, str) + and _tc_id_for_skip in _gemini_skip_tool_result_ids + ): + continue + # OpenAI's role="tool" follow-up carries the function + # result. Gemini's matching shape is a role="user" turn + # with a functionResponse part. When the caller dropped + # ``name``, recover it from the matching assistant + # tool_call so Gemini doesn't 400 on an empty name. + tool_name = msg.get("name") or msg.get("tool_name") or "" + if not tool_name: + tc_id = msg.get("tool_call_id") + if isinstance(tc_id, str) and tc_id in tool_call_names: + tool_name = tool_call_names[tc_id] + response_payload: Any + if isinstance(content, list): + # OpenAI tool messages may carry list-form content + # (`[{"type":"text","text":"..."}]`). Forwarding the + # content-part objects verbatim into Gemini's + # `functionResponse.response.result` yields + # `result:[{"type":"text","text":"..."}]` instead of + # the actual tool output text; flatten text parts so + # the result mirrors the string-content path. + _flat_parts: list[str] = [] + for _cpart in content: + if ( + isinstance(_cpart, dict) + and _cpart.get("type") == "text" + and isinstance(_cpart.get("text"), str) + ): + _flat_parts.append(_cpart["text"]) + _flat_text = "".join(_flat_parts) + try: + response_payload = _json.loads(_flat_text) + except Exception: + response_payload = {"result": _flat_text} + elif isinstance(content, str): + try: + response_payload = _json.loads(content) + except Exception: + response_payload = {"result": content} + else: + response_payload = content or {} + function_response_part: dict[str, Any] = { + "name": tool_name, + "response": ( + response_payload + if isinstance(response_payload, dict) + else {"result": response_payload} + ), + } + # Mirror tool_call_id onto functionResponse.id so + # Gemini can match the result to the originating + # functionCall when multiple parallel calls were made. + tc_id = msg.get("tool_call_id") + if isinstance(tc_id, str) and tc_id: + function_response_part["id"] = tc_id + parts = [{"functionResponse": function_response_part}] + gemini_role = "user" + if parts: + # Gemini expects parallel functionResponses (multiple + # OpenAI role="tool" messages in a row) to ride on a + # single user content with multiple functionResponse + # parts -- the docs show parallel responses grouped + # together in the next turn. Merge consecutive + # functionResponse-only user blocks so realistic + # parallel tool loops round-trip correctly. + if ( + role == "tool" + and contents + and contents[-1].get("role") == "user" + and all( + isinstance(p, dict) and "functionResponse" in p + for p in (contents[-1].get("parts") or []) + ) + ): + contents[-1]["parts"].extend(parts) + else: + contents.append({"role": gemini_role, "parts": parts}) + + body: dict[str, Any] = {"contents": contents} + if system_text_parts: + body["systemInstruction"] = { + "parts": [{"text": "\n\n".join(system_text_parts)}] + } + + # Generation config -- temperature / topP / topK / maxOutputTokens + # map straight across. The frontend capability matrix restricts + # the sliders the UI exposes for Gemini to this set. + gen_config: dict[str, Any] = {} + if temperature is not None: + gen_config["temperature"] = temperature + if top_p is not None: + gen_config["topP"] = top_p + if top_k is not None and top_k > 0: + gen_config["topK"] = top_k + # Gemini accepts ``presencePenalty`` on generationConfig with the + # same sign convention as the OpenAI knob (positive discourages + # repetition). Forward when the caller bothers to set it. + if presence_penalty: + gen_config["presencePenalty"] = presence_penalty + if max_tokens is not None: + gen_config["maxOutputTokens"] = max_tokens + + # Nano Banana image generation. Gemini only accepts + # `responseModalities: ["TEXT","IMAGE"]` on the image-capable + # model family (id contains `-image` or `nano-banana`). Text- + # only models such as `gemini-2.5-flash` 400 on the same body, + # so only force image mode when the selected model actually + # supports it -- a stale `enabled_tools=["image_generation"]` + # on a text model is silently treated as a regular turn. + # https://ai.google.dev/gemini-api/docs/image-generation + model_lc = model.lower() + is_image_picker_model = "-image" in model_lc or "nano-banana" in model_lc + # tool_choice="none" / forced-function tool_choice must also + # suppress the implicit image-generation hosted tool. Otherwise + # an explicit OpenAI-style opt-out (or an explicit user-function + # pin) still flips `responseModalities=["TEXT","IMAGE"]` on + # image-tier models and bills for image output. + _tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + _tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _hosted_builtins_allowed = ( + not _tool_choice_disabled and not _tool_choice_forced_function + ) + # Image-tier model IDs reject text-only tools (code_execution, + # user functions) and thinkingConfig regardless of whether the + # Images pill is on -- those are model-level constraints + # documented by Google. The pill only controls whether we ask + # Gemini to actually emit image output via + # `responseModalities: ["TEXT","IMAGE"]`. Decoupling the two + # avoids the case where Images is off + Code/Search is on + # forwards `tools: [{codeExecution: {}}]` plus + # `thinkingConfig` to an image model and 400s. + image_tool_requested = bool( + _hosted_builtins_allowed + and enabled_tools + and "image_generation" in enabled_tools + ) + # Strict tool / thinking strip uses the model-id check. + is_image_model_strict = is_image_picker_model + # The actual modality flip only happens when the user opted in. + is_image_model = is_image_picker_model and image_tool_requested + if is_image_model: + gen_config["responseModalities"] = ["TEXT", "IMAGE"] + elif is_image_picker_model: + # Force TEXT-only so an image-capable model with Images OFF + # doesn't still bill for image output. + gen_config["responseModalities"] = ["TEXT"] + + # Thinking control. Gemini 3 uses thinkingLevel (str), 2.5 uses + # thinkingBudget (int). Gemini 3 has no full-off; minimum is + # "minimal" on Flash, "low" on Pro. + # https://ai.google.dev/gemini-api/docs/thinking + _GEMINI3_THINKING_PREFIXES = ( + "gemini-3.5-", + "gemini-3.1-", + "gemini-3-", + "gemini-pro-latest", + "gemini-flash-latest", + "gemini-flash-lite-latest", + ) + _GEMINI3_PRO_PREFIXES = ( + "gemini-3.5-pro", + "gemini-3.1-pro", + "gemini-3-pro", + "gemini-pro-latest", + ) + _PRO_THINKING_PREFIXES = ("gemini-2.5-pro",) + is_gemini3_thinking = any( + model_lc.startswith(p) for p in _GEMINI3_THINKING_PREFIXES + ) + is_gemini3_pro = any(model_lc.startswith(p) for p in _GEMINI3_PRO_PREFIXES) + _is_pro_thinking_only = any( + model_lc == p or model_lc.startswith(p + "-") + for p in _PRO_THINKING_PREFIXES + ) + effort_lc = (reasoning_effort or "").strip().lower() + if not is_image_model_strict and is_gemini3_thinking: + # Gemini 3.x thinkingLevel matrix: + # 3.1+ Pro: low/medium/high + # 3 Pro: low/high (deprecated 2026-03-09) + # 3.x Flash*: minimal/low/medium/high + # Coerce minimal->low on Pro; medium->high on legacy 3-Pro. + _G3_LEVELS = {"minimal", "low", "medium", "high"} + level: Optional[str] = None + if effort_lc in ("none", "off"): + level = "low" if is_gemini3_pro else "minimal" + elif effort_lc == "max": + level = "high" + elif effort_lc in _G3_LEVELS: + # Coerce legacy 3-Pro (low/high only) inputs. + _is_legacy_gemini3_pro = model_lc.startswith( + ("gemini-3-pro-preview", "gemini-3-pro") + ) and not model_lc.startswith(("gemini-3.1-pro", "gemini-3.5-pro")) + if is_gemini3_pro and effort_lc == "minimal": + level = "low" + elif _is_legacy_gemini3_pro and effort_lc == "medium": + level = "high" + else: + level = effort_lc + elif enable_thinking is True: + level = "high" + elif enable_thinking is False: + level = "low" if is_gemini3_pro else "minimal" + if level is not None: + gen_config["thinkingConfig"] = {"thinkingLevel": level} + elif not is_image_model_strict: + # Gemini 2.5 / older: thinkingBudget int. Effort -> budget + # mirrors the OpenAI minimal/low/medium/high ladder so the + # existing frontend picker maps cleanly. + # NOTE: gemini-2.5-flash-lite rejects positive budgets below + # 512 with HTTP 400, so minimal=512 sits at that floor. + _EFFORT_TO_BUDGET: dict[str, int] = { + "minimal": 512, + "low": 2048, + "medium": 8192, + "high": 24576, + "xhigh": -1, + "max": -1, + } + thinking_budget: Optional[int] = None + if effort_lc == "none" or enable_thinking is False: + # Pro-tier 2.5 rejects budget=0 (400 "only works in + # thinking mode"), so coerce to a small positive value. + thinking_budget = 128 if _is_pro_thinking_only else 0 + elif effort_lc in _EFFORT_TO_BUDGET: + thinking_budget = _EFFORT_TO_BUDGET[effort_lc] + elif enable_thinking is True: + thinking_budget = -1 + if thinking_budget is not None: + gen_config["thinkingConfig"] = { + "thinkingBudget": thinking_budget, + } + + if gen_config: + body["generationConfig"] = gen_config + + # Hosted tools: googleSearch (grounding) and codeExecution. + # Image-mode rejects codeExecution; only Gemini 3 image models + # accept googleSearch. + # https://ai.google.dev/gemini-api/docs/grounding + # https://ai.google.dev/gemini-api/docs/code-execution + def _gemini_image_model_allows_google_search(_m: str) -> bool: + return ( + _m.startswith("gemini-3-pro-image") + or _m.startswith("gemini-3.1-flash-image") + or _m.startswith("nano-banana-pro") + or _m.startswith("nano-banana-2") + ) + + google_search_allowed = ( + not is_image_model_strict + or _gemini_image_model_allows_google_search(model_lc) + ) + code_execution_allowed = not is_image_model_strict + text_tools_allowed = not is_image_model_strict + # tool_choice="none" / forced-function suppresses hosted builtins + # too, matching the Anthropic / OpenRouter gates. + tools_array: list[dict[str, Any]] = [] + if ( + _hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + and google_search_allowed + ): + tools_array.append({"googleSearch": {}}) + if ( + _hosted_builtins_allowed + and enabled_tools + and "code_execution" in enabled_tools + and code_execution_allowed + ): + tools_array.append({"codeExecution": {}}) + # OpenAI-style function declarations -> Gemini functionDeclarations. + # https://ai.google.dev/gemini-api/docs/function-calling#step_1 + # Gemini's Schema accepts only the OpenAPI 3.0 subset documented + # at https://ai.google.dev/api/caching#Schema; OpenAI's strict + # tool definitions routinely include `additionalProperties`, + # `$schema`, `$defs`, `strict`, `examples`, and similar keys + # which 400 the request as INVALID_ARGUMENT. Strip them + # recursively before forwarding. + _GEMINI_ALLOWED_SCHEMA_KEYS = frozenset( + { + "type", + "format", + "title", + "description", + "nullable", + "enum", + "maxItems", + "minItems", + "properties", + "required", + "minProperties", + "maxProperties", + "items", + "minimum", + "maximum", + "minLength", + "maxLength", + "pattern", + "default", + "anyOf", + "propertyOrdering", + } + ) + + def _resolve_local_schema_ref( + root: Optional[dict[str, Any]], ref: str + ) -> Optional[Any]: + # Walk a `#/foo/bar` JSON pointer against the schema root. + # Returns None if the pointer doesn't resolve to a dict, so + # the caller can fall back to the unresolved node. + if not isinstance(root, dict) or not isinstance(ref, str): + return None + if not ref.startswith("#/"): + return None + node: Any = root + for raw_part in ref[2:].split("/"): + if not raw_part: + continue + part = raw_part.replace("~1", "/").replace("~0", "~") + if not isinstance(node, dict) or part not in node: + return None + node = node[part] + return node + + def _sanitize_gemini_schema( + node: Any, + root: Optional[dict[str, Any]] = None, + _seen_refs: Optional[frozenset[str]] = None, + ) -> Any: + # Recursively filter to Gemini's OpenAPI 3.0 subset. At a + # Schema-keyword dict layer we drop keys not in the + # allowlist; under `properties` the keys are user-defined + # field names and the values are themselves Schemas; under + # `items` / `anyOf` the values are also Schemas. + # OpenAI strict tools commonly use JSON Schema's + # `"type": ["string", "null"]` form for nullable fields; + # Gemini's OpenAPI Schema uses `"type": "string"` plus + # `"nullable": true`. Translate that here. + if root is None and isinstance(node, dict): + root = node + if _seen_refs is None: + _seen_refs = frozenset() + if isinstance(node, dict): + # Pydantic / OpenAI strict tools commonly hoist nested + # object schemas into `$defs` and reference them via + # `{"$ref": "#/$defs/Address"}`. Gemini's OpenAPI subset + # has no $ref and drops anything not in the allowlist, + # so the referenced shape would vanish if we didn't + # inline it here. Recurse into the resolved target with + # local siblings overriding the reference (normal JSON + # Schema composition), guarding against ref cycles. + _ref = node.get("$ref") + if isinstance(_ref, str): + if _ref in _seen_refs: + return {} + _target = _resolve_local_schema_ref(root, _ref) + if isinstance(_target, dict): + _merged = { + **_target, + **{k: v for k, v in node.items() if k != "$ref"}, + } + return _sanitize_gemini_schema( + _merged, root, _seen_refs | {_ref} + ) + cleaned: dict[str, Any] = {} + _nullable_from_union = False + _flattened_type: Optional[str] = None + _union_any_of: Optional[list[dict[str, Any]]] = None + _raw_type = node.get("type") + if isinstance(_raw_type, list): + _non_null = [t for t in _raw_type if t != "null"] + if len(_non_null) < len(_raw_type): + _nullable_from_union = True + if len(_non_null) == 1: + _flattened_type = _non_null[0] + elif len(_non_null) > 1: + # Preserve multi-type unions as anyOf; flattening + # to the first non-null type silently drops the + # other branches and changes the tool contract. + _union_any_of = [ + {"type": _t} for _t in _non_null if isinstance(_t, str) + ] + for _k, _v in node.items(): + if _k == "type" and isinstance(_v, list): + # Handled below via _flattened_type. + continue + if _k not in _GEMINI_ALLOWED_SCHEMA_KEYS: + continue + if _k == "properties" and isinstance(_v, dict): + cleaned[_k] = { + _name: _sanitize_gemini_schema(_subschema, root, _seen_refs) + for _name, _subschema in _v.items() + } + elif _k == "items": + cleaned[_k] = _sanitize_gemini_schema(_v, root, _seen_refs) + elif _k == "anyOf" and isinstance(_v, list): + # Optional[X] / Union[A, B, None]: Pydantic emits + # `anyOf: [..., {"type":"null"}]`. Gemini's + # OpenAPI subset rejects `"type": "null"` inside + # anyOf, so drop the null variant and surface it + # via `nullable: true`. If exactly one non-null + # branch remains, collapse it inline; otherwise + # keep the slim anyOf and mark the field + # nullable. + _saw_null = any( + isinstance(_entry, dict) and _entry.get("type") == "null" + for _entry in _v + ) + _non_null_entries = [ + _entry + for _entry in _v + if not ( + isinstance(_entry, dict) + and _entry.get("type") == "null" + ) + ] + if len(_non_null_entries) == 1 and _saw_null: + _inner = _sanitize_gemini_schema( + _non_null_entries[0], root, _seen_refs + ) + if isinstance(_inner, dict): + for _ik, _iv in _inner.items(): + cleaned.setdefault(_ik, _iv) + cleaned.setdefault("nullable", True) + else: + cleaned[_k] = [ + _sanitize_gemini_schema(_entry, root, _seen_refs) + for _entry in _non_null_entries + ] + if _saw_null: + cleaned.setdefault("nullable", True) + elif _k in ("required", "enum", "propertyOrdering"): + # Lists of plain strings; copy verbatim. + cleaned[_k] = _v + else: + cleaned[_k] = _v + if _union_any_of is not None and "anyOf" not in cleaned: + cleaned["anyOf"] = [ + _sanitize_gemini_schema(_s, root, _seen_refs) + for _s in _union_any_of + ] + elif _flattened_type is not None: + cleaned["type"] = _flattened_type + if _nullable_from_union and "nullable" not in cleaned: + cleaned["nullable"] = True + return cleaned + return node + + function_declarations: list[dict[str, Any]] = [] + if tools and text_tools_allowed and not _tool_choice_disabled: + for _tool in tools: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _fn = _tool.get("function") + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _decl: dict[str, Any] = { + "name": _fn["name"], + "description": _fn.get("description") or "", + } + _params = _fn.get("parameters") + if isinstance(_params, dict): + _decl["parameters"] = _sanitize_gemini_schema(_params) + function_declarations.append(_decl) + if function_declarations: + tools_array.append({"functionDeclarations": function_declarations}) + if tools_array: + body["tools"] = tools_array + # Tool-choice mapping: OpenAI "auto"/"none"/"required"/{name=...} + # -> Gemini toolConfig.functionCallingConfig.mode + allowedFunctionNames. + if tool_choice is not None and function_declarations and text_tools_allowed: + _mode: Optional[str] = None + _allowed: Optional[list[str]] = None + if isinstance(tool_choice, str): + _tc_lc = tool_choice.strip().lower() + if _tc_lc == "auto": + _mode = "AUTO" + elif _tc_lc == "none": + _mode = "NONE" + elif _tc_lc in ("required", "any"): + _mode = "ANY" + elif ( + isinstance(tool_choice, dict) and tool_choice.get("type") == "function" + ): + _fn_pick = tool_choice.get("function") or {} + _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None + if isinstance(_name, str) and _name: + _mode = "ANY" + _allowed = [_name] + if _mode is not None: + _fcc: dict[str, Any] = {"mode": _mode} + if _allowed: + _fcc["allowedFunctionNames"] = _allowed + body["toolConfig"] = {"functionCallingConfig": _fcc} + + # Prompt caching. The Gemini caching contract is "create a + # CachedContent resource, then pass its name on + # `cachedContent`". The cache itself is created out of band by + # the caller via POST /cachedContents; here we forward an + # explicit cache id when the dispatcher hands us one (a string + # value on enable_prompt_caching means "use this cache name"). + # https://ai.google.dev/gemini-api/docs/caching + if isinstance(enable_prompt_caching, str) and enable_prompt_caching: + body["cachedContent"] = enable_prompt_caching + + # Model id is already validated at the top of _stream_gemini so + # we never reach a path-traversed URL segment here. + url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse" + completion_id = f"chatcmpl-gemini-{model.replace('/', '-')}" + + logger.info( + "Proxying Gemini streamGenerateContent to %s (model=%s, " + "tools=%s, image=%s)", + url, + model, + [list(t.keys())[0] for t in tools_array] if tools_array else [], + is_image_model, + ) + + def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + def _text_chunk( + text: str, extra_content: Optional[dict[str, Any]] = None + ) -> str: + delta: dict[str, Any] = {"content": text} + if extra_content: + delta["extra_content"] = extra_content + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": delta, + "finish_reason": None, + } + ], + } + return f"data: {_json.dumps(chunk)}" + + def _gemini_part_extra(part: dict[str, Any]) -> Optional[dict[str, Any]]: + """Return ``{"google": {"thought_signature": ...}}`` when the + Gemini stream part carries a `thoughtSignature` we need to + replay on a follow-up turn (Gemini 3 image editing + tool + contexts both require an exact signature echo).""" + sig = part.get("thoughtSignature") or part.get("thought_signature") + if isinstance(sig, str) and sig: + return {"google": {"thought_signature": sig}} + return None + + # Gemini finish reasons -> OpenAI vocabulary. Reference: + # https://ai.google.dev/api/rest/v1beta/Candidate#FinishReason + _finish_reason_map: dict[str, Optional[str]] = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "PROHIBITED_CONTENT": "content_filter", + "BLOCKLIST": "content_filter", + "MALFORMED_FUNCTION_CALL": "stop", + "OTHER": "stop", + "FINISH_REASON_UNSPECIFIED": None, + } + + last_usage: Optional[dict[str, Any]] = None + emitted_function_call_ids: set[str] = set() + # True once any Gemini functionCall part has been emitted so the + # final finish_reason swaps STOP -> tool_calls (matches the + # OpenAI Chat Completions contract; an OAI client that sees a + # tool_calls delta followed by finish_reason="stop" never + # executes the tool). + emitted_any_function_call = False + # web_search_active drives the tool_start / tool_end envelope. + # Track on whether `googleSearch` was actually forwarded above, + # not the raw caller intent -- image-mode requests filter the + # tool out, and emitting a phantom "search complete" card on a + # turn where Gemini was never told to search confuses the UI. + web_search_active = any("googleSearch" in t for t in tools_array) + web_search_tool_id = "gemini_web_search" + web_search_tool_started = False + web_search_tool_ended = False + web_search_citations: list[dict[str, str]] = [] + # Tracks the tool_call_id minted on the most recent + # executableCode part so the matching codeExecutionResult can + # close out the same envelope. None between rounds. + gemini_code_exec_pending_id: Optional[str] = None + # The most recently emitted code_execution id + result text. Kept + # *after* the tool_end so a following inline image (matplotlib + # plot rendered by codeExecution) can attach to the same card + # via a `__IMAGES__:` marker instead of spawning a separate + # image_generation event. + last_code_exec_tool_id: Optional[str] = None + last_code_exec_result_text: str = "" + + try: + async with _http_client.stream( + "POST", + url, + json = body, + headers = self._auth_headers(), + timeout = self._stream_timeout, + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors = "replace") + logger.error( + "Gemini returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + if web_search_active: + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": web_search_tool_id, + "arguments": {}, + } + ) + web_search_tool_started = True + + # NOTE: same manual __anext__ loop pattern as the other + # streaming helpers (see stream_chat_completion for the + # Python 3.13 + httpcore 1.0.x GeneratorExit ordering). + lines_gen = response.aiter_lines().__aiter__() + final_finish_reason: Optional[str] = None + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line.strip(): + continue + if not line.startswith("data:"): + continue + data_str = line[len("data:") :].strip() + if not data_str or data_str == "[DONE]": + continue + try: + event = _json.loads(data_str) + except Exception: + logger.warning( + "Gemini: failed to parse SSE chunk: %s", + data_str[:200], + ) + continue + if not isinstance(event, dict): + continue + + # Latch usageMetadata across deltas -- the final + # fragment carries the complete totals. + usage_meta = event.get("usageMetadata") + if isinstance(usage_meta, dict): + last_usage = usage_meta + + # Prompt-level safety block: Gemini ships zero + # candidates plus a `promptFeedback.blockReason` + # (e.g. SAFETY). The downstream OAI client would + # otherwise see an empty successful assistant + # response. Surface as a content_filter error + # event so the UI can render the block reason. + prompt_feedback = event.get("promptFeedback") + if isinstance(prompt_feedback, dict) and prompt_feedback.get( + "blockReason" + ): + block_reason = str(prompt_feedback.get("blockReason")) + # Close out the synthetic web_search start so + # the UI does not show a spinner stuck on + # "searching..." after the error toast lands. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": ( + "(search aborted: Gemini blocked " + f"prompt: {block_reason})" + ), + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 400, + f"Gemini blocked prompt: {block_reason}", + self.provider_type, + ) + return + + candidates = event.get("candidates") or [] + if not isinstance(candidates, list): + continue + for cand in candidates: + if not isinstance(cand, dict): + continue + # Citations / grounding metadata. + # `groundingMetadata.groundingChunks[].web` + # carries `uri` + `title`. Collect for the + # tool_end emission at stream close. + gm = cand.get("groundingMetadata") + if isinstance(gm, dict) and web_search_active: + chunks_list = gm.get("groundingChunks") or [] + if isinstance(chunks_list, list): + for ch in chunks_list: + if not isinstance(ch, dict): + continue + web = ch.get("web") or {} + if not isinstance(web, dict): + continue + u = web.get("uri") or "" + if not u or not isinstance(u, str): + continue + if any( + c["url"] == u for c in web_search_citations + ): + continue + web_search_citations.append( + { + "url": u, + "title": (web.get("title") or u), + "snippet": "", + } + ) + + content_obj = cand.get("content") or {} + parts = ( + content_obj.get("parts") + if isinstance(content_obj, dict) + else None + ) + if isinstance(parts, list): + for part in parts: + if not isinstance(part, dict): + continue + # Text delta. Stow part-level + # `thoughtSignature` on the delta so + # Gemini 3 turns that need an exact + # signature echo round-trip cleanly. + text = part.get("text") + _part_extra = _gemini_part_extra(part) + if isinstance(text, str) and text: + yield _text_chunk( + text, + extra_content = _part_extra, + ) + elif _part_extra is not None and not any( + k in part + for k in ( + "functionCall", + "executableCode", + "codeExecutionResult", + "inlineData", + ) + ): + # Empty-content part carrying a + # thoughtSignature: emit an empty delta + # so the signature is preserved. + yield _text_chunk( + "", + extra_content = _part_extra, + ) + # functionCall -> OpenAI tool_calls + # delta envelope. + fc = part.get("functionCall") + if isinstance(fc, dict): + fc_name = fc.get("name") or "" + fc_args = fc.get("args") or {} + fc_id = ( + fc.get("id") + or f"call_{fc_name}_{time.time_ns()}" + ) + if fc_id in emitted_function_call_ids: + continue + emitted_function_call_ids.add(fc_id) + # Each distinct functionCall in an + # assistant turn needs its own + # tool_calls[*].index. Consumers + # that reassemble tool_calls by + # index collapse all calls onto + # the same slot when this is + # hardcoded to 0, breaking + # parallel/multi-tool turns. + tc_index = len(emitted_function_call_ids) - 1 + tool_call_delta: dict[str, Any] = { + "index": tc_index, + "id": fc_id, + "type": "function", + "function": { + "name": fc_name, + "arguments": _json.dumps(fc_args), + }, + } + # Gemini 3 function-calling: the + # part-level `thoughtSignature` + # must be echoed back on the + # next turn or the model rejects + # the tool-result envelope. Stow + # it on `extra_content.google` + # so the frontend can persist it + # and our outbound translator + # (below) can replay it. + thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + if isinstance(thought_sig, str) and thought_sig: + tool_call_delta["extra_content"] = { + "google": { + "thought_signature": thought_sig, + } + } + emitted_any_function_call = True + tool_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [tool_call_delta] + }, + "finish_reason": None, + } + ], + } + yield f"data: {_json.dumps(tool_chunk)}" + # executableCode + codeExecutionResult + # parts surface as the standard + # code_execution tool_start/tool_end + # envelope (same shape OpenAI and + # Anthropic emit) so the chat + # adapter can render Gemini sandbox + # output through CodeExecutionToolUI. + # https://ai.google.dev/gemini-api/docs/code-execution + exec_code = part.get("executableCode") + if isinstance(exec_code, dict): + code_str = exec_code.get("code") or "" + if code_str: + code_tool_id = ( + exec_code.get("id") + or f"gemini_code_exec_{time.time_ns()}" + ) + gemini_code_exec_pending_id = code_tool_id + # Stow the raw Gemini part so + # follow-up turns can replay + # the native `executableCode` + # (Gemini rejects a generic + # functionCall echo for code + # execution history). + _exec_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + # Per-part thoughtSignature stays + # bound to its own part (Gemini 3 + # rejects shared signatures). + _exec_part_entry: dict[str, Any] = { + "executableCode": exec_code, + } + if ( + isinstance(_exec_thought_sig, str) + and _exec_thought_sig + ): + _exec_part_entry["thoughtSignature"] = ( + _exec_thought_sig + ) + _exec_native: dict[str, Any] = { + "parts": [_exec_part_entry], + } + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": code_tool_id, + "arguments": { + "kind": "code_execution", + "language": ( + ( + exec_code.get( + "language" + ) + or "PYTHON" + ).lower() + ), + "code": code_str, + "google": { + "native_part": _exec_native, + }, + }, + } + ) + exec_result = part.get("codeExecutionResult") + if isinstance(exec_result, dict): + outcome = exec_result.get("outcome") or "" + output = exec_result.get("output") or "" + # Gemini returns + # OUTCOME_OK / OUTCOME_FAILED / + # OUTCOME_DEADLINE_EXCEEDED. Treat + # non-OK outcomes as stderr so the + # UI surfaces the error. + if outcome and outcome != "OUTCOME_OK": + result_text = ( + f"[{outcome}]\n{output}".rstrip() + ) + else: + result_text = output + # Pair tool_end with the most recent + # executableCode tool_start; fall back + # to exec_result.id then a fresh id. + pair_id = ( + gemini_code_exec_pending_id + or exec_result.get("id") + or f"gemini_code_exec_{time.time_ns()}" + ) + if gemini_code_exec_pending_id is None: + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": pair_id, + "arguments": { + "kind": "code_execution", + "code": "", + }, + } + ) + _result_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _result_part_entry: dict[str, Any] = { + "codeExecutionResult": exec_result, + } + if ( + isinstance(_result_thought_sig, str) + and _result_thought_sig + ): + _result_part_entry["thoughtSignature"] = ( + _result_thought_sig + ) + _result_native: dict[str, Any] = { + "parts": [_result_part_entry], + } + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": pair_id, + "result": result_text, + "google": { + "native_part": _result_native, + }, + } + ) + last_code_exec_tool_id = pair_id + last_code_exec_result_text = result_text + gemini_code_exec_pending_id = None + # inlineData: either a Nano Banana + # generation (own card) or a sandbox + # plot attached to the code_execution + # card via the __IMAGES__: marker. + inline = part.get("inlineData") + if isinstance(inline, dict): + b64 = inline.get("data") or "" + mime = inline.get("mimeType") or "image/png" + if b64: + image_uri = f"data:{mime};base64,{b64}" + attached_to_code_exec = ( + not is_image_model + and last_code_exec_tool_id is not None + and bool(enabled_tools) + and "code_execution" + in (enabled_tools or []) + ) + if attached_to_code_exec: + updated_result = ( + last_code_exec_result_text + + "\n__IMAGES__:" + + _json.dumps([image_uri]) + ) + # Stow inlineData so a follow-up + # turn can replay the plot with + # its per-part thoughtSignature. + _plot_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _plot_part_entry: dict[str, Any] = { + "inlineData": { + "mimeType": mime, + "data": b64, + }, + } + if ( + isinstance(_plot_thought_sig, str) + and _plot_thought_sig + ): + _plot_part_entry[ + "thoughtSignature" + ] = _plot_thought_sig + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": ( + last_code_exec_tool_id + ), + "result": updated_result, + "google": { + "native_part": { + "parts": [ + _plot_part_entry + ], + }, + }, + } + ) + last_code_exec_result_text = ( + updated_result + ) + else: + img_id = f"img_{time.time_ns()}" + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": img_id, + "arguments": { + "kind": "image", + "prompt": "", + }, + } + ) + # Gemini 3 image edit needs + # the prior thoughtSignature + # echoed on the inline image part. + _img_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _img_tool_end: dict[str, Any] = { + "type": "tool_end", + "tool_call_id": img_id, + "result": "", + "image_b64": b64, + "image_mime": mime, + } + # Stow inlineData so multi-turn + # edits replay the original + # image as native history. + _img_part_entry: dict[str, Any] = { + "inlineData": { + "mimeType": mime, + "data": b64, + }, + } + if ( + isinstance(_img_thought_sig, str) + and _img_thought_sig + ): + _img_part_entry[ + "thoughtSignature" + ] = _img_thought_sig + _img_native: dict[str, Any] = { + "parts": [_img_part_entry], + } + _img_google: dict[str, Any] = { + "native_part": _img_native, + } + if ( + isinstance(_img_thought_sig, str) + and _img_thought_sig + ): + _img_google["thought_signature"] = ( + _img_thought_sig + ) + _img_tool_end["google"] = _img_google + yield _emit_tool_event(_img_tool_end) + finish_reason = cand.get("finishReason") + if isinstance(finish_reason, str): + mapped = _finish_reason_map.get(finish_reason, "stop") + if mapped is not None: + final_finish_reason = mapped + + # End-of-stream emission order: web_search tool_end + # (with citations) -> finish_reason chunk -> usage + # chunk -> [DONE]. Matches the Anthropic / OpenAI + # helpers' contract so the frontend handler does + # not need provider-specific ordering knowledge. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + blocks: list[str] = [] + for cit in web_search_citations: + line_out = f"Title: {cit['title']}\nURL: {cit['url']}" + if cit.get("snippet"): + line_out += f"\nSnippet: {cit['snippet']}" + blocks.append(line_out) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": ( + "\n---\n".join(blocks) + if blocks + else "(search complete)" + ), + } + ) + web_search_tool_ended = True + + if final_finish_reason: + # OpenAI clients trigger tool execution when + # finish_reason="tool_calls". Gemini emits + # "STOP" even when the turn was a pure + # functionCall request, so override after the + # fact to match the OAI contract. + if emitted_any_function_call and final_finish_reason == "stop": + final_finish_reason = "tool_calls" + finish_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": final_finish_reason, + } + ], + } + yield f"data: {_json.dumps(finish_chunk)}" + + # Map Gemini usageMetadata onto OpenAI include_usage. + # thoughtsTokenCount is billed output too — fold it in + # so cost calculators don't undercount. + if isinstance(last_usage, dict): + thought_tokens = last_usage.get("thoughtsTokenCount") or 0 + candidate_tokens = last_usage.get("candidatesTokenCount") or 0 + prompt_tokens = last_usage.get("promptTokenCount") or 0 + # Gemini bills tool-call prompt slices separately + # via `toolUsePromptTokenCount`. Fold into input + # so total_tokens does not undercount tool turns. + tool_use_prompt_tokens = ( + last_usage.get("toolUsePromptTokenCount") or 0 + ) + translated_usage = { + "input_tokens": prompt_tokens + tool_use_prompt_tokens, + "output_tokens": candidate_tokens + thought_tokens, + "input_tokens_details": { + "cached_tokens": ( + last_usage.get("cachedContentTokenCount") or 0 + ), + "tool_use_prompt_tokens": tool_use_prompt_tokens, + }, + "output_tokens_details": { + "reasoning_tokens": thought_tokens, + }, + } + usage_line = _build_usage_chunk( + completion_id, "openai", translated_usage + ) + if usage_line: + yield usage_line + + yield "data: [DONE]" + finally: + # Close response first so lines_gen.aclose() becomes + # a no-op (avoids the httpcore 1.0 GeneratorExit + # path and the aclose-never-awaited RuntimeWarning). + await response.aclose() + await lines_gen.aclose() + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": f"(search aborted: connection error: {exc})", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 502, + f"Failed to connect to {self.provider_type}: {exc}", + self.provider_type, + ) + except httpx.ReadTimeout as exc: + logger.error("Read timeout from %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": "(search aborted: read timeout)", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 504, + f"Timeout waiting for {self.provider_type} response", + self.provider_type, + ) + except httpx.HTTPError as exc: + logger.error("HTTP error from %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": f"(search aborted: transport error: {exc})", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, + ) + async def _stream_openai_responses( self, messages: list[dict[str, Any]], @@ -2893,6 +4949,8 @@ class ExternalProviderClient: enable_prompt_caching: Optional[bool] = None, openai_code_exec_container_id: Optional[str] = None, compaction_threshold: Optional[int] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -2916,6 +4974,12 @@ class ExternalProviderClient: # translate user/assistant messages into the Responses input shape. instructions_parts: list[str] = [] input_items: list[dict[str, Any]] = [] + # When we drop a server-side builtin `function_call` here, the + # matching `role="tool"` follow-up must also be dropped -- + # otherwise the outbound body contains an orphan + # `function_call_output` with no matching `function_call`, which + # OpenAI Responses can reject or mis-associate. + skipped_server_builtin_call_ids: set[str] = set() openai_replay_items: list[dict[str, Any]] = [] previous_response_id: Optional[str] = None for msg in messages: @@ -2932,6 +4996,126 @@ class ExternalProviderClient: instructions_parts.append(part["text"]) continue + # OpenAI Responses uses item-shape history for function + # calling: assistant turns that invoked user tools must + # serialize each call as a `function_call` input item, and + # each role="tool" follow-up as a `function_call_output` + # item keyed by the matching `call_id`. Without this the + # second turn after a function call sends Chat Completions + # shape and Responses 400s the request. + if role == "tool": + _call_id = msg.get("tool_call_id") or "" + # If the matching assistant `function_call` was a + # server-side builtin we already dropped, drop the + # follow-up too to avoid emitting an orphan + # `function_call_output`. + if _call_id and _call_id in skipped_server_builtin_call_ids: + continue + if isinstance(content, list): + _flat_parts: list[str] = [] + for part in content: + if part.get("type") == "text" and part.get("text"): + _flat_parts.append(part["text"]) + _output_text = "".join(_flat_parts) + else: + _output_text = content if isinstance(content, str) else "" + if _call_id: + input_items.append( + { + "type": "function_call_output", + "call_id": _call_id, + "output": _output_text, + } + ) + continue + + # Assistant turns that returned tool_calls translate each + # call as a `function_call` item (carrying name + JSON + # arguments + call_id). Skip builtin server-side cards + # (canonical builtin name + `args._server_tool` marker) + # which never round-trip as user functions. We require both + # checks so a user function literally named `_server_tool` + # in its argument schema is not dropped. + _tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if role == "assistant" and isinstance(_tool_calls, list): + # Preserve the prior `response.output` ordering: the + # model's text precedes its function_call items, and + # the matching role=tool follow-up arrives AFTER the + # call. Without this guard, history replay puts + # function_call -> assistant text -> function_call_output, + # which can put the tool output after an unrelated + # assistant message and confuse multi-turn function + # calling. + if isinstance(content, str) and content: + input_items.append({"role": "assistant", "content": content}) + elif isinstance(content, list): + _asst_parts: list[dict[str, Any]] = [] + for _part in content: + if not isinstance(_part, dict): + continue + _pt = _part.get("type") + if _pt == "text" and _part.get("text"): + _asst_parts.append( + { + "type": "input_text", + "text": _part.get("text", ""), + } + ) + elif _pt == "image_url": + _u = _part.get("image_url", {}).get("url", "") + if _u: + _asst_parts.append( + {"type": "input_image", "image_url": _u} + ) + if _asst_parts: + input_items.append( + {"role": "assistant", "content": _asst_parts} + ) + + for _tc in _tool_calls: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _args_raw = _fn.get("arguments") or "" + if not isinstance(_args_raw, str): + try: + _args_raw = _json.dumps(_args_raw) + except Exception: + _args_raw = "" + _fn_name_lc = (_fn.get("name") or "").lower() + _is_server_builtin = False + if _fn_name_lc in _SERVER_SIDE_BUILTIN_TOOL_NAMES: + try: + _args_obj = _json.loads(_args_raw) if _args_raw else {} + except Exception: + _args_obj = None + if isinstance(_args_obj, dict): + if _args_obj.get("_server_tool") is True: + _is_server_builtin = True + else: + _g = _args_obj.get("google") + if isinstance(_g, dict) and isinstance( + _g.get("native_part"), dict + ): + _is_server_builtin = True + _call_id_out = _tc.get("id") or f"call_{time.time_ns()}" + if _is_server_builtin: + skipped_server_builtin_call_ids.add(_call_id_out) + continue + input_items.append( + { + "type": "function_call", + "call_id": _call_id_out, + "name": _fn["name"], + "arguments": _args_raw, + } + ) + # Assistant text already emitted above (in order) so we + # don't fall through to the generic content branches. + continue + if isinstance(content, str): input_items.append({"role": role, "content": content}) continue @@ -3121,44 +5305,15 @@ class ExternalProviderClient: if max_tokens is not None: body["max_output_tokens"] = max_tokens - # Prompt caching on /v1/responses is automatic and free, but the - # default in-memory policy only survives ~5-10 min of inactivity - # (up to ~1 hr). Opt into the 24-hour retention policy so a chat - # left idle overnight still hits the cache on the next turn. - # Pricing is identical to in_memory per OpenAI's docs. - # - # Gated on the base URL because ollama / llama.cpp / "custom" - # presets all collapse to provider_type="openai" in - # toExternalBackendProviderType, so they also land in this - # helper. Those servers expose /v1/responses-shaped routes in - # some configurations but don't implement - # prompt_cache_retention — sending the field unconditionally - # would 400 them. Match the public OpenAI host strictly so the - # field only goes to OpenAI cloud. Studio's openai model picker - # is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which - # accept this parameter (gpt-5.5+ already defaults to "24h" and - # rejects "in_memory", so it's a safe no-op there). - # OpenAI-family cloud: api.openai.com OR Azure OpenAI Foundry - # (*.openai.azure.com). Both expose the same Responses-API - # extensions used below -- prompt_cache_retention, - # context_management compaction, container shell tool -- so - # treat them uniformly. Non-cloud OpenAI-compatible servers - # (ollama / llama.cpp / vLLM / "custom" preset) hit /v1/responses - # without these extensions and would 400 on the unknown body - # fields, so they intentionally fall outside this gate. + # Opt into 24h prompt-cache retention (free, vs the default + # ~5-10 min). Gated on the OpenAI cloud host because ollama / + # llama.cpp / "custom" presets reach this code path too and + # would 400 on the unknown field. if is_openai_cloud and enable_prompt_caching is not False: body["prompt_cache_retention"] = "24h" - # OpenAI server-side context compaction — see - # https://developers.openai.com/api/docs/guides/compaction - # When `compaction_threshold` is provided on a cloud OpenAI - # request, attach `context_management: [{type:"compaction", - # compact_threshold:N}]` so the API runs server-side - # compaction when the rendered prompt crosses the threshold. - # No beta header is required; no dated version pin. The field - # is silently dropped for non-cloud backends because ollama / - # llama.cpp / "custom" presets land in this helper and would - # 400 on an unknown body field. + # Server-side context compaction (OpenAI cloud only). + # https://developers.openai.com/api/docs/guides/compaction if ( is_openai_cloud and compaction_threshold is not None @@ -3171,55 +5326,92 @@ class ExternalProviderClient: } ] - # OpenAI server-side tools — see - # https://developers.openai.com/api/docs/guides/tools - # https://developers.openai.com/api/docs/guides/tools-shell - # The frontend's Search/Code buttons map to the unified - # enabled_tools shorthand; translate that into the Responses-API - # tool schema. Other built-in tools (file_search, - # code_interpreter, image_generation, computer_use_preview) can - # be added with the same pattern when we surface their toggles. + # Map enabled_tools onto Responses-API server tools (cloud only; + # local OAI-compat backends 400 on these). + # https://developers.openai.com/api/docs/guides/tools code_execution_enabled_openai = bool( enabled_tools and "code_execution" in enabled_tools and is_openai_cloud ) - # OpenAI's image_generation tool is a Responses-API server tool. - # See https://developers.openai.com/api/docs/guides/tools-image-generation - # The model picks size / quality / background server-side and - # delegates rendering to a gpt-image-* family model; the result - # comes back inline as an `image_generation_call` output item - # with a base64 image. Available on every gpt-5.x family member - # plus gpt-4.1 / gpt-4o / o3 per the docs; restrict to cloud - # OpenAI because the local llama.cpp / ollama backends don't - # implement it and would 400. - image_generation_enabled_openai = image_generation_requested + image_generation_enabled_openai = bool( + enabled_tools and "image_generation" in enabled_tools and is_openai_cloud + ) def _openai_image_generation_tool() -> dict[str, Any]: tool: dict[str, Any] = {"type": "image_generation"} if image_generation_has_reference: - # OpenAI's Responses image tool defaults to `auto`. For - # Studio's explicit follow-up edit flow, force edit mode so - # the provider uses the previous response / call id as image - # context instead of treating the text as a fresh generation. + # Force edit mode so the prior call id is used as context. tool["action"] = "edit" return tool - if enabled_tools: - tools_array: list[dict[str, Any]] = [] - if "web_search" in enabled_tools: + # Translate Chat-Completions function tools into the Responses + # function-tool shape (flattened name/description/parameters). + responses_user_function_tools: list[dict[str, Any]] = [] + if tools: + for _tool in tools: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _fn = _tool.get("function") + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _entry: dict[str, Any] = { + "type": "function", + "name": _fn["name"], + } + if _fn.get("description"): + _entry["description"] = _fn["description"] + if isinstance(_fn.get("parameters"), dict): + _entry["parameters"] = _fn["parameters"] + responses_user_function_tools.append(_entry) + + # Translate tool_choice into the Responses shape. + _responses_tc_string: Optional[str] = None + if isinstance(tool_choice, str): + _tc_lc = tool_choice.strip().lower() + if _tc_lc in ("auto", "none", "required"): + _responses_tc_string = _tc_lc + responses_tool_choice: Optional[Any] = None + _has_responses_tools = bool(enabled_tools or responses_user_function_tools) + if _responses_tc_string is not None and _has_responses_tools: + responses_tool_choice = _responses_tc_string + elif ( + tool_choice is not None + and responses_user_function_tools + and isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + ): + _fn_pick = tool_choice.get("function") or {} + _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None + if isinstance(_name, str) and _name: + responses_tool_choice = {"type": "function", "name": _name} + + _responses_tool_choice_none = _responses_tc_string == "none" + # A pinned user function suppresses hosted builtins (privacy + + # billing), matching the Gemini / Anthropic / OpenRouter gates. + _responses_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _responses_hosted_builtins_allowed = ( + not _responses_tool_choice_none + and not _responses_tool_choice_forced_function + ) + + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: + tools_array: list[dict[str, Any]] = list(responses_user_function_tools) + if ( + _responses_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): tools_array.append({"type": "web_search"}) - if code_execution_enabled_openai: - # `container_auto` lets OpenAI auto-create a fresh - # container per request; we capture the resulting - # container_id off the SSE stream and the chat-adapter - # persists it onto the thread record. Subsequent turns - # in the same thread pass it back as - # `openai_code_exec_container_id`, which we translate to - # `container_reference` here so the model sees - # filesystem state from prior turns. Container expires - # after ~20 min of inactivity per OpenAI's default - # policy — a stale id 400s, the chat-adapter clears it - # via container_invalidated, and the next turn falls - # back to auto-create. + if _responses_hosted_builtins_allowed and code_execution_enabled_openai: + # Reuse the thread's container so filesystem state + # persists; auto-create when there isn't one yet. Stale + # ids 400 and are cleared via container_invalidated. shell_env: dict[str, Any] if openai_code_exec_container_id: shell_env = { @@ -3229,10 +5421,12 @@ class ExternalProviderClient: else: shell_env = {"type": "container_auto"} tools_array.append({"type": "shell", "environment": shell_env}) - if image_generation_enabled_openai: + if _responses_hosted_builtins_allowed and image_generation_enabled_openai: tools_array.append(_openai_image_generation_tool()) if tools_array: body["tools"] = tools_array + if responses_tool_choice is not None: + body["tool_choice"] = responses_tool_choice url = f"{self.base_url}/responses" completion_id = f"chatcmpl-openai-{model.replace('/', '-')}" @@ -3246,11 +5440,19 @@ class ExternalProviderClient: first attempt. """ attempt_body = dict(body) - if enabled_tools: - tools_array_attempt: list[dict[str, Any]] = [] - if "web_search" in enabled_tools: + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: + tools_array_attempt: list[dict[str, Any]] = list( + responses_user_function_tools + ) + if ( + _responses_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): tools_array_attempt.append({"type": "web_search"}) - if code_execution_enabled_openai: + if _responses_hosted_builtins_allowed and code_execution_enabled_openai: if container_id_for_this_attempt: env_attempt: dict[str, Any] = { "type": "container_reference", @@ -3261,12 +5463,17 @@ class ExternalProviderClient: tools_array_attempt.append( {"type": "shell", "environment": env_attempt} ) - if image_generation_enabled_openai: + if ( + _responses_hosted_builtins_allowed + and image_generation_enabled_openai + ): tools_array_attempt.append(_openai_image_generation_tool()) if tools_array_attempt: attempt_body["tools"] = tools_array_attempt else: attempt_body.pop("tools", None) + if responses_tool_choice is not None: + attempt_body["tool_choice"] = responses_tool_choice return attempt_body def _is_openai_container_expired_error(error_text: str) -> bool: @@ -3328,52 +5535,26 @@ class ExternalProviderClient: done_emitted = False reasoning_open = False reasoning_emitted = False - # Latched from response.completed / response.incomplete so - # the final log can surface input_tokens_details.cached_tokens — - # the field that proves prompt_cache_retention="24h" is - # actually hitting OpenAI's cache instead of recomputing - # the prefix every turn. + # Per-call function-tool indexing; distinct slots so + # parallel calls don't collide on delta.tool_calls[].index. + saw_function_call = False + function_call_index = 0 + # Latched from response.completed/incomplete; surfaces + # input_tokens_details.cached_tokens to prove cache hits. last_usage: Optional[dict[str, Any]] = None - # Per-call state for OpenAI's server-side web_search tool. Mapped - # back into our local _toolEvent shape so the existing chat-UI - # renderer surfaces web_search the same way it does for local - # tool calls: a "Searching…" tool-call card, then a `tool_end` - # carrying citations formatted as - # Title: …\nURL: …\nSnippet: …\n---\n… - # blocks (which the frontend's parseSourcesFromResult lifts - # into source content parts at end of stream). - # web_search_calls preserves insertion order so we can apply - # the aggregated citation list onto the *last* call's - # tool_end — that's the one the frontend's source-pill - # extraction reads (parseSourcesFromResult flatMaps every - # web_search result, so a single non-empty result is enough - # to surface all sources at message tail). - # OpenAI emits url_citation annotations on text deltas, not - # per call — there's no wire field linking a citation back - # to a specific search invocation. Hence the shared list. - # web_search_calls: { item_id -> {query} } + # web_search state. Citations are emitted on text deltas + # (not per call), so the aggregate list is shared and + # applied to the LAST web_search tool_end (parseSourcesFromResult + # flatmaps every call, one non-empty is enough). web_search_calls: dict[str, dict[str, Any]] = {} all_url_citations: list[dict[str, Any]] = [] - # Shell-tool (code execution) state. OpenAI emits - # `shell_call` items (model requesting a command list) - # paired with `shell_call_output` items (execution - # results). We mirror the Anthropic code-execution UX - # by emitting one `_toolEvent` tool_start per - # shell_call and one tool_end per shell_call_output; - # they're linked via `shell_call_output.call_id` - # matching `shell_call.id`. Items are independent of - # web_search (different keyed map). - # shell_calls: { call_id -> {commands, output} } + # shell_calls (code execution): { call_id -> {commands, output} }. + # shell_call <-> shell_call_output match by call_id; emit + # tool_start/tool_end like the Anthropic UX. shell_calls: dict[str, dict[str, Any]] = {} - # Container id captured from the response stream. When - # it differs from the inbound id, emit a synthetic - # `container_ready` _toolEvent so the frontend can - # persist it onto the thread record for the next turn. - # Where OpenAI surfaces it is documented loosely; we - # probe two known fields (response.container_id on - # response.completed, item.environment.container_id on - # shell_call output items) and latch the first one we - # see. + # Container id latched from response.container_id or + # item.environment.container_id; emit container_ready + # when it differs from the inbound id. latched_container_id: Optional[str] = None container_id_emitted = False current_openai_response_id: Optional[str] = None @@ -3453,6 +5634,7 @@ class ExternalProviderClient: return rendered def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -3775,17 +5957,9 @@ class ExternalProviderClient: f"ws_{len(web_search_calls)}" ) web_search_calls.setdefault(item_id, {"query": ""}) - # Shell-tool: register the call eagerly so - # the matching shell_call_output can link - # back even if `done` arrives out of order. - # Also probe for container_id on the - # environment field — when container_auto - # auto-creates one, this is the first place - # the new id might surface (OpenAI doesn't - # promise this in docs, but the field is - # cheap to scan and lets us emit - # container_ready earlier than - # response.completed). + # Register shell_call eagerly so out-of-order + # output links back. Probe env.container_id + # to emit container_ready before response.completed. if ( isinstance(item, dict) and item.get("type") == "shell_call" @@ -3978,24 +6152,9 @@ class ExternalProviderClient: } ) elif item.get("type") == "image_generation_call": - # OpenAI's image_generation tool returns - # a single output item with the base64 - # PNG/WebP/JPEG on `result` (sometimes - # `b64_json` depending on output_format). - # `revised_prompt` is what the gpt-image - # backbone actually used after refinement - # of the assistant's request. Emit - # tool_start + tool_end so the chat card - # renders the prompt + the generated - # image inline. The frontend chat-adapter - # decides how to render the base64 blob - # (likely an ) - # based on the `kind: "image"` hint we - # set on tool_start arguments. - # `time_ns()` (nanoseconds) instead of - # millisecond resolution so synthesised - # ids stay unique even when two image - # generations resolve in the same ms. + # Base64 image on `result` (or `b64_json`), + # `revised_prompt` for the rewritten prompt. + # ns-resolution id so concurrent gens stay unique. raw_item_id = item.get("id") item_id = raw_item_id or f"img_{time.time_ns()}" prompt_in = ( @@ -4034,6 +6193,54 @@ class ExternalProviderClient: "prompt": prompt_in, } ) + elif item.get("type") == "function_call": + # Translate to Chat-Completions delta.tool_calls. + # https://platform.openai.com/docs/guides/function-calling?api-mode=responses + fn_call_id = ( + item.get("call_id") + or item.get("id") + or f"call_{time.time_ns()}" + ) + fn_name = item.get("name") or "" + fn_args = item.get("arguments") or "" + if not isinstance(fn_args, str): + try: + fn_args = _json.dumps(fn_args) + except Exception: + fn_args = "" + _tc_index = function_call_index + function_call_index += 1 + yield ( + "data: " + + _json.dumps( + { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": _tc_index, + "id": fn_call_id, + "type": "function", + "function": { + "name": fn_name, + "arguments": ( + fn_args + ), + }, + } + ], + }, + "finish_reason": None, + } + ], + } + ) + ) + saw_function_call = True elif ( isinstance(event_type, str) @@ -4167,7 +6374,11 @@ class ExternalProviderClient: { "index": 0, "delta": {}, - "finish_reason": "stop", + "finish_reason": ( + "tool_calls" + if saw_function_call + else "stop" + ), } ], } @@ -4448,11 +6659,70 @@ class ExternalProviderClient: models = [model for model in raw_models if isinstance(model, dict)] if not models and self.provider_type == "ollama": models = await self._list_ollama_native_models() + # Gemini's native /v1beta/models returns + # {"models": [{"name": "models/gemini-2.5-flash", ...}]} + # -- repackage into the OpenAI-compatible shape the rest + # of Studio expects so dynamic model discovery works. + if not models and self.provider_type == "gemini": + models = self._parse_gemini_models(data) return models except httpx.HTTPError as exc: logger.error("Failed to list models from %s: %s", self.provider_type, exc) raise + @staticmethod + def _parse_gemini_models(payload: Any) -> list[dict[str, Any]]: + """Translate Gemini's native /v1beta/models payload to OpenAI shape. + + Native response: + {"models": [{"name": "models/gemini-2.5-flash", + "baseModelId": "gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": [...]}]} + + We only keep entries that advertise + ``generateContent`` / ``streamGenerateContent`` so the picker + does not surface embedding-only models the chat path can't + drive. + """ + if not isinstance(payload, dict): + return [] + entries = payload.get("models") or [] + if not isinstance(entries, list): + return [] + out: list[dict[str, Any]] = [] + for entry in entries: + if not isinstance(entry, dict): + continue + methods = entry.get("supportedGenerationMethods") or [] + if ( + isinstance(methods, list) + and methods + and not any( + m in methods for m in ("generateContent", "streamGenerateContent") + ) + ): + continue + base_id = entry.get("baseModelId") + name = entry.get("name") or "" + # ``name`` arrives as ``"models/gemini-2.5-flash"``; the + # chat path uses the bare id. + short_id = ( + base_id + if isinstance(base_id, str) and base_id + else (name.split("/", 1)[1] if "/" in name else name) + ) + if not short_id: + continue + out.append( + { + "id": short_id, + "owned_by": "google", + "display_name": entry.get("displayName") or short_id, + } + ) + return out + async def _list_ollama_native_models(self) -> list[dict[str, Any]]: """Fallback when Ollama's /v1/models returns an empty or null catalog.""" root = self.base_url.removesuffix("/v1").rstrip("/") @@ -4757,6 +7027,17 @@ def _build_usage_chunk( "total_tokens": prompt_tokens + completion_tokens, "prompt_tokens_details": {"cached_tokens": cached}, } + # Surface OpenAI Responses / Gemini reasoning-token detail. The + # caller pre-populates last_usage["output_tokens_details"] with + # at least {"reasoning_tokens": ...}; mirror it into the OAI + # `completion_tokens_details` shape so SDKs can render the + # hidden-thoughts slice. + out_details = last_usage.get("output_tokens_details") + if isinstance(out_details, dict) and out_details: + usage_block["completion_tokens_details"] = { + "reasoning_tokens": out_details.get("reasoning_tokens") or 0, + } + usage_block["output_tokens_details"] = out_details chunk = { "id": completion_id, diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index f18b9c177d..2688f656ee 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -65,28 +65,77 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { }, "gemini": { "display_name": "Google Gemini", - "base_url": "https://generativelanguage.googleapis.com/v1beta/openai", - # Curated lineup — Google's /v1beta/openai/models returns dozens - # of historical / experimental / embedding ids. Cap to the current - # 3.x family plus the rolling `*-latest` aliases. + # Native Gemini REST endpoint -- the Gemini API does NOT speak + # OpenAI Chat Completions on this base. Requests/responses are + # translated in `_stream_gemini` in external_provider.py. + # API reference: https://ai.google.dev/gemini-api/docs + "base_url": "https://generativelanguage.googleapis.com/v1beta", + # Curated lineup -- the live ListModels response returns dozens + # of historical / experimental / embedding ids. Cap to the + # current chat-capable Gemini families (3.5 / 3.1 / 3 Flash / + # 2.5) plus the Nano Banana image trio and the rolling + # `*-latest` aliases. Excluded on purpose: + # - `gemini-2.0-flash*` (Google retired 2026-06-01; 404 on use) + # - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects + # to `gemini-3.1-pro-preview` per Google's deprecation notice, + # so we surface 3.1 directly and skip the redirect). + # The allowlist below blocks the retired ids from re-appearing + # via the live ListModels fetch. Verified against the live + # `/v1beta/models` catalog 2026-05-24. "default_models": [ "gemini-3.1-pro-preview", + "gemini-3.5-flash", "gemini-3.1-flash-lite", "gemini-3-flash-preview", "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest", + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + "gemini-2.5-flash-image", ], "supports_streaming": True, "supports_vision": True, "supports_tool_calling": True, - "auth_header": "Authorization", - "auth_prefix": "Bearer ", - "notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.", + # The native API takes the API key on the `x-goog-api-key` + # header. An empty `auth_prefix` ensures we send the bare key. + "auth_header": "x-goog-api-key", + "auth_prefix": "", + "openai_compatible": False, + "notes": ( + "Native Gemini API. Translation lives in _stream_gemini. " + "API key from https://aistudio.google.com/apikey. " + "See https://ai.google.dev/gemini-api/docs for endpoint shapes." + ), + # Even after the regex match, drop ids that Google still + # returns from ListModels but routes via implicit redirect. + # gemini-3-pro-preview was shut down 2026-03-09 and is + # auto-aliased to gemini-3.1-pro-preview; we surface the + # canonical id only so users do not see two cards for the + # same underlying model. + "model_id_deny_exact": ("gemini-3-pro-preview",), + # Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the + # rolling *-latest aliases (which Google rolls forward as new + # generations ship). Image-tier ids (`-image`, `-image-preview`, + # `nano-banana-pro-preview`) flow through the Nano Banana + # `responseModalities` path in `_stream_gemini`. Retired 2.0 + # ids ARE NOT in this regex on purpose -- Google's ListModels + # would otherwise re-surface them and they 404 on use. "model_id_allowlist": re.compile( - r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|" - r"gemini-3\.1-pro-preview|gemini-pro-latest|" - r"gemini-flash-latest|gemini-flash-lite-latest)$" + r"^(" + r"gemini-3\.5-(?:flash|pro)(?:-preview)?|" + r"gemini-3\.1-(?:flash|pro|flash-lite)(?:-preview)?(?:-customtools)?|" + r"gemini-3\.1-flash-image-preview|" + r"gemini-3-(?:flash|pro)(?:-preview)?|" + r"gemini-3-pro-image-preview|" + r"nano-banana-pro-preview|" + r"gemini-2\.5-pro|gemini-2\.5-flash|gemini-2\.5-flash-lite|" + r"gemini-2\.5-flash-image|" + r"gemini-pro-latest|gemini-flash-latest|gemini-flash-lite-latest" + r")$" ), }, "deepseek": { diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0e9cce7c3e..35d5903826 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -632,8 +632,17 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str for *_, sockaddr in infos: ip = ipaddress.ip_address(sockaddr[0]) + # `not ip.is_global` rejects every category the denylist below + # also rejects PLUS shared address space (100.64.0.0/10 carrier- + # grade NAT) and benchmarking/documentation/exchange ranges that + # Python classifies with `is_private=False` and `is_global=False` + # (see https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global). + # The explicit predicates after it give human-readable categories + # in the error message, but a single non-global check is the + # source of truth and prevents future ranges from leaking. if ( - ip.is_private + not ip.is_global + or ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b128fb5338..b9643cac6a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3057,6 +3057,14 @@ class UnslothTrainer: logger.info("Configuring DeepSeek OCR data collator...\n") FastVisionModel.for_training(self.model) + # DeepSeek OCR's (image_size, base_size, crop_mode) is a + # coupled preset; changing image_size alone desyncs the + # per-crop pixel grid from num_queries. Use Gundam. + if training_args.get("vision_image_size") is not None: + logger.info( + "Vision image resize ignored for DeepSeek OCR " + "(uses fixed Gundam preset).\n" + ) data_collator = DeepSeekOCRDataCollator( tokenizer = self.tokenizer, model = self.model, @@ -3123,7 +3131,21 @@ class UnslothTrainer: from unsloth.trainer import UnslothVisionDataCollator FastVisionModel.for_training(self.model) - data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) + vision_image_size = training_args.get("vision_image_size") + if vision_image_size is None: + data_collator = UnslothVisionDataCollator( + self.model, self.tokenizer + ) + else: + logger.info( + f"Vision image resize: {vision_image_size} (max dimension)\n" + ) + data_collator = UnslothVisionDataCollator( + self.model, + self.tokenizer, + resize = vision_image_size, + resize_dimension = "max", + ) logger.info("Vision data collator configured\n") # ========== TRAINING CONFIGURATION ========== diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d2c2316d45..0af3349c6f 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -193,6 +193,7 @@ class TrainingBackend: "hf_token": kwargs.get("hf_token", ""), "load_in_4bit": kwargs.get("load_in_4bit", True), "max_seq_length": kwargs.get("max_seq_length", 2048), + "vision_image_size": kwargs.get("vision_image_size"), "hf_dataset": kwargs.get("hf_dataset", ""), "local_datasets": kwargs.get("local_datasets"), "local_eval_datasets": kwargs.get("local_eval_datasets"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f47a6bd599..632b38d75a 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -959,7 +959,47 @@ def _activate_transformers_version(model_name: str) -> None: activate_transformers_for_subprocess(model_name) -def _adapt_for_mlx_vlm(items): +def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]: + if width <= 0 or height <= 0 or target <= 0: + return width, height + largest_side = max(width, height) + if largest_side <= target: + return width, height + # Integer formula matches unsloth_zoo's collator (Python round() differs + # by 1px on half-pixel cases). max(1, _) avoids zero-side degenerate output. + new_w = max(1, (width * target + largest_side // 2) // largest_side) + new_h = max(1, (height * target + largest_side // 2) // largest_side) + return new_w, new_h + + +def _resize_mlx_vlm_image(image, resize): + if resize is None: + return image + try: + from PIL import Image + import numpy as np + except ImportError: + return image + if not isinstance(image, Image.Image): + return image + image = image.convert("RGB") + new_size = _mlx_vlm_max_resized_size(*image.size, int(resize)) + if new_size != image.size: + resampling = getattr(Image, "Resampling", Image).LANCZOS + image = image.resize(new_size, resampling) + # When a resize is requested, hand mlx-vlm a writable RGB ndarray so its + # PIL-path square-resize is skipped and HF processors don't warn on + # non-writable views. resize=None (Default) above keeps the original PIL. + return np.array(image, copy = True) + + +def _resize_mlx_vlm_images(value, resize): + if isinstance(value, list): + return [_resize_mlx_vlm_image(image, resize) for image in value] + return _resize_mlx_vlm_image(value, resize) + + +def _adapt_for_mlx_vlm(items, resize = None): """Adapt GPU-path VLM dataset output for mlx-vlm consumption. The GPU path embeds PIL images inside messages content as @@ -979,7 +1019,7 @@ def _adapt_for_mlx_vlm(items): if isinstance(part, dict) and part.get("type") == "image": img = part.get("image") if img is not None: - images.append(img) + images.append(_resize_mlx_vlm_image(img, resize)) new_content.append({"type": "image"}) else: new_content.append(part) @@ -990,9 +1030,9 @@ def _adapt_for_mlx_vlm(items): if images: out["image"] = images[0] if len(images) == 1 else images elif "image" in item: - out["image"] = item["image"] + out["image"] = _resize_mlx_vlm_images(item["image"], resize) elif "images" in item: - out["images"] = item["images"] + out["images"] = _resize_mlx_vlm_images(item["images"], resize) adapted.append(out) return adapted @@ -1168,6 +1208,25 @@ def _run_mlx_training(event_queue, stop_queue, config): is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False)) model._is_vlm_model = is_vlm + vision_image_size = config.get("vision_image_size") + # DeepSeek OCR uses a coupled preset tuple; skip resize like the Torch path. + _model_name_lower = str(config.get("model_name", "")).lower() + _is_deepseek_ocr = "deepseek" in _model_name_lower and "ocr" in _model_name_lower + if is_vlm and vision_image_size is not None and _is_deepseek_ocr: + _send( + "status", + status_message = ( + "MLX vision image resize ignored for DeepSeek OCR " + "(uses fixed Gundam preset)." + ), + ) + vision_image_size = None + elif is_vlm and vision_image_size is not None: + vision_image_size = int(vision_image_size) + _send( + "status", + status_message = f"MLX vision image resize: {vision_image_size} (max dimension)", + ) # ── 2. Apply LoRA / full FT ── # Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.) @@ -1302,7 +1361,10 @@ def _run_mlx_training(event_queue, stop_queue, config): progress_callback = _fmt_progress, ) if vlm_info.get("success"): - dataset = _adapt_for_mlx_vlm(vlm_info["dataset"]) + dataset = _adapt_for_mlx_vlm( + vlm_info["dataset"], + resize = vision_image_size, + ) else: errors = vlm_info.get("errors", []) raise ValueError( @@ -1317,7 +1379,10 @@ def _run_mlx_training(event_queue, stop_queue, config): dataset_name = hf_dataset or "local", ) if ev_info.get("success"): - eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"]) + eval_dataset = _adapt_for_mlx_vlm( + ev_info["dataset"], + resize = vision_image_size, + ) elif format_type: _send("status", status_message = f"Formatting dataset ({format_type})...") @@ -2248,6 +2313,7 @@ def run_training_process( eval_dataset = eval_dataset, eval_steps = eval_steps, max_seq_length = config.get("max_seq_length", 2048), + vision_image_size = config.get("vision_image_size"), optim = config.get("optim", "adamw_8bit"), lr_scheduler_type = config.get("lr_scheduler_type", "linear"), is_cpt = is_cpt, diff --git a/studio/backend/main.py b/studio/backend/main.py index d0d0cf9043..686ef5f0d6 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -18,6 +18,17 @@ _backend_dir = str(_Path(__file__).parent) if _backend_dir not in sys.path: sys.path.insert(0, _backend_dir) +# `uvicorn main:app` bypasses run.py; seed thread caps here too. +from utils.cpu_threads import configure_cpu_threads + +try: + configure_cpu_threads() +except ValueError as exc: + _raw = os.environ.get("UNSLOTH_CPU_THREADS") + raise SystemExit( + f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}" + ) from None + # Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before # any library imports that trigger attrs -> rich -> structlog -> platform crash. # See: https://github.com/python/cpython/issues/102396 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e48b0919cf..a23cd9c0c2 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -581,6 +581,14 @@ class ChatMessage(BaseModel): None, description = "OpenAI tool-result messages: name of the tool whose result this is.", ) + extra_content: Optional[dict] = Field( + None, + description = ( + "Provider-specific extra fields the translator may read. " + "Gemini reads `extra_content.google.thought_signature` " + "from assistant messages to replay text-part signatures." + ), + ) @model_validator(mode = "after") def _validate_role_shape(self) -> "ChatMessage": @@ -752,17 +760,42 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Override base URL for the external provider.", ) - enable_prompt_caching: Optional[bool] = Field( + enable_prompt_caching: Optional[Union[bool, str]] = Field( None, description = ( "[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, " - "attaches cache_control={type:ephemeral} to the system block so the " - "static prefix is reused across turns. On OpenAI cloud, caching is " - "automatic for prompts >=1024 tokens and this flag is informational. " - "Ignored for every other provider (mistral, gemini, kimi, openrouter, " - "vllm, local, etc.). Treated as enabled when omitted." + "boolean true attaches cache_control={type:ephemeral} to the system " + "block so the static prefix is reused across turns. On OpenAI cloud, " + "caching is automatic for prompts >=1024 tokens and the boolean is " + "informational. On Gemini, pass a string cache resource name such " + "as `cachedContents/abc123` to attach `cachedContent` on the native " + "request (boolean true is a no-op on Gemini because creating the " + "cache requires a separate POST /cachedContents call). Ignored for " + "every other provider. Treated as enabled when omitted." ), ) + + @field_validator("enable_prompt_caching", mode = "before") + @classmethod + def _coerce_enable_prompt_caching(cls, value: Any) -> Any: + """Preserve the pre-PR coercion: the field used to be Optional[bool], + so callers historically sent JSON strings `"true"` / `"false"` and + Pydantic v1 coerced them. Widening to Optional[Union[bool, str]] for + Gemini cache resource names lets `"false"` slip through as a truthy + string. Coerce the canonical bool literals back so explicit opt-outs + stay opt-out.""" + if isinstance(value, str): + lowered = value.strip().lower() + # Match Pydantic v1's BooleanField coercion table (yes/y/on/t/1 + # and no/n/off/f/0) so opt-outs that used to parse still parse. + # Anything else is preserved as a string for Gemini's + # cachedContent resource path. + if lowered in ("true", "t", "1", "yes", "y", "on"): + return True + if lowered in ("false", "f", "0", "no", "n", "off"): + return False + return value + prompt_cache_ttl: Optional[str] = Field( None, description = ( diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 7c53b0fee5..c6be1eff4e 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -5,10 +5,17 @@ Pydantic schemas for Training API """ +import re from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +# ASCII integer with an optional single sign. Used by _check_vision_image_size +# to reject "++512", "--256", and Unicode-digit strings ("５１２", "٥١٢") that +# would otherwise slip through str.isdigit() + int(). +_INT_RE = re.compile(r"[+-]?[0-9]+") + + _MAX_BATCH_SIZE = 4096 _MAX_GRAD_ACCUM = 4096 _MAX_STEPS = 1_000_000 @@ -18,6 +25,9 @@ _MAX_SEQ_LENGTH = 2_000_000 _MAX_LR_VALUE = 1.0 _MAX_LORA_R = 16_384 _MAX_LORA_ALPHA = 32_768 +_MIN_VISION_IMAGE_SIZE = 256 +# 2048 was the most I could get most llms to work at without getting unstable +_MAX_VISION_IMAGE_SIZE = 2048 def _parse_lr(v: Any) -> float: @@ -58,6 +68,10 @@ class TrainingStartRequest(BaseModel): hf_token: Optional[str] = Field(None, description = "HuggingFace token") load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") max_seq_length: int = Field(2048, description = "Maximum sequence length") + vision_image_size: Optional[int] = Field( + None, + description = "Optional maximum image side length for VLM training. Null uses model default.", + ) trust_remote_code: bool = Field( False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", @@ -159,6 +173,40 @@ class TrainingStartRequest(BaseModel): ) return v + @field_validator("vision_image_size", mode = "before") + @classmethod + def _check_vision_image_size(cls, v: Any) -> Optional[int]: + # mode="before" sees True/False as bool (not 1/0) for a precise error. + if v is None: + return v + if isinstance(v, bool): + raise ValueError("vision_image_size must be an integer or null") + if isinstance(v, int): + coerced = v + elif isinstance(v, str) and _INT_RE.fullmatch(v.strip()): + coerced = int(v.strip()) + elif isinstance(v, float) and v.is_integer(): + coerced = int(v) + else: + # numpy ints / Integral subclasses, without a hard numpy import. + try: + import numbers + + if isinstance(v, numbers.Integral): + coerced = int(v) + elif isinstance(v, numbers.Real) and float(v).is_integer(): + coerced = int(v) + else: + raise TypeError + except Exception: + raise ValueError("vision_image_size must be an integer or null") + if coerced < _MIN_VISION_IMAGE_SIZE or coerced > _MAX_VISION_IMAGE_SIZE: + raise ValueError( + f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, " + f"{_MAX_VISION_IMAGE_SIZE}] (got {coerced!r})" + ) + return coerced + @field_validator("warmup_steps") @classmethod def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 629ac56bf6..d06124c9a1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1705,6 +1705,7 @@ def _build_external_messages( messages: list, supports_vision: bool, provider_type: Optional[str] = None, + base_url: Optional[str] = None, ) -> list[dict]: """ Convert ChatMessage list to OpenAI-compatible dicts for external providers. @@ -1732,14 +1733,171 @@ def _build_external_messages( document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS anthropic = provider_type == "anthropic" openai = provider_type == "openai" + # `extra_content` is a Gemini-specific carrier for the assistant's + # text-part `thoughtSignature` round-trip on the native + # streamGenerateContent endpoint. Custom Gemini OpenAI-compatible + # gateways (LiteLLM etc.) route through /chat/completions where + # the field is unknown and can be rejected -- gate strictly on the + # Google-hosted Gemini base. + _native_gemini = False + if provider_type == "gemini" and base_url: + try: + from urllib.parse import urlparse as _urlparse + + _host = (_urlparse(base_url).hostname or "").lower() + _native_gemini = _host == "generativelanguage.googleapis.com" + except Exception: + _native_gemini = False + emit_extra_content = _native_gemini + + _SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} + ) + + def _is_marked_server_builtin_tool_call(tc: Any) -> bool: + """Return True iff `tc` is a synthetic provider-side tool card + with one of the canonical builtin names and either: + - the new `args._server_tool` marker stamped by the backend, or + - a Gemini `args.google.native_part` payload (durable replay + signal for code_execution / image_generation that predates + the marker). + Such cards must not be forwarded to non-native providers + because they are not real user functions and the receiving API + will reject the orphan tool history. Real user functions with + these names normally have neither signal. + """ + if not isinstance(tc, dict): + return False + fn = tc.get("function") + if not isinstance(fn, dict): + return False + name = (fn.get("name") or "").lower() + if name not in _SERVER_BUILTIN_TOOL_NAMES: + return False + raw_args = fn.get("arguments") or "" + try: + args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + except Exception: + return False + if not isinstance(args, dict): + return False + if args.get("_server_tool") is True: + return True + google = args.get("google") + return isinstance(google, dict) and isinstance(google.get("native_part"), dict) + + # When we drop a server-side builtin tool_call here, the matching + # `role="tool"` follow-up must also be dropped from the outbound + # history -- otherwise the provider receives an orphan + # tool_call_id with no matching assistant call, which OpenAI + # Responses and Anthropic both reject. + dropped_server_builtin_tool_call_ids: set[str] = set() + + def _filter_tool_calls(tool_calls: Any) -> Optional[list]: + """Sanitize assistant `tool_calls` for non-native-Gemini providers. + + Two concerns: + 1. `tool_calls[i].extra_content` carries Gemini-only + thoughtSignature metadata; strip it for providers that + cannot parse the unknown key. + 2. Marked server-side builtin cards (`_server_tool: true` on + a canonical builtin name, or a Gemini `native_part` + payload) are provider-internal Studio tool cards from a + prior native Gemini turn; forwarding them to OpenAI / + Anthropic / custom OAI-compat gateways sends an orphan + `tool_calls` entry (no matching tool declaration, often + no matching `role="tool"` reply) that can be rejected. + We record the dropped call_ids so the matching role=tool + message is also skipped below. + Native Gemini keeps both untouched so the native translator can + replay them via `native_part`. + """ + if not tool_calls: + return None + if not isinstance(tool_calls, list): + return tool_calls + if emit_extra_content: + return tool_calls + cleaned: list = [] + for _tc in tool_calls: + if _is_marked_server_builtin_tool_call(_tc): + _tc_id = _tc.get("id") if isinstance(_tc, dict) else None + if isinstance(_tc_id, str) and _tc_id: + dropped_server_builtin_tool_call_ids.add(_tc_id) + continue + if not isinstance(_tc, dict): + cleaned.append(_tc) + continue + if "extra_content" not in _tc: + cleaned.append(_tc) + continue + _stripped = {k: v for k, v in _tc.items() if k != "extra_content"} + cleaned.append(_stripped) + return cleaned + result = [] for msg in messages: + # Drop role=tool messages whose matching server-builtin + # tool_call was already filtered above. Forwarding an orphan + # tool_result with no matching tool_call would be rejected by + # OpenAI Responses and Anthropic. + if ( + msg.role == "tool" + and isinstance(msg.tool_call_id, str) + and msg.tool_call_id in dropped_server_builtin_tool_call_ids + ): + continue if isinstance(msg.content, str): - # Skip assistant messages with empty content (some providers reject them) - if msg.role == "assistant" and not msg.content.strip(): + # Drop bare assistant messages with no content AND no + # tool_calls (some providers reject empty assistant turns). + # Preserve assistant turns whose only payload is tool_calls + # so multi-turn function-call loops round-trip. + if ( + msg.role == "assistant" + and not msg.content.strip() + and not msg.tool_calls + ): continue - result.append({"role": msg.role, "content": msg.content}) - elif isinstance(msg.content, list): + out: dict[str, Any] = {"role": msg.role, "content": msg.content} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + out["tool_calls"] = _tcs + elif not msg.content.strip(): + # Every tool_call was a synthetic provider-side + # card and was dropped; the assistant turn would + # be an empty `{"role":"assistant","content":""}` + # which some providers reject. Skip it entirely. + continue + if msg.role == "tool": + if msg.tool_call_id: + out["tool_call_id"] = msg.tool_call_id + if msg.name: + out["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + out["extra_content"] = msg.extra_content + result.append(out) + continue + # Assistant messages with content=None but populated tool_calls + # are valid (post-tool-call assistant turn). Forward them so the + # provider helper can rebuild the functionCall part. + if msg.content is None and msg.role == "assistant" and msg.tool_calls: + _filtered_tcs = _filter_tool_calls(msg.tool_calls) + if not _filtered_tcs: + # Every tool_call on this turn was provider-side + # synthetic and dropped; skipping the whole message + # avoids forwarding an empty assistant turn. + continue + _assistant_only: dict[str, Any] = { + "role": "assistant", + "content": "", + "tool_calls": _filtered_tcs, + } + if emit_extra_content and msg.extra_content: + _assistant_only["extra_content"] = msg.extra_content + result.append(_assistant_only) + continue + if isinstance(msg.content, list): if supports_vision: parts = [] for part in msg.content: @@ -1797,9 +1955,27 @@ def _build_external_messages( # provider would 400 on the unknown part, so # gate by provider_type. parts.append({"type": "compaction", "content": part.content}) - if msg.role == "assistant" and not parts: + entry: dict[str, Any] = {"role": msg.role, "content": parts} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + elif not parts: + # All tool_calls were synthetic and dropped, + # and no preserved content parts survived. + # Skip rather than forward an empty assistant + # turn that downstream providers reject. + continue + elif msg.role == "assistant" and not parts: continue - result.append({"role": msg.role, "content": parts}) + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) else: # Non-vision provider: strip images / documents, keep # text, optionally keep compaction (Anthropic only -- @@ -1835,9 +2011,32 @@ def _build_external_messages( if len(preserved) == 1 and preserved[0]["type"] == "text": # Single text part collapses back to a string for # providers that don't accept content arrays. - result.append({"role": msg.role, "content": preserved[0]["text"]}) + entry = {"role": msg.role, "content": preserved[0]["text"]} else: - result.append({"role": msg.role, "content": preserved}) + entry = {"role": msg.role, "content": preserved} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + else: + # All tool_calls were synthetic and dropped; + # skip if there's no surviving content either. + _entry_content = entry.get("content") + _has_text = ( + isinstance(_entry_content, str) and _entry_content.strip() + ) or ( + isinstance(_entry_content, list) and len(_entry_content) > 0 + ) + if not _has_text: + continue + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) return result @@ -1915,6 +2114,7 @@ async def _proxy_to_external_provider( payload.messages, _supports_vision, provider_type = provider_type, + base_url = base_url, ) # Codex provider: dispatch through the local CLI / SDK instead of @@ -2011,6 +2211,14 @@ async def _proxy_to_external_provider( api_key = api_key, ) + # `top_k` defaults to 20 in ChatCompletionRequest because the local + # inference path expects an int, but the external-provider path + # should treat "field omitted from JSON" as "use provider default" + # so callers that send only model/messages do not silently get + # different sampling than before this PR. Pydantic's + # `model_fields_set` tracks explicit-vs-default per request. + _top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None + async def _stream(): gen = client.stream_chat_completion( messages = chat_messages, @@ -2019,7 +2227,7 @@ async def _proxy_to_external_provider( top_p = payload.top_p, max_tokens = payload.max_tokens, presence_penalty = payload.presence_penalty, - top_k = payload.top_k, + top_k = _top_k_explicit, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, enabled_tools = payload.enabled_tools, @@ -2028,6 +2236,8 @@ async def _proxy_to_external_provider( anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id, prompt_cache_ttl = payload.prompt_cache_ttl, compaction_threshold = payload.compaction_threshold, + tools = payload.tools, + tool_choice = payload.tool_choice, fast_mode = payload.fast_mode, stream = payload.stream, ) @@ -4571,7 +4781,17 @@ async def anthropic_messages( [m.model_dump() for m in payload.messages], payload.system, ) - openai_messages = _drop_empty_assistant_sentinels(openai_messages) + # Strip synthetic provider-side builtin tool history (web_search, + # web_fetch, code_execution, image_generation cards tagged with + # _server_tool or extra_content.google.native_part) before handing + # off to local llama-server. The local /v1/chat/completions and + # GGUF passthrough builders apply the same strip; without it an + # Anthropic /v1/messages caller replaying a prior provider-side + # tool_use forwards fake builtin tool history to a backend that + # has no matching function declarations. + openai_messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) + ) # Enforce vision guard + re-encode embedded images to PNG so the # Anthropic endpoint matches the behavior of /v1/chat/completions. @@ -5362,6 +5582,110 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]: return out +_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} +) + + +def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]: + """Drop synthetic provider-side tool_calls + matching role=tool replies + on the local-backend (llama-server / GGUF) dispatch path. + + A Gemini chat that ran code_execution / image_generation persists the + server-side tool card into thread history as an assistant tool_calls + entry tagged with ``args._server_tool`` (or a Gemini + ``args.google.native_part`` payload) plus a follow-up role=tool reply. + When the user switches the SAME thread to a local GGUF model, those + synthetic tool_calls are not real user functions, llama-server has no + matching declaration, and Gemini-only ``extra_content`` / + ``native_part`` payloads are meaningless. Forward only ordinary user + function calls; strip the matched role=tool replies too so the + backend does not see an orphan tool_call_id. + """ + dropped_ids: set[str] = set() + sanitized_assistant: list[dict] = [] + for m in messages: + if m.get("role") != "assistant": + sanitized_assistant.append(m) + continue + tool_calls = m.get("tool_calls") + if not isinstance(tool_calls, list) or not tool_calls: + # Plain text Gemini reply: still strip message-level + # `extra_content` (carries `google.thought_signature` replay + # metadata) so a text-only Gemini turn switched to a local + # GGUF backend does not leak Gemini-only fields to + # llama-server. ChatMessage previously did not have + # `extra_content`, so the field was implicitly dropped -- + # round-22 added it to ChatMessage, which is what made this + # leak possible. + if "extra_content" in m: + m = {k: v for k, v in m.items() if k != "extra_content"} + sanitized_assistant.append(m) + continue + cleaned: list[dict] = [] + for tc in tool_calls: + if not isinstance(tc, dict): + cleaned.append(tc) + continue + fn = tc.get("function") + name = "" + if isinstance(fn, dict): + name = (fn.get("name") or "").lower() + if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES: + raw_args = fn.get("arguments") if isinstance(fn, dict) else None + args_obj: Any = None + if isinstance(raw_args, str): + try: + args_obj = json.loads(raw_args) if raw_args else None + except Exception: + args_obj = None + elif isinstance(raw_args, dict): + args_obj = raw_args + is_synthetic = False + if isinstance(args_obj, dict): + if args_obj.get("_server_tool") is True: + is_synthetic = True + google = args_obj.get("google") + if isinstance(google, dict) and isinstance( + google.get("native_part"), dict + ): + is_synthetic = True + if is_synthetic: + tc_id = tc.get("id") + if isinstance(tc_id, str) and tc_id: + dropped_ids.add(tc_id) + continue + # Strip Gemini-only `extra_content` on real user tool_calls + # too — llama-server has no use for it and may pass it + # through to the model unchanged. + if "extra_content" in tc: + tc = {k: v for k, v in tc.items() if k != "extra_content"} + cleaned.append(tc) + # Drop top-level message-level `extra_content` (Gemini + # thoughtSignature replay metadata) on local dispatch. + m_clean = {k: v for k, v in m.items() if k != "extra_content"} + if cleaned: + m_clean["tool_calls"] = cleaned + else: + m_clean.pop("tool_calls", None) + if not m_clean.get("content") and not m_clean.get("tool_calls"): + continue # assistant turn now empty, drop + sanitized_assistant.append(m_clean) + + if not dropped_ids: + return sanitized_assistant + out: list[dict] = [] + for m in sanitized_assistant: + if ( + m.get("role") == "tool" + and isinstance(m.get("tool_call_id"), str) + and m["tool_call_id"] in dropped_ids + ): + continue + out.append(m) + return out + + def _openai_messages_for_passthrough(payload) -> list[dict]: """Build OpenAI-format message dicts for the /v1/chat/completions passthrough path. @@ -5378,8 +5702,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: ``image_url`` content part so vision + function-calling requests work transparently. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) if not payload.image_base64: @@ -5428,8 +5754,10 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict] all per-turn ``image_url`` parts so multi-image chat history keeps each image attached to its original turn. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) has_message_image = any( isinstance(msg.get("content"), list) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 2bb1de5366..5d4bd46e62 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -318,22 +318,45 @@ async def list_provider_models( try: models = await client.list_models() - allow_prefixes = info.get("model_id_allow_prefixes") - if allow_prefixes is not None: - prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) - if prefix_tuple: - models = [m for m in models if m.get("id", "").startswith(prefix_tuple)] - allowlist = info.get("model_id_allowlist") - if allowlist is not None: - models = [m for m in models if allowlist.match(m.get("id", ""))] - deny_exact = info.get("model_id_deny_exact") - if deny_exact is not None: - deny_ids = {str(m) for m in deny_exact if str(m)} - if deny_ids: - models = [m for m in models if m.get("id", "") not in deny_ids] - denylist = info.get("model_id_denylist") - if denylist is not None: - models = [m for m in models if not denylist.search(m.get("id", ""))] + # Registry-level model-id filters are scoped to the canonical + # native Gemini base. A custom Gemini OAI-compatible proxy + # (LiteLLM, deployment gateway) returns IDs like + # `google/gemini-2.5-flash`, `gemini/gemini-2.5-flash`, or + # team-prefixed deployment aliases; the native allowlist regex + # would strip those out and leave the picker empty even though + # the chat path now routes them via the OAI-compatible + # dispatcher (the same gate ExternalProviderClient applies for + # request building). Match the host check here so the model + # list and chat dispatch agree on what counts as "native". + apply_registry_model_filters = True + if payload.provider_type == "gemini": + try: + from urllib.parse import urlparse as _urlparse + + _host = (_urlparse(base_url).hostname or "").lower() + except Exception: + _host = "" + apply_registry_model_filters = _host == "generativelanguage.googleapis.com" + + if apply_registry_model_filters: + allow_prefixes = info.get("model_id_allow_prefixes") + if allow_prefixes is not None: + prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) + if prefix_tuple: + models = [ + m for m in models if m.get("id", "").startswith(prefix_tuple) + ] + allowlist = info.get("model_id_allowlist") + if allowlist is not None: + models = [m for m in models if allowlist.match(m.get("id", ""))] + deny_exact = info.get("model_id_deny_exact") + if deny_exact is not None: + deny_ids = {str(m) for m in deny_exact if str(m)} + if deny_ids: + models = [m for m in models if m.get("id", "") not in deny_ids] + denylist = info.get("model_id_denylist") + if denylist is not None: + models = [m for m in models if not denylist.search(m.get("id", ""))] # Apply an optional cap after filtering so registry entries with a # large remote catalog (e.g. HF Inference Providers) can stay # picker-sized. No popularity sort happens server-side, so this is diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6e2413b3e9..41a9e15562 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -194,6 +194,7 @@ async def start_training( "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, "max_seq_length": request.max_seq_length, + "vision_image_size": request.vision_image_size, "hf_dataset": request.hf_dataset or "", "local_datasets": request.local_datasets, "local_eval_datasets": request.local_eval_datasets, diff --git a/studio/backend/run.py b/studio/backend/run.py index e96e609659..5ddf404370 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -19,6 +19,16 @@ backend_dir = Path(__file__).parent if str(backend_dir) not in sys.path: sys.path.insert(0, str(backend_dir)) +from utils.cpu_threads import configure_cpu_threads + +try: + configure_cpu_threads() +except ValueError as exc: + configured = os.environ.get("UNSLOTH_CPU_THREADS") + raise SystemExit( + f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}" + ) from None + # Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before # any library imports that trigger attrs -> rich -> structlog -> platform crash. # See: https://github.com/python/cpython/issues/102396 diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py index 7f6fe58329..5c88437d17 100644 --- a/studio/backend/tests/test_anthropic_code_execution.py +++ b/studio/backend/tests/test_anthropic_code_execution.py @@ -275,7 +275,13 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): assert start["type"] == "tool_start" assert start["tool_name"] == "code_execution" assert start["tool_call_id"] == "srvtoolu_1" - assert start["arguments"] == {"kind": "bash", "command": "ls -la"} + # `_server_tool: True` marks this as a provider-side synthetic + # tool card for the frontend's history serializer. + assert start["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_1" diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index da10d679eb..88a922e7eb 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -271,7 +271,12 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): assert start["type"] == "tool_start" assert start["tool_name"] == "web_fetch" assert start["tool_call_id"] == "srvtoolu_wf1" - assert start["arguments"] == {"url": "https://example.com/article"} + # `_server_tool: True` marks this as a provider-side synthetic + # tool card for the frontend's history serializer. + assert start["arguments"] == { + "url": "https://example.com/article", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_wf1" # The source pill uses Title / URL / snippet as parseSourcesFromResult expects. diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py new file mode 100644 index 0000000000..0dbcbdb74b --- /dev/null +++ b/studio/backend/tests/test_cpu_threads.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for Studio's early CPU thread-pool configuration.""" + +import ast +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from utils.cpu_threads import _THREAD_POOL_ENV_VARS, configure_cpu_threads + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent +_RUN_PY = _BACKEND_DIR / "run.py" +_MAIN_PY = _BACKEND_DIR / "main.py" + + +# Explicit positive integers seed all four native pool env vars. +def test_cpu_thread_cap_seeds_native_pool_limits(): + env = {"UNSLOTH_CPU_THREADS": " 6 "} + + configure_cpu_threads(env) + + assert {variable: env[variable] for variable in _THREAD_POOL_ENV_VARS} == { + variable: "6" for variable in _THREAD_POOL_ENV_VARS + } + + +# Explicit per-library values win over the Studio knob via setdefault. +def test_cpu_thread_cap_preserves_runtime_specific_override(): + env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"} + + configure_cpu_threads(env) + + assert env["OMP_NUM_THREADS"] == "2" + assert env["MKL_NUM_THREADS"] == "4" + + +# Whitespace / plus-prefix / leading zero all normalise via int(). +@pytest.mark.parametrize("raw", ["+4", "007", " 4 "]) +def test_cpu_thread_cap_normalises_valid_inputs(raw): + env = {"UNSLOTH_CPU_THREADS": raw} + + configure_cpu_threads(env) + + assert env["OMP_NUM_THREADS"] == str(int(raw.strip())) + + +# Unset / empty / whitespace -> no env mutation (pure opt-in). +@pytest.mark.parametrize("raw", [None, "", " ", "\t"]) +def test_cpu_thread_cap_is_opt_in(raw): + env = {} if raw is None else {"UNSLOTH_CPU_THREADS": raw} + snapshot = dict(env) + + configure_cpu_threads(env) + + assert env == snapshot + assert all(variable not in env for variable in _THREAD_POOL_ENV_VARS) + + +# Anything that is not a positive integer raises a clear ValueError. +@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]) +def test_cpu_thread_cap_requires_positive_integer(raw): + with pytest.raises(ValueError, match="must be a positive integer"): + configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw}) + + +# env=None path uses real os.environ (production call from run.py / main.py). +def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch): + for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3") + + configure_cpu_threads() + + for variable in _THREAD_POOL_ENV_VARS: + assert os.environ[variable] == "3" + + +# Calling twice must not flip any seeded value. +def test_cpu_thread_cap_idempotent(monkeypatch): + for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5") + + configure_cpu_threads() + snapshot = {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} + configure_cpu_threads() + + assert {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} == snapshot + + +def _ast_line_of_configure_call(source: str) -> int: + tree = ast.parse(source) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "configure_cpu_threads" + ): + return node.lineno + raise AssertionError("configure_cpu_threads() call not found") + + +def _ast_line_of_platform_compat_import(source: str) -> int: + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "_platform_compat": + return node.lineno + raise AssertionError("_platform_compat import not found") + + +# AST-based ordering: configure_cpu_threads() must precede _platform_compat +# in both run.py and main.py. Robust to formatting / line shifts. +@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY]) +def test_cpu_thread_configuration_runs_before_backend_imports(entry_point): + source = entry_point.read_text() + call_line = _ast_line_of_configure_call(source) + compat_line = _ast_line_of_platform_compat_import(source) + assert call_line < compat_line, ( + f"{entry_point.name}: configure_cpu_threads() (line {call_line}) " + f"must precede import _platform_compat (line {compat_line})" + ) + + +# Invalid env -> exit 1, one-line stderr, no traceback, gated before any +# heavy import. Parametrised over both entry points. +@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY]) +def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point): + env = os.environ.copy() + env["UNSLOTH_CPU_THREADS"] = "not-a-count" + + result = subprocess.run( + [sys.executable, str(entry_point)], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 1 + assert ( + "Error: Invalid UNSLOTH_CPU_THREADS value 'not-a-count': " + "UNSLOTH_CPU_THREADS must be a positive integer" + ) in result.stderr + assert "Traceback" not in result.stderr + assert "_platform_compat" not in result.stderr diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index ee1e7fb18f..34d14b08fe 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -485,11 +485,14 @@ from typer.testing import CliRunner studio_home = Path(sys.argv[1]) real_import = builtins.__import__ -def guarded_import(name, *args, **kwargs): +def guarded_import(name, globals = None, locals = None, fromlist = (), level = 0): + # Only gate absolute imports; relative `from .utils import x` inside + # third-party packages (e.g. typer._click.decorators) hits level > 0 + # with name="utils" and must pass through. blocked = ("auth", "fastapi", "structlog", "utils") - if name in blocked or name.startswith(("auth.", "utils.")): + if level == 0 and (name in blocked or name.startswith(("auth.", "utils."))): raise ModuleNotFoundError(name) - return real_import(name, *args, **kwargs) + return real_import(name, globals, locals, fromlist, level) builtins.__import__ = guarded_import from unsloth_cli.commands import studio as studio_cli diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py new file mode 100644 index 0000000000..4dc97302e2 --- /dev/null +++ b/studio/backend/tests/test_gemini_provider.py @@ -0,0 +1,5501 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the native Gemini API translation layer. + +Gemini does NOT speak OpenAI Chat Completions on its primary endpoint +(`streamGenerateContent`). `_stream_gemini` in +`core/inference/external_provider.py` translates between the two shapes: + + Request: + OpenAI messages [{role, content}] + -> Gemini contents [{role, parts: [{text}|{inlineData}|{functionCall}|...]}] + + systemInstruction.parts[].text for role=system messages + + generationConfig.{temperature,topP,topK,maxOutputTokens} + + tools[{googleSearch:{}}] for web_search + + tools[{codeExecution:{}}] for code_execution + + responseModalities=[TEXT,IMAGE] for Nano Banana (gemini-2.5-flash-image) + + cachedContent for prompt caching + + Response: + Gemini SSE chunks { candidates:[{content:{parts:[...]}, finishReason}], + usageMetadata:{promptTokenCount, candidatesTokenCount} } + -> OpenAI chat.completion.chunk frames + (delta.content for text, delta.tool_calls for functionCall, + _toolEvent for image_b64/web_search, usage block before [DONE]) + +These tests pin the outbound body shape AND the inbound translation +using httpx.MockTransport (no live network). Mirrors the structure of +test_anthropic_cache_ttl.py and test_openai_image_generation.py. +""" + +import asyncio +import base64 +import json + +import httpx +import pytest + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +_active_mock_clients: list[httpx.AsyncClient] = [] + + +def _drive(coro): + # Create a fresh loop per drive so tests don't share asyncio state. + # Close mocked clients + shutdown async-generators inside this loop + # so Python 3.13 doesn't emit the + # `Response.aiter_*.aclose was never awaited` warning on GC. + loop = asyncio.new_event_loop() + try: + result = loop.run_until_complete(coro) + while _active_mock_clients: + mc = _active_mock_clients.pop() + loop.run_until_complete(mc.aclose()) + return result + finally: + try: + loop.run_until_complete(loop.shutdown_asyncgens()) + finally: + loop.close() + + +def _make_gemini_client( + base_url: str = "https://generativelanguage.googleapis.com/v1beta", +) -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "gemini", + base_url = base_url, + api_key = "AIza-test-key", + ) + + +def _mock_http(monkeypatch, handler): + mock_client = httpx.AsyncClient(transport = httpx.MockTransport(handler)) + monkeypatch.setattr(ep_mod, "_http_client", mock_client) + # `_drive` will aclose this at the end of the run inside the same + # event loop so we do not leak an unawaited aclose() coroutine. + _active_mock_clients.append(mock_client) + + +def _gemini_sse(events: list[dict]) -> bytes: + """Encode a list of dicts as Gemini-style SSE frames (`data:` lines).""" + chunks: list[str] = [] + for event in events: + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _capture_body(monkeypatch, **kwargs) -> dict: + """Drive a single stream and return the captured outbound request body.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + captured["url"] = str(request.url) + captured["method"] = request.method + # Minimal valid Gemini stream so the helper can complete. + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + pass + await client.close() + + _drive(run()) + return captured + + +def _collect(monkeypatch, sse_events, **kwargs) -> list[str]: + """Drive a stream with a custom set of SSE events and return raw lines.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + out.append(line) + await client.close() + + _drive(run()) + return out + + +def _parse_chunks(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if not payload or payload == "[DONE]": + continue + try: + out.append(json.loads(payload)) + except json.JSONDecodeError: + continue + return out + + +# ── request body translation ───────────────────────────────────────── + + +def test_request_body_uses_contents_and_parts_shape(monkeypatch): + """OpenAI messages must be translated to Gemini's `contents` shape.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + {"role": "user", "content": "Follow up"}, + ], + ) + body = captured["body"] + # system -> systemInstruction + assert body["systemInstruction"] == {"parts": [{"text": "Be brief."}]}, body + # user/assistant -> contents with role user/model + assert body["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]}, + {"role": "model", "parts": [{"text": "Hi there"}]}, + {"role": "user", "parts": [{"text": "Follow up"}]}, + ], body["contents"] + # generationConfig fields map across with Google's casing. + gc = body["generationConfig"] + assert gc["temperature"] == 0.7 + assert gc["topP"] == 0.95 + assert gc["maxOutputTokens"] == 64 + + +def test_request_url_targets_stream_generate_content(monkeypatch): + """Helper must POST to /v1beta/models/{model}:streamGenerateContent?alt=sse.""" + captured = _capture_body(monkeypatch, model = "gemini-2.5-pro") + url = captured["url"] + assert ":streamGenerateContent" in url, url + assert "alt=sse" in url, url + assert "/v1beta/models/gemini-2.5-pro" in url, url + assert captured["method"] == "POST" + + +def test_request_auth_header_uses_x_goog_api_key(monkeypatch): + """API key must be sent on `x-goog-api-key`, not Authorization.""" + captured = _capture_body(monkeypatch) + hdrs = captured["headers"] + assert hdrs.get("x-goog-api-key") == "AIza-test-key", hdrs + assert "authorization" not in {k.lower() for k in hdrs}, hdrs + + +def test_top_k_forwarded_only_when_positive(monkeypatch): + """top_k is opt-in; only positive integers reach the wire.""" + captured = _capture_body(monkeypatch, top_k = 40) + assert captured["body"]["generationConfig"]["topK"] == 40 + + captured = _capture_body(monkeypatch, top_k = 0) + assert "topK" not in captured["body"]["generationConfig"] + + +def test_presence_penalty_forwarded_to_generation_config(monkeypatch): + """A non-zero presence_penalty reaches generationConfig.presencePenalty.""" + captured = _capture_body(monkeypatch, presence_penalty = 0.7) + assert captured["body"]["generationConfig"]["presencePenalty"] == 0.7 + + # And the default of zero is omitted, matching top_k semantics. + captured = _capture_body(monkeypatch, presence_penalty = 0.0) + assert "presencePenalty" not in captured["body"]["generationConfig"] + + +# ── thinkingConfig translation ──────────────────────────────────────── + + +def test_gemini25_flash_thinking_disabled_sets_budget_zero(monkeypatch): + """Gemini 2.5 Flash still uses thinkingBudget; 0 = off.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": 0}, tc + + +def test_gemini3_flash_thinking_disabled_uses_minimal_level(monkeypatch): + """Gemini 3 Flash migrated to thinkingLevel; "off" maps to minimal + (Gemini 3 cannot turn thinking fully off).""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_gemini25_pro_thinking_disabled_uses_small_budget(monkeypatch): + """Gemini 2.5 Pro 400s on thinkingBudget=0 ("only works in thinking + mode"); coerce to a small positive budget.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc is not None and tc.get("thinkingBudget", 0) > 0, tc + + +def test_gemini3_pro_thinking_disabled_uses_low_level(monkeypatch): + """Gemini 3 Pro uses thinkingLevel and rejects 'minimal' (Pro tier), + so 'off' coerces to 'low' (lowest the API accepts).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-3-pro-preview", + "gemini-3.5-pro", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, (model, tc) + + +def test_gemini25_flash_effort_levels_map_to_budgets(monkeypatch): + """Gemini 2.5 Flash retains the integer thinkingBudget ladder.""" + cases = { + "minimal": 512, + "low": 2048, + "medium": 8192, + "high": 24576, + "max": -1, + "xhigh": -1, + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": expected}, (effort, tc) + + +def test_gemini3_flash_effort_levels_map_to_thinking_level(monkeypatch): + """Gemini 3 Flash thinkingLevel ladder: minimal/low/medium/high.""" + cases = { + "minimal": "minimal", + "low": "low", + "medium": "medium", + "high": "high", + "max": "high", + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": expected}, (effort, tc) + + +def test_gemini3_pro_passes_medium_through(monkeypatch): + """Gemini 3.1+ Pro accepts thinkingLevel="medium" per + https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro; + forward as-is (medium is the documented mid-tier on Gemini 3.1).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "medium", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "medium"}, (model, tc) + + +def test_gemini3_pro_minimal_effort_coerces_to_low(monkeypatch): + """Gemini 3 Pro rejects thinkingLevel="minimal"; coerce to "low".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "minimal", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_flash_effort_none_maps_to_minimal(monkeypatch): + """reasoning_effort='none' on Gemini 3 Flash -> thinkingLevel=minimal.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = "none", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_thinking_default_omits_thinking_config(monkeypatch): + """When neither knob is supplied, thinkingConfig is omitted entirely + (Google's server-side default applies).""" + captured = _capture_body(monkeypatch, model = "gemini-3.5-flash") + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, gc + + +def test_nano_banana_alias_routes_through_image_modalities(monkeypatch): + """`nano-banana-pro-preview` is an alias for the Pro image model and + must set responseModalities=[TEXT,IMAGE] when the Images pill is on + (enabled_tools includes "image_generation").""" + captured = _capture_body( + monkeypatch, + model = "nano-banana-pro-preview", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT", "IMAGE"], gc + + +def test_image_capable_model_without_image_pill_stays_text_only(monkeypatch): + """When the Images pill is off (enabled_tools has no + image_generation), an image-capable model id (gemini-2.5-flash-image) + must force responseModalities=["TEXT"]. Google's image models + default to text+image when responseModalities is omitted, so + omitting it would silently bill image output the UI says is + disabled.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT"], gc + + +def test_image_models_skip_thinking_config(monkeypatch): + """Image-tier ids do not benefit from a visible thinking knob and + must NOT forward thinkingConfig even when stale UI state still + sends `reasoning_effort` or `enable_thinking=False`.""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "high", + enable_thinking = False, + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, (model, gc) + + +def test_image_models_drop_code_execution(monkeypatch): + """All image-tier ids reject `tools: [{codeExecution: {}}]`; drop + silently. (Gemini 3 image models DO accept googleSearch -- see + test_gemini3_image_models_allow_google_search; older image models + drop everything.)""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "codeExecution" not in names, (model, tools_arr) + + +def test_gemini_35_pro_uses_thinking_level(monkeypatch): + """`gemini-3.5-pro` is part of the Gemini 3 family and uses + thinkingLevel (not thinkingBudget). "Off" maps to "low" because Pro + tier rejects "minimal".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_image_models_allow_google_search(monkeypatch): + """Google documents Search grounding on the Gemini 3 image family + (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + nano-banana-pro). codeExecution stays blocked on image mode.""" + for model in ( + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "googleSearch" in names, (model, tools_arr) + assert "codeExecution" not in names, (model, tools_arr) + + +def test_legacy_image_models_block_google_search(monkeypatch): + """Older Gemini image ids (gemini-2.5-flash-image) still 400 on + `tools: [{googleSearch: {}}]`; backend keeps stripping it.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + assert "tools" not in captured["body"], captured["body"].get("tools") + + +def test_legacy_openai_base_url_normalized(monkeypatch): + """Saved Gemini providers carrying the legacy `/v1beta/openai` base + (from the pre-PR OpenAI-compat plumbing) now point at the native + endpoint without the user re-saving the connection.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_finish_reason_swaps_to_tool_calls_when_function_call_emitted(monkeypatch): + """Gemini emits finishReason="STOP" even for pure functionCall turns; + surface as `tool_calls` so OAI clients trigger tool execution.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "lookup", "args": {"k": "v"}}} + ], + }, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if c.get("choices", [{}])[0].get("finish_reason") is not None + ] + assert finish_chunks, chunks + assert finish_chunks[-1]["choices"][0]["finish_reason"] == "tool_calls", chunks + + +def test_thought_signature_round_trips_into_gemini_function_call(monkeypatch): + """An assistant tool_call carrying `extra_content.google.thought_signature` + must echo the value back as a sibling of the Gemini functionCall part.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "lookup x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": {"google": {"thought_signature": "SIG-ABC"}}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_0", + "name": "lookup", + "content": "{}", + }, + ], + ) + contents = captured["body"]["contents"] + fc_turn = next((c for c in contents if c["role"] == "model"), None) + assert fc_turn is not None, contents + fc_part = next( + (p for p in fc_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, fc_turn + assert fc_part.get("thoughtSignature") == "SIG-ABC", fc_part + + +def test_thought_signature_emitted_in_tool_call_delta(monkeypatch): + """A Gemini functionCall part with `thoughtSignature` must surface + that signature on the outbound OpenAI tool_calls delta via + `extra_content.google.thought_signature`.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "lookup", + "args": {"k": "v"}, + "id": "call_xyz", + }, + "thoughtSignature": "SIG-FROM-GEMINI", + } + ], + }, + "finishReason": "STOP", + } + ] + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + deltas = [ + tc + for c in chunks + for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + ] + assert deltas, chunks + sig = deltas[0].get("extra_content", {}).get("google", {}).get("thought_signature") + assert sig == "SIG-FROM-GEMINI", deltas + + +def test_image_models_suppress_phantom_web_search_card(monkeypatch): + """When the image guard filters googleSearch out of the outbound + request, the inbound stream must NOT emit web_search tool_start / + tool_end (otherwise the UI shows a misleading 'Search complete' + card on a turn where Gemini never actually searched).""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "drawn"}]}, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + chunks = _parse_chunks(lines) + tool_evs = [ + ev + for c in chunks + for ev in [c.get("_toolEvent")] + if isinstance(ev, dict) and ev.get("tool_name") == "web_search" + ] + assert tool_evs == [], tool_evs + + +def test_image_generation_tool_on_image_model_drops_text_tools(monkeypatch): + """`enabled_tools=["image_generation", "web_search", "code_execution"]` + on a Gemini IMAGE model flips responseModalities to TEXT+IMAGE; in + that mode codeExecution must NOT be forwarded (Gemini rejects text + code tools alongside image responseModalities). Older image + families also drop googleSearch.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [ + "image_generation", + "web_search", + "code_execution", + ], + ) + assert "tools" not in captured["body"], captured["body"] + assert captured["body"]["generationConfig"].get("responseModalities") == [ + "TEXT", + "IMAGE", + ] + + +def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch): + """`promptFeedback.blockReason` with zero candidates must produce + an error chunk, not a silent empty assistant reply.""" + sse = [ + { + "promptFeedback": {"blockReason": "SAFETY"}, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + error_chunks = [c for c in chunks if "error" in c] + assert error_chunks, chunks + assert "SAFETY" in ( + error_chunks[0].get("error", {}).get("message") or "" + ), error_chunks + + +def test_usage_chunk_includes_thoughts_tokens(monkeypatch): + """`thoughtsTokenCount` is the hidden-reasoning slice of output; + roll it into `output_tokens` AND surface it on + `output_tokens_details.reasoning_tokens` so total_tokens reflects + the full billable spend.""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + "totalTokenCount": 35, + }, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + usage_chunk = next((c for c in chunks if isinstance(c.get("usage"), dict)), None) + assert usage_chunk is not None, chunks + usage = usage_chunk["usage"] + assert usage.get("prompt_tokens") == 10, usage + # candidates 5 + thoughts 20 = 25 output tokens; total = 35. + assert usage.get("completion_tokens") == 25, usage + assert usage.get("total_tokens") == 35, usage + + +# ── web_search forwarded as googleSearch tool ──────────────────────── + + +def test_web_search_forwarded_as_google_search_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search"], + ) + tools = captured["body"].get("tools") or [] + assert {"googleSearch": {}} in tools, tools + + +def test_code_execution_forwarded_as_code_execution_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["code_execution"], + ) + tools = captured["body"].get("tools") or [] + assert {"codeExecution": {}} in tools, tools + + +def test_omitted_tools_leaves_body_untouched(monkeypatch): + captured = _capture_body(monkeypatch, enabled_tools = []) + assert "tools" not in captured["body"], captured["body"] + + +# ── prompt caching passthrough ─────────────────────────────────────── + + +def test_cached_content_pass_through(monkeypatch): + """A string cache id on enable_prompt_caching is forwarded verbatim.""" + cache_name = "cachedContents/abc123" + captured = _capture_body( + monkeypatch, + enable_prompt_caching = cache_name, + ) + assert captured["body"].get("cachedContent") == cache_name + + +def test_boolean_caching_does_not_set_cached_content(monkeypatch): + """Studio's existing True/False signals shouldn't fabricate a cache id.""" + captured = _capture_body(monkeypatch, enable_prompt_caching = True) + assert "cachedContent" not in captured["body"] + + +# ── image generation: request modalities + response translation ────── + + +def test_image_model_sets_response_modalities(monkeypatch): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypatch): + """`enabled_tools=["image_generation"]` flips responseModalities + only when the selected model is image-capable; otherwise the + request stays plain text (text-only models 400 on + responseModalities).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_image_response_emits_image_b64_tool_event(monkeypatch): + """`inlineData` parts become a tool_end with image_b64 + image_mime.""" + fake_b64 = base64.b64encode(b"PNG-BYTES").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": fake_b64, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 0, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + image_starts = [e for e in starts if e.get("tool_name") == "image_generation"] + image_ends = [e for e in ends if e.get("image_b64")] + assert len(image_starts) == 1, tool_events + assert len(image_ends) == 1, tool_events + assert image_ends[0]["image_b64"] == fake_b64 + assert image_ends[0]["image_mime"] == "image/png" + + +# ── function calling round-trips both directions ───────────────────── + + +def test_function_call_response_translates_to_tool_calls_delta(monkeypatch): + """Gemini `functionCall` parts become OpenAI `tool_calls` delta chunks.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "get_weather", + "args": {"location": "Paris"}, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 12, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 1, chunks + tc = tool_call_chunks[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["function"]["name"] == "get_weather" + args = json.loads(tc["function"]["arguments"]) + assert args == {"location": "Paris"} + + +def test_tool_message_translates_to_function_response_part(monkeypatch): + """role=tool follow-ups are rewritten to functionResponse parts.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "name": "get_weather", + "content": json.dumps({"temp_c": 18, "summary": "Sunny"}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + # Last turn must be a functionResponse part (Gemini wraps it as a + # role=user turn carrying the result). + last = contents[-1] + assert last["role"] == "user", last + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert fr["name"] == "get_weather" + assert fr["response"] == {"temp_c": 18, "summary": "Sunny"} + # And the assistant turn carries the original functionCall so the + # model sees the round-trip context. + assistant_turn = [c for c in contents if c["role"] == "model"][0] + fc_part = next( + (p for p in assistant_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, assistant_turn + assert fc_part["functionCall"]["name"] == "get_weather" + assert fc_part["functionCall"]["args"] == {"location": "Paris"} + + +def test_parallel_function_calls_get_distinct_tool_call_indices(monkeypatch): + """Each emitted functionCall in one assistant turn needs its own + tool_calls[*].index. Hardcoding index=0 collapses parallel calls + onto a single slot in OpenAI-style reassemblers.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "id": "call_alpha", + "name": "search", + "args": {"q": "alpha"}, + } + }, + { + "functionCall": { + "id": "call_beta", + "name": "search", + "args": {"q": "beta"}, + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 2, tool_call_chunks + indices = [ + c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks + ] + assert indices == [0, 1], indices + + +def test_function_call_ids_forwarded_into_gemini_function_call_part(monkeypatch): + """OpenAI tool_call id rides functionCall.id so parallel calls disambiguate.""" + messages = [ + {"role": "user", "content": "x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_alpha", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "a"}), + }, + }, + { + "id": "call_beta", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "b"}), + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_alpha", + "content": json.dumps({"hits": ["A"]}), + }, + { + "role": "tool", + "tool_call_id": "call_beta", + "content": json.dumps({"hits": ["B"]}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + assistant_parts = next(c for c in contents if c["role"] == "model")["parts"] + call_ids = [p["functionCall"]["id"] for p in assistant_parts if "functionCall" in p] + assert call_ids == ["call_alpha", "call_beta"], assistant_parts + response_ids = [ + p["functionResponse"]["id"] + for c in contents + for p in c["parts"] + if "functionResponse" in p + ] + assert response_ids == ["call_alpha", "call_beta"], contents + + +def test_parse_gemini_models_translates_native_catalog(): + """Gemini's native /v1beta/models payload becomes OpenAI-shape entries.""" + payload = { + "models": [ + { + "name": "models/gemini-2.5-flash", + "baseModelId": "gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": [ + "generateContent", + "streamGenerateContent", + ], + }, + { + "name": "models/embedding-001", + "supportedGenerationMethods": ["embedContent"], + }, + { + "name": "models/gemini-2.5-pro", + }, + ] + } + out = ExternalProviderClient._parse_gemini_models(payload) + ids = [m["id"] for m in out] + assert "gemini-2.5-flash" in ids + assert "gemini-2.5-pro" in ids + assert "embedding-001" not in ids + flash = next(m for m in out if m["id"] == "gemini-2.5-flash") + assert flash["display_name"] == "Gemini 2.5 Flash" + assert flash["owned_by"] == "google" + + +def test_code_execution_parts_translate_to_code_execution_tool_events(monkeypatch): + """executableCode + codeExecutionResult parts emit code_execution events.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2+2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "code_execution" + ] + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and "4" in str(e.get("result", "")) + ] + assert len(code_starts) == 1, tool_events + assert code_starts[0]["arguments"]["code"] == "print(2+2)" + assert code_starts[0]["arguments"]["language"] == "python" + assert len(code_ends) == 1, tool_events + # tool_start and tool_end must share the same tool_call_id so the + # frontend pairs them onto a single CodeExecutionToolUI block. + assert code_starts[0]["tool_call_id"] == code_ends[0]["tool_call_id"] + + +def test_code_execution_failure_outcome_surfaces_in_result(monkeypatch): + """OUTCOME_FAILED is prefixed onto the result text so the UI shows it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "1/0", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_FAILED", + "output": "ZeroDivisionError", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + result_text = next( + (e["result"] for e in tool_events if e.get("type") == "tool_end"), + "", + ) + assert "OUTCOME_FAILED" in result_text + assert "ZeroDivisionError" in result_text + + +def test_tool_message_recovers_name_from_tool_call_id(monkeypatch): + """When name is omitted, recover it from the matching tool_call_id.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": json.dumps({"temp_c": 18}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + last = contents[-1] + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert ( + fr["name"] == "get_weather" + ), "name should fall back to the prior tool_call's function name" + + +# ── usage chunk surfaces promptTokenCount / candidatesTokenCount ───── + + +def test_usage_chunk_translates_gemini_token_counts(monkeypatch): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1234, + "candidatesTokenCount": 56, + "cachedContentTokenCount": 1000, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("choices") == [] and "usage" in c] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 1234 + assert usage["completion_tokens"] == 56 + assert usage["total_tokens"] == 1290 + assert usage["prompt_tokens_details"]["cached_tokens"] == 1000 + + +# ── multimodal: vision image -> inlineData ─────────────────────────── + + +def test_vision_data_url_translates_to_inline_data(monkeypatch): + fake = base64.b64encode(b"JPGBYTES").decode() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{fake}", + }, + }, + ], + } + ] + captured = _capture_body(monkeypatch, messages = messages) + parts = captured["body"]["contents"][0]["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert len(inline_parts) == 1, parts + assert inline_parts[0]["inlineData"] == { + "mimeType": "image/jpeg", + "data": fake, + } + + +# ── finish reason mapping ──────────────────────────────────────────── + + +@pytest.mark.parametrize( + "gemini_reason, openai_reason", + [ + ("STOP", "stop"), + ("MAX_TOKENS", "length"), + ("SAFETY", "content_filter"), + ("PROHIBITED_CONTENT", "content_filter"), + ], +) +def test_finish_reason_translation(monkeypatch, gemini_reason, openai_reason): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "x"}], + }, + "finishReason": gemini_reason, + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if any(ch.get("finish_reason") for ch in c.get("choices", [])) + ] + assert any( + ch["choices"][0]["finish_reason"] == openai_reason for ch in finish_chunks + ), finish_chunks + + +# ── grounding citations surface as web_search tool_end ─────────────── + + +def test_grounding_metadata_surfaces_as_tool_end_citations(monkeypatch): + """`groundingMetadata.groundingChunks[].web` -> tool_end result block.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Answer with sources."}], + }, + "groundingMetadata": { + "groundingChunks": [ + { + "web": { + "uri": "https://example.com/a", + "title": "Example A", + } + }, + { + "web": { + "uri": "https://example.com/b", + "title": "Example B", + } + }, + ] + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 3, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + web_search_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "gemini_web_search" + ] + assert len(web_search_ends) == 1, tool_events + result = web_search_ends[0]["result"] + assert "https://example.com/a" in result + assert "https://example.com/b" in result + assert "Example A" in result + assert "Example B" in result + + +# ── round 3 review follow-ups ───────────────────────────────────────── + + +def test_custom_gemini_proxy_base_url_not_rewritten(): + """Only the Google-hosted /v1beta/openai base is normalized; a + custom gateway whose path ends in /openai must be left alone.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://proxy.example.com/team/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://proxy.example.com/team/openai" + + +def test_custom_gemini_proxy_uses_openai_dispatch(): + """Any non-Google Gemini base (LiteLLM, custom OpenAI-compat + routers) must route through the OpenAI-compatible forwarder, not + the native translator. Auth uses Authorization: Bearer ..., not + x-goog-api-key.""" + for base in ( + "https://proxy.example.com/team/openai", + "https://proxy.example.com/v1", + "https://litellm.internal.example/v1", + ): + client = ExternalProviderClient( + provider_type = "gemini", + base_url = base, + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is True, base + headers = client._auth_headers() + assert "x-goog-api-key" not in {k.lower() for k in headers}, ( + base, + headers, + ) + assert headers["Authorization"] == "Bearer AIza-test-key", ( + base, + headers, + ) + + +def test_google_hosted_gemini_still_uses_native_dispatch(): + """Google-hosted Gemini keeps native dispatch + x-goog-api-key auth.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is False + headers = client._auth_headers() + assert headers.get("x-goog-api-key") == "AIza-test-key", headers + + +def test_invalid_gemini_model_id_rejected_before_request(monkeypatch): + """Path-traversal model ids must be rejected before the URL is + interpolated so the configured API key isn't sent to unintended + Gemini endpoints.""" + + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response( + 200, + content = _gemini_sse([]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + ): + out.append(line) + await client.close() + + _drive(run()) + # No outbound request should have been issued. + assert captured == [], captured + error_lines = [line for line in out if '"error"' in line] + assert error_lines, out + + +def test_top_k_omitted_when_not_explicit_default_for_gemini(monkeypatch): + """top_k=None means "use provider default"; helper must not emit + `topK` in generationConfig when the caller didn't pass it.""" + captured = _capture_body(monkeypatch, top_k = None) + assert "topK" not in captured["body"]["generationConfig"], captured["body"] + + +def test_text_model_image_generation_tool_silently_dropped(monkeypatch): + """A stale `enabled_tools=["image_generation"]` on a text-only + Gemini model (e.g. gemini-2.5-flash) must NOT switch the request + into image mode -- Google's API 400s on responseModalities for + text models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "responseModalities" not in gc, gc + + +def test_empty_text_part_with_thought_signature_emits_extra_content( + monkeypatch, +): + """Gemini 3 can ship a content-free fragment whose only payload is + `thoughtSignature`. The translator must still surface that signature + on a delta.extra_content envelope so the next turn can replay it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "answer"}, + {"thoughtSignature": "SIG-FINAL"}, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + extra_carriers = [ + c + for c in chunks + if c.get("choices") + and c["choices"][0]["delta"].get("extra_content") + == {"google": {"thought_signature": "SIG-FINAL"}} + ] + assert extra_carriers, chunks + + +def test_enable_prompt_caching_false_string_coerces_to_bool(): + """Pre-PR the field was Optional[bool]; widening to Union[bool,str] + must preserve historical coercion so callers sending `"false"` + still opt out of caching.""" + from models.inference import ChatCompletionRequest + + msg = {"role": "user", "content": "hi"} + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "false", + } + ) + assert req.enable_prompt_caching is False, req.enable_prompt_caching + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "true", + } + ) + assert req.enable_prompt_caching is True + + # An actual cache resource name passes through untouched. + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "cachedContents/abc123", + } + ) + assert req.enable_prompt_caching == "cachedContents/abc123" + + +def test_legacy_google_openai_base_url_is_rewritten(): + """The Google-hosted /v1beta/openai legacy base IS still rewritten.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_remote_image_url_downloads_and_inlines_as_base64(monkeypatch): + """Round 14: arbitrary public HTTPS image URLs cannot be sent as + Gemini fileData (that path is reserved for Files API URIs and + YouTube). The translator must fetch the bytes server-side and + inline them as base64 inlineData.""" + image_bytes = b"FAKEPNGBYTES" + + async def fake_fetch(url, fallback_mime, max_bytes = None): + assert url == "https://cdn.example.com/diagram.png" + return ("image/png", base64.b64encode(image_bytes).decode("ascii")) + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": { + "url": "https://cdn.example.com/diagram.png", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + inline = next((p for p in parts if "inlineData" in p), None) + assert inline is not None, parts + assert inline["inlineData"]["mimeType"] == "image/png" + assert inline["inlineData"]["data"] == base64.b64encode(image_bytes).decode() + assert not any("fileData" in p for p in parts), parts + + +def test_remote_image_url_dropped_when_fetch_returns_none(monkeypatch): + """Round 15: if the SSRF guard rejects the URL (private host, + non-https, oversize, non-image), the helper returns None and the + image part is silently dropped instead of forwarding raw bytes + or a fileData fallback.""" + + async def fake_fetch_reject(url, fallback_mime, max_bytes = None): + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch_reject) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "http://10.0.0.5/private.png"}, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + assert not any("fileData" in p for p in parts), parts + + +def test_safe_fetch_image_rejects_non_https(): + """SSRF guard: only https URLs may be fetched.""" + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini("http://cdn.example.com/x.png", "image/png") + ) + assert res is None + + +def test_safe_fetch_image_rejects_loopback_ip_literal(): + """SSRF guard: refuse loopback / private IP literals before any + network call.""" + for url in ( + "https://127.0.0.1/x.png", + "https://[::1]/x.png", + "https://169.254.169.254/latest/meta-data", + "https://10.0.0.5/x.png", + "https://192.168.1.1/x.png", + ): + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini(url, "image/png") + ) + assert res is None, url + + +def test_safe_fetch_image_rejects_resolved_private_host(monkeypatch): + """SSRF guard: if a hostname resolves to a private IP, refuse.""" + import socket + + def fake_getaddrinfo(host, *_args, **_kwargs): + return [(socket.AF_INET, None, None, "", ("10.0.0.5", 0))] + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini( + "https://internal.example/x.png", "image/png" + ) + ) + assert res is None + + +def test_youtube_and_files_api_uris_stay_as_file_data(monkeypatch): + """Round 14: YouTube URLs and generativelanguage.googleapis.com + Files API URIs are the documented `fileData.fileUri` paths and + must NOT be downloaded; arbitrary public URLs do get fetched.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "explain"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc123", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "https://generativelanguage.googleapis.com/v1beta/files/abc", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + parts = captured["body"]["contents"][-1]["parts"] + file_uris = [p["fileData"]["fileUri"] for p in parts if "fileData" in p] + assert "https://www.youtube.com/watch?v=abc123" in file_uris, parts + assert ( + "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris + ), parts + + +def test_tool_use_prompt_tokens_added_to_input_tokens(monkeypatch): + """`toolUsePromptTokenCount` must roll into the OpenAI prompt + total -- otherwise tool turns silently undercount input tokens.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "result"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "toolUsePromptTokenCount": 100, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 110, usage + assert usage["completion_tokens"] == 7, usage + assert usage["total_tokens"] == 117, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 2, usage + + +def test_usage_chunk_reasoning_tokens_surfaced(monkeypatch): + """thoughtsTokenCount must surface as completion_tokens_details. + reasoning_tokens in the emitted OpenAI usage chunk.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["completion_tokens"] == 25, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 20, usage + + +def test_prompt_block_pairs_web_search_tool_end(monkeypatch): + """When `promptFeedback.blockReason` triggers after the synthetic + web_search tool_start, the helper must emit a matching tool_end so + the UI does not leave a "searching..." spinner stuck on screen.""" + sse = [ + {"promptFeedback": {"blockReason": "SAFETY"}}, + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + assert len(starts) == 1, tool_events + assert len(ends) == 1, tool_events + assert ends[0]["tool_call_id"] == "gemini_web_search" + assert "aborted" in ends[0]["result"] + error_chunks = [c for c in chunks if c.get("error")] + assert error_chunks, chunks + + +def test_code_execution_tool_events_stow_native_part(monkeypatch): + """executableCode / codeExecutionResult must round-trip native ids + and thoughtSignature in google.native_part so follow-up turns can + replay Gemini's required history shape.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-CODE", + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + code_start = next( + (e for e in starts if e.get("tool_name") == "code_execution"), + None, + ) + code_end = next(iter(ends), None) + assert code_start is not None, starts + assert code_start["tool_call_id"] == "code_a", code_start + native = code_start["arguments"]["google"]["native_part"] + # Round 21: native_part now uses an ordered `parts` list so per-part + # `thoughtSignature` survives a frontend merge of executableCode + + # codeExecutionResult into one tool-call card. + start_parts = native["parts"] + assert start_parts[0]["executableCode"]["id"] == "code_a" + assert start_parts[0]["thoughtSignature"] == "SIG-CODE" + assert code_end is not None, ends + assert code_end["tool_call_id"] == "code_a", code_end + native_end = code_end["google"]["native_part"] + end_parts = native_end["parts"] + assert end_parts[0]["codeExecutionResult"]["id"] == "result_a" + + +def test_inline_image_tool_end_carries_thought_signature(monkeypatch): + """Inline image parts with thoughtSignature must persist it on the + emitted tool_end so Gemini 3 image editing can echo it back.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNG").decode(), + }, + "thoughtSignature": "SIG-IMG", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 4, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + image_ends = [ + e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64") + ] + assert image_ends, tool_events + assert image_ends[0]["google"]["thought_signature"] == "SIG-IMG" + # Multi-turn image edit must replay the original inlineData part with + # its thoughtSignature; the outbound translator reads + # google.native_part.parts[].inlineData, so stow it on the tool_end + # too. Round 21 changed native_part to an ordered parts list so a + # per-part signature stays attached to inlineData only. + native = image_ends[0]["google"]["native_part"] + image_parts = native["parts"] + assert image_parts[0]["inlineData"]["mimeType"] == "image/png" + assert image_parts[0]["inlineData"]["data"] == base64.b64encode(b"PNG").decode() + assert image_parts[0]["thoughtSignature"] == "SIG-IMG" + + +def test_code_execution_plot_attaches_inline_image_native_part(monkeypatch): + """A code_execution turn that returns a matplotlib plot must stow + the plot's inlineData on the secondary tool_end so the follow-up + turn can replay the image alongside executableCode and + codeExecutionResult.""" + plot_data = base64.b64encode(b"PLOT").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "plt.plot([0,1])", + }, + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": plot_data, + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a" + ] + # Two tool_end events on the same id: one for codeExecutionResult, + # one merging in the inlineData plot. The plot one must carry the + # native inlineData under google.native_part so the frontend + # tool_end merge union joins it with the prior executableCode and + # codeExecutionResult parts on the same card. + assert len(code_ends) == 2, code_ends + image_end = next( + (e for e in code_ends if "__IMAGES__:" in (e.get("result") or "")), + None, + ) + assert image_end is not None, code_ends + native = image_end["google"]["native_part"] + plot_parts = native["parts"] + assert plot_parts[0]["inlineData"]["mimeType"] == "image/png" + assert plot_parts[0]["inlineData"]["data"] == plot_data + + +def test_text_chunk_carries_thought_signature(monkeypatch): + """Text parts with thoughtSignature surface it on delta.extra_content + so frontend persistence can replay it on the follow-up turn.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "hello", + "thoughtSignature": "SIG-TEXT", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + text_chunks = [ + c + for c in chunks + if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello" + ] + assert text_chunks, chunks + extra = text_chunks[0]["choices"][0]["delta"].get("extra_content") + assert extra == {"google": {"thought_signature": "SIG-TEXT"}}, text_chunks + + +def test_openai_tools_translated_into_function_declarations(monkeypatch): + """Standard ChatCompletionRequest.tools must be forwarded into + Gemini's tools[].functionDeclarations envelope.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + }, + "required": ["city"], + }, + }, + } + ], + tool_choice = {"type": "function", "function": {"name": "get_weather"}}, + ) + tools_arr = captured["body"].get("tools") or [] + fn_decls = [t for t in tools_arr if "functionDeclarations" in t] + assert fn_decls, captured["body"] + decls = fn_decls[0]["functionDeclarations"] + assert decls[0]["name"] == "get_weather" + assert decls[0]["parameters"]["properties"]["city"]["type"] == "string" + tool_config = captured["body"].get("toolConfig") + assert tool_config is not None, captured["body"] + fcc = tool_config["functionCallingConfig"] + assert fcc["mode"] == "ANY" + assert fcc["allowedFunctionNames"] == ["get_weather"] + + +def test_tool_choice_auto_maps_to_function_calling_mode_auto(monkeypatch): + """tool_choice="auto" maps to toolConfig.functionCallingConfig.mode.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + tool_choice = "auto", + ) + fcc = captured["body"]["toolConfig"]["functionCallingConfig"] + assert fcc["mode"] == "AUTO" + assert "allowedFunctionNames" not in fcc + + +def test_code_exec_inline_image_attaches_to_code_execution_card(monkeypatch): + """A codeExecution sandbox plot (matplotlib) ships as an inline + image part right after the codeExecutionResult. Instead of spawning + a separate empty image_generation card, attach to the same + code_execution tool_end via the `__IMAGES__:` marker the chat + adapter already understands.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_plot", + "language": "PYTHON", + "code": "import matplotlib.pyplot as plt; plt.plot([1,2,3]); plt.savefig('out.png')", + }, + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "saved", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNGDATA").decode(), + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + # No standalone image_generation card should have been emitted. + image_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "image_generation" + ] + assert not image_starts, tool_events + # The code_execution tool_end should now carry the inline image + # via the `__IMAGES__:` marker. + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_plot" + ] + assert code_ends, tool_events + final_result = code_ends[-1]["result"] + assert "__IMAGES__:" in final_result, code_ends + assert "data:image/png;base64," in final_result, code_ends + + +def test_code_execution_tool_call_replays_native_executable_code(monkeypatch): + """An assistant tool_call with toolName=code_execution and + extra_content.google.native_part containing the originally-emitted + `executableCode` + `codeExecutionResult` must round-trip as native + Gemini parts (not a generic functionCall) on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute 2+2"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "code_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(2+2)", + }, + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + }, + "thoughtSignature": "SIG-CODE", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "what was that result"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + native_keys = [list(p.keys())[0] for p in parts if isinstance(p, dict)] + assert "executableCode" in native_keys, parts + assert "codeExecutionResult" in native_keys, parts + assert not any( + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "code_execution" + for p in parts + ), parts + exec_part = next(p for p in parts if "executableCode" in p) + assert exec_part.get("thoughtSignature") == "SIG-CODE", exec_part + + +def test_image_generation_tool_call_replays_native_inline_data(monkeypatch): + """An assistant tool_call with toolName=image_generation and + extra_content.google.native_part.inlineData must replay the prior + image as a native Gemini inlineData part (not a generic + functionCall) so multi-turn image editing keeps the image + context.""" + pixel = base64.b64encode(b"PNG").decode() + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + messages = [ + {"role": "user", "content": "make a circle"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "img_a", + "type": "function", + "function": { + "name": "image_generation", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "inlineData": { + "mimeType": "image/png", + "data": pixel, + }, + "thoughtSignature": "SIG-IMG", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "now make it blue"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert inline_parts, parts + assert inline_parts[0]["inlineData"]["mimeType"] == "image/png" + assert inline_parts[0]["inlineData"]["data"] == pixel + assert inline_parts[0].get("thoughtSignature") == "SIG-IMG", inline_parts + assert not any( + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "image_generation" + for p in parts + ), parts + + +def test_assistant_text_thought_signature_replays_on_outbound_text_part(monkeypatch): + """Assistant text with extra_content.google.thought_signature must + attach `thoughtSignature` to the LAST text part of the replayed + Gemini history. Gemini 3 strict function-calling rejects history + that drops returned signatures, so the frontend stows the latest + signed-text signature and the backend pins it on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + text_parts = [p for p in parts if "text" in p] + assert text_parts, parts + assert text_parts[-1].get("thoughtSignature") == "SIG-TEXT", text_parts + + +def test_function_declarations_strip_openai_only_schema_keys(monkeypatch): + """OpenAI strict tools commonly include `additionalProperties`, + `$schema`, `$defs`, `strict`, etc. Gemini's Schema rejects those + with INVALID_ARGUMENT, so the translator must strip them while + keeping properties..type intact.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value.", + "parameters": { + "type": "object", + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": False, + "strict": True, + "properties": { + "key": { + "type": "string", + "additionalProperties": False, + }, + }, + "required": ["key"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "additionalProperties" not in params + assert "$schema" not in params + assert "strict" not in params + assert params["type"] == "object" + assert params["properties"]["key"]["type"] == "string" + assert "additionalProperties" not in params["properties"]["key"] + assert params["required"] == ["key"] + + +def test_function_declarations_inline_local_refs_into_gemini_schema(monkeypatch): + """Round 25: Pydantic-generated tool schemas hoist nested object + shapes into `$defs` and reference them with `{"$ref": "#/$defs/..."}`. + Gemini's OpenAPI subset has no $ref, so a naive allowlist sanitizer + drops the reference and reduces the nested property to `{}`, losing + its type, fields, and required keys. The sanitizer must resolve + local `#/...` pointers and inline the referenced schema.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_user", + "description": "Persist a user.", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": { + "street": {"type": "string"}, + "zip": {"type": "string"}, + }, + "required": ["street", "zip"], + }, + }, + "properties": { + "name": {"type": "string"}, + "address": {"$ref": "#/$defs/Address"}, + }, + "required": ["name", "address"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "$defs" not in params + address = params["properties"]["address"] + assert address.get("type") == "object", address + assert address.get("properties", {}).get("street", {}).get("type") == "string" + assert address.get("properties", {}).get("zip", {}).get("type") == "string" + assert address.get("required") == ["street", "zip"] + + +def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch): + """The recursive inliner must reach through `anyOf` branches and + `items` (array element schemas) as well, not just top-level + property refs.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "bulk_set", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": {"zip": {"type": "string"}}, + "required": ["zip"], + }, + }, + "properties": { + "primary": { + "anyOf": [ + {"$ref": "#/$defs/Address"}, + {"type": "null"}, + ], + }, + "extras": { + "type": "array", + "items": {"$ref": "#/$defs/Address"}, + }, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None + params = decls[0]["parameters"] + primary = params["properties"]["primary"] + # anyOf with single non-null branch + null collapses to inline + + # nullable: true, and the inlined branch must contain the resolved + # Address shape. + assert primary.get("nullable") is True + assert primary.get("type") == "object" + assert primary.get("properties", {}).get("zip", {}).get("type") == "string" + extras = params["properties"]["extras"] + assert extras.get("type") == "array" + assert extras.get("items", {}).get("type") == "object" + assert ( + extras.get("items", {}).get("properties", {}).get("zip", {}).get("type") + == "string" + ) + + +def test_function_declarations_self_referential_schema_terminates(monkeypatch): + """Self-referential / cyclic JSON Schemas (a `Node` that contains + `children: [Node]`) must not infinite-loop. The inliner tracks the + set of refs in flight and short-circuits to `{}` on a cycle.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_tree", + "parameters": { + "type": "object", + "$defs": { + "Node": { + "type": "object", + "properties": { + "value": {"type": "string"}, + "children": { + "type": "array", + "items": {"$ref": "#/$defs/Node"}, + }, + }, + }, + }, + "properties": { + "root": {"$ref": "#/$defs/Node"}, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None + root = decls[0]["parameters"]["properties"]["root"] + assert root.get("type") == "object" + assert root.get("properties", {}).get("value", {}).get("type") == "string" + + +def test_gemini_native_skips_orphan_function_response_for_dropped_builtin( + monkeypatch, +): + """Round 26: when the assistant-side synthetic web_search/web_fetch + tool_call is dropped from native Gemini history, the matching + role="tool" follow-up must also be dropped. Otherwise the outbound + body carries an orphan functionResponse with no preceding + functionCall, which 400s the Gemini turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "web_search", contents + + +def test_gemini_native_skips_orphan_function_response_for_native_part_replay( + monkeypatch, +): + """Round 26: code_execution / image_generation tool_calls are + replayed as Gemini-native executableCode / codeExecutionResult / + inlineData parts. The matching role="tool" follow-up must NOT then + be emitted as a functionResponse named code_execution -- there is + no declared user function with that name, and Gemini's history + rules already attribute the result to the native parts above.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "plot something"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "2\n", + } + }, + ] + } + } + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_native = False + for entry in contents: + for part in entry.get("parts", []): + if "executableCode" in part or "codeExecutionResult" in part: + saw_native = True + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "code_execution", contents + assert saw_native, contents + + +def test_gemini_native_part_falls_back_to_args_google(monkeypatch): + """Round 27: a direct OpenAI-compat API caller (or imported third- + party thread) cannot use Studio's non-standard + `tool_calls[].extra_content` field, so the native_part payload + round-trips through `function.arguments` as + `{"google": {"native_part": {...}}}`. The synthetic-builtin + detector recognizes that location, but the replay branch was only + reading from `tc.extra_content.google.native_part`. Result: the + round-25 guard saw a synthetic builtin with no _native_part and + dropped the entire assistant turn, losing the prior code/image + context. The translator must fall back to args.google.native_part + and still emit the native executableCode / inlineData parts.""" + import json as _json + + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "draw a cat"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "image_generation", + "arguments": _json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "AAAA", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "user", "content": "now make it a dog"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_inline = False + for entry in contents: + for part in entry.get("parts", []): + if "inlineData" in part: + saw_inline = True + assert saw_inline, contents + + +def test_gemini_native_skips_synthetic_server_builtin_replay(monkeypatch): + """Round 25: Marked server-side builtin tool_calls (web_search / + web_fetch with `_server_tool` or `args.google.native_part`) must + not fall through to the generic Gemini `functionCall` replay path + when no replayable native part exists. Without this guard the + outbound body contains a fake `functionCall` whose name is not a + declared user function, and the Gemini turn 400s.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fc = part.get("functionCall") + if isinstance(fc, dict): + assert fc.get("name") != "web_search", contents + + +def test_chat_message_extra_content_round_trips_through_validation(): + """Round 9: ChatMessage was missing `extra_content`, so Pydantic + discarded the field during request validation and the text-part + signature replay path read nothing. The field must survive + model_validate and pass through _build_external_messages.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + assistant_msg = req.messages[1] + assert assistant_msg.extra_content == { + "google": {"thought_signature": "SIG-TEXT"}, + } + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assistant_out = built[1] + assert assistant_out["extra_content"] == { + "google": {"thought_signature": "SIG-TEXT"}, + } + # Non-Gemini providers must NOT receive extra_content; Google's + # thought_signature field is unknown to OpenAI / Mistral / etc. + built_openai = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + ) + assert "extra_content" not in built_openai[1], built_openai[1] + # Custom non-Google Gemini bases (LiteLLM / OAI-compat gateways) + # also must not receive Gemini-only extra_content because the + # backend dispatches them through /chat/completions. + built_custom = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + assert "extra_content" not in built_custom[1], built_custom[1] + + +def test_parallel_tool_results_group_into_one_user_block(monkeypatch): + """Round 14: Gemini docs show parallel functionResponses grouped + in a single subsequent user content with multiple + functionResponse parts. Consecutive OpenAI role="tool" messages + must merge into one Gemini user block, not split into separate + user turns.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "add", "arguments": '{"x":1}'}, + }, + { + "id": "call_b", + "type": "function", + "function": {"name": "mul", "arguments": '{"x":2}'}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "add", + "content": "2", + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "mul", + "content": "4", + }, + ], + ) + contents = captured["body"]["contents"] + # Initial user, model with two functionCalls, ONE user with two + # functionResponses. + tool_result_users = [ + c + for c in contents + if c.get("role") == "user" + and all( + isinstance(p, dict) and "functionResponse" in p + for p in (c.get("parts") or []) + ) + ] + assert len(tool_result_users) == 1, contents + fr_parts = tool_result_users[0]["parts"] + assert len(fr_parts) == 2, fr_parts + names = [p["functionResponse"]["name"] for p in fr_parts] + assert names == ["add", "mul"], names + + +def test_function_schema_nullable_type_array_flattens(monkeypatch): + """Round 14: OpenAI strict tools commonly use + `"type": ["string", "null"]` for optional fields. Gemini's + OpenAPI-style Schema rejects union types and expects + `"type": "string"` with `"nullable": true`. The sanitizer must + translate the union form.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "city": {"type": ["string", "null"]}, + "score": {"type": ["number", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["city"]["type"] == "string" + assert params["city"]["nullable"] is True + assert params["score"]["type"] == "number" + assert params["score"]["nullable"] is True + + +def test_image_picker_model_with_search_off_pill_strips_text_tools(monkeypatch): + """Round 11: image-tier model id rejects text-only tools and + thinkingConfig at the model level regardless of whether the Images + pill is on. Selecting gemini-2.5-flash-image + enabled_tools= + ["web_search"] with no image_generation must NOT forward + googleSearch or thinkingConfig (Gemini 400s on text tools for + legacy image ids).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["web_search"], + reasoning_effort = "high", + ) + body = captured["body"] + assert "tools" not in body, body.get("tools") + assert "thinkingConfig" not in body.get("generationConfig", {}), body[ + "generationConfig" + ] + + +def test_image_models_drop_function_declarations(monkeypatch): + """Image-mode requests cannot mix tools with responseModalities so + user-supplied function declarations must be dropped.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_safe_fetch_image_rejects_malformed_bracketed_url(): + """Round 17: bracketed IPv6 garbage like `https://[bad/x.png` makes + urlparse raise ValueError. The fetch helper must catch it and drop + the image rather than crashing the request mid-build.""" + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://[bad/x.png", "image/png")) + assert res is None + + +def test_safe_fetch_image_pins_validated_ip_no_hostname_in_request( + monkeypatch, +): + """Round 17: the fetch helper must pin the validated IP into the + outgoing request URL (with a Host header carrying the original + hostname). A second hostname-style getaddrinfo after the validate + step would be a DNS-rebinding gap, so we assert the urllib opener + is called with an IP-rewritten URL.""" + import socket + + captured: dict = {"requests": []} + + # Public IP during validate; record every getaddrinfo call. + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + captured.setdefault("dns", []).append(host) + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-type": "image/png", "content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open(self, req, timeout = None): + captured["requests"].append( + { + "url": req.full_url, + "host_header": req.get_header("Host"), + } + ) + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is not None + assert res[0] == "image/png" + # The outgoing URL must use the pinned IP literal, not the hostname. + assert any("8.8.8.8" in r["url"] for r in captured["requests"]), captured + assert all( + "cdn.example.com" not in r["url"] for r in captured["requests"] + ), captured + # Host header still carries the original hostname for vhost/SNI. + assert captured["requests"][0]["host_header"] == "cdn.example.com" + + +def test_safe_fetch_image_redirect_to_private_host_rejected(monkeypatch): + """Round 17: each redirect hop must re-validate the new host. A + public hop that redirects to an internal address must be dropped.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + if host == "internal.bad": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("10.0.0.5", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open(self, req, timeout = None): + # Simulate a 302 to a private host. + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://internal.bad/secret.png"}, + None, + ) + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is None + + +def test_files_api_substring_url_not_misclassified_as_filedata(monkeypatch): + """Round 17: a CDN URL whose path/query merely contains the Files + API substring must NOT be sent as `fileData.fileUri`; it must be + routed through the safe-fetch path. Previously the substring check + `"generativelanguage.googleapis.com/" in url.lower()` matched any + URL carrying that text anywhere.""" + captured_outbound: dict = {} + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return "image/png", base64.b64encode(b"DATA").decode("ascii") + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + captured_outbound["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": { + # Looks like a Files API URL in the path + # but the host is an attacker CDN. + "url": "https://evil.example/path/generativelanguage.googleapis.com/v1beta/files/abc.png", + }, + }, + { + "type": "image_url", + "image_url": { + # Looks YouTube-ish in the path. + "url": "https://cdn.example.com/youtube.com/cat.png", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + parts = captured_outbound["body"]["contents"][-1]["parts"] + assert not any("fileData" in p for p in parts), parts + inline_count = sum(1 for p in parts if "inlineData" in p) + assert inline_count == 2, parts + assert len(fetch_calls) == 2, fetch_calls + + +def test_function_schema_anyof_null_variant_flattens_to_nullable(monkeypatch): + """Round 17: OpenAI/Pydantic emit `anyOf: [{X}, {"type":"null"}]` + for Optional[X]. Gemini's OpenAPI subset rejects `"type":"null"` + inside anyOf. The sanitizer must collapse a singleton-plus-null + union back to the non-null branch with `nullable: true`.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "label": { + "anyOf": [ + {"type": "string"}, + {"type": "null"}, + ] + }, + "count": { + "anyOf": [ + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["label"]["type"] == "string" + assert params["label"]["nullable"] is True + assert "anyOf" not in params["label"] + assert params["count"]["type"] == "integer" + assert params["count"]["nullable"] is True + + +def test_legacy_gemini3_pro_medium_coerced_to_high(monkeypatch): + """Round 17: legacy `gemini-3-pro*` (including `-preview`, shut down + 2026-03-09) only accepted low/high. 3.1+ Pro added medium. The + backend must coerce medium → high for the legacy model so stale UI + state does not 400 the request.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "high", + } + + +def test_gemini_3_1_pro_medium_passes_through(monkeypatch): + """Round 17 regression: 3.1+ Pro accepts medium; coercion must NOT + apply when the model id is gemini-3.1-pro*.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "medium", + } + + +def test_tool_calls_extra_content_stripped_for_non_native_gemini(): + """Round 17: per-tool-call `extra_content` (Gemini thoughtSignature + carrier) must not leak through `_build_external_messages` to + non-native-Gemini providers; OpenAI / Anthropic / custom Gemini + OAI-compat gateways would 400 on the unknown key.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": { + "google": {"thought_signature": "SIG"}, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + + # Non-native providers (openai, custom Gemini OAI-compat proxy) + # must have extra_content stripped from the tool_call entry. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + assert len(result) == 1 + tc = result[0]["tool_calls"][0] + assert "extra_content" not in tc, (provider_type, tc) + + # Native Gemini still receives extra_content for the round-trip. + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + tc_native = result_native[0]["tool_calls"][0] + assert tc_native["extra_content"]["google"]["thought_signature"] == "SIG" + + +def test_user_function_named_with_server_tool_arg_not_dropped(monkeypatch): + """Round 17: the OpenAI Responses translator must NOT drop a user + function whose JSON arguments happen to contain `_server_tool: + true` UNLESS the function name is also one of the canonical + builtin names. Otherwise a user schema with an `_server_tool` field + becomes invisible to the model.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "user_function", + "arguments": json.dumps( + {"_server_tool": True, "q": "x"} + ), + }, + } + ], + }, + { + "role": "tool", + "content": "result", + "tool_call_id": "call_user", + "name": "user_function", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + # User function call must survive (matching call + output). + assert any(c.get("name") == "user_function" for c in fn_calls), items + assert len(fn_outs) == 1, items + + +def test_builtin_named_with_server_tool_marker_dropped(monkeypatch): + """Round 17 control: a builtin (web_search) tagged with + `_server_tool: true` continues to be filtered from outbound + history.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), + }, + } + ], + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + # Builtin server-side tool call must be filtered out. + assert all(c.get("name") != "web_search" for c in fn_calls), items + + +def test_gemini_tool_choice_none_disables_hosted_builtins(monkeypatch): + """Round 18: `tool_choice="none"` must drop hosted Google Search / + code execution from the outbound Gemini body, not just user + function declarations. Otherwise an API client that opted out of + tool use still triggers grounded search (privacy + billing).""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tool_choice = "none", + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_gemini_tool_choice_none_disables_function_declarations(monkeypatch): + """Round 18: `tool_choice="none"` must drop user function + declarations as well as hosted builtins from the Gemini body.""" + captured = _capture_body( + monkeypatch, + tool_choice = "none", + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_schema_anyof_multitype_with_null_keeps_anyof_and_nullable( + monkeypatch, +): + """Round 18: multi-branch unions with null (e.g. + `Union[str, int, None]`) must keep the slim anyOf without the null + branch and add `nullable: true`; Gemini rejects + `{"type":"null"}` inside anyOf.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + assert all( + not (isinstance(b, dict) and b.get("type") == "null") for b in inner + ), inner + + +def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch): + """Round 18: when the upstream 302 Location is a malformed + bracketed-IPv6 URL, the helper must return None instead of letting + a urlparse ValueError abort the chat stream.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open(self, req, timeout = None): + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://[bad/x.png"}, + None, + ) + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is None + + +def test_safe_fetch_image_malformed_port_no_crash(): + """Round 18: a URL with a non-numeric port (`https://h:bad/x.png`) + must not raise; urlparse's port property lazily ValueErrors.""" + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://example.com:bad/x.png", "image/png" + ) + ) + assert res is None + + +def test_safe_fetch_image_missing_content_type_uses_fallback(monkeypatch): + """Round 18: when the server returns image bytes but no + Content-Type header, the helper must use the caller-provided + fallback MIME (guessed from URL extension) instead of dropping the + image as `non-image content-type=`.""" + import socket + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open(self, req, timeout = None): + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/cat.png", "image/png" + ) + ) + assert res is not None + assert res[0] == "image/png" + + +def test_anthropic_translates_openai_tool_calls_into_tool_use_blocks(monkeypatch): + """Round 18: an assistant turn with OpenAI-style top-level + `tool_calls` must be translated into Anthropic native + `{type:"tool_use", id, name, input}` content blocks before being + forwarded. The OpenAI `role="tool"` follow-up must become a + `role:"user"` message with a `tool_result` content block.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + # No top-level tool_calls should remain. + assert all("tool_calls" not in m for m in msgs), msgs + # The assistant turn must now have content blocks including a + # tool_use block. + asst = [m for m in msgs if m.get("role") == "assistant"] + assert asst and isinstance(asst[0]["content"], list), asst + tool_uses = [b for b in asst[0]["content"] if b.get("type") == "tool_use"] + assert len(tool_uses) == 1, asst[0] + assert tool_uses[0]["name"] == "lookup" + assert tool_uses[0]["input"] == {"q": "x"} + # The role="tool" message must become a user/tool_result message. + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_unmarked_user_web_search_function_survives_serialization(): + """Round 18: a user-defined function literally named `web_search` + with NO `_server_tool` marker must survive `_build_external_messages` + when forwarded to a non-native provider; only marked synthetic + builtin cards may be dropped.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "x"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "web_search" + + +def test_marked_server_builtin_dropped_from_build_external_messages(): + """Round 18: when a Gemini-native turn carrying a marked + `image_generation` server-tool card is forwarded to OpenAI / a + custom Gemini OAI-compat proxy, the tool_call must be dropped, not + just have its extra_content stripped. Forwarding an orphan + `image_generation` tool_call would 400 the receiving API.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + # Non-native providers: marked builtin tool_call must be dropped + # AND if it was the only payload, the whole message disappears. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # Empty assistant turn with only synthetic tool_call dropped. + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), ( + provider_type, + result, + ) + + # Native Gemini preserves it (round-trips via extra_content). + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assert len(result_native) == 1 + assert result_native[0]["tool_calls"][0]["function"]["name"] == "image_generation" + + +def test_openai_responses_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 18: `tool_choice="none"` must also drop hosted OpenAI + Responses builtins (web_search, code execution shell, image + generation), not just user function tools.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_anthropic_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 19: tool_choice="none" must opt out of Anthropic hosted + builtins (web_search, web_fetch, code_execution) just like it does + for Gemini and OpenAI Responses.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_openrouter_tool_choice_none_drops_web_plugin(monkeypatch): + """Round 19: tool_choice="none" must drop the OpenRouter web + plugin so a request that opted out of tool use does not still + trigger hosted web search.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_tool_choice_none_skips_web_search_helper(monkeypatch): + """Round 19: when tool_choice="none" plus enabled_tools= + ["web_search"] on Kimi, the dispatcher must NOT route into + `_stream_kimi_web_search`. Falling through to the generic OAI- + compat path is the expected behavior.""" + routed_to_helper = {"called": False} + + real_helper = ExternalProviderClient._stream_kimi_web_search + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + assert routed_to_helper["called"] is False + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + real_helper, + ) + + +def test_user_code_execution_function_not_dropped(): + """Round 19: a user-declared function literally named + `code_execution` with normal `code` arguments must survive + `_build_external_messages` -- round 17's shape heuristic dropped + it, which broke function-calling round-trips.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "code_execution", + "arguments": '{"code": "print(1)"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "code_execution" + + +def test_native_part_code_execution_treated_as_server_side(): + """Round 19: a Gemini `code_execution` card persists its replay + payload at `args.google.native_part` (no `_server_tool` marker on + pre-PR cards). The backend filter must still drop it for non-native + providers because it is a synthetic card, not a real user function.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + args_with_native_part = json.dumps( + { + "google": { + "native_part": { + "executableCode": { + "language": "PYTHON", + "code": "print(1)", + } + } + } + } + ) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": { + "name": "code_execution", + "arguments": args_with_native_part, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), result + + +def test_remote_image_fetch_attempt_cap_includes_failures(monkeypatch): + """Round 19: the per-request image fetch count cap must count + ATTEMPTS, not just successes. Otherwise a request with 100 + failing/slow URLs runs 100 fetches each up to the 15s timeout.""" + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + image_parts = [ + { + "type": "image_url", + "image_url": {"url": f"https://cdn.example.com/img{idx}.png"}, + } + for idx in range(20) + ] + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + *image_parts, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert len(fetch_calls) <= 8, len(fetch_calls) + + +def test_orphan_function_call_output_dropped_when_call_skipped(monkeypatch): + """Round 19: when a marked server-side builtin `function_call` is + dropped from OpenAI Responses input items, the matching role=tool + follow-up must also be dropped to avoid an orphan + `function_call_output`.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + assert all(c.get("call_id") != "call_b" for c in fn_calls), items + assert all(o.get("call_id") != "call_b" for o in fn_outs), items + + +def test_schema_multitype_union_with_null_preserves_anyof(monkeypatch): + """Round 19: a JSON Schema `"type": ["string","integer","null"]` + must be sanitized to anyOf:[{string},{integer}] + nullable:true. + Flattening to just `{"type":"string"}` silently drops the integer + branch and changes the function contract.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": {"type": ["string", "integer", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + types = sorted( + b.get("type") for b in inner if isinstance(b, dict) and b.get("type") + ) + assert types == ["integer", "string"], inner + + +def test_invalid_gemini_model_rejected_before_image_fetch(monkeypatch): + """Round 19: invalid Gemini model IDs are rejected at the top of + `_stream_gemini`, BEFORE any user-controlled remote image fetch + runs.""" + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "image_url", + "image_url": {"url": "https://cdn.example.com/x.png"}, + }, + ], + } + ], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert fetch_calls == [], fetch_calls + + +def test_empty_assistant_turn_skipped_after_synthetic_tool_calls_dropped(): + """Round 20: when `_filter_tool_calls` drops every synthetic + server-builtin tool_call on an empty-content assistant turn, the + whole message must be skipped. Forwarding + `{"role":"assistant","content":""}` is rejected by several + providers as an empty assistant turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # The empty assistant turn (only a synthetic builtin) must + # NOT appear in the output at all. + assert result == [], (provider_type, result) + + +def test_role_tool_dropped_when_matching_synthetic_call_filtered(): + """Round 20: `_build_external_messages` drops the matching role= + tool follow-up when its tool_call was a synthetic builtin that + `_filter_tool_calls` removed. Otherwise the receiving provider + sees an orphan tool_result with no tool_call.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "query": "x"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": marked_args, + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + # Only the user "continue" message survives. + roles = [m.get("role") for m in result] + assert roles == ["user"], result + + +def test_openrouter_no_synthetic_web_search_event_on_tool_choice_none( + monkeypatch, +): + """Round 20: OpenRouter dispatcher must not emit synthetic + web_search tool_start / tool_end events when tool_choice="none"; + otherwise the chat UI shows a search card for a search that + never happened.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :].strip() + if not payload or payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + # Backend emits synthetic tool events as a top-level + # `_toolEvent` on the SSE payload (not nested inside + # `delta`). Read both shapes so a future format change + # cannot mask this regression. + evt = obj.get("_toolEvent") + if isinstance(evt, dict): + captured_events.append(evt) + for ch in obj.get("choices") or []: + delta = ch.get("delta") or {} + nested = delta.get("_toolEvent") if isinstance(delta, dict) else None + if isinstance(nested, dict): + captured_events.append(nested) + await client.close() + + _drive(run()) + # No synthetic web_search tool_start / tool_end emitted. + assert all( + e.get("tool_name") != "web_search" for e in captured_events + ), captured_events + + +def test_anthropic_role_tool_list_content_translates_to_tool_result( + monkeypatch, +): + """Round 20: an OpenAI-shape role=tool message with list content + (`content=[{"type":"text","text":"result"}]`) must be translated + into Anthropic's native tool_result block, not forwarded as an + invalid role=tool message.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": [{"type": "text", "text": "result_text"}], + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + assert all(m.get("role") != "tool" for m in msgs), msgs + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_data_url_non_image_mime_dropped(monkeypatch): + """Round 20: a `data:text/html;base64,...` image_url must be + dropped from the outbound Gemini body, not forwarded as + `inlineData.mimeType="text/html"` which Gemini rejects.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + { + "type": "image_url", + "image_url": { + "url": "data:text/html;base64,PGgxPmhpPC9oMT4=", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + + +def test_youtube_filedata_uses_video_mime(monkeypatch): + """Round 20: YouTube `fileData.fileUri` must declare a video + mimeType, not `image/jpeg` guessed from the URL path.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarise"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + yt = next((p for p in parts if "fileData" in p), None) + assert yt is not None, parts + assert yt["fileData"]["mimeType"].startswith("video/"), yt + + +def test_openai_responses_assistant_text_serialized_before_function_call( + monkeypatch, +): + """Round 20: in OpenAI Responses history, the assistant's + visible text for a turn that ALSO emitted a function_call must + serialize BEFORE the function_call item, matching the prior + response.output sequence. Otherwise function_call_output (the + role=tool follow-up) appears to follow an unrelated assistant + message.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "Let me check that.", + "tool_calls": [ + { + "id": "call_w", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "content": "sunny", + "tool_call_id": "call_w", + "name": "get_weather", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + types = [i.get("type") or i.get("role") for i in items] + # Expected order: + # user ("weather?") + # assistant ("Let me check that.") + # function_call (get_weather) + # function_call_output (sunny) + # user ("thanks") + assert types == [ + "user", + "assistant", + "function_call", + "function_call_output", + "user", + ], items + + +def test_gemini_tool_choice_none_disables_image_generation(monkeypatch): + """Round 21: `tool_choice="none"` must also flip the implicit + image-generation hosted tool off on image-tier models. Otherwise + `responseModalities=["TEXT","IMAGE"]` still rides on the outbound + body and the provider can generate (and bill for) image output + despite the explicit OpenAI tool opt-out.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = "none", + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_forced_function_tool_choice_drops_hosted_builtins(monkeypatch): + """Round 21: forced-function `tool_choice` (e.g. + `{"type":"function","function":{"name":"lookup"}}`) must suppress + hosted Google Search / code execution. Gemini's toolConfig only + constrains function declarations, not hosted tools, so leaving + `googleSearch`/`codeExecution` in `tools[]` lets them fire despite + the caller pinning a specific user function.""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + ) + body = captured["body"] + tool_kinds = [list(t.keys())[0] for t in (body.get("tools") or [])] + assert "googleSearch" not in tool_kinds, body + assert "codeExecution" not in tool_kinds, body + # User function declaration still survives. + assert "functionDeclarations" in tool_kinds, body + + +def test_gemini_forced_function_tool_choice_drops_image_generation(monkeypatch): + """Round 21: forced-function `tool_choice` must also flip the + implicit image-generation hosted tool off on image-tier models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_code_execution_native_part_list_replays_per_part_signatures( + monkeypatch, +): + """Round 21: merged code-execution history must replay per-part + `thoughtSignature`s, not fan one top-level signature across every + native subpart. Gemini 3 strict validators reject a signature + placed on the wrong part.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-EXEC", + }, + { + "codeExecutionResult": { + "id": "res_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + # Locate the assistant turn replayed as native code-exec parts. + assistant_turn = next(c for c in contents if c["role"] == "model") + parts = assistant_turn["parts"] + exec_parts = [p for p in parts if "executableCode" in p] + result_parts = [p for p in parts if "codeExecutionResult" in p] + assert exec_parts and result_parts, parts + assert exec_parts[0].get("thoughtSignature") == "SIG-EXEC", exec_parts[0] + # codeExecutionResult had no signature -- must NOT inherit one. + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_code_execution_legacy_merged_signature_only_on_executable( + monkeypatch, +): + """Round 21: backward compatibility for pre-round-21 persisted + history that stored merged `native_part` as a single object plus a + top-level `thoughtSignature`. The replay branch must attach that + signature only to `executableCode` (where Gemini 3 emits it), not + fan it across `codeExecutionResult` / `inlineData`.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_b", + "language": "PYTHON", + "code": "print(1+1)", + }, + "codeExecutionResult": { + "id": "res_b", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + "thoughtSignature": "LEGACY-SIG", + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + assistant_turn = next(c for c in contents if c["role"] == "model") + exec_parts = [p for p in assistant_turn["parts"] if "executableCode" in p] + result_parts = [p for p in assistant_turn["parts"] if "codeExecutionResult" in p] + assert exec_parts[0].get("thoughtSignature") == "LEGACY-SIG", exec_parts[0] + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_role_tool_list_content_flattens_to_result_text(monkeypatch): + """Round 21: OpenAI-shape role=tool messages may carry list content + like `[{"type":"text","text":"result"}]`. Forwarding those parts + verbatim into `functionResponse.response.result` yields a list of + content-part objects instead of the actual tool output text.""" + history = [ + {"role": "user", "content": "look up"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": json.dumps({"q": "x"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "name": "lookup", + "content": [{"type": "text", "text": "answer-text"}], + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + fn_response = None + for c in contents: + for p in c.get("parts") or []: + if isinstance(p, dict) and "functionResponse" in p: + fn_response = p["functionResponse"] + break + if fn_response: + break + assert fn_response is not None, contents + assert fn_response["response"] == {"result": "answer-text"}, fn_response + + +def test_safe_fetch_image_threads_per_request_byte_budget(monkeypatch): + """Round 21: the aggregate per-request byte cap must be passed into + `_safe_fetch_image_for_gemini` so an oversize URL is refused via + Content-Length (short-circuit) rather than fully downloaded then + discarded after the fact.""" + import socket + + captured: dict = {"reads": 0, "content_length_seen": None} + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + # Declared 5 MiB, but caller passes a 1 MiB remaining budget. + headers = { + "content-type": "image/png", + "content-length": str(5 * 1024 * 1024), + } + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + captured["reads"] += 1 + return b"\x00" * (5 * 1024 * 1024) + + class _StubOpener: + def open(self, req, timeout = None): + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/big.png", + "image/png", + max_bytes = 1 * 1024 * 1024, + ) + ) + assert res is None + # Refused via Content-Length pre-check, never read. + assert captured["reads"] == 0 + + +def test_openai_chat_delta_type_includes_tool_calls_and_extra_content(): + """Round 21: the frontend `OpenAIChatDelta` interface must expose + `tool_calls` and `extra_content` so TypeScript callers can consume + the Gemini-native stream fields without `any` casts. This test is + a static-string assertion against the .ts source; mirrors how other + frontend wire-contract tests are pinned from the backend suite.""" + import os + + here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + types_path = os.path.join( + here, "frontend", "src", "features", "chat", "types", "api.ts" + ) + with open(types_path, "r", encoding = "utf-8") as f: + src = f.read() + assert "tool_calls?: OpenAIToolCallPart[]" in src, src[:200] + assert "extra_content?: Record" in src, src[:200] + assert "boolean | string | null" in src, src[:200] + + +def test_anthropic_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 22: forced-function tool_choice must suppress Anthropic + hosted builtins the same way it does for Gemini. Pinning a user + function (`tool_choice={"type":"function","function":{"name":...}}`) + while also passing `enabled_tools=["web_search","web_fetch", + "code_execution"]` should not still fire those server-side.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + # No hosted tools should be in the body — only the caller's user- + # function declarations (which this test doesn't pass any of). + tools = body.get("tools") or [] + hosted_tool_names = {"web_search", "web_fetch", "code_execution"} + for tool in tools: + assert tool.get("name") not in hosted_tool_names, body + + +def test_openrouter_forced_function_tool_choice_drops_web_plugin(monkeypatch): + """Round 22: forced-function tool_choice must drop the OpenRouter + web plugin too — caller pinned a user function, OpenRouter must not + still attach the hosted web-search plugin.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_forced_function_tool_choice_skips_web_search_helper(monkeypatch): + """Round 22: forced-function tool_choice plus enabled_tools= + ["web_search"] on Kimi must NOT route into `_stream_kimi_web_search`. + Caller pinned a user function; hosted $web_search should be + suppressed for the same privacy/billing reason.""" + routed_to_helper = {"called": False} + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + assert not routed_to_helper["called"] + + +def test_openai_responses_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 23: forced-function tool_choice on the OpenAI Responses + path must suppress hosted builtins (web_search, shell, + image_generation) the same way it does for Gemini / Anthropic / + OpenRouter / Kimi. User-defined function tools still flow through + so the pinned function can resolve.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"event: response.completed\ndata: {}\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-openai-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tools = [ + { + "type": "function", + "function": { + "name": "lookup_record", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + tools = body.get("tools") or [] + hosted_types = {"web_search", "shell", "image_generation"} + hosted_seen = {t.get("type") for t in tools if isinstance(t, dict)} + assert not (hosted_seen & hosted_types), body + # The user function declaration must still be present so the pin + # has something to target. + user_function_seen = any( + isinstance(t, dict) and t.get("type") == "function" for t in tools + ) + assert user_function_seen, body + # And the forced-function tool_choice must be forwarded in Responses + # shape: `{type:"function", name:"..."}`. + tc = body.get("tool_choice") + assert isinstance(tc, dict) and tc.get("type") == "function", body + assert tc.get("name") == "lookup_record", body + + +def test_strip_provider_synthetic_tool_history_drops_text_only_extra_content(): + """Round 24: a plain text Gemini reply (no tool_calls) carrying + `extra_content.google.thought_signature` must still have that + metadata stripped before being forwarded to a local llama-server + backend. Without this, switching a Gemini thread mid-stream to a + local GGUF model leaks Gemini-only fields to llama-server.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "Hello!", + "extra_content": {"google": {"thought_signature": "SIG_ABC"}}, + }, + {"role": "user", "content": "now in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + # Same three turns, but the assistant's `extra_content` is gone. + assert [m["role"] for m in out] == ["user", "assistant", "user"] + assistant = out[1] + assert "extra_content" not in assistant, assistant + assert assistant["content"] == "Hello!" + + +def test_validate_and_resolve_host_blocks_shared_address_space(): + """Round 24 SSRF P1: 100.64.0.0/10 carrier-grade NAT addresses are + `is_private=False` AND `is_global=False` per Python's ipaddress + docs. The previous denylist (is_private/loopback/link_local/etc.) + missed them. Adding `not ip.is_global` as the primary gate covers + all non-public ranges, current and future.""" + import socket as _socket + from core.inference import tools as _tools + + orig_getaddrinfo = _socket.getaddrinfo + + def fake_getaddrinfo(hostname, port, *args, **kwargs): + if hostname == "shared.example": + return [ + ( + _socket.AF_INET, + _socket.SOCK_STREAM, + 0, + "", + ("100.64.0.1", port), + ), + ] + return orig_getaddrinfo(hostname, port, *args, **kwargs) + + _socket.getaddrinfo = fake_getaddrinfo + try: + ok, reason, _ip = _tools._validate_and_resolve_host("shared.example", 443) + finally: + _socket.getaddrinfo = orig_getaddrinfo + assert ok is False, (ok, reason) + assert "non-public" in reason.lower() or "100.64.0.1" in reason + + +def test_gemini_custom_oai_compat_base_skips_native_allowlist(): + """Round 24: a custom Gemini OAI-compatible base (LiteLLM/proxy) + must NOT have its model list filtered through the native Gemini + allowlist regex. A LiteLLM gateway returning + `["google/gemini-2.5-flash", "my-team/gemini", "gemini-2.5-flash"]` + should be passed through; the native filter would strip the + prefixed IDs even though the chat dispatch routes them via the + OpenAI-compatible client.""" + import asyncio as _asyncio + + from routes import providers as _providers + from routes.providers import ( + ProviderModelsRequest, + list_provider_models, + ) + + captured: dict = {"base": None} + + class _FakeClient: + def __init__(self, *, base_url, **kwargs): + captured["base"] = base_url + + async def list_models(self): + return [ + {"id": "google/gemini-2.5-flash"}, + {"id": "my-team/gemini"}, + {"id": "gemini-2.5-flash"}, + ] + + async def close(self): + return None + + orig = _providers.ExternalProviderClient + _providers.ExternalProviderClient = _FakeClient + try: + req = ProviderModelsRequest( + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + result = _asyncio.run(list_provider_models(req, current_subject = "unsloth")) + finally: + _providers.ExternalProviderClient = orig + ids = {m.id for m in result} + # All three IDs survive — the native allowlist was bypassed. + assert "google/gemini-2.5-flash" in ids, ids + assert "my-team/gemini" in ids, ids + assert "gemini-2.5-flash" in ids, ids + + +def test_strip_provider_synthetic_tool_history_drops_synthetic_only(): + """Round 22: switching a thread from native Gemini (code_execution + / image_generation tool_cards in history) to a local GGUF backend + must strip the synthetic tool_calls + matching role=tool replies + before llama-server sees them. Real user-function tool_calls and + their matching tool replies must survive.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "let me run it", + "tool_calls": [ + { + "id": "synth_ce_1", + "type": "function", + "function": { + "name": "code_execution", + "arguments": json.dumps( + { + "_server_tool": True, + "google": {"native_part": {"parts": []}}, + } + ), + }, + "extra_content": {"google": {"thought_signature": "abc"}}, + }, + { + "id": "real_lookup", + "type": "function", + "function": { + "name": "lookup_user", + "arguments": json.dumps({"id": 42}), + }, + }, + ], + "extra_content": {"google": {"thought_signature": "msglevel"}}, + }, + { + "role": "tool", + "tool_call_id": "synth_ce_1", + "content": "Gemini-only result text", + }, + { + "role": "tool", + "tool_call_id": "real_lookup", + "content": '{"name": "alice"}', + }, + ] + out = _strip_provider_synthetic_tool_history(messages) + assistant = next(m for m in out if m.get("role") == "assistant") + tcs = assistant["tool_calls"] + assert len(tcs) == 1, tcs + assert tcs[0]["id"] == "real_lookup" + assert "extra_content" not in tcs[0] + assert "extra_content" not in assistant + tool_msgs = [m for m in out if m.get("role") == "tool"] + assert len(tool_msgs) == 1 + assert tool_msgs[0]["tool_call_id"] == "real_lookup" + + +def test_strip_provider_synthetic_tool_history_drops_empty_assistant(): + """If every tool_call was synthetic and the assistant turn had no + content, the entire turn must be dropped (llama-server rejects + empty assistant messages with no tool_calls).""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "draw a sloth"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "synth_imggen", + "type": "function", + "function": { + "name": "image_generation", + "arguments": json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "Zm9v", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "tool", "tool_call_id": "synth_imggen", "content": "(image)"}, + {"role": "user", "content": "now try in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + roles = [m.get("role") for m in out] + # The synthetic assistant + its tool reply are both gone; only the + # two user turns survive. + assert roles == ["user", "user"], out + + +def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice( + monkeypatch, +): + """Round 22 sibling of the round-20 `tool_choice='none'` test: when + the caller forces a specific function via `tool_choice={"type": + "function", ...}` AND passes `enabled_tools=["web_search"]`, the + OpenRouter path must NOT synthesize a fake `web_search` tool card. + The plugin was not attached upstream so the UI must not see a + server-tool card.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = ( + b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' + b"data: [DONE]\n\n" + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + payload = line.strip().removeprefix("data: ") + if payload and payload != "[DONE]": + try: + captured_events.append(json.loads(payload)) + except Exception: + pass + await client.close() + + _drive(run()) + for evt in captured_events: + for choice in evt.get("choices") or []: + delta = choice.get("delta") or {} + extra = delta.get("extra_content") or {} + tool_event = extra.get("toolEvent") if isinstance(extra, dict) else None + if isinstance(tool_event, dict): + assert tool_event.get("tool_name") != "web_search", evt diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 98c7bdaa55..c36363b1ae 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -66,6 +66,7 @@ def _load_worker_module(): _worker = _load_worker_module() _normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer _normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler +_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size def test_mlx_studio_optimizer_aliases_are_explicit(): @@ -82,3 +83,14 @@ def test_mlx_studio_rejects_unknown_optimizer(): def test_mlx_studio_rejects_unknown_scheduler(): with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"): _normalize_mlx_studio_scheduler("linear_typo") + + +def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer(): + assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256) + assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512) + assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512) + assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128) + assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256) + # Half-pixel cases must match the Torch collator (not banker's round). + assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500) + assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167) diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py index 3d179371e3..0b1a65c69e 100644 --- a/studio/backend/tests/test_openai_code_execution.py +++ b/studio/backend/tests/test_openai_code_execution.py @@ -269,7 +269,15 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch): assert len(ends) == 1 assert starts[0]["tool_name"] == "code_execution" assert starts[0]["tool_call_id"] == "scall_1" - assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"} + # `_server_tool: True` is the synthetic-builtin marker the + # backend stamps onto every provider-side tool_start so the + # frontend serializer can distinguish hosted tools from + # user-declared functions on history replay. + assert starts[0]["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert ends[0]["tool_call_id"] == "scall_1" assert "total 24" in ends[0]["result"] diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py index 6d415316d4..f5dee2561a 100644 --- a/studio/backend/tests/test_openai_image_generation.py +++ b/studio/backend/tests/test_openai_image_generation.py @@ -207,9 +207,12 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch): ends = [e for e in image_events if e.get("type") == "tool_end"] assert len(starts) == 1, image_events assert len(ends) == 1, image_events + # `_server_tool: True` marks this as a provider-side synthetic + # tool card on the frontend's history serializer. assert starts[0]["arguments"] == { "kind": "image", "prompt": "A photorealistic cat sitting", + "_server_tool": True, "openai_image_generation_call_id": "img_abc", } assert ends[0]["image_b64"] == "AAAA" diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py index 22ccba7058..f177ed5ef3 100644 --- a/studio/backend/tests/test_openai_responses_translation.py +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -215,6 +215,254 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch): assert payloads[-1] == "[DONE]" +def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch): + """Round 12: caller-supplied function tools forwarded into /v1/responses + must have their `function_call` output items translated back into Chat + Completions delta.tool_calls, and the terminal chunk must emit + finish_reason="tool_calls" so the frontend's accumulator runs the + function instead of seeing finish_reason="stop".""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_abc", + "call_id": "call_xyz", + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "weather?"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + payloads = [ + json.loads(line[len("data:") :].strip()) + for line in lines + if line.startswith("data:") and line[len("data:") :].strip() != "[DONE]" + ] + tool_call_deltas = [ + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("delta", {}).get("tool_calls") + ] + assert tool_call_deltas, payloads + tc = tool_call_deltas[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["id"] == "call_xyz" + assert tc["function"]["name"] == "get_weather" + assert tc["function"]["arguments"] == '{"city":"SF"}' + # Final chunk reports tool_calls instead of stop. + terminal = next( + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("finish_reason") in ("stop", "tool_calls") + ) + assert terminal["choices"][0]["finish_reason"] == "tool_calls", payloads + + +def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch): + """Round 13: parallel function_call items must land on distinct + delta.tool_calls[].index slots so index-keyed clients don't + collapse the second call into the first.""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_a", + "call_id": "call_a", + "name": "lookup_a", + "arguments": "{}", + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_b", + "call_id": "call_b", + "name": "lookup_b", + "arguments": "{}", + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "lookup_a", + "parameters": {"type": "object"}, + }, + }, + { + "type": "function", + "function": { + "name": "lookup_b", + "parameters": {"type": "object"}, + }, + }, + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + indices: list[int] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + delta = (obj.get("choices") or [{}])[0].get("delta") or {} + for tc in delta.get("tool_calls") or []: + indices.append(tc.get("index")) + assert indices == [0, 1], indices + + +def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch): + """Round 13: a second turn after a Responses function call must + serialize the tool_calls history and tool result as Responses + `function_call` / `function_call_output` input items, not as + Chat Completions role="tool" content.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _responses_sse( + [ + {"type": "response.created"}, + {"type": "response.completed", "response": {}}, + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + await _collect( + client._stream_openai_responses( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": "sunny", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + + _drive(run()) + items = captured["body"]["input"] + types = [it.get("type") or it.get("role") for it in items] + assert "function_call" in types, items + assert "function_call_output" in types, items + fc = next(it for it in items if it.get("type") == "function_call") + assert fc["call_id"] == "call_xyz" + assert fc["name"] == "get_weather" + assert fc["arguments"] == '{"city":"SF"}' + fco = next(it for it in items if it.get("type") == "function_call_output") + assert fco["call_id"] == "call_xyz" + assert fco["output"] == "sunny" + + def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch): def handler(request: httpx.Request) -> httpx.Response: events = [ diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py index 7ffa9bb384..6df491610b 100644 --- a/studio/backend/tests/test_studio_train_validation.py +++ b/studio/backend/tests/test_studio_train_validation.py @@ -18,6 +18,8 @@ from models.training import ( _MAX_LORA_ALPHA, _MAX_LORA_R, _MAX_SEQ_LENGTH, + _MAX_VISION_IMAGE_SIZE, + _MIN_VISION_IMAGE_SIZE, ) @@ -62,6 +64,52 @@ class TestBatchSizeCap: _check_field("batch_size", 0) +class TestVisionImageSizeCap: + def test_none_accepts_model_default(self): + _check_field("vision_image_size", None) + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE], + ) + def test_in_range_accepts(self, value): + _check_field("vision_image_size", value) + assert _MIN_VISION_IMAGE_SIZE == 256 + assert _MAX_VISION_IMAGE_SIZE == 2048 + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True], + ) + def test_invalid_rejects(self, value): + with pytest.raises(ValidationError): + _check_field("vision_image_size", value) + + @pytest.mark.parametrize("value", [True, False]) + def test_bool_error_says_integer_not_range(self, value): + # Regression guard: bools must say "integer or null", not "in [256, 2048]". + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + @pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"]) + def test_multi_sign_string_says_integer_not_raw(self, value): + # Regression guard: multi-sign strings must not leak int()'s raw + # "invalid literal" message; precise contract is "integer or null". + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + assert "invalid literal" not in str(exc.value) + + @pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"]) + def test_unicode_digit_string_rejected(self, value): + # Full-width / Arabic-Indic / Devanagari digits must be rejected so the + # value reaching the backend equals the ASCII the user typed. + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + class TestLoraRCap: def test_at_cap_accepts(self): _check_field("lora_r", _MAX_LORA_R) diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py new file mode 100644 index 0000000000..922f95ce55 --- /dev/null +++ b/studio/backend/utils/cpu_threads.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Early CPU thread-pool configuration for Studio processes.""" + +import os +from typing import MutableMapping, Optional + + +_THREAD_POOL_ENV_VARS = ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", +) + + +def configure_cpu_threads(env: Optional[MutableMapping[str, str]] = None) -> None: + """Apply ``UNSLOTH_CPU_THREADS`` to native CPU pools when configured. + + This must run before importing libraries that initialize an OpenMP or + BLAS thread pool. Library-specific variables are left untouched so users + can override a single runtime independently. + """ + environ = os.environ if env is None else env + configured = environ.get("UNSLOTH_CPU_THREADS", "").strip() + if not configured: + return + + try: + thread_count = int(configured) + except ValueError as exc: + raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") from exc + if thread_count < 1: + raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") + + value = str(thread_count) + for variable in _THREAD_POOL_ENV_VARS: + environ.setdefault(variable, value) diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 83b1fd96f9..b43e174889 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -12,6 +12,7 @@ "lint": "eslint .", "preview": "vite preview", "typecheck": "tsc -b --pretty false", + "i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts", "biome:check": "biome check", "biome:fix": "biome check --write" }, diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index c7bc0440bd..f0a417638d 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -3,6 +3,7 @@ import { Link, createRouter, useRouterState } from "@tanstack/react-router"; import { Button } from "@/components/ui/button"; +import { useT } from "@/i18n"; import { Route as rootRoute } from "./routes/__root"; import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; @@ -31,7 +32,9 @@ const routeTree = rootRoute.addChildren([ ]); function DefaultNotFound() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); + return (

- Page not found + {t("shell.notFound.title")}

- {pathname} does not exist. + {t("shell.notFound.description", { path: pathname })}

); diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 47bff815e6..57ed233d51 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,8 +6,9 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; -import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard"; +import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; +import { useT, type TranslationKey } from "@/i18n"; import { Outlet, createRootRoute, @@ -16,24 +17,25 @@ import { useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; -import { Suspense, useEffect, useLayoutEffect, type ReactNode } from "react"; +import { Suspense, useEffect, useLayoutEffect } from "react"; import { AppProvider } from "../provider"; -// Type `staticData.title` on every route so the matched-title selector -// below stays type-safe without an inline cast. declare module "@tanstack/react-router" { interface StaticDataRouteOption { title?: string; + titleKey?: TranslationKey; } } -// Fallback while a lazy route bundle (Train/Recipes/Export) loads. -// /chat is synchronous and never hits this. -const RouteFallback: ReactNode = ( -
- Loading... -
-); +function RouteFallback() { + const t = useT(); + + return ( +
+ {t("common.loading")} +
+ ); +} const CHAT_ONLY_ALLOWED = new Set([ "/", @@ -68,6 +70,7 @@ const HIDDEN_NAVBAR_ROUTES = ["/onboarding", "/login", "/change-password"]; const DEFAULT_DOCUMENT_TITLE = "Unsloth Studio"; function RootLayout() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname); const isChatRoute = pathname.startsWith("/chat"); @@ -75,24 +78,20 @@ function RootLayout() { useTrainingUnloadGuard(); - // Walk matches deepest-first; each route declares its own title. const matchedTitle = useMatches({ select: (matches) => { for (let i = matches.length - 1; i >= 0; i--) { - const title = matches[i].staticData.title; + const { title, titleKey } = matches[i].staticData; + if (titleKey) return t(titleKey); if (title) return title; } return null; }, }); - // `/settings` redirects in `beforeLoad`, so its route never stays - // matched; surface the modal's title via the store instead. const settingsDialogOpen = useSettingsDialogStore((s) => s.open); - const documentTitle = settingsDialogOpen ? "Settings" : matchedTitle; + const documentTitle = settingsDialogOpen ? t("settings.title") : matchedTitle; - // useLayoutEffect updates the tab title before paint, avoiding a - // one-frame flash of the previous route's title on navigation. useLayoutEffect(() => { document.title = documentTitle ? `${documentTitle} - ${DEFAULT_DOCUMENT_TITLE}` @@ -116,7 +115,7 @@ function RootLayout() { {hideNavbar ? (
- + }>
@@ -142,7 +141,7 @@ function RootLayout() { transition={{ duration: 0.15 }} className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`} > - + }> diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx index 75f1a1b937..ae7f445e94 100644 --- a/studio/frontend/src/app/routes/studio.tsx +++ b/studio/frontend/src/app/routes/studio.tsx @@ -15,7 +15,7 @@ const StudioPage = lazy(() => export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/studio", - staticData: { title: "Train" }, + staticData: { titleKey: "studio.routeTitle" }, beforeLoad: () => requireAuth(), component: StudioPage, }); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index aac5f8f8a8..849e017ea8 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -90,9 +90,33 @@ import { useTrainingRuntimeStore, } from "@/features/training"; import type { TrainingRunSummary } from "@/features/training"; -import { useEffect, useRef, useState } from "react"; +import { useEffect, useRef, useState, type ReactNode } from "react"; import { toast } from "@/lib/toast"; import { ShutdownDialog } from "@/components/shutdown-dialog"; +import { translate, useT, type TranslationKey } from "@/i18n"; + +const EMPHASIS_MARKER = "__UNSLOTH_I18N_EMPHASIS_MARKER__"; + +type AppT = ReturnType; + +function renderEmphasizedTranslation( + t: AppT, + key: TranslationKey, + emphasizedValue: string, +): ReactNode { + const translated = t(key, { name: EMPHASIS_MARKER }); + const parts = translated.split(EMPHASIS_MARKER); + if (parts.length === 1) return translated; + + const nodes: ReactNode[] = []; + parts.forEach((part, index) => { + if (part.length > 0) nodes.push(part); + if (index < parts.length - 1) { + nodes.push({emphasizedValue}); + } + }); + return nodes; +} function getTourId(pathname: string): string | null { if (pathname.startsWith("/studio")) return "studio"; @@ -185,6 +209,7 @@ function NavItem({ } export function AppSidebar() { + const t = useT(); const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle(); const { pathname, search } = useRouterState({ select: (s) => ({ @@ -204,14 +229,8 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const [shutdownOpen, setShutdownOpen] = useState(false); - // Chat collapsible state — open by default, auto-expand on route entry const isChatRoute = pathname.startsWith("/chat"); const isStudioRoute = pathname === "/studio" || pathname.startsWith("/studio/"); - const [chatOpen, setChatOpen] = useState(true); - const [runsOpen, setRunsOpen] = useState(true); - - useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]); - useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]); const scrollRef = useRef(null); const [scrolled, setScrolled] = useState(false); @@ -290,7 +309,7 @@ export function AppSidebar() { try { await renameChatItem(target.item, renameTrimmed); } catch (err) { - toast.error("Failed to rename chat", { + toast.error(translate("shell.toast.failedToRenameChat"), { description: err instanceof Error ? err.message : undefined, }); } @@ -300,7 +319,7 @@ export function AppSidebar() { const updated = await renameTrainingRun(target.run.id, nextRunDisplayName); emitTrainingRunUpdated(updated); } catch (err) { - toast.error("Failed to rename run", { + toast.error(translate("shell.toast.failedToRenameRun"), { description: err instanceof Error ? err.message : undefined, }); } @@ -320,14 +339,14 @@ export function AppSidebar() { try { await handleDeleteThread(target.item); } catch (err) { - toast.error("Failed to delete chat", { + toast.error(translate("shell.toast.failedToDeleteChat"), { description: err instanceof Error ? err.message : undefined, }); } return; } if (target.run.status === "running") { - toast.error("Cannot delete a running training run"); + toast.error(t("shell.toast.cannotDeleteRunningRun")); return; } try { @@ -337,7 +356,7 @@ export function AppSidebar() { } emitTrainingRunDeleted(target.run.id); } catch (err) { - toast.error("Failed to delete run", { + toast.error(translate("shell.toast.failedToDeleteRun"), { description: err instanceof Error ? err.message : undefined, }); } @@ -366,7 +385,7 @@ export function AppSidebar() { }); }} className="flex items-center gap-[6px] select-none" - aria-label="Unsloth home" + aria-label={t("shell.aria.home")} > - BETA + {t("shell.beta")} {!isMobile && ( @@ -387,7 +406,7 @@ export function AppSidebar() { type="button" onClick={togglePinned} className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Close sidebar" + aria-label={t("shell.aria.closeSidebar")} > @@ -397,7 +416,7 @@ export function AppSidebar() { sideOffset={6} className="tooltip-compact" > - Close sidebar + {t("shell.aria.closeSidebar")} )} @@ -412,7 +431,7 @@ export function AppSidebar() { type="button" onClick={togglePinned} className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Open sidebar" + aria-label={t("shell.aria.openSidebar")} > @@ -422,7 +441,7 @@ export function AppSidebar() { sideOffset={8} className="tooltip-compact" > - Open sidebar + {t("shell.aria.openSidebar")} @@ -434,7 +453,7 @@ export function AppSidebar() { { @@ -446,7 +465,7 @@ export function AppSidebar() { /> i.id === search.compare)} disabled={chatDisabled} dataTour="chat-compare" @@ -459,7 +478,7 @@ export function AppSidebar() { /> { @@ -477,7 +496,7 @@ export function AppSidebar() { { @@ -489,7 +508,7 @@ export function AppSidebar() { { navigate({ to: "/data-recipes" }); @@ -499,7 +518,7 @@ export function AppSidebar() { { @@ -513,13 +532,16 @@ export function AppSidebar() { - {/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */} {!isStudioRoute && chatItems.length > 0 && ( - + - Recents + {t("shell.navigation.recents")} @@ -552,7 +574,7 @@ export function AppSidebar() { @@ -844,7 +872,9 @@ export function AppSidebar() { - {renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"} + {renamingTarget?.kind === "run" + ? t("shell.dialog.renameRun.title") + : t("shell.dialog.renameChat.title")} @@ -868,14 +906,14 @@ export function AppSidebar() { variant="ghost" onClick={() => setRenamingTarget(null)} > - Cancel + {t("common.cancel")} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index a5bb4029d1..51cf000863 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -802,6 +802,7 @@ const ReasoningToggle: FC = () => { { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index e9fe1d679c..c0a60d0de2 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -108,6 +108,7 @@ export const LR_DEFAULT_CPT = 5e-5; export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, + visionImageSize: null as number | null, learningRate: LR_DEFAULT_LORA, // null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT. embeddingLearningRate: null as number | null, diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 560df956c6..de43462d42 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -28,6 +28,7 @@ import { getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, + isGeminiCustomOpenAICompatBase, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebFetch, @@ -108,6 +109,33 @@ type RunMessage = RunMessages[number]; /** Tracks which user messages were sent with an audio file (messageId → filename). */ export const sentAudioNames = new Map(); +// Synthetic provider-side tool names; backend stamps args._server_tool +// so user functions with the same name aren't dropped. Mirror of +// _SERVER_SIDE_BUILTIN_TOOL_NAMES on the backend. +const SERVER_SIDE_BUILTIN_TOOL_NAMES = new Set([ + "web_search", + "web_fetch", + "code_execution", + "image_generation", +]); + +/** + * Whether a persisted tool-call part is provider-side synthetic and + * should be stripped from outbound history. Match on the + * args._server_tool marker or a Gemini native_part payload — no shape + * heuristic, because user functions can legitimately share a name. + */ +function isServerSideBuiltinToolPart( + toolNameLower: string, + _argsObj: Record | null, + hasServerToolMarker: boolean, + hasNativePart: boolean, +): boolean { + if (!SERVER_SIDE_BUILTIN_TOOL_NAMES.has(toolNameLower)) return false; + if (hasServerToolMarker) return true; + return hasNativePart; +} + /** * Match error messages that indicate the request filled or would fill * the KV cache, so the UI can show a dedicated toast pointing at the @@ -511,21 +539,200 @@ function isAnthropicRefusalMessage(message: RunMessage): boolean { return metadata?.custom?.anthropicRefusal === true; } -function toOpenAIMessage(message: RunMessage): { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; -} | null { +function collectAssistantToolCalls( + message: RunMessage, +): Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; +}> { + const out: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart & { + argsText?: string; + extra_content?: unknown; + }; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + const hasServerToolMarker = Boolean( + argsObj && (argsObj as Record)._server_tool === true, + ); + const isServerSideBuiltin = isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ); + if (isServerSideBuiltin) { + // Gemini code_execution / image_generation still need to round- + // trip the native_part payload for native replay; drop the rest. + if (!hasNativePart) continue; + } + const argumentsStr = + typeof tc.argsText === "string" && tc.argsText.length > 0 + ? tc.argsText + : JSON.stringify(tc.args ?? {}); + const entry: { + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + } = { + id: tc.toolCallId, + type: "function" as const, + function: { + name: tc.toolName ?? "", + arguments: argumentsStr, + }, + }; + // Promote args.google to extra_content.google so the backend + // native_part replay branch can find it. The backend only inspects + // extra_content, not function.arguments. + if (tc.extra_content !== undefined) { + entry.extra_content = tc.extra_content; + } else if (argsGoogle) { + entry.extra_content = { google: argsGoogle }; + } + out.push(entry); + } + return out; +} + +function collectToolResultMessages( + message: RunMessage, +): Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; +}> { + const out: Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart; + const result = (tc as { result?: unknown }).result; + // Skip provider-side builtins; see isServerSideBuiltinToolPart. + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const hasServerToolMarker = Boolean( + argsObj && argsObj._server_tool === true, + ); + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + if ( + isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ) + ) { + continue; + } + if (result === undefined || result === null) continue; + let content: string; + if (typeof result === "string") { + // Backend ChatMessage validator rejects role="tool" with empty + // content; serialise a sentinel JSON so legitimately empty tool + // outputs still round-trip the follow-up turn to the provider. + content = result.length > 0 ? result : JSON.stringify({ result: "" }); + } else { + try { + content = JSON.stringify(result); + } catch { + content = String(result); + } + } + out.push({ + role: "tool", + content, + tool_call_id: tc.toolCallId, + ...(tc.toolName ? { name: tc.toolName } : {}), + }); + } + return out; +} + +type SerializedMessage = { + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + tool_calls?: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }>; + tool_call_id?: string; + name?: string; + /** + * Gemini text-part thoughtSignature stashed during streaming on the + * last text MessagePart. Backend reads + * `extra_content.google.thought_signature` and attaches it to the + * matching Gemini text part on the outbound turn. + */ + extra_content?: unknown; +}; + +function collectAssistantTextThoughtSignature( + message: RunMessage, +): string | undefined { + if (!Array.isArray(message.content)) return undefined; + for (let i = message.content.length - 1; i >= 0; i -= 1) { + const part = message.content[i] as { type?: string } & Record< + string, + unknown + >; + if (part?.type !== "text") continue; + const sig = part._google_thought_signature; + if (typeof sig === "string" && sig) return sig; + } + return undefined; +} + +function toOpenAIMessages(message: RunMessage): SerializedMessage[] { if ( message.role !== "system" && message.role !== "user" && message.role !== "assistant" ) { - return null; + return []; } let textContent = collectTextParts(message).join("\n"); - // Strip inline audio base64 from prior assistant messages to avoid - // inflating token counts (e.g. audio-player responses with embedded WAV). if (message.role === "assistant") { textContent = textContent.replace( /data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g, @@ -534,25 +741,70 @@ function toOpenAIMessage(message: RunMessage): { if (isAnthropicRefusalMessage(message)) { // Prune refused assistant turn from outbound history; the // rendered transcript still shows the user-visible notice. - return null; + return []; } } const imageParts = collectImageParts(message); - if (imageParts.length > 0) { - return { - role: message.role, - content: [ - ...(textContent ? [{ type: "text" as const, text: textContent }] : []), - ...imageParts, - ], - }; + const toolCalls = + message.role === "assistant" ? collectAssistantToolCalls(message) : []; + const toolResults = + message.role === "assistant" ? collectToolResultMessages(message) : []; + + const base: SerializedMessage = { + role: message.role, + content: + imageParts.length > 0 + ? [{ type: "text", text: textContent }, ...imageParts] + : textContent, + }; + if (toolCalls.length > 0) { + base.tool_calls = toolCalls; + // OpenAI requires content === null on assistant turns whose + // payload is entirely tool_calls (matches the wire shape Gemini + // expects for the next functionCall replay). + if (!textContent && imageParts.length === 0) { + base.content = null; + } + } + if (message.role === "assistant") { + const sig = collectAssistantTextThoughtSignature(message); + if (sig) { + base.extra_content = { google: { thought_signature: sig } }; + } } - if (!textContent) { + return toolResults.length > 0 ? [base, ...toolResults] : [base]; +} + +// Thin singular wrapper: returns only the first serialized message +// (without tool_calls or tool follow-ups) so the OpenAI image-edit +// replay path can map a thread to flat OpenAI chat messages without +// pulling in tool history. +function toOpenAIMessage(message: RunMessage): { + role: "system" | "user" | "assistant"; + content: OpenAIMessageContent; +} | null { + const serialized = toOpenAIMessages(message); + if (serialized.length === 0) return null; + const first = serialized[0]; + if ( + first.role !== "system" && + first.role !== "user" && + first.role !== "assistant" + ) { return null; } - return { role: message.role, content: textContent }; + if (first.content === null || first.content === undefined) { + return null; + } + if (typeof first.content === "string" && !first.content) { + return null; + } + return { + role: first.role, + content: first.content as OpenAIMessageContent, + }; } function extractImageBase64(input: string): string | undefined { @@ -1087,17 +1339,25 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { clearSelectedImageEditReference(); throw new Error("Connection not found."); } - // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. - // Codex dispatches via the local CLI / SDK, no HTTP API key. + // Local providers, custom Gemini bases, and Codex (local CLI / SDK) all allow an empty key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; const externalProviderIsCodex = externalProvider ? isCodexProviderType(externalProvider.providerType) : false; - const externalProviderNeedsApiKey = - isExternalRequest && !externalProviderIsCustom && !externalProviderIsCodex; - if (externalProviderNeedsApiKey && !externalApiKey) { + const externalProviderIsGeminiCustomBase = Boolean( + externalProvider && + externalProvider.providerType === "gemini" && + isGeminiCustomOpenAICompatBase(externalProvider.baseUrl), + ); + if ( + isExternalRequest && + !externalApiKey && + !externalProviderIsCustom && + !externalProviderIsCodex && + !externalProviderIsGeminiCustomBase + ) { toast.error("Missing API key for selected connection.", { description: "Open Settings > Connections and set the API key again.", }); @@ -1105,15 +1365,38 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Missing connection API key."); } + // Image-generation flag (OpenAI cloud + Responses-capable model). + // Computed first so Gemini image mode can suppress Search/Code. + const imageGenerationEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + imageToolsEnabled && + providerSupportsBuiltinImageGeneration( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); + // Per-model Search/Code allowances live in + // providerSupportsBuiltin*; this flag just signals image-mode. + const geminiImageModeForThisTurn = + externalProvider?.providerType === "gemini" && + imageGenerationEnabledForThisTurn; const webSearchEnabledForThisTurn = Boolean( externalProvider && + externalSelection && toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType), + providerSupportsBuiltinWebSearch( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), ); const codeExecEnabledForThisTurn = Boolean( externalProvider && externalSelection && codeToolsEnabled && + !geminiImageModeForThisTurn && providerSupportsBuiltinCodeExecution( externalProvider.providerType, externalSelection.modelId, @@ -1133,20 +1416,6 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider && providerSupportsBuiltinWebFetch(externalProvider.providerType), ); - // OpenAI Responses-API image_generation server tool. Pill is - // gated on OpenAI cloud + a Responses-API model id; the backend - // additionally re-checks is_openai_cloud before appending - // {type:"image_generation"} to the request tools array. - const imageGenerationEnabledForThisTurn = Boolean( - externalProvider && - externalSelection && - imageToolsEnabled && - providerSupportsBuiltinImageGeneration( - externalProvider.providerType, - externalSelection.modelId, - externalProvider.baseUrl, - ), - ); if (selectedImageEditReference && !imageGenerationEnabledForThisTurn) { clearSelectedImageEditReference(); @@ -1157,11 +1426,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Image generation edit unavailable."); } - // Two-pass build: a refused assistant turn also drops the user - // prompt that triggered it (leaving it in context re-triggers - // the classifier). Refusal flag rides assistant - // metadata.custom.anthropicRefusal, set out-of-band from the - // backend _toolEvent. + // Drop refused assistant turns + their triggering user prompt; + // otherwise context re-triggers the classifier. const survivingMessages: RunMessage[] = []; for (const message of messages) { if (isAnthropicRefusalMessage(message)) { @@ -1174,8 +1440,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { survivingMessages.push(message); } + // toOpenAIMessages emits assistant tool_calls + role="tool" + // follow-ups; the backend Gemini translator rebuilds the + // functionCall/functionResponse parts (with thoughtSignature). const outboundMessages = survivingMessages - .map(toOpenAIMessage) + .flatMap(toOpenAIMessages) .filter((message): message is NonNullable => Boolean(message), ); @@ -1198,7 +1467,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { break; } } - outboundMessages.splice(insertAt, 0, referenceMessage); + // OpenAIChatMessage is a structural superset of SerializedMessage + // for the role/content axis the outbound pipeline consumes; cast + // through unknown since referenceMessage carries no tool_calls + // (the image_edit reference is a plain assistant turn). + outboundMessages.splice( + insertAt, + 0, + referenceMessage as unknown as SerializedMessage, + ); } const safeSystemPrompt = @@ -1280,7 +1557,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { outboundMessages[0] = { ...firstMessage, content: [ - ...firstMessage.content, + ...(Array.isArray(firstMessage.content) + ? firstMessage.content + : []), { type: "text", text: `\n\n${disabledToolGuard}` }, ], }; @@ -1491,9 +1770,26 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // because the close tag fires when the next chunk carries content // (or when the stream ends). let reasoningContentOpen = false; - // Tool call content parts — accumulated and yielded cumulatively. - // result is set directly on the tool-call part when tool_end arrives. + // Tool call parts, cumulative; result lands on tool_end. const toolCallParts: ToolCallMessagePart[] = []; + // Latest Gemini text-part thoughtSignature; pinned onto the final + // text MessagePart so next-turn replay carries it. + let latestTextThoughtSignature: string | undefined; + const pinTextThoughtSignature = ( + parts: T[], + ): T[] => { + if (!latestTextThoughtSignature || parts.length === 0) return parts; + for (let i = parts.length - 1; i >= 0; i -= 1) { + if (parts[i].type === "text") { + parts[i] = { + ...parts[i], + _google_thought_signature: latestTextThoughtSignature, + } as T; + break; + } + } + return parts; + }; const orderAssistantContent = ( textParts: ReturnType, ) => { @@ -1585,6 +1881,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { { isReasoningProvider: externalProvider.isReasoningModel === true, + baseUrl: externalProvider.baseUrl ?? null, }, ) : { @@ -1617,13 +1914,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { forceRefreshPublicKey = false, ): Promise => { if (externalSelection && externalProvider) { - // OpenAI shell-tool container reuse: pull the per-thread - // container_id (if any) so subsequent turns in the same - // thread reference the existing container instead of - // auto-creating a fresh one. Empty string / undefined → - // backend falls back to container_auto. Anthropic uses - // the parallel `anthropicCodeExecContainerId` field below - // (sent as `container` on /v1/messages). + // Per-thread container reuse; empty/undefined falls back to + // container_auto. Anthropic uses anthropicCodeExecContainerId. let openaiCodeExecContainerId: string | null = null; let anthropicCodeExecContainerId: string | null = null; if (codeExecEnabledForThisTurn && resolvedThreadId) { @@ -1637,17 +1929,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; anthropicCodeExecContainerId = null; } - // Pre-send container validation (OpenAI only). The list - // endpoint already filters status==="expired" server-side - // (studio/backend/routes/inference.py — list_openai_containers), - // so membership in this set means "OpenAI will accept it - // as container_reference". A stale id silently dropped here - // falls through to the inheritance + lazy-create logic - // below, so the user never sees "Container is expired" in - // the chat thread. On list-call failure we leave - // activeContainerIds null and skip validation — the - // backend's transparent retry path is the safety net for - // that case. + // Pre-send container validation (OpenAI). Stale ids drop + // silently and fall through to lazy-create. On list-call + // failure, skip and rely on the backend's retry path. let activeContainerIds: Set | null = null; if (externalProvider.providerType === "openai") { try { @@ -1670,15 +1954,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; } } - // Cross-thread inheritance: when the active thread has - // no container yet, default to the one most recently - // used on *any* other thread (provider-scoped). - // Matches what the Code Execution settings section - // shows in the picker, and keeps the user from getting - // a fresh container on every new thread. The picker - // can still be set to "Auto-create per thread" - // explicitly to opt into a fresh container — but - // that's done via the dropdown, not silently. + // Cross-thread inheritance: reuse the most recently used + // container from any other thread; opt-out via the picker. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1711,13 +1988,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { /* fall through to lazy-create below */ } } - // Lazy pre-create when there's no inherited container. - // We always POST /v1/containers ourselves (rather than - // letting the backend send container_auto) so every - // container shows up in the picker with a friendly - // English-word name and the user's configured TTL. - // Falls back to container_auto only if the POST fails - // — keeps the chat moving in that case. + // Pre-create our own container (vs container_auto) so it + // shows up in the picker with a friendly name and the + // configured TTL. Falls back to container_auto on failure. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1776,22 +2049,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalSelection?.modelId, ), ), - // Only forward sampling knobs the provider actually accepts; the - // backend's external-provider proxy is param-permissive and would - // surface a 400 from providers that reject unknown fields (e.g. - // OpenAI rejects top_k, Anthropic/DeepSeek reject presence_penalty). + // Forward only sampling knobs the provider accepts. ...(externalCapabilities?.topK ? { top_k: params.topK } : {}), ...(externalCapabilities?.presencePenalty ? { presence_penalty: params.presencePenalty } : {}), - // Built-in tools: Search pill maps to provider-side - // web_search (currently OpenAI / Anthropic / OpenRouter / - // Kimi); Code pill maps to Anthropic's server-side - // code_execution_20250825 tool (Anthropic is the only - // external provider that ships one today). Backend - // translates enabled_tools into each provider's tool - // schema — for Anthropic that's the entries appended to - // body["tools"] inside _stream_anthropic. + // Compose the enabled_tools list from the active pills; + // backend maps each name to the provider's tool schema. ...(webSearchEnabledForThisTurn || webFetchEnabledForThisTurn || codeExecEnabledForThisTurn || @@ -1800,14 +2064,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { enable_tools: true, enabled_tools: [ ...(webSearchEnabledForThisTurn ? ["web_search"] : []), - // web_fetch has its own Fetch pill, independent - // of Search. Anthropic-only today. ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []), ...(codeExecEnabledForThisTurn ? ["code_execution"] : []), - // OpenAI Responses-API only: `image_generation` - // returns inline image_generation_call output - // items; the backend's _stream_openai_responses - // path translates them to assistant tool events. ...(imageGenerationEnabledForThisTurn ? ["image_generation"] : []), @@ -1843,12 +2101,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.enablePromptCaching ?? true, } : {}), - // Anthropic-only: pass the cache TTL the user picked in - // Configuration → Provider. Omitted = inherit the default - // 5-minute pool. The backend's `_stream_anthropic` only - // attaches `cache_control.ttl` when the value is one of - // "5m" / "1h" (see external_provider.py near line 1375), - // so unknown values are a no-op end-to-end. + // Anthropic prompt-cache TTL; unknown values no-op on backend. ...(supportsProviderPromptCacheTtl( externalProvider.providerType, ) && @@ -2121,12 +2374,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { typeof imageB64 === "string" && imageB64 ) { - // OpenAI Responses image_generation_call: the - // backend stashes the base64 PNG/WebP/JPEG on - // separate `image_b64` / `image_mime` fields on - // the synthetic _toolEvent so the JSON result - // string stays small enough to log. Repackage as - // a structured result for the dedicated tool UI. + // Backend keeps base64 on separate image_b64 / + // image_mime fields so logs stay small; repackage. parsedResult = { image_b64: imageB64, image_mime: @@ -2153,31 +2402,107 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else { parsedResult = rawResult; } + // Merge tool_end args first, then Gemini native_part. const nextArgs = toolEvent.arguments && typeof toolEvent.arguments === "object" ? (toolEvent.arguments as ToolCallMessagePart["args"]) : undefined; - const mergedArgs = nextArgs - ? { ...(toolCallParts[idx].args ?? {}), ...nextArgs } - : toolCallParts[idx].args; + const mergedArgs: ToolCallMessagePart["args"] = { + ...(toolCallParts[idx].args ?? {}), + ...(nextArgs ?? {}), + } as ToolCallMessagePart["args"]; + // Merge tool_end native_part into args.google so the + // outbound translator replays both start (executableCode) + // and end (result / inlineData) on the same turn. + // Concatenate parts so each keeps its own thoughtSignature. + const endGoogle = ( + toolEvent as { google?: { native_part?: unknown } } + ).google; + if ( + endGoogle && + typeof endGoogle === "object" && + endGoogle.native_part && + typeof endGoogle.native_part === "object" + ) { + const argsObj = mergedArgs as Record; + const existingGoogle = (argsObj.google ?? {}) as Record< + string, + unknown + >; + const existingNative = + (existingGoogle.native_part as Record< + string, + unknown + >) ?? {}; + const endNative = endGoogle.native_part as Record< + string, + unknown + >; + // Extract part entries from either parts:[...] or + // legacy single-object native_part. Legacy + // thoughtSignature always belongs on executableCode. + const collectParts = ( + native: Record, + ): Record[] => { + if (Array.isArray(native.parts)) { + return (native.parts as unknown[]).filter( + (entry): entry is Record => + Boolean(entry) && + typeof entry === "object" && + !Array.isArray(entry), + ); + } + const out: Record[] = []; + const legacySig = + typeof native.thoughtSignature === "string" + ? native.thoughtSignature + : typeof native.thought_signature === "string" + ? (native.thought_signature as string) + : null; + for (const key of [ + "executableCode", + "codeExecutionResult", + "inlineData", + ] as const) { + const sub = native[key]; + if (sub && typeof sub === "object") { + const entry: Record = { + [key]: sub, + }; + if (key === "executableCode" && legacySig) { + entry.thoughtSignature = legacySig; + } + out.push(entry); + } + } + return out; + }; + const mergedParts = [ + ...collectParts(existingNative), + ...collectParts(endNative), + ]; + argsObj.google = { + ...existingGoogle, + native_part: { parts: mergedParts }, + }; + } toolCallParts[idx] = { ...toolCallParts[idx], args: mergedArgs, - argsText: mergedArgs - ? JSON.stringify(mergedArgs) - : toolCallParts[idx].argsText, + argsText: JSON.stringify(mergedArgs ?? {}), result: parsedResult, }; } } - // Yield cumulative state so tool UI updates. Search/code tools stay - // before the text, while generated images sit after the answer. - // ``renderFullContent()`` preserves any Codex per-tab text - // accumulated in earlier ``_toolEvent`` frames so the tabs do - // not briefly appear and then vanish when the regular content - // path arrives (the round 8 render-order fix lives in there). - const textParts = parseAssistantContent(renderFullContent()); + // Cumulative yield so tool UI updates. orderAssistantContent + // puts search / code before text and generated images after. + // renderFullContent() preserves any Codex per-tab text from + // earlier _toolEvent frames; pinTextThoughtSignature attaches + // Gemini thoughtSignature onto the final text part. + const textParts = pinTextThoughtSignature( + parseAssistantContent(renderFullContent()), + ); yield { content: orderAssistantContent(textParts), metadata: { @@ -2204,11 +2529,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } totalChunks += 1; - // OpenRouter's free router (openrouter/free) picks a different - // underlying free model per request and reports it in every - // chunk's top-level `model` field. Latch the first non-empty - // value that differs from the requested checkpoint so the - // header chip can render "openrouter/free:". + // Latch the chunk's `model` field so the openrouter/free + // chip can show the chosen underlying model. if ( isExternalRequest && externalProvider?.providerType === "openrouter" && @@ -2227,30 +2549,37 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } const rawDelta = chunk.choices?.[0]?.delta?.content; - // Providers like Mistral's magistral return delta.content as an - // array of structured parts; normalize to text (with thinking - // parts re-wrapped as inline tags) so the rest of the - // accumulator stays string-based. + // Normalize structured delta.content (mistral magistral) to text. const delta = extractDeltaText(rawDelta); - // Kimi (kimi-k2.6, kimi-k2-thinking) and DeepSeek reasoner - // stream thinking via `delta.reasoning_content` as a plain - // string field — separate from `delta.content` which carries - // the answer. Wrap reasoning chunks inline as ... - // so parseAssistantContent treats them like any - // other reasoning. The close tag fires when the next chunk - // brings content, or when the stream ends. + // Latest Gemini text-part thoughtSignature for next-turn replay. + const deltaExtraContent = ( + chunk.choices?.[0]?.delta as + | { extra_content?: unknown } + | undefined + )?.extra_content; + if ( + deltaExtraContent && + typeof deltaExtraContent === "object" + ) { + const eGoogle = (deltaExtraContent as Record) + .google; + if (eGoogle && typeof eGoogle === "object") { + const sig = (eGoogle as Record) + .thought_signature; + if (typeof sig === "string" && sig) { + latestTextThoughtSignature = sig; + } + } + } + // Kimi / DeepSeek stream thinking via delta.reasoning_content. + // Wrap inline as ... for parseAssistantContent. const rawReasoning = ( chunk.choices?.[0]?.delta as | { reasoning_content?: unknown } | undefined )?.reasoning_content; - // OpenRouter uses a third reasoning shape: a structured - // `delta.reasoning_details` array of parts (each carrying - // `text`). The router emits this regardless of which - // underlying provider it picked, so we extract here and - // merge into the same ... wrap path used - // for Kimi / DeepSeek reasoning_content. See - // https://openrouter.ai/docs/guides/best-practices/reasoning-tokens + // OpenRouter ships reasoning as delta.reasoning_details[] + // regardless of underlying provider; merge into the same wrap path. const rawReasoningDetails = ( chunk.choices?.[0]?.delta as | { reasoning_details?: unknown } @@ -2268,6 +2597,138 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const reasoning = (typeof rawReasoning === "string" ? rawReasoning : "") + reasoningFromDetails; + // OpenAI delta.tool_calls: streams fragments by index; + // accumulate into one part. extra_content carries Gemini 3 + // thoughtSignature for next-turn replay. + const rawDeltaToolCalls = ( + chunk.choices?.[0]?.delta as + | { tool_calls?: unknown } + | undefined + )?.tool_calls; + if ( + Array.isArray(rawDeltaToolCalls) && + rawDeltaToolCalls.length > 0 + ) { + for (const tc of rawDeltaToolCalls) { + if (!tc || typeof tc !== "object") continue; + const call = tc as { + id?: string; + index?: number; + function?: { name?: string; arguments?: string }; + extra_content?: unknown; + }; + const idx = + typeof call.index === "number" ? call.index : undefined; + const stableId = call.id; + // Match an existing fragment by id first (canonical), + // then by index slot. Fall back to a freshly-minted + // tool_call_ id for streams that send neither. + let existing = stableId + ? toolCallParts.find((p) => p.toolCallId === stableId) + : undefined; + if (!existing && idx !== undefined) { + existing = toolCallParts.find( + (p) => + ( + p as ToolCallMessagePart & { _delta_index?: number } + )._delta_index === idx, + ); + } + const argsFragment = call.function?.arguments ?? ""; + if (existing) { + const prevName = existing.toolName ?? ""; + const nextName = call.function?.name ?? prevName; + const merged = + (existing.argsText ?? "") + argsFragment; + let parsedArgs: + ToolCallMessagePart["args"] = existing.args ?? {}; + if (merged) { + try { + parsedArgs = JSON.parse( + merged, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: merged, + } as ToolCallMessagePart["args"]; + } + } + const prevExtra = ( + existing as ToolCallMessagePart & { + extra_content?: unknown; + } + ).extra_content; + const updated: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + ...(existing as ToolCallMessagePart), + toolName: nextName, + argsText: merged, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : prevExtra !== undefined + ? { extra_content: prevExtra } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + const replaceIdx = toolCallParts.indexOf(existing); + if (replaceIdx >= 0) { + toolCallParts[replaceIdx] = updated; + } + } else { + const callId = + stableId || + `tool_call_${idx ?? toolCallParts.length}`; + const argsText = argsFragment; + let parsedArgs: ToolCallMessagePart["args"] = {}; + if (argsText) { + try { + parsedArgs = JSON.parse( + argsText, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: argsText, + } as ToolCallMessagePart["args"]; + } + } + const fresh: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + type: "tool-call" as const, + toolCallId: callId, + toolName: call.function?.name ?? "", + argsText, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + toolCallParts.push(fresh); + } + } + yield { + content: [ + ...toolCallParts, + ...pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ), + ], + metadata: { + timing: buildTiming( + streamStartTime, + totalChunks, + firstTokenTime, + ), + custom: { reasoningDuration }, + }, + }; + continue; + } if (!delta && !reasoning) { continue; } @@ -2293,14 +2754,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } cumulativeText += delta; } - // Mistral's magistral occasionally emits a trailing - // template-literal artifact (e.g. "${response}") at the end of - // an otherwise complete answer. It is never part of a real - // reply, so strip a trailing `${...}` token from external - // provider streams. The regex anchors to end-of-string and is - // idempotent — fragments mid-stream (e.g. "${re") leave the - // string untouched and only collapse once the closing brace - // arrives. Local-model output is left alone. + // Strip a trailing ${...} template-literal artifact from + // external streams (mistral magistral occasionally emits one). if (isExternalRequest) { cumulativeText = cumulativeText.replace( /\s*\$\{[^}]*\}\s*$/, @@ -2308,8 +2763,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ); } // renderFullContent() preserves any Codex per-tab text the - // fan-out branch accumulated into codexTabBuffers. - const parts = parseAssistantContent(renderFullContent()); + // fan-out branch accumulated into codexTabBuffers; + // pinTextThoughtSignature attaches Gemini thoughtSignature. + const parts = pinTextThoughtSignature( + parseAssistantContent(renderFullContent()), + ); if ( parts.some((part) => part.type === "reasoning") && @@ -2427,10 +2885,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { content: [ // renderFullContent() ensures the Codex per-tab text is in // the FINAL message too -- otherwise the synthesis delta on - // the regular content path would have erased it. The - // orderAssistantContent step keeps the rest of the message - // (tools / images) in the right slot relative to the text. - ...orderAssistantContent(parseAssistantContent(renderFullContent())), + // the regular content path would have erased it. + // pinTextThoughtSignature attaches Gemini thoughtSignature. + ...orderAssistantContent( + pinTextThoughtSignature(parseAssistantContent(renderFullContent())), + ), ...sourceParts, ...documentCitationParts, ], diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 1783a08281..c4a5f34273 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -728,7 +728,10 @@ export function ChatPage(): ReactElement { const reasoningCaps = getExternalReasoningCapabilities( provider?.providerType, selection.modelId, - { isReasoningProvider: provider?.isReasoningModel === true }, + { + isReasoningProvider: provider?.isReasoningModel === true, + baseUrl: provider?.baseUrl ?? null, + }, ); const state = useChatRuntimeStore.getState(); const preferredEffort = state.reasoningEffort; @@ -769,6 +772,8 @@ export function ChatPage(): ReactElement { : state.reasoningEffort; const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( provider?.providerType, + selection.modelId, + provider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( provider?.providerType, @@ -968,6 +973,7 @@ export function ChatPage(): ReactElement { { isReasoningProvider: selectedProvider?.isReasoningModel === true, + baseUrl: selectedProvider?.baseUrl ?? null, }, ); const preferredEffort = store.reasoningEffort; @@ -1009,6 +1015,8 @@ export function ChatPage(): ReactElement { store.setCheckpoint(value, null); const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( selectedProvider?.providerType, + selectedExternal?.modelId, + selectedProvider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( selectedProvider?.providerType, diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index e95159e14a..c3c7d34be5 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -44,6 +44,13 @@ export interface ExternalProviderConfig { updatedAt: number; } +// Gemini supports prompt caching, but the wire flow requires a +// separate POST to /v1beta/cachedContents to create the cache before +// the generateContent call can reference it; the boolean Studio +// currently emits on enable_prompt_caching is not enough on its own. +// Until that two-step orchestration ships we keep the picker off so +// the toggle does not silently no-op for Gemini users. See +// https://ai.google.dev/gemini-api/docs/caching. const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]); export function supportsProviderPromptCaching( diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 5adc01ea2b..f550279826 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -189,7 +189,27 @@ function _inferProviderFromOpenrouterId( */ export function providerSupportsBuiltinWebSearch( providerType: string | null | undefined, + modelId?: string | null | undefined, + baseUrl?: string | null | undefined, ): boolean { + // Gemini ships grounded search via `tools: [{googleSearch: {}}]` on + // every chat-capable model. Most image-tier ids (`-image`, + // `nano-banana`) reject text-tool wiring because the + // responseModalities path is mutually exclusive with text tools, but + // Google explicitly documents Search grounding on the Gemini 3 image + // family (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + // nano-banana-pro). Allow Search on those; hide on older image ids. + // Custom Gemini OpenAI-compat proxies (non-Google bases) skip the + // native translator on the backend, so native tool envelopes never + // reach them -- hide the pill there. + if (providerType === "gemini") { + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + const normalized = modelId?.trim().toLowerCase() ?? ""; + if (normalized && isGeminiImageModel(normalized)) { + return geminiImageModelAllowsGoogleSearch(normalized); + } + return true; + } return ( providerType === "openai" || providerType === "anthropic" || @@ -319,6 +339,20 @@ export function providerSupportsBuiltinCodeExecution( normalized.startsWith(prefix), ); } + if (providerType === "gemini") { + // Gemini's `tools: [{codeExecution: {}}]` is supported on every + // chat-capable model. Image-tier ids (`-image`, `nano-banana`) + // reject text-tool wiring because the inline-image path is + // mutually exclusive with codeExecution. Custom Gemini + // OpenAI-compat proxies skip the native translator on the + // backend, so native codeExecution envelopes do not reach them. + // Wire-up lives in `_stream_gemini` on the backend; output comes + // back inline as executableCode/codeExecutionResult parts. See + // https://ai.google.dev/gemini-api/docs/code-execution. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + if (isGeminiImageModel(normalized)) return false; + return normalized.startsWith("gemini-"); + } return false; } @@ -351,12 +385,75 @@ export function providerSupportsBuiltinImageGeneration( modelId: string | null | undefined, baseUrl?: string | null, ): boolean { - if (providerType !== "openai") return false; - if (!isOpenAICloudBaseUrl(baseUrl)) return false; const normalized = modelId?.trim().toLowerCase() ?? ""; if (!normalized) return false; - return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => - normalized.startsWith(prefix), + if (providerType === "openai") { + if (!isOpenAICloudBaseUrl(baseUrl)) return false; + return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); + } + if (providerType === "gemini") { + // Gemini's Nano Banana image-output ids carry either `-image` (e.g. + // `gemini-2.5-flash-image`, `gemini-3.1-flash-image-preview`) or the + // `nano-banana` alias (`nano-banana-pro-preview`). The backend flips + // generationConfig.responseModalities to ["TEXT", "IMAGE"] when one + // is picked, and translates inlineData parts into the same image_b64 + // tool_end envelope the OpenAI path emits so the chat UI renders the + // picture inline. Custom Gemini OpenAI-compat proxies skip the + // native translator on the backend, so hide the image pill there. + // See https://ai.google.dev/gemini-api/docs/image-generation. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + return normalized.includes("-image") || normalized.includes("nano-banana"); + } + return false; +} + +/** + * Whether `modelId` is a Gemini image-output id (Nano Banana family). + * Mirrors the backend's `is_image_picker_model` guard so the frontend + * hides text-only tool pills (web_search, code_execution) for these. + */ +function isGeminiImageModel(modelId: string): boolean { + const m = modelId.toLowerCase(); + return m.includes("-image") || m.includes("nano-banana"); +} + +/** + * Whether the saved Gemini connection points at a custom + * OpenAI-compatible gateway (any non-Google host). The backend + * `_is_openai_compatible` mirrors this to route those connections + * through `/chat/completions` instead of the native translator, so + * native Gemini tool envelopes (googleSearch, codeExecution, + * responseModalities) never reach them. Hide the corresponding + * Studio pills here so the request, builder, and UI agree. + */ +export function isGeminiCustomOpenAICompatBase( + baseUrl: string | null | undefined, +): boolean { + if (!baseUrl) return false; + try { + const host = new URL(baseUrl).hostname.toLowerCase(); + return host.length > 0 && host !== "generativelanguage.googleapis.com"; + } catch { + return false; + } +} + +/** + * Whether the given Gemini image model supports `tools: [{googleSearch: {}}]`. + * Google documents Search grounding on the Gemini 3 image family + * (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + * "Nano Banana Pro"); older image ids (gemini-2.5-flash-image) reject + * it with "Search as tool is not enabled for this model". + */ +function geminiImageModelAllowsGoogleSearch(modelId: string): boolean { + const m = modelId.toLowerCase(); + return ( + m.startsWith("gemini-3-pro-image") || + m.startsWith("gemini-3.1-flash-image") || + m.startsWith("nano-banana-pro") || + m.startsWith("nano-banana-2") ); } @@ -431,7 +528,20 @@ const PROVIDER_CAPABILITIES: Record = { presencePenalty: false, }, mistral: OPENAI_COMPAT_BASE, - gemini: OPENAI_COMPAT_BASE, + // Gemini's native generationConfig accepts temperature, topP, topK and + // presencePenalty (plus a separate frequencyPenalty we do not surface + // today). minP and repetitionPenalty are not part of the contract -- + // see https://ai.google.dev/api/rest/v1beta/GenerationConfig. Backend + // request shaping lives in _stream_gemini in + // studio/backend/core/inference/external_provider.py. + gemini: { + temperature: true, + topP: true, + topK: true, + minP: false, + repetitionPenalty: false, + presencePenalty: true, + }, // Kimi k2.5/k2.6 are reasoning-class — the API locks temperature and // top_p to fixed defaults and 400s on any other value: // "invalid temperature: only 1 is allowed for this model". @@ -643,6 +753,119 @@ function resolveKimiReasoningCapabilities(modelId: string): ExternalReasoningCap return withEnableThinkingStyle(); } +// Gemini's thinking ladder. +// - Gemini 3.x (3 / 3.1 / 3.5, Pro + Flash + Flash-Lite) and the +// gemini-pro-latest / gemini-flash-latest aliases use the new +// `thinkingConfig.thinkingLevel` string field (LOW/MEDIUM/HIGH/ +// MINIMAL). Pro tier rejects MINIMAL. +// - Gemini 2.5 Flash + 2.5 Pro stay on the integer +// `thinkingConfig.thinkingBudget` (0=off on Flash, -1=dynamic, +// N>0=cap; Pro rejects 0). +// - 2.5 Flash-Lite: no native thinking surfaced; leave it off. +// - Image-tier ids (`*-image*`, `nano-banana-pro-preview`): image +// generation path -- no reasoning controls. +const GEMINI3_PRO_PREFIXES = [ + "gemini-3.5-pro", + "gemini-3.1-pro", + "gemini-3-pro-preview", + "gemini-pro-latest", +]; +const GEMINI3_FLASH_PREFIXES = [ + "gemini-3.5-flash", + "gemini-3.1-flash", + "gemini-3-flash", + "gemini-flash-latest", + "gemini-flash-lite-latest", +]; +const GEMINI25_PRO_PREFIXES = [ + "gemini-2.5-pro", +]; +const GEMINI25_FLASH_PREFIXES = [ + "gemini-2.5-flash", +]; +const GEMINI_IMAGE_HINTS = [ + "-image", + "nano-banana", +]; +function resolveGeminiReasoningCapabilities( + modelId: string, +): ExternalReasoningCapabilities { + const m = modelId.toLowerCase(); + if (GEMINI_IMAGE_HINTS.some((h) => m.includes(h))) { + // Image generation; no thinking knob. + return withEnableThinkingStyle(); + } + // Gemini 2.5 Flash-Lite supports `thinkingBudget` with `0` = off and + // a positive range starting at 512 (the backend maps "minimal" to + // that floor at external_provider._stream_gemini). Check this branch + // BEFORE the broader `gemini-2.5-flash` prefix. + // https://ai.google.dev/gemini-api/docs/thinking + if (m.startsWith("gemini-2.5-flash-lite")) { + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "minimal", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + if (GEMINI3_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3.x Pro: thinkingLevel supports low/medium/high per + // https://ai.google.dev/gemini-api/docs/thinking and + // https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro. + // Cannot fully disable thinking; "minimal" is rejected on Pro. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high"] as const, + }); + } + if (GEMINI3_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3 Flash: thinkingLevel minimal/low/medium/high. Minimal + // is the closest to "off" Google offers on Gemini 3. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: [ + "minimal", + "low", + "medium", + "high", + ] as const, + }); + } + if (GEMINI25_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Pro: thinkingBudget cannot be 0 (API rejects with + // "only works in thinking mode"); backend coerces to a small + // positive budget. The picker still hides the off switch. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high", "max"] as const, + }); + } + if (GEMINI25_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Flash: thinkingBudget supports 0 = off cleanly. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + return withEnableThinkingStyle(); +} + function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoningCapabilities { if (modelId === "magistral-medium-latest") { return withReasoningEffortStyle({ @@ -665,6 +888,8 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning export interface ExternalReasoningResolveOptions { /** vLLM connection flagged as a reasoning model in provider config. */ isReasoningProvider?: boolean; + /** Provider base URL; used to detect custom Gemini OAI-compat gateways. */ + baseUrl?: string | null; } // vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle. @@ -740,6 +965,16 @@ export function getExternalReasoningCapabilities( } if (isKimiProvider) return resolveKimiReasoningCapabilities(modelForMatching); if (isMistralProvider) return resolveMistralReasoningCapabilities(modelForMatching); + if (normalizedProvider === "gemini") { + // Custom Gemini OAI-compat gateways (LiteLLM, proxies) route + // through /chat/completions which drops the Gemini-native + // thinkingConfig payload. Hide the native thinking ladder so the + // UI does not advertise a control the backend cannot honor. + if (isGeminiCustomOpenAICompatBase(options?.baseUrl)) { + return withEnableThinkingStyle(); + } + return resolveGeminiReasoningCapabilities(modelForMatching); + } if (!isOpenAIProvider && !isAnthropicProvider) { return withEnableThinkingStyle(); } diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index ba4a897f63..a811125a20 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -396,6 +396,7 @@ export function SharedComposer({ { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; @@ -449,16 +450,36 @@ export function SharedComposer({ const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( selectedExternalProvider?.providerType, ); - const searchDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); - const codeDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); - // Images pill is only ever lit on OpenAI cloud's Responses-API models. - // No local tool runtime fallback because the only image-generation - // server tool we wire today is OpenAI's; local models cannot dispatch - // it. Hidden entirely when the active model does not advertise it so - // the pill row stays compact for providers without the capability. + // Gemini rejects codeExecution alongside image modalities. Search is + // blocked on older Gemini image ids but allowed on Gemini 3 image + // models -- supportsBuiltinWebSearch already encodes the per-model + // allowance, so we only disable Code unconditionally in Gemini + // image mode. + const isExternalGemini = selectedExternalProvider?.providerType === "gemini"; const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration; + const imageModeDisablesCode = + isExternalGemini && imageToolsEnabled && !imageDisabled; + // Image-tier Gemini models always reject codeExecution and reject + // web_search on older ids (Gemini 3.x Pro/Flash allow it -- encoded + // in supportsBuiltinWebSearch). Don't let the local `supportsTools` + // runtime flag re-enable a pill the Gemini backend will silently + // drop. Detect "external provider is Gemini AND model is image-tier" + // and gate strictly on the provider builtin support. + const isGeminiImageTier = + isExternalGemini && supportsBuiltinImageGeneration; + const searchDisabled = + !modelLoaded || + (isGeminiImageTier + ? !supportsBuiltinWebSearch + : !(supportsTools || supportsBuiltinWebSearch)); + const codeDisabled = + !modelLoaded || + (isGeminiImageTier + ? true + : !(supportsTools || supportsBuiltinCodeExecution)) || + imageModeDisablesCode; + // Images pill is only ever lit on OpenAI cloud's Responses-API models + // and Gemini Nano Banana family. No local tool runtime fallback. const showImagePill = supportsBuiltinImageGeneration; // Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209). const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch; diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 5238875b71..d313b43438 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -219,9 +219,34 @@ export type OpenAIMessageContentPart = export type OpenAIMessageContent = string | OpenAIMessageContentPart[]; +/** + * OpenAI Chat Completions tool_call shape. Assistant turns echo back + * function/tool calls as `tool_calls`; the matching tool result rides + * on a separate `role="tool"` message keyed by `tool_call_id`. + * `extra_content.google.thought_signature` is the Gemini-specific + * round-trip field the backend translator both emits (on `delta. + * tool_calls`) and consumes (when rebuilding the native functionCall + * part on the next turn). + */ +export interface OpenAIToolCallPart { + id?: string; + type?: "function"; + function?: { + name?: string; + arguments?: string; + }; + extra_content?: unknown; +} + export interface OpenAIChatMessage { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + /** Assistant tool-call deltas, when the turn invoked a function tool. */ + tool_calls?: OpenAIToolCallPart[]; + /** `role="tool"` only: id matching `assistant.tool_calls[].id`. */ + tool_call_id?: string; + /** `role="tool"` only: name of the function that produced the result. */ + name?: string; } export interface OpenAIChatCompletionsRequest { @@ -262,7 +287,14 @@ export interface OpenAIChatCompletionsRequest { external_model?: string; encrypted_api_key?: string; provider_base_url?: string | null; - enable_prompt_caching?: boolean | null; + /** + * Boolean toggle for OpenAI/Anthropic ephemeral cache_control. For + * Gemini the backend also accepts the cached-content resource name + * (`cachedContents/...`) as a string, which is forwarded as + * `generationConfig.cachedContent` on the native streamGenerateContent + * request. + */ + enable_prompt_caching?: boolean | string | null; /** * OpenAI shell-tool container id captured from the prior response in * this chat thread. When set and the Code pill is on, the backend @@ -292,7 +324,20 @@ export interface OpenAIChatCompletionsRequest { export interface OpenAIChatDelta { role?: string; - content?: string; + content?: string | null; + /** + * Streamed assistant tool calls. The Gemini and OpenAI Responses + * translators emit incremental `tool_calls` deltas (function name + + * arguments fragments) so the chat-adapter can render tool cards as + * they arrive. + */ + tool_calls?: OpenAIToolCallPart[]; + /** + * Provider-specific passthrough. Gemini ships `thoughtSignature`, + * citations, `native_part`, etc., here so the round-trip can replay + * them on follow-up turns without bleeding into other providers. + */ + extra_content?: Record; } export interface OpenAIChatChunkChoice { diff --git a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx index f30af7fcf9..ed590f226e 100644 --- a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx +++ b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx @@ -5,6 +5,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { getAuthToken } from "@/features/auth"; +import { useT } from "@/i18n"; import { toastError, toastSuccess } from "@/shared/toast"; import { Camera } from "lucide-react"; import { useMemo, useRef, useState } from "react"; @@ -37,6 +38,7 @@ function readPersistedProfile(): { displayName: string; avatarDataUrl: string | } export function ProfilePersonalizationPanel() { + const t = useT(); const displayName = useUserProfileStore((s) => s.displayName); const avatarDataUrl = useUserProfileStore((s) => s.avatarDataUrl); const setDisplayName = useUserProfileStore((s) => s.setDisplayName); @@ -60,11 +62,11 @@ export function ProfilePersonalizationPanel() { setDisplayName(trimmed); const persisted = readPersistedProfile(); if (persisted && persisted.displayName === trimmed) { - toastSuccess("Profile name saved"); + toastSuccess(t("settings.profile.nameSaved")); } else { toastError( - "Could not persist profile name", - "Name updated for this session, but may not persist after reload.", + t("settings.profile.namePersistErrorTitle"), + t("settings.profile.namePersistErrorDescription"), ); } } @@ -78,17 +80,18 @@ export function ProfilePersonalizationPanel() { setAvatarDataUrl(dataUrl); const persisted = readPersistedProfile(); if (persisted && persisted.avatarDataUrl === dataUrl) { - toastSuccess("Profile photo updated"); + toastSuccess(t("settings.profile.photoUpdated")); } else { toastError( - "Could not persist profile photo", - "Photo updated for this session, but may not persist after reload.", + t("settings.profile.photoPersistErrorTitle"), + t("settings.profile.photoPersistErrorDescription"), ); } } catch (e) { - const message = e instanceof Error ? e.message : "Could not use this image."; + const message = + e instanceof Error ? e.message : t("settings.profile.imageUseError"); setImageError(message); - toastError("Could not update profile photo", message); + toastError(t("settings.profile.photoUpdateErrorTitle"), message); } }; @@ -115,7 +118,7 @@ export function ProfilePersonalizationPanel() { type="button" onClick={() => fileInputRef.current?.click()} className="absolute right-0 bottom-0 -translate-x-[15.625%] -translate-y-[15.625%] flex size-8 items-center justify-center rounded-full border border-border bg-background text-foreground shadow-sm transition-colors hover:bg-muted focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background" - aria-label="Change profile picture" + aria-label={t("settings.profile.changePicture")} > @@ -123,7 +126,7 @@ export function ProfilePersonalizationPanel() {
diff --git a/studio/frontend/src/features/settings/components/api-key-row.tsx b/studio/frontend/src/features/settings/components/api-key-row.tsx index ced6de17d1..9a6e1ee207 100644 --- a/studio/frontend/src/features/settings/components/api-key-row.tsx +++ b/studio/frontend/src/features/settings/components/api-key-row.tsx @@ -14,30 +14,39 @@ import { MoreHorizontalIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { useT } from "@/i18n"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import type { ApiKey } from "../api/api-keys"; -function relative(iso: string | null): string { - if (!iso) return "never"; +type SettingsT = ReturnType; + +function relative(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = Date.now() - new Date(iso).getTime(); const days = Math.floor(diff / 86400000); if (days < 1) { const hours = Math.floor(diff / 3600000); - if (hours < 1) return "just now"; - return `${hours}h ago`; + if (hours < 1) return t("settings.apiKeys.relativeJustNow"); + return t("settings.apiKeys.relativeHoursAgo", { count: hours }); } - if (days < 30) return `${days}d ago`; - if (days < 365) return `${Math.floor(days / 30)}mo ago`; - return `${Math.floor(days / 365)}y ago`; + if (days < 30) return t("settings.apiKeys.relativeDaysAgo", { count: days }); + if (days < 365) { + return t("settings.apiKeys.relativeMonthsAgo", { + count: Math.floor(days / 30), + }); + } + return t("settings.apiKeys.relativeYearsAgo", { + count: Math.floor(days / 365), + }); } -function expiresText(iso: string | null): string { - if (!iso) return "never"; +function expiresText(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = new Date(iso).getTime() - Date.now(); - if (diff < 0) return "expired"; + if (diff < 0) return t("settings.apiKeys.expired"); const days = Math.floor(diff / 86400000); - if (days < 1) return "today"; - return `in ${days}d`; + if (days < 1) return t("settings.apiKeys.today"); + return t("settings.apiKeys.inDays", { count: days }); } export function ApiKeyRow({ @@ -47,6 +56,7 @@ export function ApiKeyRow({ apiKey: ApiKey; onRevoke: (key: ApiKey) => void; }) { + const t = useT(); const prefix = `sk-unsloth-${apiKey.key_prefix}…`; return (
@@ -64,11 +74,23 @@ export function ApiKeyRow({
- Created {relative(apiKey.created_at)} + + {t("settings.apiKeys.created", { + value: relative(apiKey.created_at, t), + })} + · - Used {relative(apiKey.last_used_at)} + + {t("settings.apiKeys.used", { + value: relative(apiKey.last_used_at, t), + })} + · - Expires {expiresText(apiKey.expires_at)} + + {t("settings.apiKeys.expires", { + value: expiresText(apiKey.expires_at, t), + })} +
@@ -77,7 +99,7 @@ export function ApiKeyRow({ variant="ghost" size="sm" className="size-7 p-0 opacity-0 transition-opacity group-hover:opacity-100 data-[state=open]:opacity-100 max-sm:!opacity-100 max-sm:size-9" - aria-label={`Actions for ${apiKey.name}`} + aria-label={t("settings.apiKeys.actionsFor", { name: apiKey.name })} > @@ -85,14 +107,14 @@ export function ApiKeyRow({ { await copyToClipboard(prefix); }}> - Copy prefix + {t("settings.apiKeys.copyPrefix")} onRevoke(apiKey)} className="text-destructive focus:text-destructive" > - Revoke token + {t("settings.apiKeys.revokeToken")} diff --git a/studio/frontend/src/features/settings/components/create-key-form.tsx b/studio/frontend/src/features/settings/components/create-key-form.tsx index a0f2d7f82d..93802bfad8 100644 --- a/studio/frontend/src/features/settings/components/create-key-form.tsx +++ b/studio/frontend/src/features/settings/components/create-key-form.tsx @@ -3,6 +3,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { useState } from "react"; import { createApiKey } from "../api/api-keys"; @@ -21,6 +22,7 @@ export function CreateKeyForm({ onCreated: (rawKey: string) => void; onError: (message: string) => void; }) { + const t = useT(); const [name, setName] = useState(""); const [expiry, setExpiry] = useState(null); const [loading, setLoading] = useState(false); @@ -33,8 +35,11 @@ export function CreateKeyForm({ const result = await createApiKey(name.trim(), expiry); onCreated(result.key); setName(""); - } catch (err) { - onError(err instanceof Error ? err.message : "Couldn't create access token."); + } catch { + // API helpers in ../api/api-keys.ts throw generic English Error + // messages; always use the translated message so zh-CN users do not + // see English text bleed through from internal exceptions. + onError(t("settings.apiKeys.createError")); } finally { setLoading(false); } @@ -49,9 +54,9 @@ export function CreateKeyForm({ setName(e.target.value)} - placeholder="Token name (e.g. production)" + placeholder={t("settings.apiKeys.tokenNamePlaceholder")} className="h-8 min-w-[180px] flex-1 text-sm" - aria-label="New access token name" + aria-label={t("settings.apiKeys.newAccessTokenName")} />
{EXPIRY_PRESETS.map((p) => { @@ -69,13 +74,15 @@ export function CreateKeyForm({ : "text-muted-foreground hover:text-foreground", )} > - {p.label} + {p.value === null ? t("settings.apiKeys.never") : p.label} ); })}
diff --git a/studio/frontend/src/features/settings/components/key-reveal-card.tsx b/studio/frontend/src/features/settings/components/key-reveal-card.tsx index 2b589e88fe..bdcb861c1d 100644 --- a/studio/frontend/src/features/settings/components/key-reveal-card.tsx +++ b/studio/frontend/src/features/settings/components/key-reveal-card.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Button } from "@/components/ui/button"; +import { useT } from "@/i18n"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; @@ -15,6 +16,7 @@ export function KeyRevealCard({ rawKey: string; onDone: () => void; }) { + const t = useT(); const [copied, setCopied] = useState(false); const handleCopy = async () => { @@ -32,7 +34,7 @@ export function KeyRevealCard({ className="size-3.5 text-emerald-600 dark:text-emerald-500" /> - New access token created + {t("settings.apiKeys.newTokenCreated")}

- Copy now — this won't be shown again. + {t("settings.apiKeys.copyNow")}

diff --git a/studio/frontend/src/features/settings/components/language-select.tsx b/studio/frontend/src/features/settings/components/language-select.tsx new file mode 100644 index 0000000000..9d30e06147 --- /dev/null +++ b/studio/frontend/src/features/settings/components/language-select.tsx @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { + LOCALES, + isSupportedLocale, + setLocale, + useT, + useLocale, +} from "@/i18n"; + +export function LanguageSelect() { + const t = useT(); + const locale = useLocale(); + + return ( + + ); +} diff --git a/studio/frontend/src/features/settings/components/theme-segmented.tsx b/studio/frontend/src/features/settings/components/theme-segmented.tsx index 1061995346..36e7062d8f 100644 --- a/studio/frontend/src/features/settings/components/theme-segmented.tsx +++ b/studio/frontend/src/features/settings/components/theme-segmented.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { cn } from "@/lib/utils"; +import { useT, type TranslationKey } from "@/i18n"; import { LaptopIcon, Moon02Icon, @@ -11,13 +12,18 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { motion, useReducedMotion } from "motion/react"; import { useTheme, type Theme } from "../stores/theme-store"; -const OPTIONS: { value: Theme; label: string; icon: typeof Sun02Icon }[] = [ - { value: "light", label: "Light", icon: Sun02Icon }, - { value: "dark", label: "Dark", icon: Moon02Icon }, - { value: "system", label: "System", icon: LaptopIcon }, +const OPTIONS: { + value: Theme; + labelKey: TranslationKey; + icon: typeof Sun02Icon; +}[] = [ + { value: "light", labelKey: "settings.appearance.theme.light", icon: Sun02Icon }, + { value: "dark", labelKey: "settings.appearance.theme.dark", icon: Moon02Icon }, + { value: "system", labelKey: "settings.appearance.theme.system", icon: LaptopIcon }, ]; export function ThemeSegmented() { + const t = useT(); const { theme, setTheme } = useTheme(); const reduced = useReducedMotion(); return ( @@ -49,7 +55,7 @@ export function ThemeSegmented() { /> )} - {opt.label} + {t(opt.labelKey)} ); })} diff --git a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx index e4cdccd2d7..90f66483b7 100644 --- a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx +++ b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -29,10 +30,13 @@ export type UpdateInstallSource = | "unknown"; type UpdateInstallSourceState = UpdateInstallSource | "loading"; -function getStudioUpdateInstructionLine(shell: UpdateShell): string { +function getStudioUpdateInstructionLine( + shell: UpdateShell, + t: ReturnType, +): string { return shell === "windows" - ? "Open PowerShell and run:" - : "Open Terminal and run:"; + ? t("settings.about.update.openPowerShell") + : t("settings.about.update.openTerminal"); } function isLocalInstallSource( @@ -59,6 +63,7 @@ function CopyableCommand({ command: string; copyLabel: string; }): ReactElement { + const t = useT(); const [copied, setCopied] = useState(false); const timerRef = useRef | null>(null); @@ -89,14 +94,26 @@ function CopyableCommand({ value={command} className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none" title={command} - aria-label={`${copyLabel} text`} + aria-label={t("settings.about.update.commandText", { + label: copyLabel, + })} /> ); })} @@ -125,20 +131,20 @@ export function UsageExamples() { type="button" onClick={handleCopy} className="flex items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Copy snippet" + aria-label={t("settings.apiKeys.copySnippet")} > - {copied ? "Copied" : "Copy"} + {copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")} 
           {snippets[lang]}
- Setup docs: + {t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( s.open); const activeTab = useSettingsDialogStore((s) => s.activeTab); const setActiveTab = useSettingsDialogStore((s) => s.setActiveTab); @@ -117,9 +133,9 @@ export function SettingsDialog() { "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", )} > - Settings + {t("settings.dialog.title")} - Manage your Unsloth Studio preferences. + {t("settings.dialog.description")}