diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py new file mode 100644 index 0000000000..f5b67eef70 --- /dev/null +++ b/studio/backend/core/inference/external_provider.py @@ -0,0 +1,1238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Async HTTP client for proxying chat completions to external LLM providers. + +Most registry providers expose OpenAI-compatible /v1/chat/completions endpoints; +Anthropic uses native Messages API with translation in this client. +""" + +import json as _json +import re +from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional + +import httpx +import structlog + +# Use structlog so INFO-level diagnostics actually surface in the +# studio backend's JSON log stream. The stdlib root logger defaults to +# WARNING and is not configured with handlers, so plain +# `logging.getLogger(__name__).info(...)` was being silently dropped — +# only WARNING/ERROR made it through (because they bypassed the root +# level threshold via uvicorn's stderr capture). All existing call +# sites use printf-style positional args, which structlog accepts. +logger = structlog.get_logger(__name__) + +# Claude 4.7 (Opus/Sonnet/Haiku) deprecated top_k and returns 400 +# "top_k is deprecated for this model" when it is set. 3.x and 4.5/4.6 +# still accept it. Match the 4-7 line specifically so we keep the knob +# live on every other Claude generation. +_ANTHROPIC_TOP_K_DEPRECATED = re.compile(r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)") + + +class _AnthropicThinkingSpec(NamedTuple): + prefixes: tuple[str, ...] + kind: Literal["adaptive", "manual"] + efforts: tuple[str, ...] + + +_ANTHROPIC_THINKING_SPECS = ( + _AnthropicThinkingSpec( + prefixes = ("claude-opus-4-7",), + kind = "adaptive", + efforts = ("none", "low", "medium", "high", "xhigh"), + ), + _AnthropicThinkingSpec( + prefixes = ("claude-opus-4-6", "claude-sonnet-4-6"), + kind = "adaptive", + efforts = ("none", "low", "medium", "high", "xhigh", "max"), + ), + _AnthropicThinkingSpec( + prefixes = ("claude-opus-4-5", "claude-sonnet-4-5", "claude-haiku-4-5"), + kind = "manual", + efforts = ("none", "low", "medium", "high"), + ), +) + + +def _anthropic_thinking_spec(model: str) -> Optional[_AnthropicThinkingSpec]: + for spec in _ANTHROPIC_THINKING_SPECS: + if model.startswith(spec.prefixes): + return spec + return None + + +class _MistralThinkingSpec(NamedTuple): + models: tuple[str, ...] + style: Literal["prompt_mode", "reasoning_effort", "disabled"] + efforts: tuple[str, ...] = () + + +_MISTRAL_THINKING_SPECS = ( + _MistralThinkingSpec( + models = ("magistral-medium-latest",), + style = "prompt_mode", + ), + _MistralThinkingSpec( + models = ("mistral-small-latest", "mistral-vibe-cli-latest"), + style = "reasoning_effort", + efforts = ("none", "high"), + ), +) + +_OPENROUTER_MANDATORY_REASONING_MODELS = frozenset( + { + "~google/gemini-pro-latest", + "baidu/cobuddy:free", + "inclusionai/ring-2.6-1t:free", + "deepseek/deepseek-r1", + } +) + + +def _mistral_thinking_spec(model: str) -> _MistralThinkingSpec: + for spec in _MISTRAL_THINKING_SPECS: + if model in spec.models: + return spec + return _MistralThinkingSpec(models = (), style = "disabled") + + +def _apply_mistral_reasoning_controls( + body: dict[str, Any], + model: str, + enable_thinking: Optional[bool], + reasoning_effort: Optional[str], +) -> None: + """ + Translate generic reasoning controls into Mistral's model-specific shape. + + Current contract: + - magistral-medium-latest: baseline (no extra field) or + `prompt_mode="reasoning"` for the explicit reasoning mode. + - mistral-small-latest / mistral-vibe-cli-latest: + `reasoning_effort` in {"none", "high"}. + - all other tested Mistral models: no reasoning/thinking params. + """ + model_for_matching = model.rsplit("/", 1)[-1].strip().lower() + spec = _mistral_thinking_spec(model_for_matching) + body.pop("prompt_mode", None) + body.pop("reasoning_effort", None) + + if spec.style == "prompt_mode": + # Magistral baseline is already reasoning-capable. The explicit + # prompt_mode path is only used for the "high" UI selection. + if enable_thinking is True or reasoning_effort == "high": + body["prompt_mode"] = "reasoning" + return + + if spec.style == "reasoning_effort": + if reasoning_effort in spec.efforts: + body["reasoning_effort"] = reasoning_effort + elif enable_thinking is False: + body["reasoning_effort"] = "none" + elif enable_thinking is True: + body["reasoning_effort"] = "high" + + +# Shared client reused across all requests for HTTP connection pooling. +# Auth headers and timeouts are passed per-request, so a single client +# handles every provider without storing credentials. +_http_client = httpx.AsyncClient() + + +class ExternalProviderClient: + """Async proxy for OpenAI-compatible external LLM APIs.""" + + def __init__( + self, + provider_type: str, + base_url: str, + api_key: str, + timeout: float = 120.0, + ): + self.provider_type = provider_type + self.base_url = base_url.rstrip("/") + 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. + self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None) + + def _auth_headers(self) -> dict[str, str]: + """Build authentication headers using the provider's registry config.""" + from core.inference.providers import get_provider_info + + provider_info = get_provider_info(self.provider_type) or {} + auth_header = provider_info.get("auth_header", "Authorization") + auth_prefix = provider_info.get("auth_prefix", "Bearer ") + + headers = { + "Content-Type": "application/json", + auth_header: f"{auth_prefix}{self.api_key}", + } + # Merge any provider-specific extra headers (e.g. anthropic-version, OpenRouter attribution) + headers.update(provider_info.get("extra_headers", {})) + return headers + + def _is_openai_compatible(self) -> bool: + """Return False for providers that need request/response translation (e.g. Anthropic).""" + from core.inference.providers import get_provider_info + + info = get_provider_info(self.provider_type) or {} + return info.get("openai_compatible", True) + + async def stream_chat_completion( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float = 0.7, + top_p: float = 0.95, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + top_k: Optional[int] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + stream: bool = True, + ) -> AsyncGenerator[str, None]: + """ + Yield OpenAI-format SSE lines from the external provider. + + For OpenAI-compatible providers, lines are forwarded verbatim. + For Anthropic, the native Messages API SSE is translated to OpenAI format. + + ``top_k`` and ``presence_penalty`` are forwarded only when the caller + supplies a value the provider accepts — the frontend's + provider-capability map already filters these per provider, so we + treat them as opt-in here. + """ + if not self._is_openai_compatible(): + async for line in self._stream_anthropic( + messages, + model, + temperature, + top_p, + max_tokens, + top_k, + enable_thinking, + reasoning_effort, + ): + yield line + return + + # OpenAI moved their flagship models (gpt-5.x) off /v1/chat/completions + # — those endpoints return 404 with "This is not a chat model" for the + # new families. Route all OpenAI traffic through /v1/responses instead; + # we translate the Responses SSE back into Chat Completions chunks so + # the frontend stays endpoint-agnostic. + if self.provider_type == "openai": + async for line in self._stream_openai_responses( + messages, + model, + temperature, + top_p, + max_tokens, + enable_thinking, + reasoning_effort, + ): + yield line + return + + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": stream, + "temperature": temperature, + "top_p": top_p, + "presence_penalty": presence_penalty, + } + if max_tokens is not None: + # OpenAI newer models (gpt-4o, gpt-5.x) reject max_tokens + if self.provider_type == "openai": + body["max_completion_tokens"] = max_tokens + 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. + 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. + 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. + pass + elif enable_thinking: + body["thinking"] = {"type": "enabled", "keep": "all"} + else: + body["thinking"] = {"type": "disabled"} + elif self.provider_type == "mistral": + _apply_mistral_reasoning_controls( + body, model, enable_thinking, reasoning_effort + ) + + # 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". + if self.provider_type == "openrouter": + normalized_or_model = model.strip().lower() + if reasoning_effort in ("low", "medium", "high"): + body["reasoning"] = {"effort": reasoning_effort} + elif enable_thinking is True: + body["reasoning"] = {"enabled": True} + elif enable_thinking is False: + if normalized_or_model in _OPENROUTER_MANDATORY_REASONING_MODELS: + body.pop("reasoning", None) + else: + body["reasoning"] = {"enabled": False} + + url = f"{self.base_url}/chat/completions" + logger.info( + "Proxying chat completion to %s (provider=%s, model=%s)", + url, + self.provider_type, + model, + ) + + 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( + "External provider returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + 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. + 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. + event_counts: dict[str, int] = {} + chosen_model: Optional[str] = None + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line.strip(): + continue + if line.startswith("data:"): + data_str = line[len("data:") :].strip() + if data_str == "[DONE]": + event_counts["done"] = event_counts.get("done", 0) + 1 + elif data_str: + try: + parsed = _json.loads(data_str) + except Exception: + parsed = None + if isinstance(parsed, dict): + # Mid-stream provider error event. OpenRouter + # in particular returns 200 then surfaces the + # actual failure as an SSE error event. + if "error" in parsed: + event_counts["error"] = ( + event_counts.get("error", 0) + 1 + ) + logger.warning( + "%s SSE error event: %s", + self.provider_type, + parsed.get("error"), + ) + else: + event_counts["delta"] = ( + event_counts.get("delta", 0) + 1 + ) + # OpenRouter (and most OAI-compat providers) + # report the underlying model that handled + # the request in every chunk's `model` field. + # Latch the first non-empty value so the + # router-picked model surfaces in logs and + # is available to the proxy caller. + if chosen_model is None and isinstance( + parsed.get("model"), str + ): + chosen_model = parsed["model"] + yield line + except GeneratorExit: + await response.aclose() # set PoolByteStream._closed=True FIRST + await lines_gen.aclose() # now safe — aclose() is a no-op + raise + finally: + logger.info( + "%s stream complete (model=%s, chosen=%s, events=%s)", + self.provider_type, + model, + chosen_model, + event_counts, + ) + await response.aclose() + await lines_gen.aclose() + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + 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) + 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) + yield _error_sse_line( + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, + ) + + async def _stream_anthropic( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float, + top_p: float, + max_tokens: Optional[int], + top_k: Optional[int] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + ) -> AsyncGenerator[str, None]: + """ + Call the Anthropic Messages API and translate its SSE to OpenAI format. + + Anthropic SSE event types: + content_block_delta → OpenAI chunk with delta.content + message_delta → OpenAI chunk with finish_reason + message_stop → data: [DONE] + (all others skipped) + """ + import json as _json + + # Extract system prompt and translate image_url parts to Anthropic format + system: Optional[str] = None + filtered: list[dict[str, Any]] = [] + for msg in messages: + if msg.get("role") == "system": + content = msg.get("content", "") + system = ( + content + if isinstance(content, str) + else "\n".join( + p["text"] for p in content if p.get("type") == "text" + ) + ) + continue + + content = msg.get("content") + if isinstance(content, list): + # Translate OpenAI image_url parts → Anthropic native image format + anthropic_parts: list[dict[str, Any]] = [] + for part in content: + if part.get("type") == "text": + anthropic_parts.append({"type": "text", "text": part["text"]}) + elif part.get("type") == "image_url": + url = part.get("image_url", {}).get("url", "") + if url.startswith("data:"): + # data:image/png;base64, → split header and data + header, _, b64data = url.partition(",") + media_type = ( + header.split(";")[0].replace("data:", "") + or "image/jpeg" + ) + anthropic_parts.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": media_type, + "data": b64data, + }, + } + ) + else: + # Remote URL — Anthropic supports url source type natively. + # See: https://docs.anthropic.com/en/docs/build-with-claude/vision#url-based-images + anthropic_parts.append( + { + "type": "image", + "source": { + "type": "url", + "url": url, + }, + } + ) + filtered.append({"role": msg["role"], "content": anthropic_parts}) + else: + filtered.append(msg) + + body: dict[str, Any] = { + "model": model, + "messages": filtered, + "max_tokens": max_tokens or 1024, # required by Anthropic + "temperature": temperature, + "stream": True, + } + # top_k is deprecated on Claude 4.7 (Opus/Sonnet/Haiku) — the API + # returns 400 "top_k is deprecated for this model" when it is set. + # 3.x and 4.5/4.6 still accept it, so gate strictly on the 4.7 ids. + if ( + top_k is not None + and top_k > 0 + and not _ANTHROPIC_TOP_K_DEPRECATED.match(model) + ): + body["top_k"] = top_k + if system: + body["system"] = system + thinking_spec = _anthropic_thinking_spec(model) + allowed_efforts = ( + thinking_spec.efforts + if thinking_spec + else ("none", "low", "medium", "high") + ) + effort = reasoning_effort if reasoning_effort in allowed_efforts else None + # Claude 4.6 Opus/Sonnet accept top-tier adaptive effort as "max" only; + # "xhigh" is rejected (supported on Claude 4.7). Map our shared "xhigh" + # semantic to "max" for 4.6 outbound requests while still accepting + # both in ``allowed_efforts`` for persisted / cross-provider UI state. + if effort == "xhigh" and model.startswith( + ("claude-opus-4-6", "claude-sonnet-4-6") + ): + effort = "max" + if effort is None: + if enable_thinking is False: + effort = "none" + elif enable_thinking is True: + effort = "medium" + # Normalize one semantic Thinking control into Anthropic's two model-era + # APIs: adaptive effort on Claude 4.6/4.7, manual budget_tokens on 4.5. + if effort and effort != "none": + # Anthropic rejects top_k whenever thinking is enabled. + body.pop("top_k", None) + # Anthropic requires temperature=1 whenever thinking is enabled, + # AND forbids top_p in the same request: setting both produces + # "temperature and top_p cannot both be specified for this + # model. Please use only one." + # The base body never sets top_p, but pop defensively in case + # an upstream edit ever adds it before this branch runs. + body["temperature"] = 1 + body.pop("top_p", None) + if thinking_spec and thinking_spec.kind == "adaptive": + # `display` defaults to "omitted" on Claude Opus 4.7 (per the + # adaptive-thinking docs) — without an explicit opt-in the + # API emits an empty thinking block plus a signature_delta, + # so our SSE handler would surface a stray + # and the reasoning panel would stay blank. Force + # "summarized" so 4.7 streams thinking_delta events like + # 4.6 does. On 4.6 / Sonnet 4.6 this is the default, so + # setting it explicitly is harmless. + body["thinking"] = {"type": "adaptive", "display": "summarized"} + # Per the Messages API reference, the effort knob for + # adaptive thinking lives under `output_config.effort` — + # NOT as a top-level field. Sending `effort: ...` directly + # produces a 400 "effort: Extra inputs are not permitted". + # Allowed values: low | medium | high | xhigh | max. See: + # https://platform.claude.com/docs/en/api/messages + body["output_config"] = {"effort": effort} + elif thinking_spec and thinking_spec.kind == "manual": + budget_tokens = {"low": 1024, "medium": 2048, "high": 4096}[effort] + body["thinking"] = { + "type": "enabled", + "budget_tokens": budget_tokens, + } + # Anthropic requires max_tokens to be strictly greater than + # thinking.budget_tokens on the manual-thinking path. + if body.get("max_tokens", 0) <= budget_tokens: + body["max_tokens"] = budget_tokens + 1024 + + 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. + 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, " + "max_tokens=%s)", + model, + "thinking" in body, + body.get("thinking"), + body.get("output_config"), + body.get("temperature"), + "top_p" in body, + "top_k" in body, + body.get("max_tokens"), + ) + + _finish_reason_map = { + "end_turn": "stop", + "max_tokens": "length", + "stop_sequence": "stop", + } + + logger.info("Proxying Anthropic Messages API to %s (model=%s)", url, model) + + 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( + "Anthropic returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + # NOTE: same manual __anext__ loop as stream_chat_completion — see comment there. + lines_gen = response.aiter_lines().__aiter__() + thinking_open = False + # Diagnostic counters for the next time the user reports + # "no thinking content" — distinguishes "Anthropic never sent + # thinking_delta" from "frontend didn't render the chunks". + event_counts: dict[str, int] = {} + + def _content_chunk(text: str) -> str: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"content": text}, + "finish_reason": None, + } + ], + } + return f"data: {_json.dumps(chunk)}" + + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line or line.startswith("event:"): + continue + if not line.startswith("data:"): + continue + + data_str = line[len("data:") :].strip() + if not data_str: + continue + + try: + event = _json.loads(data_str) + except _json.JSONDecodeError: + continue + + event_type = event.get("type") + if event_type == "content_block_delta": + delta_kind = (event.get("delta") or {}).get("type") + key = f"{event_type}:{delta_kind}" + else: + key = event_type or "" + event_counts[key] = event_counts.get(key, 0) + 1 + + if event_type == "content_block_delta": + 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. + thinking_text = delta.get("thinking", "") + if thinking_text: + if not thinking_open: + thinking_text = f"{thinking_text}" + thinking_open = True + yield _content_chunk(thinking_text) + elif delta_type == "text_delta": + # First text after a thinking block closes the + # tag we opened above. Anthropic emits + # a content_block_stop between blocks, but + # closing on the text_delta transition is more + # forgiving if events arrive out of order. + if thinking_open: + yield _content_chunk("") + thinking_open = False + text = delta.get("text", "") + if text: + yield _content_chunk(text) + # signature_delta and any other delta types are + # intentionally skipped — they carry trust / + # verification metadata, not user-visible content. + + elif event_type == "content_block_stop": + # Close the tag when the thinking block + # ends, in case no text_delta follows (e.g. + # display=omitted on Claude 4.7, or thinking-only + # turns). + if thinking_open: + yield _content_chunk("") + thinking_open = False + + elif event_type == "message_delta": + stop_reason = event.get("delta", {}).get("stop_reason") + if stop_reason: + if thinking_open: + yield _content_chunk("") + thinking_open = False + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": _finish_reason_map.get( + stop_reason, "stop" + ), + } + ], + } + yield f"data: {_json.dumps(chunk)}" + + elif event_type == "message_stop": + if thinking_open: + yield _content_chunk("") + thinking_open = False + yield "data: [DONE]" + await ( + response.aclose() + ) # set PoolByteStream._closed=True FIRST + break + except GeneratorExit: + await response.aclose() # set PoolByteStream._closed=True FIRST + await lines_gen.aclose() # now safe — aclose() is a no-op + raise + finally: + # Surface per-event-type counts so reports of "no + # reasoning panel content" can be triaged at a glance: + # zero `content_block_delta:thinking_delta` entries + # means Anthropic skipped thinking for this prompt + # (adaptive can choose to); non-zero means thinking + # arrived and we wrapped it — any visual gap is then + # on the frontend. + logger.info( + "Anthropic stream event counts (model=%s): %s", + model, + event_counts, + ) + await response.aclose() + await lines_gen.aclose() + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + 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) + 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) + 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]], + model: str, + temperature: float, + top_p: float, + max_tokens: Optional[int], + enable_thinking: Optional[bool], + reasoning_effort: Optional[str], + ) -> AsyncGenerator[str, None]: + """ + Call OpenAI's /v1/responses endpoint and translate its SSE stream back + into OpenAI Chat Completions chunk format. + + The Responses API uses a different request shape (``input`` instead of + ``messages``, ``instructions`` for system prompts, ``max_output_tokens`` + for the budget) and emits event-typed SSE frames (e.g. + ``response.output_text.delta``) rather than chat-completion chunks. + ``presence_penalty`` / ``top_k`` are not part of the Responses contract + and are dropped here intentionally. + """ + import json as _json + + # Split system messages out into a single `instructions` string and + # translate user/assistant messages into the Responses input shape. + instructions_parts: list[str] = [] + input_items: list[dict[str, Any]] = [] + for msg in messages: + role = msg.get("role") + content = msg.get("content", "") + + if role == "system": + if isinstance(content, str): + if content: + instructions_parts.append(content) + elif isinstance(content, list): + for part in content: + if part.get("type") == "text" and part.get("text"): + instructions_parts.append(part["text"]) + continue + + if isinstance(content, str): + input_items.append({"role": role, "content": content}) + continue + + if isinstance(content, list): + translated_parts: list[dict[str, Any]] = [] + for part in content: + part_type = part.get("type") + if part_type == "text": + translated_parts.append( + {"type": "input_text", "text": part.get("text", "")} + ) + elif part_type == "image_url": + url = part.get("image_url", {}).get("url", "") + if url: + # Responses takes image_url as a flat string (both + # https:// URLs and data: URLs are accepted). + translated_parts.append( + {"type": "input_image", "image_url": url} + ) + if translated_parts: + input_items.append({"role": role, "content": translated_parts}) + + # NOTE: gpt-5.x / o3 / gpt-4.5 are reasoning-class models. They reject + # temperature and top_p with `Unsupported parameter` 400s on + # /v1/responses (and on /v1/chat/completions for the same families). + # The PROVIDER_REGISTRY['openai'] model_id_allowlist already scopes + # the picker to those families, so we never need to send sampling + # knobs here. ``reasoning.effort`` defaults to "medium" server-side + # if omitted — surface it in a future commit if a knob is wanted. + del temperature, top_p # explicit drop — params are accepted for + # API symmetry with the other stream methods but not forwarded. + + body: dict[str, Any] = { + "model": model, + "input": input_items, + "stream": True, + } + # `summary: "auto"` is what makes /v1/responses emit reasoning + # summary events — without it OpenAI returns no thinking text on + # most reasoning models, the SSE handler has no … + # to wrap, and the chat reasoning panel stays blank. Always pair + # an explicit effort with summary except for the explicit "off" + # case (effort: "none"), where summaries are pointless. + if reasoning_effort in ( + "minimal", + "low", + "medium", + "high", + "max", + "xhigh", + ): + body["reasoning"] = {"effort": reasoning_effort, "summary": "auto"} + elif reasoning_effort == "none" or enable_thinking is False: + body["reasoning"] = {"effort": "none"} + elif enable_thinking is True: + body["reasoning"] = {"effort": "medium", "summary": "auto"} + if instructions_parts: + body["instructions"] = "\n\n".join(instructions_parts) + if max_tokens is not None: + body["max_output_tokens"] = max_tokens + + url = f"{self.base_url}/responses" + completion_id = f"chatcmpl-openai-{model.replace('/', '-')}" + + logger.info("Proxying OpenAI Responses API to %s (model=%s)", url, model) + + 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( + "OpenAI Responses returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + # NOTE: same manual __anext__ loop as stream_chat_completion — + # see comment there for the GeneratorExit / aclose ordering. + lines_gen = response.aiter_lines().__aiter__() + done_emitted = False + reasoning_open = False + reasoning_emitted = False + + def _extract_reasoning_text(payload: Any) -> str: + if payload is None: + return "" + if isinstance(payload, str): + return payload + if isinstance(payload, list): + out: list[str] = [] + for item in payload: + text = _extract_reasoning_text(item) + if text: + out.append(text) + return "".join(out) + if isinstance(payload, dict): + # OpenAI responses may carry reasoning summaries in + # different envelope fields across event variants. + for key in ("text", "delta", "content", "summary"): + if key in payload: + text = _extract_reasoning_text(payload.get(key)) + if text: + return text + if payload.get("type") == "summary_text": + return _extract_reasoning_text(payload.get("text")) + return "" + + def _chunk_with_text(text: str) -> str: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"content": text}, + "finish_reason": None, + } + ], + } + return f"data: {_json.dumps(chunk)}" + + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line or line.startswith("event:"): + continue + if not line.startswith("data:"): + continue + + data_str = line[len("data:") :].strip() + if not data_str: + continue + if data_str == "[DONE]": + if not done_emitted: + yield "data: [DONE]" + done_emitted = True + break + + try: + event = _json.loads(data_str) + except _json.JSONDecodeError: + continue + + event_type = event.get("type") + + if event_type == "response.output_text.delta": + delta_text = event.get("delta", "") + if delta_text: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(delta_text) + + elif event_type == "response.output_item.done": + item = event.get("item", {}) + if ( + isinstance(item, dict) + and item.get("type") == "reasoning" + ): + summary_text = _extract_reasoning_text( + item.get("summary") + ) + if summary_text and not reasoning_emitted: + if not reasoning_open: + summary_text = f"{summary_text}" + reasoning_open = True + yield _chunk_with_text(summary_text) + reasoning_emitted = True + + elif isinstance(event_type, str) and "reasoning" in event_type: + reasoning_delta = _extract_reasoning_text(event) + if reasoning_delta: + if not reasoning_open: + reasoning_delta = f"{reasoning_delta}" + reasoning_open = True + yield _chunk_with_text(reasoning_delta) + reasoning_emitted = True + + elif event_type == "response.completed": + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], + } + yield f"data: {_json.dumps(chunk)}" + + elif event_type == "response.incomplete": + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "length", + } + ], + } + yield f"data: {_json.dumps(chunk)}" + + elif event_type in ("response.failed", "error"): + # Surface the failure to the client; let the + # outer route emit [DONE] as part of its cleanup. + error_payload = event.get("response", {}).get( + "error", {} + ) or { + "message": event.get("message", "Unknown error"), + "code": event.get("code"), + } + yield _error_sse_line( + 502, + _json.dumps(error_payload), + self.provider_type, + ) + break + except GeneratorExit: + await response.aclose() + await lines_gen.aclose() + raise + finally: + await response.aclose() + await lines_gen.aclose() + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + 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) + 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) + yield _error_sse_line( + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, + ) + + async def chat_completion( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float = 0.7, + top_p: float = 0.95, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + ) -> dict[str, Any]: + """Non-streaming chat completion. Returns the full response dict. + + Note: only valid for OpenAI-compatible providers. Anthropic requires its + own Messages API; use stream_chat_completion (with stream=False) instead + if a non-streaming Anthropic path is needed in the future. + """ + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": False, + "temperature": temperature, + "top_p": top_p, + "presence_penalty": presence_penalty, + } + if max_tokens is not None: + if self.provider_type == "openai": + body["max_completion_tokens"] = max_tokens + else: + body["max_tokens"] = max_tokens + + response = await _http_client.post( + f"{self.base_url}/chat/completions", + json = body, + headers = self._auth_headers(), + timeout = self._timeout, + ) + response.raise_for_status() + return response.json() + + async def list_models(self) -> list[dict[str, Any]]: + """ + Call GET /models on the provider to discover available models. + + Returns a list of model dicts with at least 'id' and optionally + 'created', 'owned_by', etc. + + All supported providers expose a /models endpoint: + - OpenAI-compatible: standard {"data": [...]} response + - Anthropic: https://api.anthropic.com/v1/models — same {"data": [...]} shape + """ + try: + response = await _http_client.get( + f"{self.base_url}/models", + headers = self._auth_headers(), + timeout = self._timeout, + ) + response.raise_for_status() + data = response.json() + # OpenAI format: {"data": [{"id": "...", ...}, ...]} + models = data.get("data", []) + return models + except httpx.HTTPError as exc: + logger.error("Failed to list models from %s: %s", self.provider_type, exc) + raise + + async def verify_models_endpoint_lightweight(self) -> None: + """ + Confirm GET /models returns 200 without buffering the full response body. + + Used for providers with enormous catalogs (e.g. OpenRouter, Hugging Face router) + where downloading the full JSON would be prohibitive. + """ + url = f"{self.base_url}/models" + try: + async with _http_client.stream( + "GET", + url, + headers = self._auth_headers(), + timeout = self._timeout, + ) as response: + if response.status_code != 200: + response.raise_for_status() + async for _chunk in response.aiter_bytes(chunk_size = 2048): + break + except httpx.HTTPError as exc: + logger.error( + "Lightweight /models check failed for %s: %s", + self.provider_type, + exc, + ) + raise + + async def close(self) -> None: + """No-op — the underlying client is shared across requests.""" + + +def _error_sse_line(status_code: int, message: str, provider_type: str) -> str: + """Format an error as an SSE data line in OpenAI error format.""" + import json + + error_obj = { + "error": { + "message": message, + "type": "provider_error", + "code": str(status_code), + "provider": provider_type, + } + } + return f"data: {json.dumps(error_obj)}" diff --git a/studio/backend/core/inference/key_exchange.py b/studio/backend/core/inference/key_exchange.py new file mode 100644 index 0000000000..f43bb16cf6 --- /dev/null +++ b/studio/backend/core/inference/key_exchange.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +RSA key pair for encrypting API keys in transit. + +The frontend encrypts API keys with the server's public key before +including them in requests. The backend decrypts with its private key +before forwarding to external providers. + +The key pair is generated at server startup and lives only in memory — +it is regenerated on each restart. The frontend fetches the public key +via GET /api/providers/public-key on load. +""" + +import base64 +import hashlib +import logging + +from cryptography.hazmat.primitives.asymmetric import rsa, padding +from cryptography.hazmat.primitives import serialization, hashes + +logger = logging.getLogger(__name__) + +_private_key: rsa.RSAPrivateKey | None = None +_public_key_pem: str | None = None +_public_key_fingerprint: str | None = None + + +def _compute_fingerprint(pem: str) -> str: + """SHA256 of the PEM bytes, truncated for log compactness.""" + return hashlib.sha256(pem.encode("utf-8")).hexdigest()[:16] + + +def init_key_pair() -> None: + """Generate an RSA-2048 key pair. Called once at server startup.""" + global _private_key, _public_key_pem, _public_key_fingerprint + if _private_key is not None: + # Re-entry is suspicious — every fresh keypair invalidates all + # in-flight ciphertext encrypted against the previous public key. + # Log loudly so a regression that calls init twice is visible. + logger.warning( + "init_key_pair called again — replacing existing RSA keypair " + "(previous fingerprint=%s). Any frontend that cached the old " + "public key will start hitting decryption failures.", + _public_key_fingerprint, + ) + _private_key = rsa.generate_private_key( + public_exponent = 65537, + key_size = 2048, + ) + _public_key_pem = ( + _private_key.public_key() + .public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + .decode("utf-8") + ) + _public_key_fingerprint = _compute_fingerprint(_public_key_pem) + logger.info( + "RSA key pair generated for API key encryption (fingerprint=%s)", + _public_key_fingerprint, + ) + + +def get_public_key_fingerprint() -> str | None: + """Short SHA256 of the current public key PEM; None before init.""" + return _public_key_fingerprint + + +def get_public_key_pem() -> str: + """Return the PEM-encoded public key for the frontend.""" + if _public_key_pem is None: + raise RuntimeError("Key pair not initialized. Call init_key_pair() first.") + return _public_key_pem + + +def decrypt_api_key(encrypted_b64: str) -> str: + """ + Decrypt an API key that was encrypted with the public key. + + Args: + encrypted_b64: Base64-encoded RSA-OAEP ciphertext. + + Returns: + The plaintext API key string. + """ + if _private_key is None: + raise RuntimeError("Key pair not initialized. Call init_key_pair() first.") + + try: + ciphertext = base64.b64decode(encrypted_b64) + except Exception as exc: + logger.warning( + "decrypt_api_key: base64 decode failed (input_len=%d, fingerprint=%s): %s: %s", + len(encrypted_b64), + _public_key_fingerprint, + type(exc).__name__, + exc, + ) + raise + + try: + plaintext = _private_key.decrypt( + ciphertext, + padding.OAEP( + mgf = padding.MGF1(algorithm = hashes.SHA256()), + algorithm = hashes.SHA256(), + label = None, + ), + ) + except Exception as exc: + # Surface enough state to distinguish key mismatch (wrong public key + # used on encrypt) from a padding/algo mismatch or corrupted bytes. + # Expected ciphertext length for RSA-2048 is exactly 256 bytes. + logger.warning( + "decrypt_api_key: RSA decrypt failed (ciphertext_len=%d, expected=256, " + "fingerprint=%s, exc=%s): %s", + len(ciphertext), + _public_key_fingerprint, + type(exc).__name__, + exc, + ) + raise + + return plaintext.decode("utf-8") diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py new file mode 100644 index 0000000000..4b6d7d6b17 --- /dev/null +++ b/studio/backend/core/inference/providers.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Static registry of supported external LLM providers. + +All providers expose OpenAI-compatible /v1/chat/completions endpoints +with Bearer token authentication and SSE streaming support. +""" + +import re +from typing import Any + +PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { + "openai": { + "display_name": "OpenAI", + "base_url": "https://api.openai.com/v1", + "default_models": [ + "gpt-5.5", + "gpt-5.4", + "gpt-5.4-mini", + "o3", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + # Keep the model picker scoped to the current generation. The remote + # /v1/models listing returns dozens of historical snapshots, fine-tunes + # and non-chat models (embeddings, TTS, image, moderation) that we + # never want to surface in the chat UI. Filtering here so backend + # is the single source of truth. + "model_id_allowlist": re.compile(r"^(gpt-5\.[345]|gpt-4\.5|o3)(?:[-.]|$)"), + # Hide dated snapshots and the retired plain gpt-5.3 id. + "model_id_denylist": re.compile(r"^(gpt-5\.3)$|-\d{4}-\d{2}-\d{2}$"), + }, + "anthropic": { + "display_name": "Anthropic", + "base_url": "https://api.anthropic.com/v1", + "default_models": [ + "claude-opus-4-7", + "claude-opus-4-6", + "claude-sonnet-4-6", + "claude-opus-4-5", + "claude-sonnet-4-5", + "claude-haiku-4-5", + ], + # Anthropic /v1/models returns dated snapshot ids alongside the + # canonical names (e.g. claude-3-5-sonnet-20241022). Hide the + # YYYYMMDD-suffixed variants from the picker — same intent as the + # OpenAI denylist, just a different date format (no dashes between + # year/month/day). + "model_id_denylist": re.compile(r"-\d{8}$"), + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": False, + "auth_header": "x-api-key", + "auth_prefix": "", + "extra_headers": { + "anthropic-version": "2023-06-01", + }, + "openai_compatible": False, + "notes": "Native Anthropic Messages API. Uses x-api-key header and /v1/messages endpoint with SSE translation.", + }, + "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. + "default_models": [ + "gemini-3.1-pro-preview", + "gemini-3.1-flash-lite", + "gemini-3-flash-preview", + "gemini-pro-latest", + "gemini-flash-latest", + "gemini-flash-lite-latest", + ], + "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.", + "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)$" + ), + }, + "deepseek": { + "display_name": "DeepSeek", + "base_url": "https://api.deepseek.com/v1", + "default_models": [ + "deepseek-chat", + "deepseek-reasoner", + ], + "supports_streaming": True, + "supports_vision": False, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "OpenAI-compatible API. deepseek-chat = V3, deepseek-reasoner = R1 thinking mode.", + }, + "mistral": { + "display_name": "Mistral AI", + "base_url": "https://api.mistral.ai/v1", + "default_models": [ + "codestral-latest", + "devstral-latest", + "devstral-medium-latest", + "magistral-medium-latest", + "ministral-14b-latest", + "ministral-3b-latest", + "ministral-8b-latest", + "mistral-large-latest", + "mistral-medium-latest", + "mistral-small-latest", + "mistral-tiny-latest", + "mistral-vibe-cli-latest", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "model_id_allowlist": re.compile( + r"^(codestral-latest|devstral-latest|devstral-medium-latest|" + r"magistral-medium-latest|ministral-(?:14b|3b|8b)-latest|" + r"mistral-(?:large|medium|small|tiny)-latest|" + r"mistral-vibe-cli-latest)$" + ), + }, + "kimi": { + "display_name": "Kimi", + "base_url": "https://api.moonshot.ai/v1", + # Current Kimi model lineup per the official docs: + # https://platform.kimi.ai/docs/models + # Listing/overview endpoints used to enumerate them: + # https://platform.kimi.ai/docs/api/list-models + # https://platform.kimi.ai/docs/api/overview + # kimi-k2.6 and kimi-k2.5 are the two SoTA multimodal models we + # surface in the picker; everything else (moonshot-v1-*, dated + # k2 previews) is filtered out by model_id_allowlist below. + "default_models": [ + "kimi-k2.6", + "kimi-k2.5", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "Moonshot API key. China: use base URL https://api.moonshot.cn/v1", + "model_id_allowlist": re.compile(r"^kimi-k2\.[56]$"), + # Both k2.6 and k2.5 are reasoning-class. The API rejects custom + # sampling: "invalid temperature: only 1 is allowed for this model" + # (and the same shape for top_p). Strip both fields from the + # outbound body so the server falls back to its required defaults. + "body_omit": ("temperature", "top_p"), + }, + "qwen": { + "display_name": "Qwen", + "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + "default_models": [ + "qwen-plus", + "qwen-turbo", + "qwen-max", + "qwen2.5-72b-instruct", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "DashScope API key. China mainland: override base URL to https://dashscope.aliyuncs.com/compatible-mode/v1", + }, + "huggingface": { + "display_name": "Hugging Face", + "base_url": "https://router.huggingface.co/v1", + # Seed the picker with a few popular ids so something is selectable + # before the live /v1/models call resolves. The remote listing is + # the source of truth — see model_list_mode below. + "default_models": [ + "openai/gpt-oss-120b", + "deepseek-ai/DeepSeek-V3", + "meta-llama/Llama-3.3-70B-Instruct", + "Qwen/Qwen2.5-72B-Instruct", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": ( + "HF token from huggingface.co/settings/tokens. Uses the " + "OpenAI-compatible router at /v1/chat/completions; /v1/models " + "returns the cross-provider chat catalog. See " + "https://huggingface.co/docs/inference-providers/index." + ), + # /v1/models works on the HF router and returns the full chat-model + # catalog (state.org/model[:policy] ids). Switch to remote so users + # see live availability — the picker has a search box, and + # loadModels() merges defaults so default_models entries remain + # visible if the remote call fails. + "model_list_mode": "remote", + # Scope the catalog to first-party org repos we trust as primary + # sources. The HF /v1/models response is otherwise hundreds of + # ids long (community fine-tunes, mirrors, fp8 variants, etc.). + "model_id_allowlist": re.compile( + r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|" + r"mistralai|zai-org)/" + ), + # Cap the post-filter list. /v1/models has no server-side limit + # or popularity sort, so this is just "first N matches" — pair it + # with the default_models seed so the most useful flagship ids + # are always among the top regardless of the API's order. + "model_id_limit": 15, + }, + "openrouter": { + "display_name": "OpenRouter", + "base_url": "https://openrouter.ai/api/v1", + # Curated list for Studio's picker (explicitly locked, not live /models). + "default_models": [ + "openrouter/free", + "openai/gpt-4o", + "anthropic/claude-sonnet-4-5", + "google/gemini-2.5-flash", + "mistralai/mistral-large-2411", + "deepseek/deepseek-r1", + "mistralai/mistral-small-3.1-24b-instruct", + "perceptron/perceptron-mk1", + "inclusionai/ring-2.6-1t:free", + "google/gemini-3.1-flash-lite", + "baidu/cobuddy:free", + "openai/gpt-chat-latest", + "x-ai/grok-4.3", + "ibm-granite/granite-4.1-8b", + "openrouter/owl-alpha", + "poolside/laguna-xs.2:free", + "~google/gemini-pro-latest", + "~moonshotai/kimi-latest", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "extra_headers": { + "HTTP-Referer": "https://unsloth.ai", + "X-Title": "Unsloth Studio", + }, + "notes": "Unified gateway to 300+ models across all major providers. HTTP-Referer and X-Title headers sent for attribution.", + "model_list_mode": "curated", + }, +} + + +def get_provider_info(provider_type: str) -> dict[str, Any] | None: + """Return the registry entry for a provider type, or None if unknown.""" + return PROVIDER_REGISTRY.get(provider_type) + + +def get_base_url(provider_type: str) -> str | None: + """Return the default base URL for a provider type.""" + info = PROVIDER_REGISTRY.get(provider_type) + return info["base_url"] if info else None + + +def list_available_providers() -> list[dict[str, Any]]: + """Return all registered providers (for the /registry endpoint).""" + result = [] + for provider_type, info in PROVIDER_REGISTRY.items(): + result.append( + { + "provider_type": provider_type, + "display_name": info["display_name"], + "base_url": info["base_url"], + "default_models": info["default_models"], + "supports_streaming": info["supports_streaming"], + "supports_vision": info.get("supports_vision", False), + "supports_tool_calling": info.get("supports_tool_calling", False), + "model_list_mode": info.get("model_list_mode", "remote"), + } + ) + return result diff --git a/studio/backend/main.py b/studio/backend/main.py index 4955e988e6..c1c9ed1d90 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -120,6 +120,7 @@ from routes import ( inference_router, inference_studio_router, models_router, + providers_router, training_history_router, training_router, ) @@ -222,6 +223,11 @@ async def lifespan(app: FastAPI): threading.Thread(target = _precache, daemon = True).start() + # Initialize RSA key pair for API key encryption (external providers) + from core.inference.key_exchange import init_key_pair + + init_key_pair() + if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() app.state.bootstrap_password = bootstrap_pw @@ -474,6 +480,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = [" # so external tools (Open WebUI, SillyTavern, etc.) can use the # standard /v1/chat/completions path. app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) +app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 746ac8bbc2..013328f6c6 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -531,9 +531,11 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models", ) - reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field( + reasoning_effort: Optional[ + Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] + ] = Field( None, - description = "[x-unsloth] Reasoning effort level ('low'|'medium'|'high') for Harmony-style reasoning models (e.g. gpt-oss). Overrides enable_thinking when the active model uses reasoning_effort style.", + description = "[x-unsloth] Reasoning effort level ('none'|'minimal'|'low'|'medium'|'high'|'max'|'xhigh'). OpenAI `/v1/responses` accepts model-dependent subsets; Anthropic adaptive thinking uses `max` as the top tier on Claude 4.6 Opus/Sonnet (inbound `xhigh` is mapped to `max`) and `xhigh` on Claude 4.7 Opus; local Harmony/gpt-oss templates support low|medium|high.", ) preserve_thinking: Optional[bool] = Field( None, @@ -570,6 +572,28 @@ class ChatCompletionRequest(BaseModel): description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.", ) + # ── External provider routing (x-unsloth extensions) ────────── + provider_id: Optional[str] = Field( + None, + description = "[x-unsloth] Saved provider config ID. If set with encrypted_api_key, routes to external LLM.", + ) + provider_type: Optional[str] = Field( + None, + description = "[x-unsloth] Provider type (e.g. 'openai', 'mistral'). Used if provider_id is not set.", + ) + external_model: Optional[str] = Field( + None, + description = "[x-unsloth] Model ID at the external provider.", + ) + encrypted_api_key: Optional[str] = Field( + None, + description = "[x-unsloth] RSA-encrypted, base64-encoded API key for the external provider.", + ) + provider_base_url: Optional[str] = Field( + None, + description = "[x-unsloth] Override base URL for the external provider.", + ) + # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py new file mode 100644 index 0000000000..5678e69f62 --- /dev/null +++ b/studio/backend/models/providers.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Pydantic schemas for the external LLM providers API. +""" + +from typing import Literal, Optional + +from pydantic import BaseModel, Field + + +# ── Registry (static provider info) ─────────────────────────────── + + +class ProviderRegistryEntry(BaseModel): + """A supported provider type with its default configuration.""" + + provider_type: str = Field( + ..., description = "Provider identifier (e.g. 'openai', 'mistral')" + ) + display_name: str = Field(..., description = "Human-readable provider name") + base_url: str = Field(..., description = "Default API base URL") + default_models: list[str] = Field( + default_factory = list, description = "Well-known model IDs for this provider" + ) + supports_streaming: bool = Field( + True, description = "Whether this provider supports SSE streaming" + ) + supports_vision: bool = Field( + False, description = "Whether this provider supports vision/image input" + ) + supports_tool_calling: bool = Field( + False, description = "Whether this provider supports tool/function calling" + ) + model_list_mode: Literal["remote", "curated"] = Field( + "remote", + description = "remote = fetch /models; curated = huge catalogs — UI uses defaults + manual IDs only", + ) + + +# ── Provider config CRUD ────────────────────────────────────────── + + +class ProviderCreate(BaseModel): + """Request to create a saved provider configuration.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + display_name: str = Field( + ..., description = "User-chosen label (e.g. 'My OpenAI Key')" + ) + base_url: Optional[str] = Field( + None, + description = "Custom base URL (overrides registry default). Omit to use the default.", + ) + + +class ProviderUpdate(BaseModel): + """Request to update a saved provider configuration.""" + + display_name: Optional[str] = Field(None, description = "New display name") + base_url: Optional[str] = Field(None, description = "New base URL") + is_enabled: Optional[bool] = Field( + None, description = "Enable or disable this provider" + ) + + +class ProviderResponse(BaseModel): + """A saved provider configuration (returned by list/get endpoints).""" + + id: str = Field(..., description = "Unique provider config ID") + provider_type: str = Field(..., description = "Provider type (e.g. 'openai')") + display_name: str = Field(..., description = "User-chosen label") + base_url: str = Field(..., description = "API base URL") + is_enabled: bool = Field(True, description = "Whether this provider is enabled") + created_at: str = Field(..., description = "ISO 8601 creation timestamp") + updated_at: str = Field(..., description = "ISO 8601 last-update timestamp") + + +# ── Model listing ───────────────────────────────────────────────── + + +class ProviderModelInfo(BaseModel): + """A model available from an external provider.""" + + id: str = Field(..., description = "Model ID as expected by the provider API") + display_name: str = Field("", description = "Human-readable model name") + context_length: Optional[int] = Field( + None, description = "Maximum context length in tokens" + ) + owned_by: Optional[str] = Field(None, description = "Model owner/organization") + + +class ProviderModelsRequest(BaseModel): + """Request to list models from an external provider.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + encrypted_api_key: str = Field( + ..., description = "RSA-encrypted, base64-encoded API key" + ) + base_url: Optional[str] = Field( + None, description = "Custom base URL (overrides registry default)" + ) + + +# ── Connection testing ──────────────────────────────────────────── + + +class ProviderTestRequest(BaseModel): + """Request to test connectivity to an external provider.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + encrypted_api_key: str = Field( + ..., description = "RSA-encrypted, base64-encoded API key" + ) + base_url: Optional[str] = Field( + None, description = "Custom base URL (overrides registry default)" + ) + + +class ProviderTestResult(BaseModel): + """Result of a provider connectivity test.""" + + success: bool = Field(..., description = "Whether the test succeeded") + message: str = Field(..., description = "Human-readable result message") + models_count: Optional[int] = Field( + None, description = "Number of models found (if test succeeded)" + ) diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 1bf751c368..96f8816b57 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -16,3 +16,5 @@ huggingface-hub==0.36.2 structlog>=24.1.0 diceware ddgs +cryptography>=42.0.0 +httpx>=0.27.0 diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index cf4586281b..62320b9084 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -14,6 +14,7 @@ from routes.auth import router as auth_router from routes.data_recipe import router as data_recipe_router from routes.export import router as export_router from routes.training_history import router as training_history_router +from routes.providers import router as providers_router __all__ = [ "training_router", @@ -25,4 +26,5 @@ __all__ = [ "data_recipe_router", "export_router", "training_history_router", + "providers_router", ] diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7102e12bf8..59928be3cf 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -204,6 +204,11 @@ from core.inference.anthropic_compat import ( ) from auth.authentication import get_current_subject +from core.inference.key_exchange import decrypt_api_key +from core.inference.providers import get_provider_info, get_base_url +from core.inference.external_provider import ExternalProviderClient +from storage import providers_db + import io import wave import base64 @@ -1464,6 +1469,161 @@ def _extract_content_parts( return system_prompt, chat_messages, first_image_b64 +# ── External provider proxy ────────────────────────────────────── + + +def _build_external_messages( + messages: list, + supports_vision: bool, +) -> list[dict]: + """ + Convert ChatMessage list to OpenAI-compatible dicts for external providers. + + - Vision providers: preserve multimodal content arrays (image_url parts intact). + - Non-vision providers: flatten to text-only (images silently dropped). + """ + result = [] + for msg in messages: + if isinstance(msg.content, str): + # Skip assistant messages with empty content (some providers reject them) + if msg.role == "assistant" and not msg.content.strip(): + continue + result.append({"role": msg.role, "content": msg.content}) + elif isinstance(msg.content, list): + if supports_vision: + parts = [] + for part in msg.content: + if part.type == "text": + parts.append({"type": "text", "text": part.text}) + elif part.type == "image_url": + parts.append( + { + "type": "image_url", + "image_url": {"url": part.image_url.url}, + } + ) + result.append({"role": msg.role, "content": parts}) + else: + # Non-vision provider — strip images, keep text only + text = "\n".join(p.text for p in msg.content if p.type == "text") + result.append({"role": msg.role, "content": text}) + return result + + +async def _proxy_to_external_provider( + payload: ChatCompletionRequest, + request: Request, +) -> StreamingResponse: + """ + Proxy a chat completion request to an external LLM provider. + + Resolves provider config (from DB or registry), decrypts the API key, + and streams the response back in OpenAI SSE format. + """ + # Resolve provider type and base URL + provider_type = payload.provider_type + base_url = payload.provider_base_url + + if payload.provider_id: + config = providers_db.get_provider(payload.provider_id) + if config is None: + raise HTTPException( + status_code = 404, + detail = f"Provider config not found: {payload.provider_id}", + ) + if not config["is_enabled"]: + raise HTTPException( + status_code = 400, + detail = f"Provider '{config['display_name']}' is disabled.", + ) + provider_type = provider_type or config["provider_type"] + base_url = base_url or config["base_url"] + + if not provider_type: + raise HTTPException( + status_code = 400, + detail = "Either provider_id or provider_type is required for external provider routing.", + ) + + # Fall back to registry default base URL + if not base_url: + base_url = get_base_url(provider_type) + if not base_url: + raise HTTPException( + status_code = 400, + detail = f"Unknown provider type: {provider_type}", + ) + + # Decrypt the API key + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("external_provider.decrypt_failed", error = str(exc)) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.", + ) + + model = payload.external_model or payload.model + if model == "default": + raise HTTPException( + status_code = 400, + detail = "external_model is required when using an external provider.", + ) + + # Build messages preserving multimodal content for vision-capable providers + from core.inference.providers import get_provider_info as _get_provider_info + + _pinfo = _get_provider_info(provider_type) or {} + _supports_vision = _pinfo.get("supports_vision", False) + chat_messages = _build_external_messages(payload.messages, _supports_vision) + + client = ExternalProviderClient( + provider_type = provider_type, + base_url = base_url, + api_key = api_key, + ) + + async def _stream(): + gen = client.stream_chat_completion( + messages = chat_messages, + model = model, + temperature = payload.temperature, + top_p = payload.top_p, + max_tokens = payload.max_tokens, + presence_penalty = payload.presence_penalty, + top_k = payload.top_k, + enable_thinking = payload.enable_thinking, + reasoning_effort = payload.reasoning_effort, + stream = payload.stream, + ) + try: + sent_done = False + async for line in gen: + yield f"{line}\n\n" + if "[DONE]" in line: + sent_done = True + if not sent_done: + yield "data: [DONE]\n\n" + except Exception as exc: + logger.error("external_provider.stream_error", error = str(exc)) + finally: + try: + await gen.aclose() + except RuntimeError: + pass # suppress httpcore asyncgen cleanup error (Python 3.13 + httpcore 1.0.x) + await client.close() + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, @@ -1483,6 +1643,10 @@ async def openai_chat_completions( - GGUF models → llama-server via LlamaCppBackend - Other models → Unsloth/transformers via InferenceBackend """ + # ── External provider routing ──────────────────────────────── + if payload.encrypted_api_key and (payload.provider_id or payload.provider_type): + return await _proxy_to_external_provider(payload, request) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py new file mode 100644 index 0000000000..e21985d60b --- /dev/null +++ b/studio/backend/routes/providers.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +API routes for external LLM provider management. + +Provides endpoints for: + - Discovering available provider types (registry) + - CRUD for saved provider configurations (no API keys stored) + - Fetching the RSA public key for API key encryption + - Testing provider connectivity + - Listing models from a provider +""" + +import uuid +import structlog +from fastapi import APIRouter, Depends, HTTPException + +from auth.authentication import get_current_subject +from core.inference.key_exchange import ( + decrypt_api_key, + get_public_key_fingerprint, + get_public_key_pem, +) +from core.inference.providers import ( + get_base_url, + get_provider_info, + list_available_providers, +) +from core.inference.external_provider import ExternalProviderClient +from models.providers import ( + ProviderCreate, + ProviderModelsRequest, + ProviderModelInfo, + ProviderResponse, + ProviderRegistryEntry, + ProviderTestRequest, + ProviderTestResult, + ProviderUpdate, +) +from storage import providers_db + +logger = structlog.get_logger(__name__) + +router = APIRouter() + + +# ── Public key for API key encryption ───────────────────────────── + + +@router.get("/public-key") +async def get_public_key( + current_subject: str = Depends(get_current_subject), +): + """Return the RSA public key PEM for client-side API key encryption. + + The ``fingerprint`` field is a short SHA256 of the PEM and is meant + purely for diagnostics — a mismatch between what the frontend + captured at encrypt time and what the server reports here is a + clear signal that the keypair rotated mid-flight (e.g. the server + re-ran ``init_key_pair`` for any reason). + """ + return { + "public_key": get_public_key_pem(), + "fingerprint": get_public_key_fingerprint(), + } + + +# ── Provider registry (static) ─────────────────────────────────── + + +@router.get("/registry", response_model = list[ProviderRegistryEntry]) +async def list_registry( + current_subject: str = Depends(get_current_subject), +): + """List all supported provider types with their default configurations.""" + return list_available_providers() + + +# ── Provider config CRUD ────────────────────────────────────────── + + +@router.get("/", response_model = list[ProviderResponse]) +async def list_provider_configs( + current_subject: str = Depends(get_current_subject), +): + """List all saved provider configurations.""" + rows = providers_db.list_providers() + return [ + ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + for row in rows + ] + + +@router.post("/", response_model = ProviderResponse, status_code = 201) +async def create_provider_config( + payload: ProviderCreate, + current_subject: str = Depends(get_current_subject), +): + """Create a new saved provider configuration (no API key stored).""" + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}. " + f"Use GET /api/providers/registry to see available types.", + ) + + provider_id = uuid.uuid4().hex[:16] + base_url = payload.base_url or info["base_url"] + + providers_db.create_provider( + id = provider_id, + provider_type = payload.provider_type, + display_name = payload.display_name, + base_url = base_url, + ) + + row = providers_db.get_provider(provider_id) + return ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + +@router.put("/{provider_id}", response_model = ProviderResponse) +async def update_provider_config( + provider_id: str, + payload: ProviderUpdate, + current_subject: str = Depends(get_current_subject), +): + """Update a saved provider configuration.""" + existing = providers_db.get_provider(provider_id) + if not existing: + raise HTTPException(status_code = 404, detail = "Provider not found") + + updated = providers_db.update_provider( + id = provider_id, + display_name = payload.display_name, + base_url = payload.base_url, + is_enabled = payload.is_enabled, + ) + if not updated: + raise HTTPException(status_code = 400, detail = "No fields to update") + + row = providers_db.get_provider(provider_id) + return ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + +@router.delete("/{provider_id}", status_code = 204) +async def delete_provider_config( + provider_id: str, + current_subject: str = Depends(get_current_subject), +): + """Delete a saved provider configuration.""" + deleted = providers_db.delete_provider(provider_id) + if not deleted: + raise HTTPException(status_code = 404, detail = "Provider not found") + + +# ── Test connectivity ───────────────────────────────────────────── + + +@router.post("/test", response_model = ProviderTestResult) +async def test_provider( + payload: ProviderTestRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Test connectivity to an external provider. + + Makes a lightweight GET /models call to verify the API key works. + The encrypted_api_key is decrypted server-side and never stored. + """ + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}", + ) + + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) + + base_url = payload.base_url or info["base_url"] + client = ExternalProviderClient( + provider_type = payload.provider_type, + base_url = base_url, + api_key = api_key, + timeout = 15.0, + ) + + try: + if info.get("model_list_mode") == "curated": + await client.verify_models_endpoint_lightweight() + return ProviderTestResult( + success = True, + message = ( + "Connected successfully. Full model list is not fetched for this provider — " + "use suggestions and manual model IDs in the dialog." + ), + models_count = None, + ) + models = await client.list_models() + return ProviderTestResult( + success = True, + message = f"Connected successfully. Found {len(models)} model(s).", + models_count = len(models), + ) + except Exception as exc: + logger.warning("Provider test failed for %s: %s", payload.provider_type, exc) + return ProviderTestResult( + success = False, + message = f"Connection failed: {exc}", + models_count = None, + ) + finally: + await client.close() + + +# ── List models from provider ───────────────────────────────────── + + +@router.post("/models", response_model = list[ProviderModelInfo]) +async def list_provider_models( + payload: ProviderModelsRequest, + current_subject: str = Depends(get_current_subject), +): + """ + List models available from an external provider. + + The encrypted_api_key is decrypted server-side and never stored. + """ + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}", + ) + + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) + + if info.get("model_list_mode") == "curated": + return [ + ProviderModelInfo( + id = m, + display_name = m, + context_length = None, + owned_by = None, + ) + for m in info.get("default_models", []) + ] + + base_url = payload.base_url or info["base_url"] + client = ExternalProviderClient( + provider_type = payload.provider_type, + base_url = base_url, + api_key = api_key, + timeout = 15.0, + ) + + 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", ""))] + # 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 + # "first N matches" — pair with default_models for any must-have + # flagship ids. + limit = info.get("model_id_limit") + if isinstance(limit, int) and limit > 0: + models = models[:limit] + return [ + ProviderModelInfo( + id = m.get("id", ""), + display_name = m.get("id", ""), + context_length = m.get("context_length") or m.get("context_window"), + owned_by = m.get("owned_by"), + ) + for m in models + ] + except Exception as exc: + logger.error("Failed to list models from %s: %s", payload.provider_type, exc) + raise HTTPException( + status_code = 502, + detail = f"Failed to list models from {payload.provider_type}: {exc}", + ) + finally: + await client.close() diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py new file mode 100644 index 0000000000..ca47fcbd80 --- /dev/null +++ b/studio/backend/storage/providers_db.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +SQLite storage for external LLM provider configurations. + +Follows the same pattern as studio_db.py — module-level functions, +raw sqlite3, WAL mode, per-function connections. + +NOTE: API keys are NOT stored here. They live only in the browser +(localStorage) and are sent encrypted per-request. +""" + +import logging +import sqlite3 +import threading +from datetime import datetime, timezone +from typing import Optional + +logger = logging.getLogger(__name__) + +from utils.paths import studio_db_path, ensure_dir + +_schema_lock = threading.Lock() +_schema_ready = False + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + """Create the llm_providers table if it doesn't exist. Called once per process.""" + conn.execute("PRAGMA journal_mode=WAL") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS llm_providers ( + id TEXT NOT NULL PRIMARY KEY, + provider_type TEXT NOT NULL, + display_name TEXT NOT NULL, + base_url TEXT NOT NULL, + is_enabled INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ) + """ + ) + + +def get_connection() -> sqlite3.Connection: + """Open studio.db with WAL mode, create table once per process.""" + global _schema_ready + db_path = studio_db_path() + ensure_dir(db_path.parent) + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + if not _schema_ready: + with _schema_lock: + if not _schema_ready: + try: + _ensure_schema(conn) + _schema_ready = True + except Exception: + conn.close() + raise + return conn + + +def create_provider( + id: str, + provider_type: str, + display_name: str, + base_url: str, +) -> None: + """Insert a new provider configuration.""" + now = datetime.now(timezone.utc).isoformat() + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + (id, provider_type, display_name, base_url, now, now), + ) + conn.commit() + finally: + conn.close() + + +def update_provider( + id: str, + display_name: Optional[str] = None, + base_url: Optional[str] = None, + is_enabled: Optional[bool] = None, +) -> bool: + """Update fields on an existing provider. Returns True if a row was updated.""" + updates = [] + params = [] + if display_name is not None: + updates.append("display_name = ?") + params.append(display_name) + if base_url is not None: + updates.append("base_url = ?") + params.append(base_url) + if is_enabled is not None: + updates.append("is_enabled = ?") + params.append(1 if is_enabled else 0) + if not updates: + return False + updates.append("updated_at = ?") + params.append(datetime.now(timezone.utc).isoformat()) + params.append(id) + + conn = get_connection() + try: + cursor = conn.execute( + f"UPDATE llm_providers SET {', '.join(updates)} WHERE id = ?", + params, + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def delete_provider(id: str) -> bool: + """Delete a provider by ID. Returns True if a row was deleted.""" + conn = get_connection() + try: + cursor = conn.execute("DELETE FROM llm_providers WHERE id = ?", (id,)) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def get_provider(id: str) -> Optional[dict]: + """Fetch a single provider by ID.""" + conn = get_connection() + try: + row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone() + return dict(row) if row else None + finally: + conn.close() + + +def list_providers() -> list[dict]: + """List all provider configurations, ordered by creation time.""" + conn = get_connection() + try: + rows = conn.execute( + "SELECT * FROM llm_providers ORDER BY created_at" + ).fetchall() + return [dict(row) for row in rows] + finally: + conn.close() diff --git a/studio/backend/tests/test_anthropic_thinking_translation.py b/studio/backend/tests/test_anthropic_thinking_translation.py new file mode 100644 index 0000000000..14f261ae6b --- /dev/null +++ b/studio/backend/tests/test_anthropic_thinking_translation.py @@ -0,0 +1,404 @@ +# 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 Anthropic extended-thinking translation in +external_provider. + +Covers: +- Adaptive-mode request body nests effort under + ``output_config: {effort: ""}`` per the Messages API + reference (a top-level ``effort`` field 400s with + "effort: Extra inputs are not permitted"). +- Streaming SSE: ``content_block_delta`` with + ``delta.type == "thinking_delta"`` is translated into inline + ``...`` chat-completion chunks so the frontend's + reasoning-panel pipeline lifts it correctly. +- The ```` tag closes when the first ``text_delta`` arrives, + on ``content_block_stop``, on ``message_delta``, or on + ``message_stop``. +- Thinking is paired with ``temperature=1`` and no ``top_p`` / + ``top_k`` on the wire (Anthropic extended-thinking contract). +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _anthropic_sse(events: list[dict]) -> bytes: + """Serialize a list of Messages-API event dicts as an SSE byte stream.""" + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _payloads_from_lines(lines: list[str]) -> list: + out = [] + for line in lines: + if not line.startswith("data:"): + continue + raw = line[len("data:") :].strip() + if not raw: + continue + if raw == "[DONE]": + out.append("[DONE]") + else: + out.append(json.loads(raw)) + return out + + +def test_adaptive_thinking_body_uses_output_config_effort_shape(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = None, + reasoning_effort = "medium", + ): + pass + await client.close() + + _drive(run()) + + body = captured["body"] + # display=summarized is set explicitly so Opus 4.7 (which defaults to + # "omitted") still emits thinking_delta events for the reasoning panel. + assert body["thinking"] == {"type": "adaptive", "display": "summarized"} + # Documented shape: effort is nested under output_config. + # A top-level `effort` field produces a 400: + # "effort: Extra inputs are not permitted". + assert body["output_config"] == {"effort": "medium"} + assert "effort" not in body + # Extended-thinking contract: temperature=1, no top_p / top_k. + assert body["temperature"] == 1 + assert "top_p" not in body + assert "top_k" not in body + + +def test_adaptive_thinking_maps_xhigh_to_max_on_claude_4_6(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = None, + reasoning_effort = "xhigh", + ): + pass + await client.close() + + _drive(run()) + + assert captured["body"]["output_config"] == {"effort": "max"} + + +def test_adaptive_thinking_keeps_max_on_claude_4_6(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = None, + reasoning_effort = "max", + ): + pass + await client.close() + + _drive(run()) + + assert captured["body"]["output_config"] == {"effort": "max"} + + +def test_adaptive_thinking_keeps_xhigh_on_claude_4_7(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = None, + reasoning_effort = "xhigh", + ): + pass + await client.close() + + _drive(run()) + + body = captured["body"] + assert body["output_config"] == {"effort": "xhigh"} + assert "effort" not in body + + +def test_manual_thinking_body_uses_budget_tokens_on_4_5(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _anthropic_sse([{"type": "message_stop"}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 1024, + top_k = None, + enable_thinking = None, + reasoning_effort = "high", + ): + pass + await client.close() + + _drive(run()) + + body = captured["body"] + assert body["thinking"] == {"type": "enabled", "budget_tokens": 4096} + # max_tokens must be strictly greater than budget_tokens; we shipped 1024 + # and budget is 4096, so the wrapper should bump max_tokens. + assert body["max_tokens"] > body["thinking"]["budget_tokens"] + # Manual-thinking path does not use output_config / effort — those are + # the adaptive-mode controls (Claude 4.6 / 4.7). + assert "effort" not in body + assert "output_config" not in body + + +def test_thinking_delta_wrapped_in_think_tags(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + events = [ + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "First "}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "I plan."}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "abc123"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "Answer."}, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}}, + {"type": "message_stop"}, + ] + return httpx.Response( + 200, + content = _anthropic_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = True, + reasoning_effort = None, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + payloads = _payloads_from_lines(lines) + + combined = "".join( + p["choices"][0]["delta"].get("content", "") + for p in payloads + if isinstance(p, dict) and p["choices"][0]["delta"] + ) + + # Reasoning text should be wrapped in ..., followed by the + # answer text, and the stream should terminate with [DONE]. + assert "First I plan." in combined + assert combined.endswith("Answer.") + # signature_delta is intentionally dropped — no leaked signature text. + assert "abc123" not in combined + assert "[DONE]" in payloads + + +def test_thinking_only_turn_closes_tag_without_text_delta(monkeypatch): + """display=omitted on Claude 4.7 emits a signature_delta and no text. + + The open is still triggered by the (synthetic) thinking_delta; + we want content_block_stop to close it cleanly so the tag never leaks + into the next chunk.""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "internal"}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}}, + {"type": "message_stop"}, + ] + return httpx.Response( + 200, + content = _anthropic_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_anthropic( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + top_k = None, + enable_thinking = True, + reasoning_effort = None, + ) + ) + await client.close() + return lines + + payloads = _payloads_from_lines(_drive(run())) + combined = "".join( + p["choices"][0]["delta"].get("content", "") + for p in payloads + if isinstance(p, dict) and p["choices"][0]["delta"] + ) + assert combined == "internal" diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index a7201ac433..913c3cc355 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -437,6 +437,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): inference_router = APIRouter(), inference_studio_router = APIRouter(), models_router = APIRouter(), + providers_router = APIRouter(), training_history_router = APIRouter(), training_router = APIRouter(), ) diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py new file mode 100644 index 0000000000..4ad6a19ea9 --- /dev/null +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -0,0 +1,432 @@ +# 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 OpenAI `/v1/responses` translation in external_provider. + +Covers: +- Request body shape: system messages collapse into `instructions`, user/ + assistant messages go into `input`, sampling knobs Responses does not + support (presence_penalty, top_k) are not forwarded. +- SSE translation: `response.output_text.delta` events become OpenAI Chat + Completions chunks, `response.completed` emits a `finish_reason: stop` + chunk, the stream terminates with `data: [DONE]`. +- Image parts in user content are rewritten from Chat Completions + `{type: image_url, image_url: {url}}` into Responses + `{type: input_image, image_url: }`. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + + +def _responses_sse(events: list[dict]) -> bytes: + """Serialize a list of Responses-API event dicts as an SSE byte stream.""" + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + chunks.append("data: [DONE]") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def test_responses_request_body_uses_input_and_instructions(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _responses_sse([{"type": "response.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [ + {"role": "system", "content": "You are concise."}, + {"role": "user", "content": "Hi"}, + ], + model = "gpt-5.5", + temperature = 0.5, + top_p = 0.9, + max_tokens = 512, + enable_thinking = None, + reasoning_effort = None, + ): + pass + await client.close() + + _drive(run()) + + assert captured["url"] == "https://api.openai.com/v1/responses" + body = captured["body"] + assert body["model"] == "gpt-5.5" + assert body["instructions"] == "You are concise." + assert body["input"] == [{"role": "user", "content": "Hi"}] + assert body["max_output_tokens"] == 512 + assert body["stream"] is True + # Responses API on reasoning-class models (gpt-5.x / o3 / gpt-4.5 — the + # only OpenAI ids the registry allowlist exposes) rejects these as + # `Unsupported parameter`. Make sure we never silently forward them. + assert "temperature" not in body + assert "top_p" not in body + assert "presence_penalty" not in body + assert "frequency_penalty" not in body + assert "top_k" not in body + assert "messages" not in body + + +def test_responses_translates_image_parts(monkeypatch): + 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.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAA"}, + }, + ], + } + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ): + pass + await client.close() + + _drive(run()) + + parts = captured["body"]["input"][0]["content"] + assert parts[0] == {"type": "input_text", "text": "What is this?"} + assert parts[1] == { + "type": "input_image", + "image_url": "data:image/png;base64,AAA", + } + # No max_output_tokens key when caller passes max_tokens=None. + assert "max_output_tokens" not in captured["body"] + + +def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + {"type": "response.output_text.delta", "delta": "Hello"}, + {"type": "response.output_text.delta", "delta": ", world"}, + {"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": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + + # Drop empty / non-data lines for assertion clarity. + data_lines = [line for line in lines if line.startswith("data:")] + payloads = [] + for line in data_lines: + raw = line[len("data:") :].strip() + if raw == "[DONE]": + payloads.append("[DONE]") + else: + payloads.append(json.loads(raw)) + + # Two text deltas, one terminal chunk, then [DONE]. + assert payloads[0]["choices"][0]["delta"]["content"] == "Hello" + assert payloads[0]["choices"][0]["finish_reason"] is None + assert payloads[1]["choices"][0]["delta"]["content"] == ", world" + assert payloads[2]["choices"][0]["delta"] == {} + assert payloads[2]["choices"][0]["finish_reason"] == "stop" + assert payloads[-1] == "[DONE]" + + +def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.output_text.delta", "delta": "partial"}, + {"type": "response.incomplete", "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": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + finish_reasons = [ + json.loads(line[len("data:") :].strip())["choices"][0]["finish_reason"] + for line in lines + if line.startswith("data:") + and line[len("data:") :].strip() not in ("", "[DONE]") + ] + assert "length" in finish_reasons + + +def test_responses_reasoning_effort_included_when_requested(monkeypatch): + 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.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = "high", + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "high", "summary": "auto"} + + +def test_responses_reasoning_effort_none_omits_summary(monkeypatch): + 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.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = "none", + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "none"} + + +def test_responses_reasoning_effort_xhigh_passthrough(monkeypatch): + 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.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = "xhigh", + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "xhigh", "summary": "auto"} + + +def test_responses_enable_thinking_false_maps_to_reasoning_none(monkeypatch): + 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.completed", "response": {}}]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + async for _ in client._stream_openai_responses( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = False, + reasoning_effort = None, + ): + pass + await client.close() + + _drive(run()) + assert captured["body"]["reasoning"] == {"effort": "none"} + + +def test_responses_reasoning_summary_wrapped_in_think_tags(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + events = [ + { + "type": "response.output_item.done", + "item": { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "plan"}], + }, + }, + {"type": "response.output_text.delta", "delta": "answer"}, + {"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": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + data_lines = [ + line[len("data:") :].strip() + for line in lines + if line.startswith("data:") + and line[len("data:") :].strip() not in ("", "[DONE]") + ] + payloads = [json.loads(raw) for raw in data_lines] + combined = "".join( + payload["choices"][0]["delta"].get("content", "") + for payload in payloads + if payload["choices"][0]["delta"] + ) + assert "plananswer" in combined diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py new file mode 100644 index 0000000000..0e668944f4 --- /dev/null +++ b/studio/backend/tests/test_providers_api.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Integration tests for the external providers API. + +Requires a running Unsloth Studio server. Configure via environment variables: + + export STUDIO_TEST_URL="http://localhost:8888" # default + export STUDIO_TEST_USER="unsloth" # default + export STUDIO_TEST_PASSWORD="..." # required — see .bootstrap_password + + # Provider API keys — any left unset will have their tests automatically skipped + export OPENAI_API_KEY="sk-..." + export MISTRAL_API_KEY="..." + export GOOGLE_API_KEY="..." + export TOGETHER_API_KEY="..." + export FIREWORKS_API_KEY="..." + export PERPLEXITY_API_KEY="..." + +Run: + cd studio/backend + pytest tests/test_providers_api.py -v -s +""" + +import base64 +import json +import os + +import pytest +import requests +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import padding + +# ── Configuration ───────────────────────────────────────────────── + +BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8000") +USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth") +PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "") + +# These tests require a live Studio server reachable at BASE_URL with a known +# bootstrap password. Skip the whole module when that environment is missing +# (e.g. on CI runners) so pytest discovery does not error out. +pytestmark = pytest.mark.skipif( + not PASSWORD, + reason = "Integration test requires a running Studio server; set STUDIO_TEST_PASSWORD to enable.", +) + +# Map provider_type → (env var name, model to use for inference test) +_PROVIDER_CONFIGS: dict[str, tuple[str, str]] = { + "openai": ("OPENAI_API_KEY", "gpt-4o-mini"), + "mistral": ("MISTRAL_API_KEY", "mistral-small-2506"), + "gemini": ("GEMINI_API_KEY", "gemini-3-flash-preview"), + "openrouter": ("OPENROUTER_API_KEY", "openai/gpt-4o-mini"), + "anthropic": ("ANTHROPIC_API_KEY", "claude-haiku-4-5"), + "deepseek": ("DEEPSEEK_API_KEY", "deepseek-chat"), + "huggingface": ("HUGGINGFACE_API_KEY", "meta-llama/Llama-3.3-70B-Instruct"), + "kimi": ("MOONSHOT_API_KEY", "moonshot-v1-8k"), + "qwen": ("DASHSCOPE_API_KEY", "qwen-turbo"), +} + +PROVIDER_KEYS: dict[str, str] = { + ptype: os.getenv(env_var, "") for ptype, (env_var, _) in _PROVIDER_CONFIGS.items() +} + +EXPECTED_PROVIDER_TYPES = set(_PROVIDER_CONFIGS.keys()) + +# ── Helpers ──────────────────────────────────────────────────────── + + +def _url(path: str) -> str: + return f"{BASE_URL}/{path.lstrip('/')}" + + +def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: + """ + Read a streaming SSE response and return (assembled_text, saw_done). + + Each chunk is a JSON object with choices[0].delta.content. + The stream ends with `data: [DONE]`. + """ + reply_parts: list[str] = [] + saw_done = False + + for raw_line in response.iter_lines(): + if isinstance(raw_line, bytes): + raw_line = raw_line.decode("utf-8") + if not raw_line.startswith("data:"): + continue + data = raw_line[len("data:") :].strip() + if data == "[DONE]": + saw_done = True + break + try: + chunk = json.loads(data) + # Handle both error payloads and normal chunks + if "error" in chunk: + raise RuntimeError(f"Provider error in stream: {chunk['error']}") + delta = chunk.get("choices", [{}])[0].get("delta", {}) + content = delta.get("content") or "" + if content: + reply_parts.append(content) + except (json.JSONDecodeError, IndexError, KeyError): + pass # skip malformed lines + + return "".join(reply_parts), saw_done + + +# ── Session-scoped fixtures ──────────────────────────────────────── + + +@pytest.fixture(scope = "session") +def auth_headers() -> dict[str, str]: + """ + Log in once per session and return auth headers. + + On a fresh Studio install the bootstrap password triggers a forced password + change (must_change_password=True). Any subsequent API call using that token + returns 403 "Password change required". This fixture detects that state, + automatically completes the change-password flow, and re-logs in so all other + tests get a fully usable token. + + The new password used during auto-change is: + STUDIO_TEST_NEW_PASSWORD (env var, optional) + or PASSWORD + "-test" (derived default) + + On the second run, set STUDIO_TEST_PASSWORD to the new password. + """ + assert PASSWORD, ( + "STUDIO_TEST_PASSWORD is not set.\n" + "Run: export STUDIO_TEST_PASSWORD=$(cat studio/backend/.bootstrap_password)" + ) + + resp = requests.post( + _url("/api/auth/login"), + json = {"username": USERNAME, "password": PASSWORD}, + timeout = 10, + ) + assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" + body = resp.json() + token = body["access_token"] + assert token, "access_token is empty" + + if body.get("must_change_password"): + # Bootstrap token is restricted — only /api/auth/change-password works with it. + # Auto-complete the forced change so the rest of the tests get a full token. + new_password = os.getenv("STUDIO_TEST_NEW_PASSWORD") or f"{PASSWORD}-test" + change_resp = requests.post( + _url("/api/auth/change-password"), + headers = {"Authorization": f"Bearer {token}"}, + json = {"current_password": PASSWORD, "new_password": new_password}, + timeout = 10, + ) + assert ( + change_resp.status_code == 200 + ), f"Auto password-change failed ({change_resp.status_code}): {change_resp.text}" + token = change_resp.json()["access_token"] + + return {"Authorization": f"Bearer {token}"} + + +@pytest.fixture(scope = "session") +def public_key_pem(auth_headers: dict[str, str]) -> str: + """Fetch RSA public key PEM once per session.""" + resp = requests.get( + _url("/api/providers/public-key"), + headers = auth_headers, + timeout = 10, + ) + assert resp.status_code == 200, f"Public key fetch failed: {resp.text}" + pem = resp.json().get("public_key", "") + assert pem.startswith("-----BEGIN PUBLIC KEY-----"), "Not a valid PEM public key" + return pem + + +@pytest.fixture(scope = "session") +def vision_image_data_url() -> str: + """ + Download the sloth image once per session and return it as a base64 data URI. + + Using a data URI instead of a remote URL ensures every provider receives + the image inline — Gemini's OpenAI-compatible layer does not fetch external + HTTP URLs, so raw image_url links silently produce empty replies for Gemini. + """ + resp = requests.get(_VISION_IMAGE_URL, timeout = 30) + resp.raise_for_status() + content_type = resp.headers.get("Content-Type", "image/jpeg").split(";")[0].strip() + b64 = base64.b64encode(resp.content).decode("utf-8") + return f"data:{content_type};base64,{b64}" + + +@pytest.fixture(scope = "session") +def encrypt_key(public_key_pem: str): + """ + Return a callable encrypt_key(plaintext: str) -> str (base64 RSA-OAEP ciphertext). + Uses the backend's RSA public key — mirrors what the frontend does. + """ + # Decode PEM → load RSA public key + pem_bytes = public_key_pem.encode("utf-8") + rsa_pub = serialization.load_pem_public_key(pem_bytes) + + def _encrypt(plaintext: str) -> str: + ciphertext = rsa_pub.encrypt( + plaintext.encode("utf-8"), + padding.OAEP( + mgf = padding.MGF1(algorithm = hashes.SHA256()), + algorithm = hashes.SHA256(), + label = None, + ), + ) + return base64.b64encode(ciphertext).decode("utf-8") + + return _encrypt + + +# ── TestAuth ──────────────────────────────────────────────────────── + + +class TestAuth: + def test_login_returns_token(self): + """POST /api/auth/login returns a non-empty access_token.""" + assert PASSWORD, "STUDIO_TEST_PASSWORD not set" + resp = requests.post( + _url("/api/auth/login"), + json = {"username": USERNAME, "password": PASSWORD}, + timeout = 10, + ) + assert ( + resp.status_code == 200 + ), f"Login failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert body.get("access_token"), "access_token is missing or empty" + assert body.get("token_type") == "bearer" + + +# ── TestPublicKey ──────────────────────────────────────────────────── + + +class TestPublicKey: + def test_public_key_is_valid_pem( + self, auth_headers: dict[str, str], public_key_pem: str + ): + """GET /api/providers/public-key returns an importable RSA PEM key.""" + pem_bytes = public_key_pem.encode("utf-8") + key = serialization.load_pem_public_key(pem_bytes) + key_size = key.key_size # type: ignore[attr-defined] + assert key_size >= 2048, f"Key size too small: {key_size}" + print(f"\n RSA-{key_size} public key OK") + + +# ── TestRegistry ──────────────────────────────────────────────────── + + +class TestRegistry: + def test_registry_returns_all_providers(self, auth_headers: dict[str, str]): + """GET /api/providers/registry returns all supported providers.""" + resp = requests.get( + _url("/api/providers/registry"), + headers = auth_headers, + timeout = 10, + ) + assert resp.status_code == 200, f"Registry failed: {resp.text}" + providers = resp.json() + assert ( + len(providers) == 9 + ), f"Expected 9 providers, got {len(providers)}: {providers}" + print(f"\n {'Provider':<12} {'Base URL'}") + print(f" {'-'*12} {'-'*45}") + for p in providers: + print(f" {p['provider_type']:<12} {p['base_url']}") + + def test_registry_has_expected_types(self, auth_headers: dict[str, str]): + """All expected provider_type values are present in the registry.""" + resp = requests.get( + _url("/api/providers/registry"), + headers = auth_headers, + timeout = 10, + ) + assert resp.status_code == 200 + returned_types = {p["provider_type"] for p in resp.json()} + missing = EXPECTED_PROVIDER_TYPES - returned_types + assert not missing, f"Missing provider types: {missing}" + + def test_registry_entries_have_required_fields(self, auth_headers: dict[str, str]): + """Each registry entry has provider_type, display_name, base_url, default_models.""" + resp = requests.get( + _url("/api/providers/registry"), headers = auth_headers, timeout = 10 + ) + assert resp.status_code == 200 + for entry in resp.json(): + for field in ( + "provider_type", + "display_name", + "base_url", + "default_models", + "model_list_mode", + ): + assert field in entry, f"Missing field '{field}' in entry: {entry}" + assert entry["model_list_mode"] in ("remote", "curated") + assert isinstance(entry["default_models"], list) + assert len(entry["default_models"]) > 0 + + +# ── TestProviderCRUD ──────────────────────────────────────────────── + + +class TestProviderCRUD: + """ + These tests run sequentially within the class and share state via class variables. + They create, read, update, and delete a single test provider config. + """ + + _created_id: str = "" + + def test_create_provider(self, auth_headers: dict[str, str]): + """POST /api/providers/ creates a provider config and returns 201.""" + resp = requests.post( + _url("/api/providers/"), + headers = auth_headers, + json = {"provider_type": "openai", "display_name": "Test OpenAI (pytest)"}, + timeout = 10, + ) + assert ( + resp.status_code == 201 + ), f"Create failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert body.get("id"), "No id in response" + assert body["provider_type"] == "openai" + assert body["display_name"] == "Test OpenAI (pytest)" + assert body["is_enabled"] is True + TestProviderCRUD._created_id = body["id"] + print(f"\n created id={body['id']}") + + def test_list_includes_created(self, auth_headers: dict[str, str]): + """GET /api/providers/ includes the newly created config.""" + assert ( + TestProviderCRUD._created_id + ), "No created_id (run test_create_provider first)" + resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10) + assert resp.status_code == 200 + ids = [p["id"] for p in resp.json()] + assert ( + TestProviderCRUD._created_id in ids + ), f"Created id {TestProviderCRUD._created_id!r} not found in list: {ids}" + print(f"\n found id={TestProviderCRUD._created_id} in list of {len(ids)}") + + def test_update_display_name(self, auth_headers: dict[str, str]): + """PUT /api/providers/{id} updates the display_name.""" + assert TestProviderCRUD._created_id, "No created_id" + new_name = "Test OpenAI (pytest updated)" + resp = requests.put( + _url(f"/api/providers/{TestProviderCRUD._created_id}"), + headers = auth_headers, + json = {"display_name": new_name}, + timeout = 10, + ) + assert ( + resp.status_code == 200 + ), f"Update failed ({resp.status_code}): {resp.text}" + assert resp.json()["display_name"] == new_name + print(f"\n updated display_name to '{new_name}'") + + def test_delete_provider(self, auth_headers: dict[str, str]): + """DELETE /api/providers/{id} removes the config (204) and it's gone from list.""" + assert TestProviderCRUD._created_id, "No created_id" + resp = requests.delete( + _url(f"/api/providers/{TestProviderCRUD._created_id}"), + headers = auth_headers, + timeout = 10, + ) + assert ( + resp.status_code == 204 + ), f"Delete failed ({resp.status_code}): {resp.text}" + + # Confirm gone from list + list_resp = requests.get( + _url("/api/providers/"), headers = auth_headers, timeout = 10 + ) + ids = [p["id"] for p in list_resp.json()] + assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list" + print(f"\n deleted id={TestProviderCRUD._created_id} confirmed gone") + + +# ── TestProviderInference ──────────────────────────────────────────── + + +# Build parametrize list: (provider_type, model, api_key) for configured providers only +_INFERENCE_PARAMS = [ + pytest.param( + ptype, + model, + PROVIDER_KEYS.get(ptype, ""), + id = ptype, + marks = pytest.mark.skipif( + not PROVIDER_KEYS.get(ptype, ""), + reason = f"no {env_var} set", + ), + ) + for ptype, (env_var, model) in _PROVIDER_CONFIGS.items() +] + + +class TestProviderInference: + """ + Live inference tests — one parametrized set per provider. + Each test is automatically skipped when the provider's API key env var is not set. + """ + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_connection( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /api/providers/test → success: true.""" + encrypted = encrypt_key(api_key) + resp = requests.post( + _url("/api/providers/test"), + headers = auth_headers, + json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, + timeout = 30, + ) + assert ( + resp.status_code == 200 + ), f"Request failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert ( + body["success"] is True + ), f"Connection test failed for {provider_type}: {body.get('message')}" + print(f"\n [{provider_type}] connection OK — {body['message']}") + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_list_models( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /api/providers/models → non-empty list, print first 3.""" + encrypted = encrypt_key(api_key) + resp = requests.post( + _url("/api/providers/models"), + headers = auth_headers, + json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, + timeout = 30, + ) + assert ( + resp.status_code == 200 + ), f"Request failed ({resp.status_code}): {resp.text}" + models = resp.json() + assert isinstance(models, list), f"Expected list, got {type(models)}" + assert len(models) > 0, f"No models returned for {provider_type}" + preview = [m["id"] for m in models[:3]] + print(f"\n [{provider_type}] {len(models)} models — first 3: {preview}") + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_chat_inference( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /v1/chat/completions with provider fields → streamed reply.""" + encrypted = encrypt_key(api_key) + payload = { + "messages": [{"role": "user", "content": "Say hello in one sentence."}], + "stream": True, + "temperature": 0.7, + "max_tokens": 64, + "provider_type": provider_type, + "external_model": model, + "encrypted_api_key": encrypted, + } + with requests.post( + _url("/v1/chat/completions"), + headers = {**auth_headers, "Content-Type": "application/json"}, + json = payload, + stream = True, + timeout = 60, + ) as resp: + assert ( + resp.status_code == 200 + ), f"Chat completions failed ({resp.status_code}): {resp.text[:500]}" + reply, saw_done = _parse_sse_stream(resp) + + assert reply.strip(), f"Empty reply from {provider_type}/{model}" + assert saw_done, f"Stream did not end with [DONE] for {provider_type}/{model}" + print(f'\n [{provider_type}/{model}] reply: "{reply.strip()}"') + + +# ── TestVisionInference ───────────────────────────────────────────── + +# Sloth photo — used to test vision routing across providers +_VISION_IMAGE_URL = ( + "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" +) + +_VISION_PARAMS = [ + pytest.param( + ptype, + model, + PROVIDER_KEYS.get(ptype, ""), + id = ptype, + marks = pytest.mark.skipif( + not PROVIDER_KEYS.get(ptype, ""), + reason = f"no key for {ptype}", + ), + ) + for ptype, (_, model) in _PROVIDER_CONFIGS.items() + if ptype in {"openai", "mistral", "gemini", "anthropic", "openrouter"} +] + + +class TestVisionInference: + """ + Send a 1×1 white PNG alongside a text question to each vision-capable provider. + Verifies that image content parts survive the proxy and the provider replies. + """ + + @pytest.mark.parametrize("provider_type,model,api_key", _VISION_PARAMS) + def test_vision_chat_inference( + self, + auth_headers: dict[str, str], + encrypt_key, + vision_image_data_url: str, + provider_type: str, + model: str, + api_key: str, + ): + """Image URL + text message → non-empty streamed reply.""" + encrypted = encrypt_key(api_key) + payload = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Which animal is in this image? Reply in one word.", + }, + { + "type": "image_url", + "image_url": {"url": vision_image_data_url}, + }, + ], + } + ], + "stream": True, + "max_tokens": 215, + "provider_type": provider_type, + "external_model": model, + "encrypted_api_key": encrypted, + } + with requests.post( + _url("/v1/chat/completions"), + headers = {**auth_headers, "Content-Type": "application/json"}, + json = payload, + stream = True, + timeout = 60, + ) as resp: + assert ( + resp.status_code == 200 + ), f"Vision request failed ({resp.status_code}): {resp.text[:300]}" + reply, saw_done = _parse_sse_stream(resp) + + assert reply.strip(), f"Empty reply from {provider_type}/{model}" + assert saw_done, f"Stream did not end with [DONE] for {provider_type}/{model}" + print(f"\n [{provider_type}/{model}] vision reply: {reply.strip()!r}") + + +# ── TestLocalInferenceUnaffected ──────────────────────────────────── + + +class TestLocalInferenceUnaffected: + def test_chat_without_provider(self, auth_headers: dict[str, str]): + """ + POST /v1/chat/completions without provider fields must not return 422 or 500. + + 200 = a local model is loaded and responded. + 503 = no model loaded (expected in test environment — that's fine). + Any other 4xx/5xx (except 503) = regression in request handling. + """ + resp = requests.post( + _url("/v1/chat/completions"), + headers = {**auth_headers, "Content-Type": "application/json"}, + json = { + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + timeout = 15, + ) + allowed = {200, 400, 503} + assert resp.status_code in allowed, ( + f"Unexpected status {resp.status_code} for local inference path: {resp.text[:300]}\n" + f"This likely means the provider fields broke the base request schema." + ) + status_label = ( + "local model responded" + if resp.status_code == 200 + else "no model loaded (expected)" + ) + print(f"\n status={resp.status_code} ({status_label}) — local path unaffected") diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index 21d31d81e6..464f47c09c 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -58,6 +58,7 @@ "motion": "^12.34.0", "next": "^16.1.6", "next-themes": "^0.4.6", + "node-forge": "^1.4.0", "radix-ui": "^1.4.3", "react": "^19.2.4", "react-day-picker": "^9.13.2", @@ -80,6 +81,7 @@ "@eslint/js": "^9.39.1", "@types/js-yaml": "^4.0.9", "@types/node": "^25.5.2", + "@types/node-forge": "^1.3.14", "@types/react": "^19.2.5", "@types/react-dom": "^19.2.3", "@vitejs/plugin-react": "^6.0.1", @@ -7377,6 +7379,16 @@ "undici-types": "~7.19.0" } }, + "node_modules/@types/node-forge": { + "version": "1.3.14", + "resolved": "https://registry.npmjs.org/@types/node-forge/-/node-forge-1.3.14.tgz", + "integrity": "sha512-mhVF2BnD4BO+jtOp7z1CdzaK4mbuK0LLQYAvdOLqHTavxFNq4zA1EmYkpnFjP8HOUzedfQkRnp0E2ulSAYSzAw==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@types/react": { "version": "19.2.14", "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.14.tgz", @@ -13285,6 +13297,15 @@ "url": "https://opencollective.com/node-fetch" } }, + "node_modules/node-forge": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/node-forge/-/node-forge-1.4.0.tgz", + "integrity": "sha512-LarFH0+6VfriEhqMMcLX2F7SwSXeWwnEAJEsYm5QKWchiVYVvJyV9v7UDvUv+w5HO23ZpQTXDv/GxdDdMyOuoQ==", + "license": "(BSD-3-Clause OR GPL-2.0)", + "engines": { + "node": ">= 6.13.0" + } + }, "node_modules/node-releases": { "version": "2.0.38", "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.38.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 3a02bb926a..c69b2fdf3e 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -66,6 +66,7 @@ "motion": "^12.34.0", "next": "^16.1.6", "next-themes": "^0.4.6", + "node-forge": "^1.4.0", "radix-ui": "^1.4.3", "react": "^19.2.4", "react-day-picker": "^9.13.2", @@ -92,6 +93,7 @@ "@biomejs/biome": "^1.9.4", "@eslint/js": "^9.39.1", "@types/js-yaml": "^4.0.9", + "@types/node-forge": "^1.3.14", "@types/node": "^25.5.2", "@types/react": "^19.2.5", "@types/react-dom": "^19.2.3", diff --git a/studio/frontend/public/provider-logos/anthropic.svg b/studio/frontend/public/provider-logos/anthropic.svg new file mode 100644 index 0000000000..7545cc8f3e --- /dev/null +++ b/studio/frontend/public/provider-logos/anthropic.svg @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/deepseek.svg b/studio/frontend/public/provider-logos/deepseek.svg new file mode 100644 index 0000000000..d1ba06b942 --- /dev/null +++ b/studio/frontend/public/provider-logos/deepseek.svg @@ -0,0 +1,14 @@ + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/gemini.svg b/studio/frontend/public/provider-logos/gemini.svg new file mode 100644 index 0000000000..9090dfb68e --- /dev/null +++ b/studio/frontend/public/provider-logos/gemini.svg @@ -0,0 +1,72 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/studio/frontend/public/provider-logos/huggingface.svg b/studio/frontend/public/provider-logos/huggingface.svg new file mode 100644 index 0000000000..ab959d165f --- /dev/null +++ b/studio/frontend/public/provider-logos/huggingface.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/studio/frontend/public/provider-logos/kimi.jpg b/studio/frontend/public/provider-logos/kimi.jpg new file mode 100644 index 0000000000..956a5b58b1 Binary files /dev/null and b/studio/frontend/public/provider-logos/kimi.jpg differ diff --git a/studio/frontend/public/provider-logos/misc/meta.svg b/studio/frontend/public/provider-logos/misc/meta.svg new file mode 100644 index 0000000000..9fa656bd6b --- /dev/null +++ b/studio/frontend/public/provider-logos/misc/meta.svg @@ -0,0 +1,19 @@ + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/microsoft.svg b/studio/frontend/public/provider-logos/misc/microsoft.svg new file mode 100644 index 0000000000..5334aa7ca6 --- /dev/null +++ b/studio/frontend/public/provider-logos/misc/microsoft.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/minimax.png b/studio/frontend/public/provider-logos/misc/minimax.png new file mode 100644 index 0000000000..e9472c676d Binary files /dev/null and b/studio/frontend/public/provider-logos/misc/minimax.png differ diff --git a/studio/frontend/public/provider-logos/misc/nvidia.svg b/studio/frontend/public/provider-logos/misc/nvidia.svg new file mode 100644 index 0000000000..ae65b09a2b --- /dev/null +++ b/studio/frontend/public/provider-logos/misc/nvidia.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/perplexity.png b/studio/frontend/public/provider-logos/misc/perplexity.png new file mode 100644 index 0000000000..9845765c7f Binary files /dev/null and b/studio/frontend/public/provider-logos/misc/perplexity.png differ diff --git a/studio/frontend/public/provider-logos/misc/xai.svg b/studio/frontend/public/provider-logos/misc/xai.svg new file mode 100644 index 0000000000..0c83eb3d9b --- /dev/null +++ b/studio/frontend/public/provider-logos/misc/xai.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/z-ai.svg b/studio/frontend/public/provider-logos/misc/z-ai.svg new file mode 100644 index 0000000000..28ca7280a1 --- /dev/null +++ b/studio/frontend/public/provider-logos/misc/z-ai.svg @@ -0,0 +1,215 @@ + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/mistral.svg b/studio/frontend/public/provider-logos/mistral.svg new file mode 100644 index 0000000000..40c2591b31 --- /dev/null +++ b/studio/frontend/public/provider-logos/mistral.svg @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + + diff --git a/studio/frontend/public/provider-logos/openai.svg b/studio/frontend/public/provider-logos/openai.svg new file mode 100644 index 0000000000..74d9b1b44b --- /dev/null +++ b/studio/frontend/public/provider-logos/openai.svg @@ -0,0 +1,5 @@ + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/openrouter.svg b/studio/frontend/public/provider-logos/openrouter.svg new file mode 100644 index 0000000000..4a4968b639 --- /dev/null +++ b/studio/frontend/public/provider-logos/openrouter.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/qwen.png b/studio/frontend/public/provider-logos/qwen.png new file mode 100644 index 0000000000..67d2258f40 Binary files /dev/null and b/studio/frontend/public/provider-logos/qwen.png differ diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 22bd7412ab..b4f7dd08d2 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -13,21 +13,71 @@ import { usePlatformStore } from "@/config/env"; import { cn } from "@/lib/utils"; import { ArrowDown01Icon, + CloudIcon, FolderSearchIcon, Logout01Icon, + Search01Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMemo, useState } from "react"; import type { DeletedModelRef, + ExternalModelOption, LoraModelOption, ModelOption, ModelSelectorChangeMeta, } from "./model-selector/types"; import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; +import { Input } from "../ui/input"; + +const PROVIDER_LOGO_EXT: Record = { + openai: "svg", + mistral: "svg", + gemini: "svg", + anthropic: "svg", + deepseek: "svg", + huggingface: "svg", + kimi: "jpg", + qwen: "png", + openrouter: "svg", +}; + +function providerLogoSrc(providerType: string | undefined): string | undefined { + if (!providerType) return undefined; + const ext = PROVIDER_LOGO_EXT[providerType]; + if (!ext) return undefined; + return `${import.meta.env.BASE_URL}provider-logos/${providerType}.${ext}`; +} + +function ExternalProviderLogo({ + providerType, + className, + title, +}: { + providerType: string | undefined; + className?: string; + title?: string; +}) { + const src = providerLogoSrc(providerType); + if (!src) return null; + return ( + + ); +} export type { DeletedModelRef, + ExternalModelOption, LoraModelOption, ModelOption, ModelSelectorChangeMeta, @@ -36,6 +86,7 @@ export type { interface ModelSelectorProps { models: ModelOption[]; loraModels?: LoraModelOption[]; + externalModels?: ExternalModelOption[]; value?: string; defaultValue?: string; activeGgufVariant?: string | null; @@ -53,11 +104,13 @@ interface ModelSelectorProps { onOpenChange?: (open: boolean) => void; triggerDataTour?: string; contentDataTour?: string; + showCloudIndicator?: boolean; } function ModelSelectorTrigger({ currentModel, isLoaded, + showCloudIndicator = false, variant = "outline", size = "default", className, @@ -65,6 +118,7 @@ function ModelSelectorTrigger({ }: { currentModel?: ModelOption; isLoaded: boolean; + showCloudIndicator?: boolean; variant?: "outline" | "ghost" | "muted"; size?: "sm" | "default" | "lg"; className?: string; @@ -90,12 +144,27 @@ function ModelSelectorTrigger({ {isLoaded && ( )} - - + {currentModel?.icon ? ( + {currentModel.icon} + ) : null} + + {currentModel?.name ?? "Select model"} + {showCloudIndicator ? ( + + ) : null} {currentModel?.description && ( - + {currentModel.description} )} @@ -115,6 +184,7 @@ function ModelSelectorTrigger({ function ModelSelectorContent({ models, loraModels, + externalModels, value, onSelect, onEject, @@ -127,6 +197,7 @@ function ModelSelectorContent({ }: { models: ModelOption[]; loraModels: LoraModelOption[]; + externalModels: ExternalModelOption[]; value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; @@ -139,6 +210,20 @@ function ModelSelectorContent({ }) { const hasSelection = Boolean(value); const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const hasExternal = externalModels.length > 0; + const chatOnlyTabsDefault = useMemo( + () => (value && externalModels.some((model) => model.id === value) ? "external" : "hub"), + [externalModels, value], + ); + const studioTabsDefault = useMemo((): "hub" | "lora" | "external" => { + if (value && externalModels.some((model) => model.id === value)) { + return "external"; + } + if (value && loraModels.some((model) => model.id === value)) { + return "lora"; + } + return "hub"; + }, [externalModels, loraModels, value]); return ( {chatOnly ? ( - + hasExternal ? ( + + + Hub models + External + + + + + + + + + ) : ( + + ) ) : ( - + Hub models Fine-tuned + {hasExternal ? External : null} @@ -171,6 +276,16 @@ function ModelSelectorContent({ deleteDisabled={deleteDisabled} /> + + {hasExternal ? ( + + + + ) : null} )} @@ -207,6 +322,7 @@ function ModelSelectorContent({ export function ModelSelector({ models, loraModels = [], + externalModels = [], value, defaultValue, activeGgufVariant, @@ -224,6 +340,7 @@ export function ModelSelector({ onOpenChange, triggerDataTour, contentDataTour, + showCloudIndicator = false, }: ModelSelectorProps) { const [uncontrolledOpen, setUncontrolledOpen] = useState(false); const open = controlledOpen ?? uncontrolledOpen; @@ -266,8 +383,21 @@ export function ModelSelector({ description: tag, }); } + for (const externalModel of externalModels) { + all.set(externalModel.id, { + ...externalModel, + description: externalModel.providerName, + icon: ( + + ), + }); + } return all; - }, [loraModels, models]); + }, [externalModels, loraModels, models]); const currentModel = useMemo(() => { if (!selected) return undefined; @@ -303,6 +433,7 @@ export function ModelSelector({ void; +}) { + const [query, setQuery] = useState(""); + const grouped = useMemo(() => { + const needle = normalizeForSearch(query.trim()); + const byProvider = new Map< + string, + { providerName: string; models: ExternalModelOption[] } + >(); + for (const model of externalModels) { + const searchText = normalizeForSearch( + `${model.name} ${model.providerName} ${model.id}`, + ); + if (needle && !searchText.includes(needle)) continue; + const prev = byProvider.get(model.providerId); + if (prev) { + prev.models.push(model); + } else { + byProvider.set(model.providerId, { + providerName: model.providerName, + models: [model], + }); + } + } + return [...byProvider.entries()] + .map(([providerId, group]) => ({ + providerId, + providerName: group.providerName, + models: group.models.sort((a, b) => a.name.localeCompare(b.name)), + })) + .sort((a, b) => a.providerName.localeCompare(b.providerName)); + }, [externalModels, query]); + + return ( +
+
+ + setQuery(event.target.value)} + placeholder="Search external models" + className="h-9 pl-8" + /> +
+
+
+ {grouped.length === 0 ? ( +
+ No external models configured. +
+ ) : ( + grouped.map((group) => ( +
+
+ + {group.providerName} +
+ {group.models.map((model) => ( + + ))} +
+ )) + )} +
+
+
+ ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index 3da75b4d4e..4cc5d779ce 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -18,8 +18,15 @@ export interface LoraModelOption extends ModelOption { exportType?: "lora" | "merged" | "gguf"; } +export interface ExternalModelOption extends ModelOption { + providerId: string; + providerName: string; + /** Registry key (e.g. openai, gemini) for provider branding. */ + providerType: string; +} + export interface ModelSelectorChangeMeta { - source: "hub" | "lora" | "exported" | "local"; + source: "hub" | "lora" | "exported" | "local" | "external"; isLora: boolean; ggufVariant?: string; isDownloaded?: boolean; diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 417637801c..d75a8cce10 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -31,6 +31,9 @@ import { DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { sentAudioNames } from "@/features/chat/api/chat-adapter"; +import { parseExternalModelId } from "@/features/chat/external-providers"; +import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities"; +import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params"; import { isTauri } from "@/lib/api-base"; @@ -474,15 +477,69 @@ const ReasoningToggle: FC = () => { const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, ); + const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning); + const reasoningAlwaysOn = useChatRuntimeStore((s) => s.reasoningAlwaysOn); const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled); const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled); const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle); const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort); + const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff); + const reasoningEffortLevels = useChatRuntimeStore((s) => s.reasoningEffortLevels); const setReasoningEffort = useChatRuntimeStore((s) => s.setReasoningEffort); - const disabled = !(modelLoaded && supportsReasoning); + const lastOpenRouterChosenModel = useChatRuntimeStore( + (s) => s.lastOpenRouterChosenModel, + ); + const externalProviders = useExternalProvidersStore((s) => s.providers); + const externalSelection = parseExternalModelId(checkpoint); + const selectedExternalProvider = + externalSelection != null + ? externalProviders.find((p) => p.id === externalSelection.providerId) + : undefined; + const effectiveExternalModelId = + selectedExternalProvider?.providerType === "openrouter" && + externalSelection?.modelId === "openrouter/free" && + lastOpenRouterChosenModel + ? lastOpenRouterChosenModel + : externalSelection?.modelId; + const externalReasoningCaps = + externalSelection != null + ? getExternalReasoningCapabilities( + selectedExternalProvider?.providerType, + effectiveExternalModelId, + ) + : null; + const effectiveReasoningStyle = + externalReasoningCaps?.reasoningStyle ?? reasoningStyle; + const effectiveReasoningAlwaysOn = + externalReasoningCaps?.reasoningAlwaysOn ?? reasoningAlwaysOn; + const effectiveSupportsReasoningOff = + externalReasoningCaps?.supportsReasoningOff ?? supportsReasoningOff; + const effectiveReasoningEffortLevels = + externalReasoningCaps?.reasoningEffortLevels ?? reasoningEffortLevels; + const effectiveSupportsReasoning = + externalReasoningCaps?.supportsReasoning ?? supportsReasoning; + const reasoningLockedOn = + effectiveSupportsReasoning && + (effectiveReasoningAlwaysOn || !effectiveSupportsReasoningOff); + const effectiveReasoningEnabled = reasoningLockedOn ? true : reasoningEnabled; + const effectiveReasoningVisualEnabled = + effectiveReasoningEnabled && reasoningEffort !== "none"; + const disabled = !(modelLoaded && effectiveSupportsReasoning); + const formatEffortLabel = (level: typeof reasoningEffort): string => { + if (level !== "xhigh") return level.charAt(0).toUpperCase() + level.slice(1); + const normalized = externalSelection?.modelId?.trim().toLowerCase() ?? ""; + if ( + normalized.startsWith("claude-opus-4-6") || + normalized.startsWith("claude-sonnet-4-6") + ) { + return "Max"; + } + return "Extra High"; + }; + const effortLabel = formatEffortLabel(reasoningEffort); - if (reasoningStyle === "reasoning_effort") { + if (effectiveReasoningStyle === "reasoning_effort") { return ( @@ -493,26 +550,47 @@ const ReasoningToggle: FC = () => { "flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors", disabled ? "cursor-not-allowed opacity-40" - : "bg-primary/10 text-primary hover:bg-primary/20", + : effectiveReasoningVisualEnabled + ? "bg-primary/10 text-primary hover:bg-primary/20" + : "text-muted-foreground hover:bg-muted-foreground/15", )} aria-label={`Reasoning effort: ${reasoningEffort}`} > - + {effectiveReasoningVisualEnabled ? ( + + ) : ( + + )} - Think:{" "} - {reasoningEffort.charAt(0).toUpperCase() + - reasoningEffort.slice(1)} + Think: {effectiveReasoningVisualEnabled ? effortLabel : "None"} - {(["low", "medium", "high"] as const).map((level) => ( + {effectiveSupportsReasoningOff && ( + { + setReasoningEnabled(false); + applyQwenThinkingParams(false); + }} + > + None + {!effectiveReasoningVisualEnabled ? " \u2713" : ""} + + )} + {effectiveReasoningEffortLevels + .filter((level) => level !== "none") + .map((level) => ( setReasoningEffort(level)} + onSelect={() => { + setReasoningEffort(level); + setReasoningEnabled(true); + applyQwenThinkingParams(true); + }} > - {level.charAt(0).toUpperCase() + level.slice(1)} - {reasoningEffort === level ? " \u2713" : ""} + {formatEffortLabel(level)} + {effectiveReasoningVisualEnabled && reasoningEffort === level ? " \u2713" : ""} ))} @@ -523,17 +601,34 @@ const ReasoningToggle: FC = () => { return ( +
+ + Cloud + + + + {editingProviderId ? "Edit" : "New"} + +
+ + +
+
+
+
+
+ +

+ Supported registry or Custom. +

+
+ +
+ +
+
+ +

+ Stored locally. +

+
+
+ setApiKey(event.target.value)} + placeholder="Enter API key" + className="h-9 pr-9 text-sm" + /> + +
+
+ + {isCustomProvider ? ( +
+ + + setCustomProviderName(event.target.value) + } + placeholder="Custom" + className="h-9 text-sm" + /> +
+ ) : null} + + {isCustomProvider ? ( +
+
+ +

+ OpenAI-compatible endpoint. +

+
+ setBaseUrlDraft(event.target.value)} + placeholder="https://my-vllm-server.com/v1" + className="h-9 text-sm" + /> +
+ ) : null} +
+
+ +
+ + +
+
+ +

+ {modelStatusLabel} +

+
+ +
+ {isCustomProvider ? ( +
+
+ +