* Studio: add API key authentication for programmatic access External users want to hit the Studio API (chat completions with tool calling, training, export, etc.) without going through the browser login flow. This adds sk-unsloth- prefixed API keys that work as a drop-in replacement for JWTs in the Authorization: Bearer header. Backend: - New api_keys table in SQLite (storage.py) - create/list/revoke/validate functions with SHA-256 hashed storage - API key detection in _get_current_subject before the JWT path - POST/GET/DELETE /api/auth/api-keys endpoints on the auth router Frontend: - /api-keys page with create form, one-time key reveal, keys table - API Keys link in desktop and mobile navbar - Route registered with requireAuth guard Zero changes to any existing route handler -- every endpoint that uses Depends(get_current_subject) automatically works with API keys. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use actual origin in API key usage examples The examples on /api-keys were hardcoded to localhost:8888 which is wrong for remote users. Use window.location.origin so the examples show the correct URL regardless of where the user is connecting from. * Add `unsloth studio run` CLI command for one-liner model serving Adds a `run` subcommand that starts Studio, loads a model, creates an API key, and prints a ready-to-use curl command -- similar to `ollama run` or `vllm serve`. Usage: unsloth studio run -m unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add end-to-end tests for `unsloth studio run` and API key usage Tests the 4 usage examples from the API Keys page: 1. curl basic (non-streaming) chat completions 2. curl streaming (SSE) chat completions 3. OpenAI Python SDK streaming completions 4. curl with tools (web_search + python) Also tests --help output, invalid key rejection, and no-key rejection. All 7 tests pass against Qwen3-1.7B-GGUF. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add /v1/completions, /v1/embeddings, /v1/responses endpoints and --parallel support - llama_cpp.py: accept n_parallel param, pass to llama-server --parallel - run.py: plumb llama_parallel_slots through to app.state - inference.py: add /completions and /embeddings as transparent proxies to llama-server, add /responses as application-level endpoint that converts to ChatCompletionRequest; thread n_parallel through load_model - studio.py: set llama_parallel_slots=4 for `unsloth studio run` path * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make /v1/responses endpoint match OpenAI Responses API format The existing /v1/responses shim returned Chat Completions format, which broke OpenAI SDK clients using openai.responses.create(). This commit replaces the endpoint with a proper implementation that: - Returns `output` array with `output_text` content parts instead of `choices` with `message` - Uses `input_tokens`/`output_tokens` instead of `prompt_tokens`/ `completion_tokens` in usage - Sets `object: "response"` and `id: "resp_..."` - Emits named SSE events for streaming (response.created, response.output_text.delta, response.completed, etc.) - Accepts all OpenAI Responses API fields (tools, store, metadata, previous_response_id) without erroring -- silently ignored - Maps `developer` role to `system` and `input_text`/`input_image` content parts to the internal Chat format Adds Pydantic schemas for request/response models and 23 unit tests covering schema validation, input normalisation, and response format. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add Anthropic-compatible /v1/messages endpoint (#4981) * Add Anthropic-compatible /v1/messages endpoint with tool support Translate Anthropic Messages API format to/from internal OpenAI format and reuse the existing server-side agentic tool loop. Supports streaming SSE (message_start, content_block_delta, etc.) and non-streaming JSON. Includes offline unit tests and e2e tests in test_studio_run.py. * Add enable_tools, enabled_tools, session_id to /v1/messages endpoint Support the same shorthand as /v1/chat/completions: enable_tools=true with an optional enabled_tools list uses built-in server tools without requiring full Anthropic tool definitions. session_id is passed through for sandbox isolation. max_tokens is now optional. * Strip leaked tool-call XML from Anthropic endpoint content Apply _TOOL_XML_RE to content events in both streaming and non-streaming tool paths, matching the OpenAI endpoint behavior. * Emit custom tool_result SSE event in Anthropic stream Adds a non-standard tool_result event between the tool_use block close and the next text block, so clients can see server-side tool execution results. Anthropic SDKs ignore unknown event types. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Split /v1/messages into server-side and client-side tool paths enable_tools=true runs the existing server-side agentic loop with built-in tools (web_search/python/terminal). A bare tools=[...] field now triggers a client-side pass-through: client-provided tools are forwarded to llama-server and any tool_use output is returned to the caller with stop_reason=tool_use for client execution. This fixes Claude Code (and any Anthropic SDK client) which sends tools=[...] expecting client-side execution but was previously routed through execute_tool() and failing with 'Unknown tool'. Adds AnthropicPassthroughEmitter to convert llama-server OpenAI SSE chunks into Anthropic SSE events, plus unit tests covering text blocks, tool_use blocks, mixed, stop reasons, and usage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix httpcore GeneratorExit in /v1/messages passthrough stream Explicitly aclose aiter_lines() before the surrounding async with blocks unwind, mirroring the prior fix in external_provider.py (a41160d3) and cc757b78's RuntimeError suppression. * Wire stop_sequences through /v1/messages; warn on tool_choice Plumb payload.stop_sequences to all three code paths (server-side tool loop, no-tool plain, client-side passthrough) so Anthropic SDK clients setting stop_sequences get the behavior they expect. The llama_cpp backend already accepted `stop` on both generate_chat_ completion and generate_chat_completion_with_tools; the Anthropic handler simply wasn't passing it. tool_choice remains declared on the request model for Anthropic SDK compatibility (the SDK often sets it by default) but is not yet honored. Log a structured warning on each request carrying a non- null tool_choice so the silent drop is visible to operators. * Wire min_p / repetition_penalty / presence_penalty through /v1/messages Align the Anthropic endpoint's sampling surface with /v1/chat/completions. Adds the three fields as x-unsloth extensions on AnthropicMessagesRequest and threads them through all three code paths: server-side tool loop, no-tool plain, and client-side passthrough. The passthrough builder emits "repeat_penalty" (not "repetition_penalty") because that is llama-server's field name; the backend methods already apply the same rename internally. * Fix block ordering and prev_text reset in non-streaming tool path _anthropic_tool_non_streaming was building the response by appending all tool_use blocks first, then a single concatenated text block at the end — losing generation order and merging pre-tool and post-tool text into one block. It also never reset prev_text between synthesis turns, so the first N characters of each post-tool turn were dropped (where N = length of the prior turn's final cumulative text). Rewrite to build content_blocks incrementally in generation order, matching the streaming emitter's behavior: deltas within a turn are merged into the trailing text block, tool_use blocks interrupt the text sequence, and prev_text is reset on tool_end so turn N+1 diffs against an empty baseline. Caught by gemini-code-assist[bot] review on #4981. * Make test_studio_run.py e2e tests pytest-compatible Add a hybrid session-scoped studio_server fixture in conftest.py that feeds base_url / api_key into the existing e2e test functions. Three invocation modes are now supported: 1. Script mode (unchanged) — python tests/test_studio_run.py 2. Pytest + external server — point at a running instance via UNSLOTH_E2E_BASE_URL / UNSLOTH_E2E_API_KEY env vars, no per-run GGUF load cost 3. Pytest + fixture-managed server — pytest drives _start_server / _kill_server itself via --unsloth-model / --unsloth-gguf-variant, CI-friendly The existing _start_server / _kill_server helpers and main() stay untouched so the script entry point keeps working exactly as before. Test function signatures are unchanged — the (base_url, api_key) parameters now resolve via the new fixtures when running under pytest. * Rename test_studio_run.py -> test_studio_api.py The file is entirely about HTTP API endpoint testing (OpenAI-compatible /v1/chat/completions, Anthropic-compatible /v1/messages, API key auth, plus a CLI --help sanity check on the command that runs the API). None of its tests cover training, export, chat-UI, or internal-Python-API concerns. The old name misleadingly suggested "tests for the unsloth studio run CLI subcommand" — the new name reflects the actual scope. Updates: - git mv the file (rename tracked, history preserved) - Rewrite opening docstring to state the API surface focus and call out what is explicitly out of scope - Update all 4 Usage-block path references to the new filename - LOG_FILE renamed to test_studio_api.log - conftest.py fixture import rewritten from test_studio_run to test_studio_api, plus 7 docstring/comment references updated No functional changes to test logic, signatures, or main(). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Fix httpcore asyncgen cleanup in /v1/messages and /v1/completions The earlier fix in985e92a9was incomplete: it closed aiter_lines() explicitly but still used `async with httpx.AsyncClient()` / `async with client.stream()` inside the generator. When the generator is orphaned (e.g. client disconnects mid-stream and Starlette drops the StreamingResponse iterator without explicitly calling aclose()), Python's asyncgen finalizer runs the cleanup in a DIFFERENT task than the one that originally entered the httpx context managers. The `async with` exits then trigger httpcore's HTTP11ConnectionByteStream .aclose(), which enters anyio.CancelScope.__exit__ with a mismatched task and raises RuntimeError("Attempted to exit cancel scope in a different task"). That error escapes any user-owned try/except because it happens during GC finalization. Replace `async with` with manual client/response lifecycle in both /v1/messages passthrough and /v1/completions proxy. Close the response and client in a finally block wrapped in `try: ... except Exception: pass`. This suppresses RuntimeError (and other Exception subclasses) from the anyio cleanup noise while letting GeneratorExit (a BaseException, not Exception) propagate cleanly so the generator terminates as Python expects. Traceback observed in user report: File ".../httpcore/_async/connection_pool.py", line 404, in __aiter__ yield part RuntimeError: async generator ignored GeneratorExit ... File ".../anyio/_backends/_asyncio.py", line 455, in __exit__ raise RuntimeError( RuntimeError: Attempted to exit cancel scope in a different task * Expand unsloth studio run banner with SDK base URL and more curl examples Add an explicit "OpenAI / Anthropic SDK base URL" line inside the info box so SDK users don't accidentally copy the bare server URL (without /v1) into their OpenAI/Anthropic SDK constructors and hit 404s. Replace the single /v1/chat/completions curl example with three labeled blocks: chat/completions, Anthropic /messages, and OpenAI Responses. The Anthropic example includes max_tokens (Anthropic SDKs require it even though Studio accepts None). All examples derived from a computed sdk_base_url so the /v1 prefix stays in sync if the public path ever changes. * Hash API keys with HMAC-SHA256 + persistent server secret Stores the HMAC secret in a new app_secrets singleton table. Fixes CodeQL py/weak-sensitive-data-hashing alert on storage.py:74-76, 394-395. Refresh tokens stay on plain SHA-256 (unchanged _hash_token) so existing user sessions survive upgrade — API keys are new on this branch so there is no migration. * Use PBKDF2 for API key hashing per CodeQL recommendation HMAC-SHA256 was still flagged by py/weak-sensitive-data-hashing. Switch to hashlib.pbkdf2_hmac, which is in CodeQL's recommended allowlist (Argon2/scrypt/bcrypt/PBKDF2). Persistent server-side salt stays in app_secrets for defense-in-depth. 100k iterations to match auth/hashing.py's password hasher. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2942 lines
113 KiB
Python
2942 lines
113 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Inference API routes for model loading and text generation.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
|
from fastapi.responses import StreamingResponse, JSONResponse, Response
|
|
from typing import Optional
|
|
import json
|
|
import httpx
|
|
import structlog
|
|
from loggers import get_logger
|
|
import asyncio
|
|
import threading
|
|
|
|
|
|
import re as _re
|
|
|
|
# Model size extraction (shared with core/inference/llama_cpp.py)
|
|
from utils.models import extract_model_size_b as _extract_model_size_b
|
|
|
|
|
|
def _friendly_error(exc: Exception) -> str:
|
|
"""Extract a user-friendly message from known llama-server errors."""
|
|
msg = str(exc)
|
|
m = _re.search(
|
|
r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)",
|
|
msg,
|
|
)
|
|
if m:
|
|
return (
|
|
f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token "
|
|
f"context window. Try increasing the Context Length in Model settings, "
|
|
f"or shorten the conversation."
|
|
)
|
|
if "Lost connection to llama-server" in msg:
|
|
return "Lost connection to the model server. It may have crashed -- try reloading the model."
|
|
return "An internal error occurred"
|
|
|
|
|
|
# Add backend directory to path
|
|
backend_path = Path(__file__).parent.parent.parent
|
|
if str(backend_path) not in sys.path:
|
|
sys.path.insert(0, str(backend_path))
|
|
|
|
# Import backend functions
|
|
try:
|
|
from core.inference import get_inference_backend
|
|
from core.inference.llama_cpp import LlamaCppBackend
|
|
from utils.models import ModelConfig
|
|
from utils.inference import load_inference_config
|
|
from utils.models.model_config import load_model_defaults
|
|
except ImportError:
|
|
parent_backend = backend_path.parent / "backend"
|
|
if str(parent_backend) not in sys.path:
|
|
sys.path.insert(0, str(parent_backend))
|
|
from core.inference import get_inference_backend
|
|
from core.inference.llama_cpp import LlamaCppBackend
|
|
from utils.models import ModelConfig
|
|
from utils.inference import load_inference_config
|
|
from utils.models.model_config import load_model_defaults
|
|
|
|
from models.inference import (
|
|
LoadRequest,
|
|
UnloadRequest,
|
|
GenerateRequest,
|
|
LoadResponse,
|
|
UnloadResponse,
|
|
InferenceStatusResponse,
|
|
ChatCompletionRequest,
|
|
ChatCompletionChunk,
|
|
ChatCompletion,
|
|
ChatMessage,
|
|
ChunkChoice,
|
|
ChoiceDelta,
|
|
CompletionChoice,
|
|
CompletionMessage,
|
|
CompletionUsage,
|
|
ValidateModelRequest,
|
|
ValidateModelResponse,
|
|
TextContentPart,
|
|
ImageContentPart,
|
|
ImageUrl,
|
|
ResponsesRequest,
|
|
ResponsesInputMessage,
|
|
ResponsesInputTextPart,
|
|
ResponsesInputImagePart,
|
|
ResponsesOutputTextContent,
|
|
ResponsesOutputMessage,
|
|
ResponsesUsage,
|
|
ResponsesResponse,
|
|
AnthropicMessagesRequest,
|
|
AnthropicMessagesResponse,
|
|
AnthropicResponseTextBlock,
|
|
AnthropicResponseToolUseBlock,
|
|
AnthropicUsage,
|
|
)
|
|
from core.inference.anthropic_compat import (
|
|
anthropic_messages_to_openai,
|
|
anthropic_tools_to_openai,
|
|
AnthropicStreamEmitter,
|
|
AnthropicPassthroughEmitter,
|
|
)
|
|
from auth.authentication import get_current_subject
|
|
|
|
import io
|
|
import wave
|
|
import base64
|
|
import numpy as np
|
|
from datetime import date as _date
|
|
|
|
router = APIRouter()
|
|
|
|
# Appended to tool-use nudge to discourage plan-without-action
|
|
_TOOL_ACTION_NUDGE = (
|
|
" IMPORTANT: Always call tools directly -- never write code yourself."
|
|
" Never describe what you plan to do -- just call the tool immediately."
|
|
" For any code request, call the python tool. For any factual question, call web_search."
|
|
" Do NOT output code blocks -- use the python tool instead."
|
|
)
|
|
|
|
# Regex for stripping leaked tool-call XML from assistant messages/stream
|
|
_TOOL_XML_RE = _re.compile(
|
|
r"<tool_call>.*?</tool_call>|<function=\w+>.*?</function>",
|
|
_re.DOTALL,
|
|
)
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# GGUF inference backend (llama-server)
|
|
_llama_cpp_backend = LlamaCppBackend()
|
|
|
|
|
|
def get_llama_cpp_backend() -> LlamaCppBackend:
|
|
return _llama_cpp_backend
|
|
|
|
|
|
@router.post("/load", response_model = LoadResponse)
|
|
async def load_model(
|
|
request: LoadRequest,
|
|
fastapi_request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Load a model for inference.
|
|
|
|
The model_path should be a clean identifier from GET /models/list.
|
|
Returns inference configuration parameters (temperature, top_p, top_k, min_p)
|
|
from the model's YAML config, falling back to default.yaml for missing values.
|
|
|
|
GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth.
|
|
"""
|
|
try:
|
|
# Version switching is handled automatically by the subprocess-based
|
|
# inference backend — no need for ensure_transformers_version() here.
|
|
|
|
# ── Already-loaded check: skip reload if the exact model is active ──
|
|
backend = get_inference_backend()
|
|
llama_backend = get_llama_cpp_backend()
|
|
|
|
if request.gguf_variant:
|
|
if (
|
|
llama_backend.is_loaded
|
|
and llama_backend.hf_variant
|
|
and llama_backend.hf_variant.lower() == request.gguf_variant.lower()
|
|
and llama_backend.model_identifier
|
|
and llama_backend.model_identifier.lower() == request.model_path.lower()
|
|
):
|
|
logger.info(
|
|
f"Model already loaded (GGUF): {request.model_path} variant={request.gguf_variant}, skipping reload"
|
|
)
|
|
inference_config = load_inference_config(llama_backend.model_identifier)
|
|
from utils.models import is_audio_input_type
|
|
|
|
_gguf_audio = (
|
|
llama_backend._audio_type
|
|
if hasattr(llama_backend, "_audio_type")
|
|
else None
|
|
)
|
|
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
|
|
return LoadResponse(
|
|
status = "already_loaded",
|
|
model = llama_backend.model_identifier,
|
|
display_name = llama_backend.model_identifier,
|
|
is_vision = llama_backend._is_vision,
|
|
is_lora = False,
|
|
is_gguf = True,
|
|
is_audio = _gguf_is_audio,
|
|
audio_type = _gguf_audio,
|
|
has_audio_input = is_audio_input_type(_gguf_audio)
|
|
if _gguf_audio
|
|
else False,
|
|
inference = inference_config,
|
|
requires_trust_remote_code = bool(
|
|
inference_config.get("trust_remote_code", False)
|
|
),
|
|
context_length = llama_backend.context_length,
|
|
max_context_length = llama_backend.max_context_length,
|
|
native_context_length = llama_backend.native_context_length,
|
|
supports_reasoning = llama_backend.supports_reasoning,
|
|
reasoning_always_on = llama_backend.reasoning_always_on,
|
|
chat_template = llama_backend.chat_template,
|
|
speculative_type = llama_backend.speculative_type,
|
|
)
|
|
else:
|
|
if (
|
|
backend.active_model_name
|
|
and backend.active_model_name.lower() == request.model_path.lower()
|
|
):
|
|
logger.info(
|
|
f"Model already loaded (Unsloth): {request.model_path}, skipping reload"
|
|
)
|
|
inference_config = load_inference_config(backend.active_model_name)
|
|
_model_info = backend.models.get(backend.active_model_name, {})
|
|
_chat_template = None
|
|
try:
|
|
_tpl_info = _model_info.get("chat_template_info", {})
|
|
_chat_template = _tpl_info.get("template")
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Could not retrieve chat template for {backend.active_model_name}: {e}"
|
|
)
|
|
return LoadResponse(
|
|
status = "already_loaded",
|
|
model = backend.active_model_name,
|
|
display_name = backend.active_model_name,
|
|
is_vision = _model_info.get("is_vision", False),
|
|
is_lora = _model_info.get("is_lora", False),
|
|
is_gguf = False,
|
|
is_audio = _model_info.get("is_audio", False),
|
|
audio_type = _model_info.get("audio_type"),
|
|
has_audio_input = _model_info.get("has_audio_input", False),
|
|
inference = inference_config,
|
|
requires_trust_remote_code = bool(
|
|
inference_config.get("trust_remote_code", False)
|
|
),
|
|
chat_template = _chat_template,
|
|
)
|
|
|
|
# Create config using clean factory method
|
|
# is_lora is auto-detected from adapter_config.json on disk/HF
|
|
config = ModelConfig.from_identifier(
|
|
model_id = request.model_path,
|
|
hf_token = request.hf_token,
|
|
gguf_variant = request.gguf_variant,
|
|
)
|
|
|
|
if not config:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Invalid model identifier: {request.model_path}",
|
|
)
|
|
|
|
# Normalize gpu_ids: empty list means auto-selection, same as None
|
|
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
|
|
|
|
# ── GGUF path: load via llama-server ──────────────────────
|
|
if config.is_gguf:
|
|
if effective_gpu_ids is not None:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "gpu_ids is not supported for GGUF models yet.",
|
|
)
|
|
|
|
llama_backend = get_llama_cpp_backend()
|
|
unsloth_backend = get_inference_backend()
|
|
|
|
# Unload any active Unsloth model first to free VRAM
|
|
if unsloth_backend.active_model_name:
|
|
logger.info(
|
|
f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
|
|
)
|
|
unsloth_backend.unload_model(unsloth_backend.active_model_name)
|
|
|
|
# Route to HF mode or local mode based on config
|
|
# Run in a thread so the event loop stays free for progress
|
|
# polling and other requests during the (potentially long)
|
|
# GGUF download + llama-server startup.
|
|
_n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1)
|
|
|
|
if config.gguf_hf_repo:
|
|
# HF mode: download via huggingface_hub then start llama-server
|
|
success = await asyncio.to_thread(
|
|
llama_backend.load_model,
|
|
hf_repo = config.gguf_hf_repo,
|
|
hf_variant = config.gguf_variant,
|
|
hf_token = request.hf_token,
|
|
model_identifier = config.identifier,
|
|
is_vision = config.is_vision,
|
|
n_ctx = request.max_seq_length,
|
|
chat_template_override = request.chat_template_override,
|
|
cache_type_kv = request.cache_type_kv,
|
|
speculative_type = request.speculative_type,
|
|
n_parallel = _n_parallel,
|
|
)
|
|
else:
|
|
# Local mode: llama-server loads via -m <path>
|
|
success = await asyncio.to_thread(
|
|
llama_backend.load_model,
|
|
gguf_path = config.gguf_file,
|
|
mmproj_path = config.gguf_mmproj_file,
|
|
model_identifier = config.identifier,
|
|
is_vision = config.is_vision,
|
|
n_ctx = request.max_seq_length,
|
|
chat_template_override = request.chat_template_override,
|
|
cache_type_kv = request.cache_type_kv,
|
|
speculative_type = request.speculative_type,
|
|
n_parallel = _n_parallel,
|
|
)
|
|
|
|
if not success:
|
|
raise HTTPException(
|
|
status_code = 500,
|
|
detail = f"Failed to load GGUF model: {config.display_name}",
|
|
)
|
|
|
|
logger.info(f"Loaded GGUF model via llama-server: {config.identifier}")
|
|
|
|
# Detect TTS audio by probing the loaded model's vocabulary
|
|
from utils.models import is_audio_input_type
|
|
|
|
_gguf_audio = llama_backend.detect_audio_type()
|
|
_gguf_is_audio = _gguf_audio in ("snac", "bicodec", "dac")
|
|
llama_backend._is_audio = _gguf_is_audio
|
|
llama_backend._audio_type = _gguf_audio
|
|
if _gguf_is_audio:
|
|
logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}")
|
|
await asyncio.to_thread(llama_backend.init_audio_codec, _gguf_audio)
|
|
|
|
inference_config = load_inference_config(config.identifier)
|
|
|
|
return LoadResponse(
|
|
status = "loaded",
|
|
model = config.identifier,
|
|
display_name = config.display_name,
|
|
is_vision = config.is_vision,
|
|
is_lora = False,
|
|
is_gguf = True,
|
|
is_audio = _gguf_is_audio,
|
|
audio_type = _gguf_audio,
|
|
has_audio_input = is_audio_input_type(_gguf_audio),
|
|
inference = inference_config,
|
|
requires_trust_remote_code = bool(
|
|
inference_config.get("trust_remote_code", False)
|
|
),
|
|
context_length = llama_backend.context_length,
|
|
max_context_length = llama_backend.max_context_length,
|
|
native_context_length = llama_backend.native_context_length,
|
|
supports_reasoning = llama_backend.supports_reasoning,
|
|
reasoning_always_on = llama_backend.reasoning_always_on,
|
|
supports_tools = llama_backend.supports_tools,
|
|
cache_type_kv = llama_backend.cache_type_kv,
|
|
chat_template = llama_backend.chat_template,
|
|
speculative_type = llama_backend.speculative_type,
|
|
)
|
|
|
|
# ── Standard path: load via Unsloth/transformers ──────────
|
|
backend = get_inference_backend()
|
|
|
|
# Unload any active GGUF model first
|
|
llama_backend = get_llama_cpp_backend()
|
|
if llama_backend.is_loaded:
|
|
logger.info("Unloading GGUF model before loading Unsloth model")
|
|
llama_backend.unload_model()
|
|
|
|
# Shut down any export subprocess to free VRAM
|
|
try:
|
|
from core.export import get_export_backend
|
|
|
|
exp_backend = get_export_backend()
|
|
if exp_backend.current_checkpoint:
|
|
logger.info(
|
|
"Shutting down export subprocess to free GPU memory for inference"
|
|
)
|
|
exp_backend._shutdown_subprocess()
|
|
exp_backend.current_checkpoint = None
|
|
exp_backend.is_vision = False
|
|
exp_backend.is_peft = False
|
|
except Exception as e:
|
|
logger.warning("Could not shut down export subprocess: %s", e)
|
|
|
|
# Auto-detect quantization for LoRA adapters from adapter_config.json
|
|
# The training pipeline patches this file with "unsloth_training_method"
|
|
# which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False.
|
|
load_in_4bit = request.load_in_4bit
|
|
if config.is_lora and config.path:
|
|
import json
|
|
from pathlib import Path
|
|
|
|
adapter_cfg_path = Path(config.path) / "adapter_config.json"
|
|
if adapter_cfg_path.exists():
|
|
try:
|
|
with open(adapter_cfg_path) as f:
|
|
adapter_cfg = json.load(f)
|
|
training_method = adapter_cfg.get("unsloth_training_method")
|
|
if training_method == "lora" and load_in_4bit:
|
|
logger.info(
|
|
f"adapter_config.json says unsloth_training_method='lora' — "
|
|
f"setting load_in_4bit=False to match 16-bit training"
|
|
)
|
|
load_in_4bit = False
|
|
elif training_method == "qlora" and not load_in_4bit:
|
|
logger.info(
|
|
f"adapter_config.json says unsloth_training_method='qlora' — "
|
|
f"setting load_in_4bit=True to match QLoRA training"
|
|
)
|
|
load_in_4bit = True
|
|
elif training_method:
|
|
logger.info(
|
|
f"Training method: {training_method}, load_in_4bit={load_in_4bit}"
|
|
)
|
|
else:
|
|
# No unsloth_training_method — fallback to base model name
|
|
if (
|
|
config.base_model
|
|
and "-bnb-4bit" not in config.base_model.lower()
|
|
and load_in_4bit
|
|
):
|
|
logger.info(
|
|
f"No unsloth_training_method in adapter_config.json. "
|
|
f"Base model '{config.base_model}' has no -bnb-4bit suffix — "
|
|
f"setting load_in_4bit=False"
|
|
)
|
|
load_in_4bit = False
|
|
except Exception as e:
|
|
logger.warning(f"Could not read adapter_config.json: {e}")
|
|
|
|
# Load the model in a thread so the event loop stays free
|
|
# for download progress polling and other requests.
|
|
success = await asyncio.to_thread(
|
|
backend.load_model,
|
|
config = config,
|
|
max_seq_length = request.max_seq_length,
|
|
load_in_4bit = load_in_4bit,
|
|
hf_token = request.hf_token,
|
|
trust_remote_code = request.trust_remote_code,
|
|
gpu_ids = effective_gpu_ids,
|
|
)
|
|
|
|
if not success:
|
|
# Check if YAML says this model needs trust_remote_code
|
|
if not request.trust_remote_code:
|
|
model_defaults = load_model_defaults(config.identifier)
|
|
yaml_trust = model_defaults.get("inference", {}).get(
|
|
"trust_remote_code", False
|
|
)
|
|
if yaml_trust:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = (
|
|
f"Model '{config.display_name}' requires trust_remote_code to be enabled. "
|
|
f"Please enable 'Trust remote code' in Chat Settings and try again."
|
|
),
|
|
)
|
|
raise HTTPException(
|
|
status_code = 500, detail = f"Failed to load model: {config.display_name}"
|
|
)
|
|
|
|
logger.info(f"Loaded model: {config.identifier}")
|
|
|
|
# Load inference configuration parameters
|
|
inference_config = load_inference_config(config.identifier)
|
|
|
|
# Get chat template from tokenizer
|
|
_chat_template = None
|
|
try:
|
|
_model_info = backend.models.get(config.identifier, {})
|
|
_tpl_info = _model_info.get("chat_template_info", {})
|
|
_chat_template = _tpl_info.get("template")
|
|
except Exception:
|
|
pass
|
|
|
|
return LoadResponse(
|
|
status = "loaded",
|
|
model = config.identifier,
|
|
display_name = config.display_name,
|
|
is_vision = config.is_vision,
|
|
is_lora = config.is_lora,
|
|
is_gguf = False,
|
|
is_audio = config.is_audio,
|
|
audio_type = config.audio_type,
|
|
has_audio_input = config.has_audio_input,
|
|
inference = inference_config,
|
|
requires_trust_remote_code = bool(
|
|
inference_config.get("trust_remote_code", False)
|
|
),
|
|
chat_template = _chat_template,
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except ValueError as e:
|
|
logger.warning("Rejected inference GPU selection: %s", e)
|
|
raise HTTPException(status_code = 400, detail = str(e))
|
|
except Exception as e:
|
|
logger.error(f"Error loading model: {e}", exc_info = True)
|
|
msg = str(e)
|
|
# Surface a friendlier message for models that Unsloth cannot load
|
|
not_supported_hints = [
|
|
"No config file found",
|
|
"not yet supported",
|
|
"is not supported",
|
|
"does not support",
|
|
]
|
|
if any(h.lower() in msg.lower() for h in not_supported_hints):
|
|
msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
|
|
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
|
|
|
|
|
|
@router.post("/validate", response_model = ValidateModelResponse)
|
|
async def validate_model(
|
|
request: ValidateModelRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Lightweight validation endpoint for model identifiers.
|
|
|
|
This checks that ModelConfig.from_identifier() can resolve the given
|
|
model_path, but it does NOT actually load model weights into GPU memory.
|
|
"""
|
|
try:
|
|
config = ModelConfig.from_identifier(
|
|
model_id = request.model_path,
|
|
hf_token = request.hf_token,
|
|
gguf_variant = request.gguf_variant,
|
|
)
|
|
|
|
if not config:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Invalid model identifier: {request.model_path}",
|
|
)
|
|
|
|
return ValidateModelResponse(
|
|
valid = True,
|
|
message = "Model identifier is valid.",
|
|
identifier = config.identifier,
|
|
display_name = getattr(config, "display_name", config.identifier),
|
|
is_gguf = getattr(config, "is_gguf", False),
|
|
is_lora = getattr(config, "is_lora", False),
|
|
is_vision = getattr(config, "is_vision", False),
|
|
requires_trust_remote_code = bool(
|
|
load_inference_config(config.identifier).get("trust_remote_code", False)
|
|
),
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error validating model identifier '{request.model_path}': {e}",
|
|
exc_info = True,
|
|
)
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Invalid model: {str(e)}",
|
|
)
|
|
|
|
|
|
@router.post("/unload", response_model = UnloadResponse)
|
|
async def unload_model(
|
|
request: UnloadRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Unload a model from memory.
|
|
Routes to the correct backend (llama-server for GGUF, Unsloth otherwise).
|
|
"""
|
|
try:
|
|
# Check if the GGUF backend has this model loaded or is loading it
|
|
llama_backend = get_llama_cpp_backend()
|
|
if llama_backend.is_active and (
|
|
llama_backend.model_identifier == request.model_path
|
|
or not llama_backend.is_loaded
|
|
):
|
|
llama_backend.unload_model()
|
|
logger.info(f"Unloaded GGUF model: {request.model_path}")
|
|
return UnloadResponse(status = "unloaded", model = request.model_path)
|
|
|
|
# Otherwise, unload from Unsloth backend
|
|
backend = get_inference_backend()
|
|
backend.unload_model(request.model_path)
|
|
logger.info(f"Unloaded model: {request.model_path}")
|
|
return UnloadResponse(status = "unloaded", model = request.model_path)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error unloading model: {e}", exc_info = True)
|
|
raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}")
|
|
|
|
|
|
@router.post("/generate/stream")
|
|
async def generate_stream(
|
|
request: GenerateRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Generate a chat response with Server-Sent Events (SSE) streaming.
|
|
|
|
For vision models, provide image_base64 with the base64-encoded image.
|
|
"""
|
|
backend = get_inference_backend()
|
|
|
|
if not backend.active_model_name:
|
|
raise HTTPException(
|
|
status_code = 400, detail = "No model loaded. Call POST /inference/load first."
|
|
)
|
|
|
|
# Decode image if provided (for vision models)
|
|
image = None
|
|
if request.image_base64:
|
|
try:
|
|
import base64
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
|
|
# Check if current model supports vision
|
|
model_info = backend.models.get(backend.active_model_name, {})
|
|
if not model_info.get("is_vision"):
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Image provided but current model is text-only. Load a vision model.",
|
|
)
|
|
|
|
image_data = base64.b64decode(request.image_base64)
|
|
image = Image.open(BytesIO(image_data))
|
|
image = backend.resize_image(image)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code = 400, detail = f"Failed to decode image: {str(e)}"
|
|
)
|
|
|
|
async def stream():
|
|
try:
|
|
for chunk in backend.generate_chat_response(
|
|
messages = request.messages,
|
|
system_prompt = request.system_prompt,
|
|
image = image,
|
|
temperature = request.temperature,
|
|
top_p = request.top_p,
|
|
top_k = request.top_k,
|
|
max_new_tokens = request.max_new_tokens,
|
|
repetition_penalty = request.repetition_penalty,
|
|
):
|
|
yield f"data: {json.dumps({'content': chunk})}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
except Exception as e:
|
|
backend.reset_generation_state()
|
|
logger.error(f"Error during generation: {e}", exc_info = True)
|
|
yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n"
|
|
|
|
return StreamingResponse(
|
|
stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
},
|
|
)
|
|
|
|
|
|
@router.get("/status", response_model = InferenceStatusResponse)
|
|
async def get_status(
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Get current inference backend status.
|
|
Reports whichever backend (Unsloth or llama-server) is currently active.
|
|
"""
|
|
try:
|
|
llama_backend = get_llama_cpp_backend()
|
|
|
|
# If a GGUF model is loaded via llama-server, report that
|
|
if llama_backend.is_loaded:
|
|
_model_id = llama_backend.model_identifier
|
|
_inference_cfg = load_inference_config(_model_id) if _model_id else None
|
|
return InferenceStatusResponse(
|
|
active_model = _model_id,
|
|
is_vision = llama_backend.is_vision,
|
|
is_gguf = True,
|
|
gguf_variant = llama_backend.hf_variant,
|
|
is_audio = getattr(llama_backend, "_is_audio", False),
|
|
audio_type = getattr(llama_backend, "_audio_type", None),
|
|
loading = [],
|
|
loaded = [_model_id],
|
|
inference = _inference_cfg,
|
|
requires_trust_remote_code = bool(
|
|
(_inference_cfg or {}).get("trust_remote_code", False)
|
|
),
|
|
supports_reasoning = llama_backend.supports_reasoning,
|
|
reasoning_always_on = llama_backend.reasoning_always_on,
|
|
supports_tools = llama_backend.supports_tools,
|
|
context_length = llama_backend.context_length,
|
|
max_context_length = llama_backend.max_context_length,
|
|
native_context_length = llama_backend.native_context_length,
|
|
speculative_type = llama_backend.speculative_type,
|
|
)
|
|
|
|
# Otherwise, report Unsloth backend status
|
|
backend = get_inference_backend()
|
|
|
|
is_vision = False
|
|
is_audio = False
|
|
audio_type = None
|
|
has_audio_input = False
|
|
if backend.active_model_name:
|
|
model_info = backend.models.get(backend.active_model_name, {})
|
|
is_vision = model_info.get("is_vision", False)
|
|
is_audio = model_info.get("is_audio", False)
|
|
audio_type = model_info.get("audio_type")
|
|
has_audio_input = model_info.get("has_audio_input", False)
|
|
|
|
# gpt-oss safetensors models support reasoning via harmony channels
|
|
supports_reasoning = False
|
|
if backend.active_model_name and hasattr(backend, "_is_gpt_oss_model"):
|
|
supports_reasoning = backend._is_gpt_oss_model()
|
|
inference_config = (
|
|
load_inference_config(backend.active_model_name)
|
|
if backend.active_model_name
|
|
else None
|
|
)
|
|
|
|
return InferenceStatusResponse(
|
|
active_model = backend.active_model_name,
|
|
is_vision = is_vision,
|
|
is_gguf = False,
|
|
is_audio = is_audio,
|
|
audio_type = audio_type,
|
|
has_audio_input = has_audio_input,
|
|
loading = list(getattr(backend, "loading_models", set())),
|
|
loaded = list(backend.models.keys()),
|
|
inference = inference_config,
|
|
requires_trust_remote_code = bool(
|
|
(inference_config or {}).get("trust_remote_code", False)
|
|
),
|
|
supports_reasoning = supports_reasoning,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error getting status: {e}", exc_info = True)
|
|
raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}")
|
|
|
|
|
|
# =====================================================================
|
|
# Audio (TTS) Generation (/audio/generate)
|
|
# =====================================================================
|
|
|
|
|
|
@router.post("/audio/generate")
|
|
async def generate_audio(
|
|
payload: ChatCompletionRequest,
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Generate audio (TTS) from the latest user message.
|
|
Returns a JSON response with base64-encoded WAV audio.
|
|
Works with both GGUF (llama-server) and Unsloth/transformers backends.
|
|
"""
|
|
import base64
|
|
|
|
# Extract text from the last user message
|
|
_, chat_messages, _ = _extract_content_parts(payload.messages)
|
|
if not chat_messages:
|
|
raise HTTPException(status_code = 400, detail = "No messages provided.")
|
|
last_user_msg = next(
|
|
(m for m in reversed(chat_messages) if m["role"] == "user"), None
|
|
)
|
|
if not last_user_msg:
|
|
raise HTTPException(status_code = 400, detail = "No user message found.")
|
|
text = last_user_msg["content"]
|
|
|
|
# Pick backend — both return (wav_bytes, sample_rate)
|
|
llama_backend = get_llama_cpp_backend()
|
|
if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False):
|
|
model_name = llama_backend.model_identifier
|
|
gen = lambda: llama_backend.generate_audio_response(
|
|
text = text,
|
|
audio_type = llama_backend._audio_type,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_new_tokens = payload.max_tokens or 2048,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
)
|
|
else:
|
|
backend = get_inference_backend()
|
|
if not backend.active_model_name:
|
|
raise HTTPException(status_code = 400, detail = "No model loaded.")
|
|
model_info = backend.models.get(backend.active_model_name, {})
|
|
if not model_info.get("is_audio"):
|
|
raise HTTPException(
|
|
status_code = 400, detail = "Active model is not an audio model."
|
|
)
|
|
model_name = backend.active_model_name
|
|
gen = lambda: backend.generate_audio_response(
|
|
text = text,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_new_tokens = payload.max_tokens or 2048,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
use_adapter = payload.use_adapter,
|
|
)
|
|
|
|
try:
|
|
wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(
|
|
None, gen
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Audio generation error: {e}", exc_info = True)
|
|
raise HTTPException(status_code = 500, detail = str(e))
|
|
|
|
audio_b64 = base64.b64encode(wav_bytes).decode("ascii")
|
|
return JSONResponse(
|
|
content = {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
|
|
"object": "chat.completion.audio",
|
|
"model": model_name,
|
|
"audio": {"data": audio_b64, "format": "wav", "sample_rate": sample_rate},
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": f'[Generated audio from: "{text[:100]}"]',
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# OpenAI-Compatible Chat Completions (/chat/completions)
|
|
# =====================================================================
|
|
|
|
|
|
def _decode_audio_base64(b64: str) -> np.ndarray:
|
|
"""Decode base64 audio (any format) → float32 numpy array at 16kHz."""
|
|
import torch
|
|
import torchaudio
|
|
import tempfile
|
|
import os
|
|
from utils.paths import ensure_dir, tmp_root
|
|
|
|
raw = base64.b64decode(b64)
|
|
# torchaudio.load needs a file path or file-like object with format hint
|
|
# Write to a temp file so torchaudio can auto-detect the format
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix = ".audio",
|
|
delete = False,
|
|
dir = str(ensure_dir(tmp_root())),
|
|
) as tmp:
|
|
tmp.write(raw)
|
|
tmp_path = tmp.name
|
|
try:
|
|
waveform, sr = torchaudio.load(tmp_path)
|
|
finally:
|
|
os.unlink(tmp_path)
|
|
|
|
# Convert to mono if stereo
|
|
if waveform.shape[0] > 1:
|
|
waveform = waveform.mean(dim = 0, keepdim = True)
|
|
|
|
# Resample to 16kHz if needed
|
|
if sr != 16000:
|
|
resampler = torchaudio.transforms.Resample(orig_freq = sr, new_freq = 16000)
|
|
waveform = resampler(waveform)
|
|
|
|
return waveform.squeeze(0).numpy()
|
|
|
|
|
|
def _extract_content_parts(
|
|
messages: list,
|
|
) -> tuple[str, list[dict], "Optional[str]"]:
|
|
"""
|
|
Parse OpenAI-format messages into components the inference backend expects.
|
|
|
|
Handles both plain-string ``content`` and multimodal content-part arrays
|
|
(``[{type: "text", ...}, {type: "image_url", ...}]``).
|
|
|
|
Returns:
|
|
system_prompt: The system message text (empty string if none provided).
|
|
chat_messages: Non-system messages with content flattened to strings.
|
|
image_base64: Base64 data of the *first* image found, or ``None``.
|
|
"""
|
|
system_prompt = ""
|
|
chat_messages: list[dict] = []
|
|
first_image_b64: Optional[str] = None
|
|
|
|
for msg in messages:
|
|
# ── System messages → extract as system_prompt ────────
|
|
if msg.role == "system":
|
|
if isinstance(msg.content, str):
|
|
system_prompt = msg.content
|
|
elif isinstance(msg.content, list):
|
|
# Unlikely but handle: join text parts
|
|
system_prompt = "\n".join(
|
|
p.text for p in msg.content if p.type == "text"
|
|
)
|
|
continue
|
|
|
|
# ── User / assistant messages ─────────────────────────
|
|
if isinstance(msg.content, str):
|
|
# Plain string content — pass through
|
|
chat_messages.append({"role": msg.role, "content": msg.content})
|
|
elif isinstance(msg.content, list):
|
|
# Multimodal content parts
|
|
text_parts: list[str] = []
|
|
for part in msg.content:
|
|
if part.type == "text":
|
|
text_parts.append(part.text)
|
|
elif part.type == "image_url" and first_image_b64 is None:
|
|
url = part.image_url.url
|
|
if url.startswith("data:"):
|
|
# data:image/png;base64,<DATA> → extract <DATA>
|
|
first_image_b64 = url.split(",", 1)[1] if "," in url else None
|
|
else:
|
|
logger.warning(
|
|
f"Remote image URLs not yet supported: {url[:80]}..."
|
|
)
|
|
combined_text = "\n".join(text_parts) if text_parts else ""
|
|
chat_messages.append({"role": msg.role, "content": combined_text})
|
|
|
|
return system_prompt, chat_messages, first_image_b64
|
|
|
|
|
|
@router.post("/chat/completions")
|
|
async def openai_chat_completions(
|
|
payload: ChatCompletionRequest,
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
OpenAI-compatible chat completions endpoint.
|
|
|
|
Supports multimodal messages: ``content`` may be a plain string or a
|
|
list of content parts (``text`` / ``image_url``).
|
|
|
|
Streaming (default): returns SSE chunks matching OpenAI's format.
|
|
Non-streaming: returns a single ChatCompletion JSON object.
|
|
|
|
Automatically routes to the correct backend:
|
|
- GGUF models → llama-server via LlamaCppBackend
|
|
- Other models → Unsloth/transformers via InferenceBackend
|
|
"""
|
|
llama_backend = get_llama_cpp_backend()
|
|
using_gguf = llama_backend.is_loaded
|
|
|
|
# ── Determine which backend is active ─────────────────────
|
|
if using_gguf:
|
|
model_name = llama_backend.model_identifier or payload.model
|
|
if getattr(llama_backend, "_is_audio", False):
|
|
return await generate_audio(payload, request)
|
|
else:
|
|
backend = get_inference_backend()
|
|
if not backend.active_model_name:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "No model loaded. Call POST /inference/load first.",
|
|
)
|
|
model_name = backend.active_model_name or payload.model
|
|
|
|
# ── Audio TTS path: auto-route to audio generation ────
|
|
# (Whisper is ASR not TTS — handled below in audio input path)
|
|
model_info = backend.models.get(backend.active_model_name, {})
|
|
if model_info.get("is_audio") and model_info.get("audio_type") != "whisper":
|
|
return await generate_audio(payload, request)
|
|
|
|
# ── Whisper without audio: return clear error ──
|
|
if model_info.get("audio_type") == "whisper" and not payload.audio_base64:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Whisper models require audio input. Please upload an audio file.",
|
|
)
|
|
|
|
# ── Audio INPUT path: decode WAV and route to audio input generation ──
|
|
if payload.audio_base64 and model_info.get("has_audio_input"):
|
|
audio_array = _decode_audio_base64(payload.audio_base64)
|
|
system_prompt, chat_messages, _ = _extract_content_parts(payload.messages)
|
|
cancel_event = threading.Event()
|
|
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
|
created = int(time.time())
|
|
|
|
def audio_input_generate():
|
|
if model_info.get("audio_type") == "whisper":
|
|
return backend.generate_whisper_response(
|
|
audio_array = audio_array,
|
|
cancel_event = cancel_event,
|
|
)
|
|
return backend.generate_audio_input_response(
|
|
messages = chat_messages,
|
|
system_prompt = system_prompt,
|
|
audio_array = audio_array,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_new_tokens = payload.max_tokens or 2048,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
cancel_event = cancel_event,
|
|
)
|
|
|
|
if payload.stream:
|
|
|
|
async def audio_input_stream():
|
|
try:
|
|
first_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(role = "assistant"),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
for chunk_text in audio_input_generate():
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
return
|
|
if chunk_text:
|
|
chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(content = chunk_text),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
final_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")
|
|
],
|
|
)
|
|
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
except asyncio.CancelledError:
|
|
cancel_event.set()
|
|
raise
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error during audio input streaming: {e}", exc_info = True
|
|
)
|
|
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
|
|
|
|
return StreamingResponse(
|
|
audio_input_stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
else:
|
|
full_text = "".join(audio_input_generate())
|
|
response = ChatCompletion(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
CompletionChoice(
|
|
message = CompletionMessage(content = full_text),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
return JSONResponse(content = response.model_dump())
|
|
|
|
# ── Parse messages (handles multimodal content parts) ─────
|
|
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
|
|
payload.messages
|
|
)
|
|
|
|
if not chat_messages:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "At least one non-system message is required.",
|
|
)
|
|
|
|
# ── GGUF path: proxy to llama-server /v1/chat/completions ──
|
|
if using_gguf:
|
|
# Reject images if this GGUF model doesn't support vision
|
|
image_b64 = extracted_image_b64 or payload.image_base64
|
|
if image_b64 and not llama_backend.is_vision:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Image provided but current GGUF model does not support vision.",
|
|
)
|
|
|
|
# Convert image to PNG for llama-server (stb_image has limited format support)
|
|
if image_b64:
|
|
try:
|
|
import base64 as _b64
|
|
from io import BytesIO as _BytesIO
|
|
from PIL import Image as _Image
|
|
|
|
raw = _b64.b64decode(image_b64)
|
|
img = _Image.open(_BytesIO(raw))
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
buf = _BytesIO()
|
|
img.save(buf, format = "PNG")
|
|
image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii")
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code = 400, detail = f"Failed to process image: {e}"
|
|
)
|
|
|
|
# Build message list with system prompt prepended
|
|
gguf_messages = []
|
|
if system_prompt:
|
|
gguf_messages.append({"role": "system", "content": system_prompt})
|
|
gguf_messages.extend(chat_messages)
|
|
|
|
cancel_event = threading.Event()
|
|
|
|
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
|
created = int(time.time())
|
|
|
|
# ── Tool-calling path (agentic loop) ──────────────────
|
|
use_tools = (
|
|
payload.enable_tools and llama_backend.supports_tools and not image_b64
|
|
)
|
|
|
|
if use_tools:
|
|
from core.inference.tools import ALL_TOOLS
|
|
|
|
if payload.enabled_tools is not None:
|
|
tools_to_use = [
|
|
t
|
|
for t in ALL_TOOLS
|
|
if t["function"]["name"] in payload.enabled_tools
|
|
]
|
|
else:
|
|
tools_to_use = ALL_TOOLS
|
|
|
|
# ── Tool-use system prompt nudge ──────────────────────
|
|
_tool_names = {t["function"]["name"] for t in tools_to_use}
|
|
_has_web = "web_search" in _tool_names
|
|
_has_code = "python" in _tool_names or "terminal" in _tool_names
|
|
|
|
_date_line = f"The current date is {_date.today().isoformat()}."
|
|
|
|
# Small models (<9B) struggle with multi-step search plans,
|
|
# so simplify the web tips to avoid plan-then-stall behavior.
|
|
_model_size_b = _extract_model_size_b(model_name)
|
|
_is_small_model = _model_size_b is not None and _model_size_b < 9
|
|
|
|
if _is_small_model:
|
|
_web_tips = "Do not repeat the same search query."
|
|
else:
|
|
_web_tips = (
|
|
"When you search and find a relevant URL in the results, "
|
|
"fetch its full content by calling web_search with the url parameter. "
|
|
"Do not repeat the same search query. If a search returns "
|
|
"no useful results, try rephrasing or fetching a result URL directly."
|
|
)
|
|
_code_tips = (
|
|
"Use code execution for math, calculations, data processing, "
|
|
"or to parse and analyze information from tool results."
|
|
)
|
|
|
|
if _has_web and _has_code:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"tools rather than answering from memory. "
|
|
+ _web_tips
|
|
+ " "
|
|
+ _code_tips
|
|
)
|
|
elif _has_code:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"code execution rather than answering from memory. " + _code_tips
|
|
)
|
|
elif _has_web:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"web search for up-to-date or uncertain factual "
|
|
"information rather than answering from memory. " + _web_tips
|
|
)
|
|
else:
|
|
_nudge = ""
|
|
|
|
if _nudge:
|
|
_nudge += _TOOL_ACTION_NUDGE
|
|
# Append nudge to system prompt (preserve user's prompt)
|
|
if system_prompt:
|
|
system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
|
|
else:
|
|
system_prompt = _nudge
|
|
# Rebuild gguf_messages with updated system prompt
|
|
gguf_messages = []
|
|
if system_prompt:
|
|
gguf_messages.append({"role": "system", "content": system_prompt})
|
|
gguf_messages.extend(chat_messages)
|
|
|
|
# ── Strip stale tool-call XML from conversation history ─
|
|
for _msg in gguf_messages:
|
|
if _msg.get("role") == "assistant" and isinstance(
|
|
_msg.get("content"), str
|
|
):
|
|
_msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
|
|
|
|
def gguf_generate_with_tools():
|
|
return llama_backend.generate_chat_completion_with_tools(
|
|
messages = gguf_messages,
|
|
tools = tools_to_use,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_tokens = payload.max_tokens,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
presence_penalty = payload.presence_penalty,
|
|
cancel_event = cancel_event,
|
|
enable_thinking = payload.enable_thinking,
|
|
auto_heal_tool_calls = payload.auto_heal_tool_calls
|
|
if payload.auto_heal_tool_calls is not None
|
|
else True,
|
|
max_tool_iterations = payload.max_tool_calls_per_message
|
|
if payload.max_tool_calls_per_message is not None
|
|
else 25,
|
|
tool_call_timeout = payload.tool_call_timeout
|
|
if payload.tool_call_timeout is not None
|
|
else 300,
|
|
session_id = payload.session_id,
|
|
)
|
|
|
|
_tool_sentinel = object()
|
|
|
|
async def gguf_tool_stream():
|
|
try:
|
|
first_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(role = "assistant"),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
# Iterate the synchronous generator in a thread so
|
|
# the event loop stays free for disconnect detection.
|
|
gen = gguf_generate_with_tools()
|
|
prev_text = ""
|
|
_stream_usage = None
|
|
_stream_timings = None
|
|
while True:
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
return
|
|
|
|
event = await asyncio.to_thread(next, gen, _tool_sentinel)
|
|
if event is _tool_sentinel:
|
|
break
|
|
|
|
if event["type"] == "status":
|
|
# Empty status marks an iteration boundary
|
|
# in the GGUF tool loop (e.g. after a
|
|
# re-prompt). Reset the cumulative cursor
|
|
# so the next assistant turn streams cleanly.
|
|
if not event["text"]:
|
|
prev_text = ""
|
|
# Emit tool status as a custom SSE event
|
|
# (including empty ones to clear UI badges)
|
|
status_data = json.dumps(
|
|
{
|
|
"type": "tool_status",
|
|
"content": event["text"],
|
|
}
|
|
)
|
|
yield f"data: {status_data}\n\n"
|
|
continue
|
|
|
|
if event["type"] in ("tool_start", "tool_end"):
|
|
if event["type"] == "tool_start":
|
|
prev_text = ""
|
|
yield f"data: {json.dumps(event)}\n\n"
|
|
continue
|
|
|
|
if event["type"] == "metadata":
|
|
_stream_usage = event.get("usage")
|
|
_stream_timings = event.get("timings")
|
|
continue
|
|
|
|
# "content" type -- cumulative text
|
|
# Sanitize the full cumulative then diff against
|
|
# the last sanitized snapshot so cross-chunk XML
|
|
# tags are handled correctly.
|
|
raw_cumulative = event.get("text", "")
|
|
clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
|
|
new_text = clean_cumulative[len(prev_text) :]
|
|
prev_text = clean_cumulative
|
|
if not new_text:
|
|
continue
|
|
chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(content = new_text),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
final_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
# Usage chunk (OpenAI-standard: choices=[], usage populated)
|
|
if _stream_usage or _stream_timings:
|
|
usage_obj = CompletionUsage(
|
|
prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0),
|
|
completion_tokens = (_stream_usage or {}).get(
|
|
"completion_tokens", 0
|
|
),
|
|
total_tokens = (_stream_usage or {}).get("total_tokens", 0),
|
|
)
|
|
usage_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [],
|
|
usage = usage_obj,
|
|
timings = _stream_timings,
|
|
)
|
|
yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
except asyncio.CancelledError:
|
|
cancel_event.set()
|
|
raise
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
tb = traceback.format_exc()
|
|
logger.error(f"Error during GGUF tool streaming: {e}\n{tb}")
|
|
error_chunk = {
|
|
"error": {
|
|
"message": _friendly_error(e),
|
|
"type": "server_error",
|
|
},
|
|
}
|
|
yield f"data: {json.dumps(error_chunk)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
gguf_tool_stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
# ── Standard GGUF path (no tools) ─────────────────────
|
|
|
|
def gguf_generate():
|
|
return llama_backend.generate_chat_completion(
|
|
messages = gguf_messages,
|
|
image_b64 = image_b64,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_tokens = payload.max_tokens,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
presence_penalty = payload.presence_penalty,
|
|
cancel_event = cancel_event,
|
|
enable_thinking = payload.enable_thinking,
|
|
)
|
|
|
|
_gguf_sentinel = object()
|
|
|
|
if payload.stream:
|
|
|
|
async def gguf_stream_chunks():
|
|
try:
|
|
# First chunk: role
|
|
first_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(role = "assistant"),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
# Iterate the synchronous generator in a thread so
|
|
# the event loop stays free for disconnect detection.
|
|
gen = gguf_generate()
|
|
prev_text = ""
|
|
_stream_usage = None
|
|
_stream_timings = None
|
|
while True:
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
return
|
|
cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel)
|
|
if cumulative is _gguf_sentinel:
|
|
break
|
|
# Capture server metadata for final usage chunk
|
|
if isinstance(cumulative, dict):
|
|
if cumulative.get("type") == "metadata":
|
|
_stream_usage = cumulative.get("usage")
|
|
_stream_timings = cumulative.get("timings")
|
|
else:
|
|
logger.warning(
|
|
"gguf_stream_chunks: unexpected dict event: %s",
|
|
{
|
|
k: v
|
|
for k, v in cumulative.items()
|
|
if k != "timings"
|
|
},
|
|
)
|
|
continue
|
|
new_text = cumulative[len(prev_text) :]
|
|
prev_text = cumulative
|
|
if not new_text:
|
|
continue
|
|
chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(content = new_text),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
# Final chunk
|
|
final_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
# Usage chunk (OpenAI-standard: choices=[], usage populated)
|
|
if _stream_usage or _stream_timings:
|
|
usage_obj = CompletionUsage(
|
|
prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0),
|
|
completion_tokens = (_stream_usage or {}).get(
|
|
"completion_tokens", 0
|
|
),
|
|
total_tokens = (_stream_usage or {}).get("total_tokens", 0),
|
|
)
|
|
usage_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [],
|
|
usage = usage_obj,
|
|
timings = _stream_timings,
|
|
)
|
|
yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
except asyncio.CancelledError:
|
|
cancel_event.set()
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error during GGUF streaming: {e}", exc_info = True)
|
|
error_chunk = {
|
|
"error": {
|
|
"message": _friendly_error(e),
|
|
"type": "server_error",
|
|
},
|
|
}
|
|
yield f"data: {json.dumps(error_chunk)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
gguf_stream_chunks(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
else:
|
|
try:
|
|
full_text = ""
|
|
for token in gguf_generate():
|
|
if isinstance(token, dict):
|
|
continue # skip metadata dict in non-streaming path
|
|
full_text = token
|
|
|
|
response = ChatCompletion(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
CompletionChoice(
|
|
message = CompletionMessage(content = full_text),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
return JSONResponse(content = response.model_dump())
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during GGUF completion: {e}", exc_info = True)
|
|
raise HTTPException(status_code = 500, detail = str(e))
|
|
|
|
# ── Standard Unsloth path ─────────────────────────────────
|
|
|
|
# Decode image (from content parts OR legacy field)
|
|
image_b64 = extracted_image_b64 or payload.image_base64
|
|
image = None
|
|
|
|
if image_b64:
|
|
try:
|
|
import base64
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
|
|
model_info = backend.models.get(backend.active_model_name, {})
|
|
if not model_info.get("is_vision"):
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Image provided but current model is text-only. Load a vision model.",
|
|
)
|
|
|
|
image_data = base64.b64decode(image_b64)
|
|
image = Image.open(BytesIO(image_data))
|
|
image = backend.resize_image(image)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}")
|
|
|
|
# Shared generation kwargs
|
|
gen_kwargs = dict(
|
|
messages = chat_messages,
|
|
system_prompt = system_prompt,
|
|
image = image,
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
top_k = payload.top_k,
|
|
min_p = payload.min_p,
|
|
max_new_tokens = payload.max_tokens or 2048,
|
|
repetition_penalty = payload.repetition_penalty,
|
|
)
|
|
|
|
# Choose generation path (adapter-controlled or standard)
|
|
cancel_event = threading.Event()
|
|
|
|
if payload.use_adapter is not None:
|
|
|
|
def generate():
|
|
return backend.generate_with_adapter_control(
|
|
use_adapter = payload.use_adapter,
|
|
cancel_event = cancel_event,
|
|
**gen_kwargs,
|
|
)
|
|
else:
|
|
|
|
def generate():
|
|
return backend.generate_chat_response(
|
|
cancel_event = cancel_event, **gen_kwargs
|
|
)
|
|
|
|
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
|
created = int(time.time())
|
|
|
|
# ── Streaming response ────────────────────────────────────────
|
|
if payload.stream:
|
|
|
|
async def stream_chunks():
|
|
try:
|
|
first_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(role = "assistant"),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
prev_text = ""
|
|
# Run sync generator in thread pool to avoid blocking
|
|
# the event loop. Critical for compare mode: two SSE
|
|
# requests arrive concurrently but the orchestrator
|
|
# serializes them via _gen_lock. Without run_in_executor
|
|
# the second request's blocking lock acquisition would
|
|
# freeze the entire event loop, stalling both streams.
|
|
_DONE = object() # sentinel for generator exhaustion
|
|
loop = asyncio.get_event_loop()
|
|
gen = generate()
|
|
while True:
|
|
# next(gen, _DONE) returns _DONE instead of raising
|
|
# StopIteration — StopIteration cannot propagate
|
|
# through asyncio futures (Python limitation).
|
|
cumulative = await loop.run_in_executor(None, next, gen, _DONE)
|
|
if cumulative is _DONE:
|
|
break
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
backend.reset_generation_state()
|
|
return
|
|
new_text = cumulative[len(prev_text) :]
|
|
prev_text = cumulative
|
|
if not new_text:
|
|
continue
|
|
chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(content = new_text),
|
|
finish_reason = None,
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
|
|
final_chunk = ChatCompletionChunk(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
ChunkChoice(
|
|
delta = ChoiceDelta(),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
except asyncio.CancelledError:
|
|
cancel_event.set()
|
|
backend.reset_generation_state()
|
|
raise
|
|
except Exception as e:
|
|
backend.reset_generation_state()
|
|
logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
|
|
error_chunk = {
|
|
"error": {
|
|
"message": _friendly_error(e),
|
|
"type": "server_error",
|
|
},
|
|
}
|
|
yield f"data: {json.dumps(error_chunk)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
stream_chunks(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
# ── Non-streaming response ────────────────────────────────────
|
|
else:
|
|
try:
|
|
full_text = ""
|
|
for token in generate():
|
|
full_text = token
|
|
|
|
response = ChatCompletion(
|
|
id = completion_id,
|
|
created = created,
|
|
model = model_name,
|
|
choices = [
|
|
CompletionChoice(
|
|
message = CompletionMessage(content = full_text),
|
|
finish_reason = "stop",
|
|
)
|
|
],
|
|
)
|
|
return JSONResponse(content = response.model_dump())
|
|
|
|
except Exception as e:
|
|
backend.reset_generation_state()
|
|
logger.error(f"Error during OpenAI completion: {e}", exc_info = True)
|
|
raise HTTPException(status_code = 500, detail = str(e))
|
|
|
|
|
|
# =====================================================================
|
|
# Sandbox file serving (/sandbox/{session_id}/{filename})
|
|
# =====================================================================
|
|
|
|
_SANDBOX_MEDIA_TYPES = {
|
|
".png": "image/png",
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".gif": "image/gif",
|
|
".webp": "image/webp",
|
|
".bmp": "image/bmp",
|
|
}
|
|
|
|
|
|
@router.get("/sandbox/{session_id}/{filename}")
|
|
async def serve_sandbox_file(
|
|
session_id: str,
|
|
filename: str,
|
|
request: Request,
|
|
token: Optional[str] = None,
|
|
):
|
|
"""
|
|
Serve image files created by Python tool execution.
|
|
|
|
Accepts auth via Authorization header OR ?token= query param
|
|
(needed because <img src> cannot send custom headers).
|
|
"""
|
|
from fastapi.responses import FileResponse
|
|
|
|
# ── Authentication (header or query param) ──────────────────
|
|
auth_header = request.headers.get("authorization")
|
|
if auth_header and auth_header.lower().startswith("bearer "):
|
|
jwt_token = auth_header[7:]
|
|
elif token:
|
|
jwt_token = token
|
|
else:
|
|
raise HTTPException(
|
|
status_code = status.HTTP_401_UNAUTHORIZED,
|
|
detail = "Missing authentication token",
|
|
)
|
|
from fastapi.security import HTTPAuthorizationCredentials
|
|
|
|
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token)
|
|
await get_current_subject(creds)
|
|
|
|
# ── Filename sanitization ───────────────────────────────────
|
|
safe_filename = os.path.basename(filename)
|
|
if not safe_filename or safe_filename in (".", ".."):
|
|
raise HTTPException(status_code = 404, detail = "Not found")
|
|
|
|
# ── Extension allowlist ─────────────────────────────────────
|
|
ext = os.path.splitext(safe_filename)[1].lower()
|
|
media_type = _SANDBOX_MEDIA_TYPES.get(ext)
|
|
if not media_type:
|
|
raise HTTPException(
|
|
status_code = status.HTTP_403_FORBIDDEN,
|
|
detail = "File type not allowed",
|
|
)
|
|
|
|
# ── Path containment check ──────────────────────────────────
|
|
home = os.path.expanduser("~")
|
|
sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox"))
|
|
safe_session = os.path.basename(session_id.replace("..", ""))
|
|
if not safe_session:
|
|
raise HTTPException(status_code = 404, detail = "Not found")
|
|
|
|
file_path = os.path.realpath(
|
|
os.path.join(sandbox_root, safe_session, safe_filename)
|
|
)
|
|
if not file_path.startswith(sandbox_root + os.sep):
|
|
raise HTTPException(
|
|
status_code = status.HTTP_403_FORBIDDEN,
|
|
detail = "Access denied",
|
|
)
|
|
|
|
if not os.path.isfile(file_path):
|
|
raise HTTPException(status_code = 404, detail = "Not found")
|
|
|
|
return FileResponse(
|
|
path = file_path,
|
|
media_type = media_type,
|
|
headers = {
|
|
"Cache-Control": "private, no-store",
|
|
"X-Content-Type-Options": "nosniff",
|
|
},
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# OpenAI-Compatible Models Listing (/models → /v1/models)
|
|
# =====================================================================
|
|
|
|
|
|
@router.get("/models")
|
|
async def openai_list_models(
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
OpenAI-compatible model listing endpoint.
|
|
|
|
Returns the currently loaded model in the format expected by
|
|
OpenAI-compatible clients (``GET /v1/models``).
|
|
"""
|
|
models = []
|
|
|
|
# Check GGUF backend
|
|
llama_backend = get_llama_cpp_backend()
|
|
if llama_backend.is_loaded:
|
|
models.append(
|
|
{
|
|
"id": llama_backend.model_identifier,
|
|
"object": "model",
|
|
"owned_by": "local",
|
|
}
|
|
)
|
|
|
|
# Check Unsloth backend
|
|
backend = get_inference_backend()
|
|
if backend.active_model_name:
|
|
models.append(
|
|
{
|
|
"id": backend.active_model_name,
|
|
"object": "model",
|
|
"owned_by": "local",
|
|
}
|
|
)
|
|
|
|
return {"object": "list", "data": models}
|
|
|
|
|
|
# =====================================================================
|
|
# OpenAI-Compatible Completions Proxy (/completions → /v1/completions)
|
|
# =====================================================================
|
|
|
|
|
|
@router.post("/completions")
|
|
async def openai_completions(
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
OpenAI-compatible text completions endpoint (non-chat).
|
|
|
|
Transparently proxies to the running llama-server's ``/v1/completions``.
|
|
Only available when a GGUF model is loaded.
|
|
"""
|
|
llama_backend = get_llama_cpp_backend()
|
|
if not llama_backend.is_loaded:
|
|
raise HTTPException(
|
|
status_code = 503,
|
|
detail = "No GGUF model loaded. Load a GGUF model first.",
|
|
)
|
|
|
|
body = await request.json()
|
|
target_url = f"{llama_backend.base_url}/v1/completions"
|
|
is_stream = body.get("stream", False)
|
|
|
|
if is_stream:
|
|
|
|
async def _stream():
|
|
# Manual httpx client/response lifecycle — see
|
|
# _anthropic_passthrough_stream for the full rationale. Briefly:
|
|
# `async with` inside an async generator causes
|
|
# "Attempted to exit cancel scope in a different task" /
|
|
# "async generator ignored GeneratorExit" on Python 3.13 +
|
|
# httpcore 1.0.x when the generator is orphaned and finalized
|
|
# by GC. Closing via a finally block that catches Exception
|
|
# (but not BaseException) suppresses the anyio cleanup noise
|
|
# while letting GeneratorExit propagate cleanly.
|
|
client = httpx.AsyncClient(timeout = 600)
|
|
resp = None
|
|
try:
|
|
req = client.build_request("POST", target_url, json = body)
|
|
resp = await client.send(req, stream = True)
|
|
async for chunk in resp.aiter_bytes():
|
|
yield chunk
|
|
except Exception as e:
|
|
logger.error("openai_completions stream error: %s", e)
|
|
finally:
|
|
if resp is not None:
|
|
try:
|
|
await resp.aclose()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
await client.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
return StreamingResponse(_stream(), media_type = "text/event-stream")
|
|
else:
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(target_url, json = body, timeout = 600)
|
|
return Response(
|
|
content = resp.content,
|
|
status_code = resp.status_code,
|
|
media_type = "application/json",
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# OpenAI-Compatible Embeddings Proxy (/embeddings → /v1/embeddings)
|
|
# =====================================================================
|
|
|
|
|
|
@router.post("/embeddings")
|
|
async def openai_embeddings(
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
OpenAI-compatible embeddings endpoint.
|
|
|
|
Transparently proxies to the running llama-server's ``/v1/embeddings``.
|
|
Only available when a GGUF model is loaded.
|
|
Note: the loaded model must support pooling; otherwise llama-server
|
|
will return an error (expected).
|
|
"""
|
|
llama_backend = get_llama_cpp_backend()
|
|
if not llama_backend.is_loaded:
|
|
raise HTTPException(
|
|
status_code = 503,
|
|
detail = "No GGUF model loaded. Load a GGUF model first.",
|
|
)
|
|
|
|
body = await request.json()
|
|
target_url = f"{llama_backend.base_url}/v1/embeddings"
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(target_url, json = body, timeout = 600)
|
|
return Response(
|
|
content = resp.content,
|
|
status_code = resp.status_code,
|
|
media_type = "application/json",
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# OpenAI Responses API (/responses → /v1/responses)
|
|
# =====================================================================
|
|
|
|
|
|
def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]:
|
|
"""Convert a ResponsesRequest into a list of ChatMessage for the completions backend."""
|
|
messages: list[ChatMessage] = []
|
|
|
|
# System / developer instructions
|
|
if payload.instructions:
|
|
messages.append(ChatMessage(role = "system", content = payload.instructions))
|
|
|
|
# Simple string input
|
|
if isinstance(payload.input, str):
|
|
if payload.input:
|
|
messages.append(ChatMessage(role = "user", content = payload.input))
|
|
return messages
|
|
|
|
# List of ResponsesInputMessage
|
|
for msg in payload.input:
|
|
role = "system" if msg.role == "developer" else msg.role
|
|
|
|
if isinstance(msg.content, str):
|
|
messages.append(ChatMessage(role = role, content = msg.content))
|
|
else:
|
|
# Convert Responses content parts -> Chat content parts
|
|
parts = []
|
|
for part in msg.content:
|
|
if isinstance(part, ResponsesInputTextPart):
|
|
parts.append(TextContentPart(type = "text", text = part.text))
|
|
elif isinstance(part, ResponsesInputImagePart):
|
|
parts.append(
|
|
ImageContentPart(
|
|
type = "image_url",
|
|
image_url = ImageUrl(url = part.image_url, detail = part.detail),
|
|
)
|
|
)
|
|
messages.append(ChatMessage(role = role, content = parts if parts else ""))
|
|
|
|
return messages
|
|
|
|
|
|
def _build_chat_request(
|
|
payload: ResponsesRequest, messages: list[ChatMessage], stream: bool
|
|
) -> ChatCompletionRequest:
|
|
"""Build a ChatCompletionRequest from a ResponsesRequest."""
|
|
chat_kwargs = dict(
|
|
model = payload.model,
|
|
messages = messages,
|
|
stream = stream,
|
|
)
|
|
if payload.temperature is not None:
|
|
chat_kwargs["temperature"] = payload.temperature
|
|
if payload.top_p is not None:
|
|
chat_kwargs["top_p"] = payload.top_p
|
|
if payload.max_output_tokens is not None:
|
|
chat_kwargs["max_tokens"] = payload.max_output_tokens
|
|
return ChatCompletionRequest(**chat_kwargs)
|
|
|
|
|
|
async def _responses_non_streaming(
|
|
payload: ResponsesRequest,
|
|
messages: list[ChatMessage],
|
|
request: Request,
|
|
) -> JSONResponse:
|
|
"""Handle a non-streaming Responses API call."""
|
|
chat_req = _build_chat_request(payload, messages, stream = False)
|
|
result = await openai_chat_completions(chat_req, request)
|
|
|
|
# openai_chat_completions returns a JSONResponse for non-streaming
|
|
if isinstance(result, JSONResponse):
|
|
body = json.loads(result.body.decode())
|
|
elif isinstance(result, Response):
|
|
body = json.loads(result.body.decode())
|
|
else:
|
|
body = result
|
|
|
|
# Extract content and usage from the Chat Completions response
|
|
choices = body.get("choices", [])
|
|
text = ""
|
|
if choices:
|
|
msg = choices[0].get("message", {})
|
|
text = msg.get("content", "") or ""
|
|
|
|
usage_data = body.get("usage", {})
|
|
input_tokens = usage_data.get("prompt_tokens", 0)
|
|
output_tokens = usage_data.get("completion_tokens", 0)
|
|
|
|
resp_id = f"resp_{uuid.uuid4().hex[:12]}"
|
|
msg_id = f"msg_{uuid.uuid4().hex[:12]}"
|
|
|
|
response = ResponsesResponse(
|
|
id = resp_id,
|
|
created_at = int(time.time()),
|
|
status = "completed",
|
|
model = body.get("model", payload.model),
|
|
output = [
|
|
ResponsesOutputMessage(
|
|
id = msg_id,
|
|
status = "completed",
|
|
role = "assistant",
|
|
content = [
|
|
ResponsesOutputTextContent(text = text),
|
|
],
|
|
),
|
|
],
|
|
usage = ResponsesUsage(
|
|
input_tokens = input_tokens,
|
|
output_tokens = output_tokens,
|
|
total_tokens = input_tokens + output_tokens,
|
|
),
|
|
temperature = payload.temperature,
|
|
top_p = payload.top_p,
|
|
max_output_tokens = payload.max_output_tokens,
|
|
instructions = payload.instructions,
|
|
)
|
|
return JSONResponse(content = response.model_dump())
|
|
|
|
|
|
async def _responses_stream(
|
|
payload: ResponsesRequest,
|
|
messages: list[ChatMessage],
|
|
request: Request,
|
|
):
|
|
"""Handle a streaming Responses API call, emitting named SSE events."""
|
|
resp_id = f"resp_{uuid.uuid4().hex[:12]}"
|
|
msg_id = f"msg_{uuid.uuid4().hex[:12]}"
|
|
item_id = f"item_{uuid.uuid4().hex[:12]}"
|
|
created_at = int(time.time())
|
|
|
|
chat_req = _build_chat_request(payload, messages, stream = True)
|
|
result = await openai_chat_completions(chat_req, request)
|
|
|
|
async def event_generator():
|
|
full_text = ""
|
|
input_tokens = 0
|
|
output_tokens = 0
|
|
|
|
# ── Preamble events ──
|
|
yield f"event: response.created\ndata: {json.dumps({'type': 'response.created', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'in_progress', 'model': payload.model, 'output': [], 'usage': {'input_tokens': 0, 'output_tokens': 0, 'total_tokens': 0}}})}\n\n"
|
|
|
|
# output_item.added
|
|
output_item = {
|
|
"type": "message",
|
|
"id": msg_id,
|
|
"status": "in_progress",
|
|
"role": "assistant",
|
|
"content": [],
|
|
}
|
|
yield f"event: response.output_item.added\ndata: {json.dumps({'type': 'response.output_item.added', 'output_index': 0, 'item': output_item})}\n\n"
|
|
|
|
# content_part.added
|
|
content_part = {"type": "output_text", "text": "", "annotations": []}
|
|
yield f"event: response.content_part.added\ndata: {json.dumps({'type': 'response.content_part.added', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': content_part})}\n\n"
|
|
|
|
# ── Stream delta events from the inner chat completions stream ──
|
|
if isinstance(result, StreamingResponse):
|
|
async for raw_chunk in result.body_iterator:
|
|
if isinstance(raw_chunk, bytes):
|
|
raw_chunk = raw_chunk.decode("utf-8", errors = "replace")
|
|
|
|
for line in raw_chunk.split("\n"):
|
|
line = line.strip()
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data_str = line[6:]
|
|
if data_str == "[DONE]":
|
|
continue
|
|
try:
|
|
chunk_data = json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
choices = chunk_data.get("choices", [])
|
|
if not choices:
|
|
# Check for usage in final chunk
|
|
usage = chunk_data.get("usage")
|
|
if usage:
|
|
input_tokens = usage.get("prompt_tokens", input_tokens)
|
|
output_tokens = usage.get(
|
|
"completion_tokens", output_tokens
|
|
)
|
|
continue
|
|
|
|
delta = choices[0].get("delta", {})
|
|
content = delta.get("content")
|
|
if content:
|
|
full_text += content
|
|
delta_event = {
|
|
"type": "response.output_text.delta",
|
|
"item_id": msg_id,
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": content,
|
|
}
|
|
yield f"event: response.output_text.delta\ndata: {json.dumps(delta_event)}\n\n"
|
|
|
|
# Check for usage in chunk
|
|
usage = chunk_data.get("usage")
|
|
if usage:
|
|
input_tokens = usage.get("prompt_tokens", input_tokens)
|
|
output_tokens = usage.get("completion_tokens", output_tokens)
|
|
|
|
# ── Closing events ──
|
|
# output_text.done
|
|
yield f"event: response.output_text.done\ndata: {json.dumps({'type': 'response.output_text.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'text': full_text})}\n\n"
|
|
|
|
# content_part.done
|
|
yield f"event: response.content_part.done\ndata: {json.dumps({'type': 'response.content_part.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': {'type': 'output_text', 'text': full_text, 'annotations': []}})}\n\n"
|
|
|
|
# output_item.done
|
|
yield f"event: response.output_item.done\ndata: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': {'type': 'message', 'id': msg_id, 'status': 'completed', 'role': 'assistant', 'content': [{'type': 'output_text', 'text': full_text, 'annotations': []}]}})}\n\n"
|
|
|
|
# response.completed
|
|
total_tokens = input_tokens + output_tokens
|
|
completed_response = {
|
|
"type": "response.completed",
|
|
"response": {
|
|
"id": resp_id,
|
|
"object": "response",
|
|
"created_at": created_at,
|
|
"status": "completed",
|
|
"model": payload.model,
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": msg_id,
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": full_text,
|
|
"annotations": [],
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"usage": {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": total_tokens,
|
|
},
|
|
},
|
|
}
|
|
yield f"event: response.completed\ndata: {json.dumps(completed_response)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
@router.post("/responses")
|
|
async def openai_responses(
|
|
payload: ResponsesRequest,
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
OpenAI Responses API endpoint.
|
|
|
|
Accepts the Responses-format request, converts it to a
|
|
ChatCompletionRequest internally, and returns a response
|
|
matching the OpenAI Responses API schema (output array,
|
|
input_tokens/output_tokens, named SSE events for streaming).
|
|
"""
|
|
messages = _normalise_responses_input(payload)
|
|
if not messages:
|
|
raise HTTPException(status_code = 400, detail = "No input provided.")
|
|
|
|
if payload.stream:
|
|
return await _responses_stream(payload, messages, request)
|
|
return await _responses_non_streaming(payload, messages, request)
|
|
|
|
|
|
# =====================================================================
|
|
# Anthropic-Compatible Messages API (/messages → /v1/messages)
|
|
# =====================================================================
|
|
|
|
|
|
@router.post("/messages")
|
|
async def anthropic_messages(
|
|
payload: AnthropicMessagesRequest,
|
|
request: Request,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Anthropic-compatible Messages API endpoint.
|
|
|
|
Translates Anthropic message format to internal OpenAI format, runs
|
|
through the existing agentic tool loop when tools are provided, and
|
|
returns responses in Anthropic Messages API format (streaming SSE or
|
|
non-streaming JSON).
|
|
"""
|
|
llama_backend = get_llama_cpp_backend()
|
|
if not llama_backend.is_loaded:
|
|
raise HTTPException(
|
|
status_code = 503,
|
|
detail = "No GGUF model loaded. Load a GGUF model first.",
|
|
)
|
|
|
|
model_name = getattr(llama_backend, "model_identifier", None) or payload.model
|
|
message_id = f"msg_{uuid.uuid4().hex[:24]}"
|
|
|
|
# ── Translate Anthropic → OpenAI ──────────────────────────
|
|
openai_messages = anthropic_messages_to_openai(
|
|
[m.model_dump() for m in payload.messages],
|
|
payload.system,
|
|
)
|
|
|
|
temperature = payload.temperature if payload.temperature is not None else 0.6
|
|
top_p = payload.top_p if payload.top_p is not None else 0.95
|
|
top_k = payload.top_k if payload.top_k is not None else 20
|
|
min_p = payload.min_p if payload.min_p is not None else 0.01
|
|
repetition_penalty = (
|
|
payload.repetition_penalty if payload.repetition_penalty is not None else 1.0
|
|
)
|
|
presence_penalty = (
|
|
payload.presence_penalty if payload.presence_penalty is not None else 0.0
|
|
)
|
|
stop = payload.stop_sequences or None
|
|
|
|
# tool_choice is declared on AnthropicMessagesRequest for Anthropic SDK
|
|
# compatibility (the SDK often sets it by default), but it is not
|
|
# currently honored by Unsloth's backend. Warn once per request so the
|
|
# silent drop is visible to operators instead of looking like a model
|
|
# quality issue to clients.
|
|
if payload.tool_choice is not None:
|
|
logger.warning(
|
|
"anthropic_messages.tool_choice_ignored",
|
|
tool_choice = payload.tool_choice,
|
|
note = (
|
|
"tool_choice is accepted for Anthropic SDK compatibility but not "
|
|
"honored by Unsloth. Use enable_tools / enabled_tools (server-side "
|
|
"built-in tools) or restrict the `tools` array (client-side) to "
|
|
"control which tools the model sees."
|
|
),
|
|
)
|
|
|
|
cancel_event = threading.Event()
|
|
|
|
# ── Tool routing ──────────────────────────────────────────
|
|
# Three paths:
|
|
# 1. enable_tools=true → server-side execution of built-in tools (Unsloth shorthand)
|
|
# 2. tools=[...] only → client-side pass-through (standard Anthropic behavior)
|
|
# 3. neither → plain chat
|
|
server_tools = payload.enable_tools and llama_backend.supports_tools
|
|
client_tools = (
|
|
not server_tools
|
|
and payload.tools
|
|
and len(payload.tools) > 0
|
|
and llama_backend.supports_tools
|
|
)
|
|
|
|
# ── Client-side pass-through path ─────────────────────────
|
|
if client_tools:
|
|
openai_tools = anthropic_tools_to_openai(payload.tools)
|
|
|
|
if payload.stream:
|
|
return await _anthropic_passthrough_stream(
|
|
request,
|
|
cancel_event,
|
|
llama_backend,
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
payload.max_tokens,
|
|
message_id,
|
|
model_name,
|
|
stop = stop,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
return await _anthropic_passthrough_non_streaming(
|
|
llama_backend,
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
payload.max_tokens,
|
|
message_id,
|
|
model_name,
|
|
stop = stop,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
|
|
if server_tools:
|
|
from core.inference.tools import ALL_TOOLS
|
|
|
|
if payload.enabled_tools is not None:
|
|
openai_tools = [
|
|
t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools
|
|
]
|
|
else:
|
|
openai_tools = ALL_TOOLS
|
|
|
|
# Build tool-use system prompt nudge (same logic as /chat/completions)
|
|
_tool_names = {t["function"]["name"] for t in openai_tools}
|
|
_has_web = "web_search" in _tool_names
|
|
_has_code = "python" in _tool_names or "terminal" in _tool_names
|
|
|
|
_date_line = f"The current date is {_date.today().isoformat()}."
|
|
_model_size_b = _extract_model_size_b(model_name)
|
|
_is_small_model = _model_size_b is not None and _model_size_b < 9
|
|
|
|
if _is_small_model:
|
|
_web_tips = "Do not repeat the same search query."
|
|
else:
|
|
_web_tips = (
|
|
"When you search and find a relevant URL in the results, "
|
|
"fetch its full content by calling web_search with the url parameter. "
|
|
"Do not repeat the same search query. If a search returns "
|
|
"no useful results, try rephrasing or fetching a result URL directly."
|
|
)
|
|
_code_tips = (
|
|
"Use code execution for math, calculations, data processing, "
|
|
"or to parse and analyze information from tool results."
|
|
)
|
|
|
|
if _has_web and _has_code:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"tools rather than answering from memory. "
|
|
+ _web_tips
|
|
+ " "
|
|
+ _code_tips
|
|
)
|
|
elif _has_code:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"code execution rather than answering from memory. " + _code_tips
|
|
)
|
|
elif _has_web:
|
|
_nudge = (
|
|
_date_line + " "
|
|
"You have access to tools. When appropriate, prefer using "
|
|
"web search for up-to-date or uncertain factual "
|
|
"information rather than answering from memory. " + _web_tips
|
|
)
|
|
else:
|
|
_nudge = ""
|
|
|
|
if _nudge:
|
|
_nudge += _TOOL_ACTION_NUDGE
|
|
# Inject into system prompt
|
|
if openai_messages and openai_messages[0].get("role") == "system":
|
|
openai_messages[0]["content"] = (
|
|
openai_messages[0]["content"].rstrip() + "\n\n" + _nudge
|
|
)
|
|
else:
|
|
openai_messages.insert(0, {"role": "system", "content": _nudge})
|
|
|
|
# Strip stale tool-call XML from conversation
|
|
for _msg in openai_messages:
|
|
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
|
|
_msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
|
|
|
|
def _run_tool_gen():
|
|
return llama_backend.generate_chat_completion_with_tools(
|
|
messages = openai_messages,
|
|
tools = openai_tools,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
max_tokens = payload.max_tokens,
|
|
stop = stop,
|
|
cancel_event = cancel_event,
|
|
max_tool_iterations = 25,
|
|
auto_heal_tool_calls = True,
|
|
tool_call_timeout = 300,
|
|
session_id = payload.session_id,
|
|
)
|
|
|
|
if payload.stream:
|
|
return await _anthropic_tool_stream(
|
|
request,
|
|
cancel_event,
|
|
_run_tool_gen,
|
|
message_id,
|
|
model_name,
|
|
)
|
|
return await _anthropic_tool_non_streaming(
|
|
_run_tool_gen,
|
|
message_id,
|
|
model_name,
|
|
)
|
|
|
|
# ── No-tool path ──────────────────────────────────────────
|
|
def _run_plain_gen():
|
|
return llama_backend.generate_chat_completion(
|
|
messages = openai_messages,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
max_tokens = payload.max_tokens,
|
|
stop = stop,
|
|
cancel_event = cancel_event,
|
|
)
|
|
|
|
if payload.stream:
|
|
return await _anthropic_plain_stream(
|
|
request,
|
|
cancel_event,
|
|
_run_plain_gen,
|
|
message_id,
|
|
model_name,
|
|
)
|
|
return await _anthropic_plain_non_streaming(
|
|
_run_plain_gen,
|
|
message_id,
|
|
model_name,
|
|
)
|
|
|
|
|
|
async def _anthropic_tool_stream(
|
|
request,
|
|
cancel_event,
|
|
run_gen,
|
|
message_id,
|
|
model_name,
|
|
):
|
|
"""Streaming response for the tool-calling path."""
|
|
_sentinel = object()
|
|
|
|
async def _stream():
|
|
emitter = AnthropicStreamEmitter()
|
|
for line in emitter.start(message_id, model_name):
|
|
yield line
|
|
|
|
gen = run_gen()
|
|
try:
|
|
while True:
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
return
|
|
event = await asyncio.to_thread(next, gen, _sentinel)
|
|
if event is _sentinel:
|
|
break
|
|
# Strip leaked tool-call XML from content events
|
|
if event.get("type") == "content":
|
|
event = dict(event)
|
|
event["text"] = _TOOL_XML_RE.sub("", event["text"])
|
|
for line in emitter.feed(event):
|
|
yield line
|
|
except Exception as e:
|
|
logger.error("anthropic_messages stream error: %s", e)
|
|
|
|
for line in emitter.finish("end_turn"):
|
|
yield line
|
|
|
|
return StreamingResponse(
|
|
_stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
async def _anthropic_plain_stream(
|
|
request,
|
|
cancel_event,
|
|
run_gen,
|
|
message_id,
|
|
model_name,
|
|
):
|
|
"""Streaming response for the no-tool path."""
|
|
_sentinel = object()
|
|
|
|
async def _stream():
|
|
emitter = AnthropicStreamEmitter()
|
|
for line in emitter.start(message_id, model_name):
|
|
yield line
|
|
|
|
gen = run_gen()
|
|
try:
|
|
while True:
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
return
|
|
cumulative = await asyncio.to_thread(next, gen, _sentinel)
|
|
if cumulative is _sentinel:
|
|
break
|
|
if isinstance(cumulative, dict):
|
|
if cumulative.get("type") == "metadata":
|
|
for line in emitter.feed(cumulative):
|
|
yield line
|
|
continue
|
|
# Plain generator yields cumulative text strings
|
|
for line in emitter.feed({"type": "content", "text": cumulative}):
|
|
yield line
|
|
except Exception as e:
|
|
logger.error("anthropic_messages stream error: %s", e)
|
|
|
|
for line in emitter.finish("end_turn"):
|
|
yield line
|
|
|
|
return StreamingResponse(
|
|
_stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
async def _anthropic_tool_non_streaming(run_gen, message_id, model_name):
|
|
"""Non-streaming response for the tool-calling path.
|
|
|
|
Builds ``content_blocks`` in generation order (text → tool_use → text →
|
|
tool_use → ...), mirroring the streaming emitter's behavior. Deltas
|
|
within a single synthesis turn are merged into the trailing text block;
|
|
tool_use blocks interrupt the text sequence and open a new text block on
|
|
the next content event.
|
|
|
|
``prev_text`` is reset on ``tool_end`` because
|
|
``generate_chat_completion_with_tools`` yields cumulative content *per
|
|
turn* — the first content event of turn N+1 must diff against an empty
|
|
baseline, not against turn N's final length.
|
|
"""
|
|
content_blocks: list = []
|
|
usage = {}
|
|
prev_text = ""
|
|
|
|
for event in run_gen():
|
|
etype = event.get("type", "")
|
|
if etype == "content":
|
|
# Strip leaked tool-call XML
|
|
clean = _TOOL_XML_RE.sub("", event["text"])
|
|
new = clean[len(prev_text) :]
|
|
prev_text = clean
|
|
if new:
|
|
if content_blocks and isinstance(
|
|
content_blocks[-1], AnthropicResponseTextBlock
|
|
):
|
|
content_blocks[-1].text += new
|
|
else:
|
|
content_blocks.append(AnthropicResponseTextBlock(text = new))
|
|
elif etype == "tool_start":
|
|
content_blocks.append(
|
|
AnthropicResponseToolUseBlock(
|
|
id = event["tool_call_id"],
|
|
name = event["tool_name"],
|
|
input = event.get("arguments", {}),
|
|
)
|
|
)
|
|
elif etype == "tool_end":
|
|
prev_text = ""
|
|
elif etype == "metadata":
|
|
usage = event.get("usage", {})
|
|
|
|
resp = AnthropicMessagesResponse(
|
|
id = message_id,
|
|
model = model_name,
|
|
content = content_blocks,
|
|
stop_reason = "end_turn",
|
|
usage = AnthropicUsage(
|
|
input_tokens = usage.get("prompt_tokens", 0),
|
|
output_tokens = usage.get("completion_tokens", 0),
|
|
),
|
|
)
|
|
return JSONResponse(content = resp.model_dump())
|
|
|
|
|
|
async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
|
|
"""Non-streaming response for the no-tool path."""
|
|
text_parts = []
|
|
usage = {}
|
|
prev_text = ""
|
|
|
|
for cumulative in run_gen():
|
|
if isinstance(cumulative, dict):
|
|
if cumulative.get("type") == "metadata":
|
|
usage = cumulative.get("usage", {})
|
|
continue
|
|
new = cumulative[len(prev_text) :]
|
|
prev_text = cumulative
|
|
if new:
|
|
text_parts.append(new)
|
|
|
|
full_text = "".join(text_parts)
|
|
content_blocks = []
|
|
if full_text:
|
|
content_blocks.append(AnthropicResponseTextBlock(text = full_text))
|
|
|
|
resp = AnthropicMessagesResponse(
|
|
id = message_id,
|
|
model = model_name,
|
|
content = content_blocks,
|
|
stop_reason = "end_turn",
|
|
usage = AnthropicUsage(
|
|
input_tokens = usage.get("prompt_tokens", 0),
|
|
output_tokens = usage.get("completion_tokens", 0),
|
|
),
|
|
)
|
|
return JSONResponse(content = resp.model_dump())
|
|
|
|
|
|
# =====================================================================
|
|
# Client-side tool pass-through (Anthropic-native tools field)
|
|
# =====================================================================
|
|
|
|
|
|
def _build_passthrough_payload(
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
max_tokens,
|
|
stream,
|
|
stop = None,
|
|
min_p = None,
|
|
repetition_penalty = None,
|
|
presence_penalty = None,
|
|
):
|
|
body = {
|
|
"messages": openai_messages,
|
|
"tools": openai_tools,
|
|
"tool_choice": "auto",
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"stream": stream,
|
|
}
|
|
if stream:
|
|
body["stream_options"] = {"include_usage": True}
|
|
if max_tokens is not None:
|
|
body["max_tokens"] = max_tokens
|
|
if stop:
|
|
body["stop"] = stop
|
|
if min_p is not None:
|
|
body["min_p"] = min_p
|
|
if repetition_penalty is not None:
|
|
# llama-server's field is "repeat_penalty", not "repetition_penalty"
|
|
body["repeat_penalty"] = repetition_penalty
|
|
if presence_penalty is not None:
|
|
body["presence_penalty"] = presence_penalty
|
|
return body
|
|
|
|
|
|
async def _anthropic_passthrough_stream(
|
|
request,
|
|
cancel_event,
|
|
llama_backend,
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
max_tokens,
|
|
message_id,
|
|
model_name,
|
|
stop = None,
|
|
min_p = None,
|
|
repetition_penalty = None,
|
|
presence_penalty = None,
|
|
):
|
|
"""Streaming client-side pass-through: forward tools to llama-server and
|
|
translate its streaming response to Anthropic SSE without executing anything."""
|
|
target_url = f"{llama_backend.base_url}/v1/chat/completions"
|
|
body = _build_passthrough_payload(
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
max_tokens,
|
|
True,
|
|
stop = stop,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
|
|
async def _stream():
|
|
emitter = AnthropicPassthroughEmitter()
|
|
for line in emitter.start(message_id, model_name):
|
|
yield line
|
|
|
|
# Manage the httpx client and response MANUALLY — no `async with`.
|
|
#
|
|
# On Python 3.13 + httpcore 1.0.x, an orphaned async generator (e.g.
|
|
# when the client disconnects mid-stream and Starlette drops the
|
|
# StreamingResponse iterator without explicitly calling aclose())
|
|
# is finalized by Python's asyncgen GC hook in a DIFFERENT asyncio
|
|
# task than the one that originally entered the httpx context
|
|
# managers. When `async with` exits run in the wrong task, httpcore's
|
|
# internal `HTTP11ConnectionByteStream.aclose()` hits
|
|
# `anyio.CancelScope.__exit__` with a mismatched task and raises
|
|
# RuntimeError("Attempted to exit cancel scope in a different task"),
|
|
# which escapes as "Exception ignored in:" because it happens during
|
|
# GC finalization outside any user-owned try/except.
|
|
#
|
|
# The fix: do not use `async with` for the client/response. Close
|
|
# them in a finally block wrapped in `try: ... except Exception: pass`.
|
|
# This narrowly suppresses RuntimeError / other Exception subclasses
|
|
# from the anyio cleanup noise while letting GeneratorExit (a
|
|
# BaseException, not Exception) propagate through cleanly so the
|
|
# generator terminates as Python expects.
|
|
client = httpx.AsyncClient(timeout = 600)
|
|
resp = None
|
|
try:
|
|
req = client.build_request("POST", target_url, json = body)
|
|
resp = await client.send(req, stream = True)
|
|
|
|
async for raw_line in resp.aiter_lines():
|
|
if await request.is_disconnected():
|
|
cancel_event.set()
|
|
break
|
|
if not raw_line or not raw_line.startswith("data: "):
|
|
continue
|
|
data_str = raw_line[6:]
|
|
if data_str.strip() == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
for line in emitter.feed_chunk(chunk):
|
|
yield line
|
|
except Exception as e:
|
|
logger.error("anthropic_messages passthrough stream error: %s", e)
|
|
finally:
|
|
if resp is not None:
|
|
try:
|
|
await resp.aclose()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
await client.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
for line in emitter.finish():
|
|
yield line
|
|
|
|
return StreamingResponse(
|
|
_stream(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
async def _anthropic_passthrough_non_streaming(
|
|
llama_backend,
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
max_tokens,
|
|
message_id,
|
|
model_name,
|
|
stop = None,
|
|
min_p = None,
|
|
repetition_penalty = None,
|
|
presence_penalty = None,
|
|
):
|
|
"""Non-streaming client-side pass-through."""
|
|
target_url = f"{llama_backend.base_url}/v1/chat/completions"
|
|
body = _build_passthrough_payload(
|
|
openai_messages,
|
|
openai_tools,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
max_tokens,
|
|
False,
|
|
stop = stop,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(target_url, json = body, timeout = 600)
|
|
|
|
if resp.status_code != 200:
|
|
raise HTTPException(
|
|
status_code = resp.status_code,
|
|
detail = f"llama-server error: {resp.text[:500]}",
|
|
)
|
|
|
|
data = resp.json()
|
|
choice = (data.get("choices") or [{}])[0]
|
|
message = choice.get("message") or {}
|
|
finish_reason = choice.get("finish_reason")
|
|
|
|
content_blocks = []
|
|
text = message.get("content") or ""
|
|
if text:
|
|
text = _TOOL_XML_RE.sub("", text).strip()
|
|
if text:
|
|
content_blocks.append(AnthropicResponseTextBlock(text = text))
|
|
|
|
tool_calls = message.get("tool_calls") or []
|
|
for tc in tool_calls:
|
|
fn = tc.get("function") or {}
|
|
try:
|
|
args = json.loads(fn.get("arguments", "{}"))
|
|
except json.JSONDecodeError:
|
|
args = {}
|
|
content_blocks.append(
|
|
AnthropicResponseToolUseBlock(
|
|
id = tc.get("id", ""),
|
|
name = fn.get("name", ""),
|
|
input = args,
|
|
)
|
|
)
|
|
|
|
if tool_calls:
|
|
stop_reason = "tool_use"
|
|
elif finish_reason == "length":
|
|
stop_reason = "max_tokens"
|
|
else:
|
|
stop_reason = "end_turn"
|
|
|
|
usage = data.get("usage") or {}
|
|
resp_obj = AnthropicMessagesResponse(
|
|
id = message_id,
|
|
model = model_name,
|
|
content = content_blocks,
|
|
stop_reason = stop_reason,
|
|
usage = AnthropicUsage(
|
|
input_tokens = usage.get("prompt_tokens", 0),
|
|
output_tokens = usage.get("completion_tokens", 0),
|
|
),
|
|
)
|
|
return JSONResponse(content = resp_obj.model_dump())
|