| import httpx |
| import json |
| import uuid |
| import time |
| import os |
| import random |
| import string |
| import logging |
| import asyncio |
|
|
| from contextlib import asynccontextmanager |
| from collections import defaultdict |
|
|
| from fastapi import FastAPI, Request, HTTPException, Depends |
| from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials |
|
|
| from config import ( |
| DEFAULT_UPSTREAM_URL, UPSTREAMS, UPSTREAM_API_KEYS, |
| GITHUB_TOKEN, GIST_ID, MODELS_GIST_ID, |
| HF_TOKEN, HF_DATASET, |
| MASTER_KEY, |
| AI_GATEWAY_URL, UPSTREAM_14448_URL, |
| BLOCKED_IPS, |
| CHATGPT_IMAGE_API_URL, GEMINI_HUB_URL, |
| TURNSTILE_SITE_KEY, TURNSTILE_SECRET_KEY, BALANCE_ENDPOINTS, |
| ) |
| from pydantic import BaseModel |
| import ipaddress |
| from templates import templates |
|
|
| |
| PROXYCHECK_URL = "https://proxycheck.io/v2/{ip}?vpn=1&asn=1" |
|
|
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger("apiarium") |
|
|
| |
|
|
| FALLBACK_UPSTREAM = os.environ.get("UPSTREAM_URL", DEFAULT_UPSTREAM_URL) |
| FALLBACK_API_KEY = os.environ.get("UPSTREAM_API_KEY", "") |
| MAX_CONTINUATIONS = 0 |
|
|
| |
|
|
| RPM_LIMIT = 4 |
| KEYS_FILE = os.path.join(os.path.dirname(__file__), "data", "api_keys.json") |
| PROVIDERS_FILE = os.path.join(os.path.dirname(__file__), "data", "providers.json") |
|
|
| |
| |
| |
| |
| |
| GEMINI_HUB_VIDEO_MODELS = {"veo-3", "veo-3-fast", "veo-3.1-fast"} |
| GEMINI_HUB_IMAGE_MODELS = {"gemini-3-image", "gemini-3-flash-image", "gemini-3.1-flash-image"} |
| GEMINI_HUB_TEXT_MODELS = {"gemini-3", "gemini-3-flash", "gemini-3-pro"} |
| GEMINI_HUB_MODELS = GEMINI_HUB_VIDEO_MODELS | GEMINI_HUB_IMAGE_MODELS | GEMINI_HUB_TEXT_MODELS |
|
|
| _rate_limits: dict = defaultdict(list) |
|
|
| |
| _dynamic_providers: dict = {} |
|
|
|
|
| |
|
|
| _github_token = os.environ.get("GITHUB_TOKEN", GITHUB_TOKEN) |
| _gist_id = os.environ.get("GIST_ID", GIST_ID) |
| |
| |
| |
| _models_gist_id = os.environ.get("MODELS_GIST_ID", MODELS_GIST_ID) or _gist_id |
| _keys_cache: dict = {} |
| _keys_loaded = False |
|
|
| |
| _shared_client: httpx.AsyncClient | None = None |
|
|
|
|
| def _get_shared_client() -> httpx.AsyncClient: |
| global _shared_client |
| if _shared_client is None or _shared_client.is_closed: |
| _shared_client = httpx.AsyncClient( |
| timeout=httpx.Timeout(connect=30.0, read=300.0, write=60.0, pool=30.0), |
| limits=httpx.Limits(max_connections=50, max_keepalive_connections=10), |
| ) |
| return _shared_client |
|
|
|
|
| def _load_keys() -> dict: |
| """Load keys from Gist (if configured) or local file.""" |
| global _keys_cache, _keys_loaded |
|
|
| |
| if _github_token and _gist_id: |
| try: |
| resp = httpx.get( |
| f"https://api.github.com/gists/{_gist_id}", |
| headers={"Authorization": f"Bearer {_github_token}"}, |
| timeout=10, |
| ) |
| if resp.status_code == 200: |
| gist = resp.json() |
| content = gist.get("files", {}).get("api_keys.json", {}).get("content", "{}") |
| _keys_cache = json.loads(content) |
| _keys_loaded = True |
| return _keys_cache |
| except Exception as e: |
| logger.warning("Gist load failed, falling back to local: %s", e) |
|
|
| |
| if not _keys_loaded: |
| try: |
| with open(KEYS_FILE, "r") as f: |
| _keys_cache = json.load(f) |
| except (FileNotFoundError, json.JSONDecodeError): |
| _keys_cache = {} |
| _keys_loaded = True |
|
|
| return _keys_cache |
|
|
|
|
| def _save_keys(data: dict): |
| """Save keys to Gist (if configured) and local file.""" |
| global _keys_cache |
| _keys_cache = data |
|
|
| |
| os.makedirs(os.path.dirname(KEYS_FILE), exist_ok=True) |
| with open(KEYS_FILE, "w") as f: |
| json.dump(data, f, indent=2) |
|
|
| |
| if _github_token and _gist_id: |
| try: |
| resp = httpx.patch( |
| f"https://api.github.com/gists/{_gist_id}", |
| headers={ |
| "Authorization": f"Bearer {_github_token}", |
| "Content-Type": "application/json", |
| }, |
| json={"files": {"api_keys.json": {"content": json.dumps(data, indent=2)}}}, |
| timeout=10, |
| ) |
| if resp.status_code == 200: |
| logger.info("Keys saved to Gist %s", _gist_id) |
| else: |
| logger.warning("Gist save failed: %s %s", resp.status_code, resp.text) |
| except Exception as e: |
| logger.warning("Gist save error: %s", e) |
|
|
|
|
| |
|
|
| def generate_api_key() -> str: |
| chars = string.ascii_lowercase + string.digits |
| suffix = "".join(random.choices(chars, k=15)) |
| return f"api-{suffix}" |
|
|
|
|
| def get_client_ip(request: Request) -> str: |
| forwarded = request.headers.get("x-forwarded-for") |
| if forwarded: |
| return forwarded.split(",")[0].strip() |
| real_ip = request.headers.get("x-real-ip") |
| if real_ip: |
| return real_ip.strip() |
| return request.client.host if request.client else "unknown" |
|
|
|
|
| async def check_vpn(ip: str) -> dict: |
| """Check if IP is VPN/proxy using proxycheck.io (free, no key).""" |
| try: |
| async with httpx.AsyncClient(timeout=5) as client: |
| resp = await client.get(PROXYCHECK_URL.format(ip=ip)) |
| if resp.status_code == 200: |
| data = resp.json() |
| ip_data = data.get(ip, {}) |
| is_proxy = ip_data.get("proxy", "no") == "yes" |
| provider = ip_data.get("provider", "unknown") |
| return {"vpn": is_proxy, "provider": provider} |
| except Exception as e: |
| logger.warning("VPN check failed: %s", e) |
| return {"vpn": False, "provider": "unknown"} |
|
|
|
|
| def find_key_by_ip(ip: str) -> str | None: |
| """Find existing key for a given IP.""" |
| keys = _load_keys() |
| for key, info in keys.items(): |
| if info.get("ip") == ip: |
| return key |
| return None |
|
|
|
|
| def verify_key_ip(api_key: str, client_ip: str) -> tuple: |
| if api_key == MASTER_KEY: |
| return True, "master" |
|
|
| keys = _load_keys() |
| if api_key not in keys: |
| return False, "Invalid API key" |
|
|
| stored_ip = keys[api_key]["ip"] |
| if stored_ip != client_ip: |
| return False, "IP mismatch: this key is not valid from your network" |
|
|
| return True, "ok" |
|
|
|
|
| def check_rate_limit(api_key: str) -> tuple: |
| if api_key == MASTER_KEY: |
| return True, "master", 0 |
|
|
| now = time.time() |
| window_start = now - 60 |
|
|
| _rate_limits[api_key] = [t for t in _rate_limits[api_key] if t > window_start] |
|
|
| if len(_rate_limits[api_key]) >= RPM_LIMIT: |
| oldest = min(_rate_limits[api_key]) |
| cooldown = int(60 - (now - oldest)) + 1 |
| return False, f"Rate limit exceeded ({RPM_LIMIT} RPM). Cooldown: {cooldown}s", cooldown |
|
|
| _rate_limits[api_key].append(now) |
| return True, "ok", 0 |
|
|
|
|
| |
|
|
| _chat_buffer: list = [] |
| _chat_buffer_lock = False |
| CHAT_LOG_DIR = os.path.join(os.path.dirname(__file__), "data", "chat_logs") |
|
|
|
|
| def _get_today() -> str: |
| from datetime import datetime |
| return datetime.utcnow().strftime("%Y-%m-%d") |
|
|
|
|
| def log_chat(ip: str, model: str, messages: list, response_content: str, api_key: str = ""): |
| """Log a chat request with full input/output to HF Dataset.""" |
| global _chat_buffer |
| from datetime import datetime |
|
|
| |
| clean_messages = [] |
| for msg in messages: |
| clean_messages.append({ |
| "role": msg.get("role", "unknown"), |
| "content": msg.get("content", ""), |
| }) |
|
|
| |
| key_type = "master" if api_key == MASTER_KEY else "regular" |
|
|
| entry = { |
| "timestamp": datetime.utcnow().isoformat(), |
| "ip": ip, |
| "model": model, |
| "key_type": key_type, |
| "input": clean_messages, |
| "output": response_content or "", |
| } |
|
|
| |
| _chat_buffer.append(entry) |
|
|
| |
| today = _get_today() |
| os.makedirs(CHAT_LOG_DIR, exist_ok=True) |
| log_file = os.path.join(CHAT_LOG_DIR, f"{today}.jsonl") |
| with open(log_file, "a") as f: |
| f.write(json.dumps(entry) + "\n") |
|
|
| |
| try: |
| loop = asyncio.get_event_loop() |
| if loop.is_running(): |
| asyncio.create_task(_flush_to_hf_async(today)) |
| else: |
| _flush_to_hf(today) |
| except RuntimeError: |
| _flush_to_hf(today) |
|
|
|
|
| async def _flush_to_hf_async(date_str: str): |
| """Run HF flush in background thread to avoid blocking response.""" |
| loop = asyncio.get_event_loop() |
| await loop.run_in_executor(None, _flush_to_hf, date_str) |
|
|
|
|
| def _flush_to_hf(date_str: str): |
| """Upload ONLY new entries (buffer) to Hugging Face Dataset.""" |
| global _chat_buffer |
| if not _chat_buffer or not HF_TOKEN: |
| return |
|
|
| try: |
| from huggingface_hub import HfApi, hf_hub_download |
| api = HfApi(token=HF_TOKEN) |
| repo_path = f"logs/{date_str}.jsonl" |
|
|
| |
| new_lines = "\n".join([json.dumps(e) for e in _chat_buffer]) + "\n" |
|
|
| |
| existing_content = "" |
| try: |
| existing_path = hf_hub_download( |
| HF_DATASET, repo_path, |
| repo_type="dataset", token=HF_TOKEN |
| ) |
| with open(existing_path, "r") as f: |
| existing_content = f.read() |
| except Exception: |
| pass |
|
|
| |
| temp_file = os.path.join(CHAT_LOG_DIR, f"{date_str}_upload.jsonl") |
| with open(temp_file, "w") as f: |
| f.write(existing_content + new_lines) |
|
|
| api.upload_file( |
| path_or_fileobj=temp_file, |
| path_in_repo=repo_path, |
| repo_id=HF_DATASET, |
| repo_type="dataset", |
| ) |
|
|
| |
| os.remove(temp_file) |
|
|
| logger.info("Flushed %d chat logs to HF Dataset", len(_chat_buffer)) |
| _chat_buffer = [] |
| except Exception as e: |
| logger.warning("HF upload failed: %s", e) |
|
|
|
|
| |
|
|
| from itertools import cycle |
|
|
| _round_robin_iters: dict = {} |
|
|
|
|
| def _get_next_upstream(model: str) -> dict: |
| """Pick next upstream for model. Returns {url, model} dict.""" |
| upstreams = UPSTREAMS.get(model) |
| if not upstreams: |
| return {"url": FALLBACK_UPSTREAM, "model": model} |
| if len(upstreams) == 1: |
| return upstreams[0] |
| if model not in _round_robin_iters: |
| _round_robin_iters[model] = cycle(upstreams) |
| return next(_round_robin_iters[model]) |
|
|
|
|
| |
|
|
| |
| _available_models: list = [] |
| _upstream_models: dict = {} |
|
|
|
|
| def _build_aliases() -> dict: |
| """Auto-generate aliases from UPSTREAMS keys.""" |
| aliases = {} |
| for key in UPSTREAMS: |
| aliases[key] = key |
| return aliases |
|
|
|
|
| async def _fetch_upstream_models(): |
| """Fetch available models from each unique upstream URL.""" |
| global _available_models, _upstream_models |
|
|
| |
| unique_urls = set() |
| for upstreams in UPSTREAMS.values(): |
| for u in upstreams: |
| unique_urls.add(u["url"]) |
|
|
| client = httpx.AsyncClient(timeout=10) |
| for url in unique_urls: |
| if url == AI_GATEWAY_URL: |
| _upstream_models[url] = [] |
| logger.info("Skipping /models fetch for AI gateway %s", url) |
| continue |
| models_url = url.rsplit("/chat/completions", 1)[0] + "/models" |
| try: |
| resp = await client.get(models_url, headers=build_upstream_headers(url)) |
| if resp.status_code == 200: |
| data = resp.json() |
| model_ids = [m["id"] for m in data.get("data", [])] |
| _upstream_models[url] = model_ids |
| logger.info("Fetched %d models from %s", len(model_ids), models_url) |
| else: |
| logger.warning("Failed to fetch models from %s: %s", models_url, resp.status_code) |
| _upstream_models[url] = [] |
| except Exception as e: |
| logger.warning("Error fetching models from %s: %s", models_url, e) |
| _upstream_models[url] = [] |
| await client.aclose() |
|
|
| |
| _available_models = list(UPSTREAMS.keys()) + ["gpt-image-1", "gpt-image-2"] + sorted(GEMINI_HUB_MODELS) |
| logger.info("Available models: %s", _available_models) |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(application: FastAPI): |
| |
| _load_keys() |
| logger.info("Preloaded %d API keys from Gist", len(_keys_cache)) |
| _load_providers() |
| await _fetch_upstream_models() |
| yield |
|
|
| app = FastAPI( |
| title="π APIarium", |
| version="0.4.0", |
| description=""" |
| # π APIarium |
| |
| **OpenAI-compatible LLM API gateway** with multi-provider routing, encrypted configuration, and admin management. |
| |
| ## β¨ Features |
| |
| - π **Multi-Model Routing** β Route requests to multiple LLM providers through a unified OpenAI-compatible API |
| - π **Encrypted Config** β Sensitive values stored as AES-256-GCM encrypted blobs |
| - π‘οΈ **Cloudflare Turnstile** β Bot protection for admin endpoints |
| - π **VPN Detection** β Free VPN/proxy detection via proxycheck.io |
| - π« **IP Blocking** β Built-in IP blocklist support |
| - πΎ **GitHub Gist Storage** β API key persistence via encrypted GitHub Gists |
| - π **HuggingFace Dataset Logging** β Request/response logging to HF Datasets |
| - β‘ **Dynamic Providers** β Add/remove upstream providers at runtime |
| |
| ## π Authentication |
| |
| All proxy endpoints require a Bearer token: |
| |
| ``` |
| Authorization: Bearer <your-api-key> |
| ``` |
| |
| Admin endpoints require the master key. |
| |
| ## π Endpoints |
| |
| ### Proxy (OpenAI-compatible) |
| - `POST /v1/chat/completions` β Chat completions (streaming supported) |
| - `POST /v1/completions` β Text completions |
| - `GET /v1/models` β List available model aliases |
| |
| ### Admin (Turnstile protected) |
| - `GET /admin` β Dashboard |
| - `GET /admin/keys` β API key management |
| - `GET /admin/providers` β Provider management |
| - `GET /admin/balance` β Balance monitoring |
| - `POST /admin/verify-turnstile` β Turnstile token verification |
| |
| ### Health |
| - `GET /health` β Health check |
| |
| ## π Links |
| |
| - **Space:** [huggingface.co/spaces/rnilkyway/APIarium](https://huggingface.co/spaces/rnilkyway/APIarium) |
| - **Redoc:** [/redoc](/redoc) |
| - **OpenAPI JSON:** [/openapi.json](/openapi.json) |
| """, |
| |
| |
| docs_url="/openapi-docs", |
| redoc_url="/redoc", |
| lifespan=lifespan, |
| ) |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_credentials=True, |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| security = HTTPBearer(auto_error=False) |
|
|
|
|
| |
|
|
| async def verify_request(request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)): |
| path = request.url.path |
| client_ip = get_client_ip(request) |
|
|
| |
| if client_ip in BLOCKED_IPS and path != "/": |
| raise HTTPException(status_code=403, detail="Access denied: your IP is blocked") |
|
|
| |
| if path in ("/", "/health", "/v1/models", "/v1/keys/me", "/v1/dash", "/v1/dashboard"): |
| return |
|
|
| |
| if credentials is None: |
| raise HTTPException(status_code=401, detail="Missing Authorization header. Use: Bearer <key>") |
|
|
| api_key = credentials.credentials |
| client_ip = get_client_ip(request) |
|
|
| valid, detail = verify_key_ip(api_key, client_ip) |
| if not valid: |
| raise HTTPException(status_code=403, detail=detail) |
|
|
| if api_key != MASTER_KEY: |
| allowed, rl_detail, cooldown = check_rate_limit(api_key) |
| if not allowed: |
| raise HTTPException( |
| status_code=429, |
| detail={ |
| "error": { |
| "message": rl_detail, |
| "type": "rate_limit_error", |
| "code": "rate_limit_exceeded", |
| }, |
| "content": f"You are ratelimited, wait for {cooldown} seconds", |
| } |
| ) |
|
|
|
|
| |
|
|
| def make_client() -> httpx.AsyncClient: |
| return httpx.AsyncClient( |
| timeout=httpx.Timeout(connect=30.0, read=300.0, write=60.0, pool=30.0), |
| limits=httpx.Limits(max_connections=50, max_keepalive_connections=10), |
| ) |
|
|
|
|
| def normalize_model(model: str) -> str: |
| aliases = _build_aliases() |
| return aliases.get(model, model) |
|
|
|
|
| def parse_provider_prefix(model: str): |
| """Parse 'prefix/model_id' format. Returns (prefix, model_id) or (None, model) if no prefix.""" |
| if "/" in model: |
| parts = model.split("/", 1) |
| return parts[0], parts[1] |
| return None, model |
|
|
|
|
| def _find_provider_by_prefix(prefix: str): |
| """Return (pid, prov) for a given prefix (case-insensitive), or (None, None).""" |
| if not prefix: |
| return None, None |
| pl = prefix.lower() |
| for pid, prov in _dynamic_providers.items(): |
| if (prov.get("prefix") or "").lower() == pl: |
| return pid, prov |
| |
| if (prov.get("name") or "").lower() == pl: |
| return pid, prov |
| return None, None |
|
|
|
|
| def _resolve_provider_model(prov: dict, alias: str) -> str: |
| """Resolve user-facing alias to real upstream model id using provider's model_aliases map.""" |
| aliases = prov.get("model_aliases") or {} |
| |
| |
| for upstream_id, exposed in aliases.items(): |
| if exposed == alias: |
| return upstream_id |
| |
| if alias in (prov.get("available_models") or []): |
| return alias |
| return alias |
|
|
|
|
| def _get_upstream_for_model(model: str): |
| """Find upstream URL and API key for a model, handling provider prefix routing.""" |
| prefix, alias = parse_provider_prefix(model) |
|
|
| if prefix: |
| pid, prov = _find_provider_by_prefix(prefix) |
| if prov: |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
| keys = prov.get("keys", []) |
| api_key = keys[0] if keys else "" |
| real_model = _resolve_provider_model(prov, alias) |
| return chat_url, real_model, api_key |
| |
|
|
| |
| upstreams = UPSTREAMS.get(alias, []) |
| if not upstreams: |
| upstreams = UPSTREAMS.get(model, []) |
| if not upstreams: |
| |
| return DEFAULT_UPSTREAM_URL, model, "" |
|
|
| upstream = upstreams[0] |
| url = upstream.get("url", DEFAULT_UPSTREAM_URL) |
| api_key = UPSTREAM_API_KEYS.get(url, "") |
| return url, alias, api_key |
|
|
|
|
| def parse_upstream_delta(clean: str) -> dict: |
| if not clean or clean == "[DONE]" or clean.startswith(":"): |
| return {} |
| try: |
| data = json.loads(clean) |
| if not isinstance(data, dict): |
| return {} |
| choices = data.get("choices") |
| if not isinstance(choices, list) or len(choices) == 0: |
| if isinstance(data.get("content"), str): |
| return {"content": data["content"]} |
| if isinstance(data.get("delta"), str): |
| return {"content": data["delta"]} |
| if isinstance(data.get("delta"), dict): |
| return {"content": data["delta"].get("content", "") or ""} |
| if isinstance(data.get("message"), str): |
| return {"content": data["message"]} |
| if isinstance(data.get("message"), dict): |
| return {"content": data["message"].get("content", "") or ""} |
| if data.get("finish") is True: |
| return {"finish_reason": data.get("reason") or "stop"} |
| return {} |
|
|
| delta = choices[0].get("delta", {}) or {} |
| message = choices[0].get("message", {}) or {} |
| result = {} |
|
|
| if delta.get("content"): |
| result["content"] = delta["content"] |
| elif message.get("content"): |
| result["content"] = message["content"] |
|
|
| if delta.get("tool_calls"): |
| result["tool_calls"] = delta["tool_calls"] |
| elif message.get("tool_calls"): |
| result["tool_calls"] = message["tool_calls"] |
|
|
| for field in ("reasoning_content", "thinking", "reasoning"): |
| val = delta.get(field) or message.get(field) |
| if val: |
| result["reasoning_content"] = val |
| break |
|
|
| fr = choices[0].get("finish_reason") |
| if fr: |
| result["finish_reason"] = fr |
|
|
| return result |
| except json.JSONDecodeError: |
| return {"content": clean} |
|
|
|
|
| def looks_incomplete(text: str) -> bool: |
| text = text.rstrip() |
| if not text: |
| return True |
| if text.count("```") % 2 != 0: |
| return True |
| if text.endswith("|"): |
| return True |
| last_char = text[-1] |
| if last_char in '.!?;"\n*#>`)}]-': |
| return False |
| return True |
|
|
|
|
| def merge_tool_calls(acc: list, new_calls: list) -> None: |
| for tc in new_calls: |
| idx = tc.get("index", 0) |
| while len(acc) <= idx: |
| acc.append({"id": "", "type": "function", "function": {"name": "", "arguments": ""}}) |
| entry = acc[idx] |
| if tc.get("id"): |
| entry["id"] = tc["id"] |
| if tc.get("type"): |
| entry["type"] = tc["type"] |
| fn = tc.get("function") |
| if isinstance(fn, dict): |
| if fn.get("name"): |
| entry["function"]["name"] += fn["name"] |
| if fn.get("arguments"): |
| entry["function"]["arguments"] += fn["arguments"] |
|
|
|
|
| def build_upstream_headers(url: str) -> dict: |
| headers = {"Content-Type": "application/json"} |
| api_key = UPSTREAM_API_KEYS.get(url, FALLBACK_API_KEY) |
| if api_key: |
| headers["Authorization"] = f"Bearer {api_key}" |
| return headers |
|
|
|
|
| def _extract_text_from_content(content) -> str: |
| if isinstance(content, str): |
| return content |
| if isinstance(content, list): |
| parts = [] |
| for item in content: |
| if isinstance(item, dict): |
| if item.get("type") == "text": |
| parts.append(str(item.get("text", ""))) |
| elif "text" in item: |
| parts.append(str(item.get("text", ""))) |
| elif item is not None: |
| parts.append(str(item)) |
| return "\n".join(p for p in parts if p).strip() |
| if content is None: |
| return "" |
| return str(content) |
|
|
|
|
| def _build_gateway_message(messages: list) -> str: |
| parts = [] |
| for msg in messages or []: |
| if not isinstance(msg, dict): |
| continue |
| role = msg.get("role", "user") |
| if role == "tool": |
| tool_name = msg.get("name") or "tool" |
| content = _extract_text_from_content(msg.get("content")) |
| parts.append(f"tool:{tool_name}: {content}".strip()) |
| continue |
| content = _extract_text_from_content(msg.get("content")) |
| tool_calls = msg.get("tool_calls") |
| if tool_calls: |
| try: |
| content = (content + "\n" if content else "") + json.dumps(tool_calls, ensure_ascii=False) |
| except Exception: |
| pass |
| parts.append(f"{role}: {content}".strip()) |
| return "\n".join(p for p in parts if p).strip() |
|
|
|
|
| def _build_gateway_tool_instruction(tools, tool_choice) -> str: |
| tool_lines = [] |
| for tool in tools or []: |
| if not isinstance(tool, dict) or tool.get("type") != "function": |
| continue |
| fn = tool.get("function") or {} |
| tool_lines.append({ |
| "name": fn.get("name", ""), |
| "description": fn.get("description", ""), |
| "parameters": fn.get("parameters", {}), |
| }) |
|
|
| policy = "If a tool is needed, respond ONLY with JSON matching the schema below. If no tool is needed, answer normally." |
| if tool_choice == "required": |
| policy = "You MUST respond ONLY with JSON matching the schema below and choose the best tool. Do not answer normally." |
| elif isinstance(tool_choice, dict): |
| forced = ((tool_choice.get("function") or {}).get("name") if tool_choice.get("type") == "function" else None) |
| if forced: |
| policy = f'You MUST call the tool "{forced}" and respond ONLY with JSON matching the schema below. Do not answer normally.' |
|
|
| schema = { |
| "tool_calls": [ |
| { |
| "name": "function_name", |
| "arguments": {"example": "value"} |
| } |
| ] |
| } |
|
|
| return ( |
| "TOOL MODE ENABLED.\n" |
| f"{policy}\n" |
| "Available tools:\n" |
| f"{json.dumps(tool_lines, ensure_ascii=False)}\n" |
| "When calling tools, output ONLY valid JSON with this exact shape:\n" |
| f"{json.dumps(schema, ensure_ascii=False)}\n" |
| "Do not wrap JSON in markdown. Do not add commentary before or after the JSON." |
| ) |
|
|
|
|
| def _parse_gateway_tool_calls(text: str): |
| if not text: |
| return None |
| text = text.strip() |
| candidates = [text] |
| if "```json" in text: |
| candidates.append(text.split("```json", 1)[1].split("```", 1)[0].strip()) |
| if "```" in text: |
| chunks = text.split("```") |
| for chunk in chunks: |
| chunk = chunk.strip() |
| if chunk and chunk != "json": |
| candidates.append(chunk) |
|
|
| for candidate in candidates: |
| try: |
| data = json.loads(candidate) |
| except Exception: |
| continue |
| calls = data.get("tool_calls") if isinstance(data, dict) else None |
| if not isinstance(calls, list) or not calls: |
| continue |
| out = [] |
| for i, call in enumerate(calls): |
| if not isinstance(call, dict): |
| continue |
| name = call.get("name") or ((call.get("function") or {}).get("name")) |
| arguments = call.get("arguments") or ((call.get("function") or {}).get("arguments")) or {} |
| if not name: |
| continue |
| if isinstance(arguments, str): |
| arg_str = arguments |
| else: |
| arg_str = json.dumps(arguments, ensure_ascii=False) |
| out.append({ |
| "id": f"call_{uuid.uuid4().hex[:24]}", |
| "type": "function", |
| "function": { |
| "name": name, |
| "arguments": arg_str, |
| }, |
| }) |
| if out: |
| return out |
| return None |
|
|
|
|
| class _NoopClient: |
| async def aclose(self): |
| return None |
|
|
|
|
| class _BufferedResponse: |
| def __init__(self, status_code: int, body: bytes, upstream_url: str = ""): |
| self.status_code = status_code |
| self._body = body |
| self.upstream_url = upstream_url |
|
|
| async def aread(self): |
| return self._body |
|
|
| async def aclose(self): |
| return None |
|
|
| async def aiter_lines(self): |
| text = self._body.decode(errors="replace") |
| for line in text.splitlines(): |
| yield line |
| if text and "\n" not in text: |
| return |
|
|
|
|
| async def _send_via_curl(url: str, payload: dict, api_key: str = ""): |
| headers = build_upstream_headers(url) |
| curl_args = [ |
| "curl", "-sS", "-X", "POST", url, |
| "-H", "Content-Type: application/json", |
| "-H", "Accept: application/json", |
| "--data", json.dumps(payload), |
| "-w", "\n__APIARIUM_HTTP_STATUS__:%{http_code}", |
| ] |
| |
| auth = f"Bearer {api_key}" if api_key else headers.get("Authorization") |
| if auth: |
| curl_args[7:7] = ["-H", f"Authorization: {auth}"] |
|
|
| proc = await asyncio.create_subprocess_exec( |
| *curl_args, |
| stdout=asyncio.subprocess.PIPE, |
| stderr=asyncio.subprocess.PIPE, |
| ) |
| stdout, stderr = await proc.communicate() |
| output = stdout.decode(errors="replace") |
|
|
| marker = "\n__APIARIUM_HTTP_STATUS__:" |
| if marker in output: |
| body_text, status_text = output.rsplit(marker, 1) |
| try: |
| status_code = int(status_text.strip()) |
| except ValueError: |
| status_code = 502 |
| body = body_text.encode() |
| else: |
| status_code = 502 |
| err = stderr.decode(errors="replace") if stderr else output |
| body = err.encode() |
|
|
| return _BufferedResponse(status_code, body, url), _NoopClient() |
|
|
|
|
| def get_upstream(model: str) -> dict: |
| """Returns {url, model} for the next upstream.""" |
| return _get_next_upstream(model) |
|
|
|
|
| async def send_upstream_request(body: dict, model: str): |
| """Forward entire request body to upstream, overriding model name.""" |
| url, upstream_model, api_key = _get_upstream_for_model(model) |
| logger.info("Routing model=%s β %s (upstream model: %s)", model, url, upstream_model) |
|
|
| payload = {**body, "model": upstream_model} |
|
|
| if "max_tokens" in payload: |
| try: |
| payload["max_tokens"] = max(int(payload["max_tokens"]), 1024) |
| except Exception: |
| payload["max_tokens"] = 1024 |
| else: |
| payload["max_tokens"] = 1024 |
|
|
| if url == UPSTREAM_14448_URL: |
| payload["max_output_tokens"] = payload.pop("max_tokens") |
|
|
| if url == AI_GATEWAY_URL: |
| messages = payload.get("messages") or [] |
| message_text = _build_gateway_message(messages) |
| tools = payload.get("tools") |
| tool_choice = payload.get("tool_choice") |
| if tools: |
| tool_instruction = _build_gateway_tool_instruction(tools, tool_choice) |
| message_text = f"system: {tool_instruction}\n{message_text}" if message_text else f"system: {tool_instruction}" |
| gateway_payload = { |
| "message": message_text or "user: Hello", |
| "model": upstream_model, |
| } |
| return await _send_via_curl(url, gateway_payload, api_key) |
|
|
| curl_payload = {**payload, "stream": False} |
| return await _send_via_curl(url, curl_payload, api_key) |
|
|
|
|
| |
|
|
| def landing_html() -> str: |
| return templates.load('landing.html') |
|
|
|
|
|
|
| |
|
|
| @app.get("/") |
| async def root(): |
| return HTMLResponse(landing_html()) |
|
|
|
|
| @app.get("/docs", include_in_schema=False) |
| async def custom_docs(): |
| """Branded API documentation page (replaces the default Swagger UI). |
| |
| Swagger UI is still available at ``/openapi-docs`` for users who prefer |
| the interactive playground. |
| """ |
| return HTMLResponse( |
| templates.load("docs.html"), |
| headers={"Cache-Control": "no-store, no-cache, must-revalidate"}, |
| ) |
|
|
|
|
| @app.get("/health") |
| async def health(): |
| return {"status": "ok", "timestamp": int(time.time())} |
|
|
|
|
| @app.get("/v1/models") |
| async def list_models(): |
| """List all available models. |
| |
| Built-in aliases are listed by their alias id. |
| Dynamic provider models are listed as ``{prefix}/{custom_alias}`` if the |
| provider has a prefix configured; otherwise the raw model id is used. |
| """ |
| seen = set() |
| data = [] |
| for m in _available_models: |
| if m in seen: |
| continue |
| seen.add(m) |
| data.append({"id": m, "object": "model", "owned_by": "apiarium"}) |
|
|
| |
| for pid, prov in _dynamic_providers.items(): |
| prov_name = prov.get("name", pid) |
| prefix = (prov.get("prefix") or "").strip() |
| aliases = prov.get("model_aliases") or {} |
| |
| for upstream_id, exposed in aliases.items(): |
| if not exposed: |
| continue |
| full_id = f"{prefix}/{exposed}" if prefix else exposed |
| if full_id in seen: |
| continue |
| seen.add(full_id) |
| data.append({"id": full_id, "object": "model", "owned_by": prov_name}) |
| return {"object": "list", "data": data} |
|
|
|
|
| @app.get("/v1/dashboard") |
| async def dashboard_usage(request: Request): |
| """Usage dashboard with charts and API keys. Auth handled client-side via /admin/keys.""" |
|
|
| |
| chat_logs = [] |
| try: |
| from huggingface_hub import HfApi, hf_hub_download |
| if HF_TOKEN: |
| api = HfApi(token=HF_TOKEN) |
| files = api.list_repo_files(HF_DATASET, repo_type="dataset") |
| log_files = [f for f in files if f.startswith("logs/") and f.endswith(".jsonl")] |
| for lf in sorted(log_files, reverse=True): |
| try: |
| path = hf_hub_download(HF_DATASET, lf, repo_type="dataset", token=HF_TOKEN) |
| with open(path) as f: |
| for line in f: |
| try: |
| chat_logs.append(json.loads(line)) |
| except json.JSONDecodeError: |
| pass |
| except Exception: |
| pass |
| except Exception as e: |
| logger.warning("Failed to load chat logs: %s", e) |
|
|
| |
| keys = _load_keys() |
| ip_to_key = {} |
| for key, data in keys.items(): |
| ip = data.get("ip", "unknown") |
| if ip not in ip_to_key: |
| ip_to_key[ip] = [] |
| ip_to_key[ip].append(key) |
|
|
| |
| ip_stats = {} |
| model_stats = {} |
| master_ips = set() |
|
|
| for log in chat_logs: |
| ip = log.get("ip", "unknown") |
| model = log.get("model", "unknown") |
| key_type = log.get("key_type", "regular") |
|
|
| |
| if key_type == "master" and ip and ip != "unknown": |
| master_ips.add(ip) |
|
|
| |
| if ip and ip != "unknown": |
| if ip not in ip_stats: |
| ip_stats[ip] = {"total": 0, "models": {}} |
| ip_stats[ip]["total"] += 1 |
| if model and model != "unknown": |
| if model not in ip_stats[ip]["models"]: |
| ip_stats[ip]["models"][model] = 0 |
| ip_stats[ip]["models"][model] += 1 |
|
|
| |
| if model and model != "unknown": |
| if model not in model_stats: |
| model_stats[model] = 0 |
| model_stats[model] += 1 |
|
|
| |
| ip_labels = json.dumps(list(ip_stats.keys())) |
| ip_totals = json.dumps([v["total"] for v in ip_stats.values()]) |
| model_labels = json.dumps(list(model_stats.keys())) |
| model_totals = json.dumps(list(model_stats.values())) |
|
|
| |
| ip_rows = "" |
| for ip, stats in sorted(ip_stats.items(), key=lambda x: x[1]["total"], reverse=True): |
| keys_list = ip_to_key.get(ip, ["N/A"]) |
| keys_str = " ".join([f'<code style="background:var(--bg-primary);padding:4px 8px;border-radius:6px;font-size:0.75rem;font-family:JetBrains Mono,monospace">{k}</code>' for k in keys_list[:3]]) |
| model_breakdown = " ".join([f'<span style="background:rgba(0,255,136,0.1);color:#fff;padding:4px 10px;border-radius:100px;font-size:0.75rem;font-weight:500">{m}: {c}</span>' for m, c in stats["models"].items()]) |
| god_badge = '<span style="background:linear-gradient(135deg,#fbbf24,#f59e0b);color:#000;padding:4px 10px;border-radius:100px;font-size:0.7rem;font-weight:700;margin-left:8px">π GOD</span>' if ip in master_ips else '' |
| ip_rows += f""" |
| <tr> |
| <td style="font-family:JetBrains Mono,monospace;font-size:0.9rem">{ip}{god_badge}</td> |
| <td>{keys_str}</td> |
| <td style="font-weight:700;color:#fff;font-size:1.1rem">{stats['total']}</td> |
| <td>{model_breakdown}</td> |
| </tr>""" |
|
|
| html = templates.load('dashboard.html') |
|
|
| total_requests = len(chat_logs) |
| unique_ips = len(ip_stats) |
| total_models = len(model_stats) |
| total_keys = len(keys) |
|
|
| html = html.replace("TOTAL_REQUESTS", str(total_requests)) |
| html = html.replace("TOTAL_IPS", str(unique_ips)) |
| html = html.replace("TOTAL_MODELS", str(total_models)) |
| html = html.replace("TOTAL_KEYS", str(total_keys)) |
| html = html.replace("GOD_IPS", str(len(master_ips))) |
| html = html.replace("IP_LABELS", ip_labels) |
| html = html.replace("IP_TOTALS", ip_totals) |
| html = html.replace("MODEL_LABELS", model_labels) |
| html = html.replace("MODEL_TOTALS", model_totals) |
| html = html.replace("IP_ROWS", ip_rows) |
|
|
| return HTMLResponse(html, headers={"Cache-Control": "no-store, no-cache, must-revalidate", "Pragma": "no-cache"}) |
|
|
|
|
| @app.get("/v1/upstream-models") |
| async def upstream_models(): |
| """Show what models each upstream reports.""" |
| return { |
| "configured": {k: [{"url": u["url"], "model": u["model"]} for u in v] for k, v in UPSTREAMS.items()}, |
| "upstream_reported": _upstream_models, |
| "aliases": _build_aliases(), |
| } |
|
|
|
|
| |
|
|
| class AdminKeyRequest(BaseModel): |
| ip: str |
| label: str = "" |
|
|
| @app.get("/admin/keys") |
| async def admin_list_keys(request: Request): |
| """List all admin-generated API keys. Requires MASTER_KEY auth.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| |
| keys = _load_keys() |
| admin_keys = [] |
| for key, data in keys.items(): |
| if data.get("admin_created"): |
| admin_keys.append({ |
| "key": key, |
| "ip": data.get("ip", ""), |
| "label": data.get("label", ""), |
| "created_at": data.get("created_at", ""), |
| "status": data.get("status", "active"), |
| }) |
| return {"keys": admin_keys} |
|
|
| @app.get("/admin/stats") |
| async def admin_stats(request: Request): |
| """Return aggregated stats for dashboard charts. Requires MASTER_KEY auth.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| from datetime import datetime, timedelta |
| |
| chat_logs = [] |
| try: |
| if os.path.isdir(CHAT_LOG_DIR): |
| for fname in sorted(os.listdir(CHAT_LOG_DIR)): |
| if fname.endswith(".jsonl") and not fname.endswith("_upload.jsonl"): |
| fpath = os.path.join(CHAT_LOG_DIR, fname) |
| try: |
| with open(fpath) as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| try: |
| chat_logs.append(json.loads(line)) |
| except json.JSONDecodeError: |
| pass |
| except Exception: |
| pass |
| except Exception as e: |
| logger.warning("Failed to load local chat logs for stats: %s", e) |
|
|
| |
| today = datetime.utcnow().date() |
| daily_counts = [0] * 7 |
| model_counts = {} |
| rate_limited = 0 |
| for log in chat_logs: |
| ts = log.get("timestamp", "") |
| try: |
| log_date = datetime.fromisoformat(ts.replace("Z", "+00:00")).date() if ts else None |
| except Exception: |
| log_date = None |
| if log_date: |
| days_ago = (today - log_date).days |
| if 0 <= days_ago < 7: |
| daily_counts[6 - days_ago] += 1 |
| model = log.get("model", "unknown") |
| model_counts[model] = model_counts.get(model, 0) + 1 |
| if log.get("rate_limited"): |
| rate_limited += 1 |
|
|
| |
| top_models = sorted(model_counts.items(), key=lambda x: x[1], reverse=True)[:5] |
| |
| while len(top_models) < 5: |
| top_models.append(("other", 0)) |
|
|
| return { |
| "total_requests": len(chat_logs), |
| "rate_limited": rate_limited, |
| "requests_by_day": daily_counts, |
| "models": [{"name": m[0], "count": m[1]} for m in top_models], |
| } |
|
|
| @app.post("/admin/keys") |
| async def admin_create_key(request: Request, body: AdminKeyRequest): |
| """Create a new IP-bound API key. Requires MASTER_KEY auth.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| |
| |
| try: |
| ip_obj = ipaddress.ip_address(body.ip) |
| if "/" in body.ip: |
| raise HTTPException(status_code=400, detail={"error": {"message": "CIDR notation not allowed"}}) |
| except ValueError: |
| raise HTTPException(status_code=400, detail={"error": {"message": "Invalid IP address format"}}) |
| |
| |
| keys = _load_keys() |
| for k, data in keys.items(): |
| if data.get("ip") == body.ip and data.get("admin_created"): |
| raise HTTPException(status_code=409, detail={"error": {"message": f"IP {body.ip} already has key {k}"}}) |
| |
| |
| new_key = generate_api_key() |
| keys[new_key] = { |
| "ip": body.ip, |
| "label": body.label, |
| "created_at": int(time.time()), |
| "status": "active", |
| "admin_created": True, |
| } |
| _save_keys(keys) |
| |
| return { |
| "key": new_key, |
| "ip": body.ip, |
| "label": body.label, |
| "created_at": keys[new_key]["created_at"], |
| } |
|
|
| @app.get("/admin/models") |
| async def admin_models(request: Request): |
| """Return all currently routed models with provider info. Requires MASTER_KEY auth.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| models = [] |
|
|
| |
| dynamic_models = set() |
| for pid, prov in _dynamic_providers.items(): |
| prov_name = prov.get("name", pid) |
| prefix = (prov.get("prefix") or "").strip() |
| aliases = prov.get("model_aliases") or {} |
| for upstream_id, exposed in aliases.items(): |
| if not exposed: |
| continue |
| full_id = f"{prefix}/{exposed}" if prefix else exposed |
| models.append({ |
| "name": full_id, |
| "alias": exposed, |
| "upstream": upstream_id, |
| "prefix": prefix, |
| "provider": prov_name, |
| "provider_id": pid, |
| "type": "dynamic", |
| }) |
| |
| |
| dynamic_models.add(upstream_id) |
|
|
| |
| |
| for vid in sorted(GEMINI_HUB_VIDEO_MODELS): |
| models.append({"name": vid, "provider": "GeminiHub", "type": "video-backend"}) |
| dynamic_models.add(vid) |
| for iid in sorted(GEMINI_HUB_IMAGE_MODELS): |
| models.append({"name": iid, "provider": "GeminiHub", "type": "image-backend"}) |
| dynamic_models.add(iid) |
| for tid in sorted(GEMINI_HUB_TEXT_MODELS): |
| models.append({"name": tid, "provider": "GeminiHub", "type": "text-backend"}) |
| dynamic_models.add(tid) |
| |
| def _provider_from_url(url: str) -> str: |
| """Extract provider name from upstream URL.""" |
| try: |
| from urllib.parse import urlparse |
| domain = urlparse(url).hostname or "" |
| if "openai" in domain: |
| return "OpenAI" |
| elif "anthropic" in domain: |
| return "Anthropic" |
| elif "deepseek" in domain: |
| return "DeepSeek" |
| elif "openrouter" in domain: |
| return "OpenRouter" |
| elif "google" in domain or "generativelanguage" in domain: |
| return "Google" |
| elif "mistral" in domain: |
| return "Mistral" |
| elif "together" in domain: |
| return "Together" |
| elif "groq" in domain: |
| return "Groq" |
| elif "x.ai" in domain or "xai" in domain: |
| return "xAI" |
| elif "cohere" in domain: |
| return "Cohere" |
| elif "perplexity" in domain: |
| return "Perplexity" |
| else: |
| |
| parts = domain.split(".") |
| if len(parts) >= 2: |
| return parts[-2].capitalize() |
| return domain.capitalize() if domain else "Unknown" |
| except Exception: |
| return "Unknown" |
| |
| for model_id in sorted(UPSTREAMS.keys()): |
| if model_id in dynamic_models: |
| continue |
| upstreams = UPSTREAMS[model_id] |
| |
| first_url = upstreams[0]["url"] if upstreams else "" |
| provider = _provider_from_url(first_url) if first_url else "Unknown" |
| |
| if len(upstreams) > 1: |
| all_providers = list(dict.fromkeys(_provider_from_url(u["url"]) for u in upstreams)) |
| if len(all_providers) > 1: |
| provider = ", ".join(all_providers) |
| models.append({ |
| "name": model_id, |
| "provider": provider, |
| "type": "built-in", |
| }) |
|
|
| return {"models": models} |
|
|
|
|
| |
| def _legacy_admin_models_logic(): |
| """Old logic kept for reference - no longer used.""" |
| prov_lookup = {} |
| for pid, prov in _dynamic_providers.items(): |
| base_url = prov["base_url"].rstrip("/") |
| prov_lookup[f"{base_url}/chat/completions"] = prov["name"] |
| for m in prov.get("models", []): |
| prov_lookup.setdefault("model:" + m, prov["name"]) |
|
|
| models = [] |
| for model_id in sorted(UPSTREAMS.keys()): |
| upstreams = UPSTREAMS[model_id] |
| urls = [u["url"] for u in upstreams] |
| provider = "built-in" |
| for u in urls: |
| if u in prov_lookup: |
| provider = prov_lookup[u] |
| break |
| if f"model:{model_id}" in prov_lookup: |
| provider = prov_lookup[f"model:{model_id}"] |
| break |
| models.append({ |
| "id": model_id, |
| "upstreams": len(upstreams), |
| "provider": provider, |
| }) |
|
|
| for m_id in ["gpt-image-1", "gpt-image-2"]: |
| if not any(m["id"] == m_id for m in models): |
| models.append({"id": m_id, "upstreams": 1, "provider": "image-backend"}) |
|
|
| return {"models": models, "total": len(models)} |
|
|
| @app.delete("/admin/keys/{key_id}") |
| async def admin_revoke_key(request: Request, key_id: str): |
| """Revoke an admin-created API key. Requires MASTER_KEY auth.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| |
| keys = _load_keys() |
| if key_id not in keys: |
| raise HTTPException(status_code=404, detail={"error": {"message": "Key not found"}}) |
| if not keys[key_id].get("admin_created"): |
| raise HTTPException(status_code=403, detail={"error": {"message": "Cannot revoke non-admin keys"}}) |
| |
| del keys[key_id] |
| _save_keys(keys) |
| return {"success": True} |
|
|
|
|
| |
|
|
| class ProviderImportRequest(BaseModel): |
| name: str |
| base_url: str |
| keys: list |
| prefix: str = "" |
| |
| |
| |
| model_aliases: dict = {} |
|
|
|
|
| class ProviderUpdateRequest(BaseModel): |
| name: str | None = None |
| prefix: str | None = None |
| keys: list | None = None |
| model_aliases: dict | None = None |
|
|
|
|
| MODELS_CONFIG_FILENAME = "models_config.json" |
|
|
|
|
| def _load_providers(): |
| """Load dynamic providers from the GitHub Gist (file: ``models_config.json``). |
| |
| HF Space disks are ephemeral, so the Gist is the source of truth. A local |
| JSON copy under ``app/data/providers.json`` is used as a best-effort fallback |
| when the Gist is unreachable. |
| |
| On success, every loaded provider is also registered into UPSTREAM_API_KEYS |
| via :func:`_register_provider_in_upstreams` so that routing works immediately |
| after start-up. |
| """ |
| global _dynamic_providers |
| data: dict = {} |
|
|
| |
| if _github_token and _models_gist_id: |
| try: |
| resp = httpx.get( |
| f"https://api.github.com/gists/{_models_gist_id}", |
| headers={"Authorization": f"Bearer {_github_token}"}, |
| timeout=10, |
| ) |
| if resp.status_code == 200: |
| gist = resp.json() |
| content = gist.get("files", {}).get(MODELS_CONFIG_FILENAME, {}).get("content") |
| if content: |
| try: |
| parsed = json.loads(content) |
| if isinstance(parsed, dict): |
| data = parsed |
| except json.JSONDecodeError as e: |
| logger.warning("models_config.json in gist is invalid JSON: %s", e) |
| else: |
| logger.info( |
| "Models gist %s has no %s yet (cold start)", |
| _models_gist_id, MODELS_CONFIG_FILENAME, |
| ) |
| else: |
| logger.warning( |
| "Models gist providers load failed: HTTP %s", resp.status_code |
| ) |
| except Exception as e: |
| logger.warning("Models gist providers load error: %s", e) |
|
|
| |
| if not data: |
| try: |
| with open(PROVIDERS_FILE, "r") as f: |
| parsed = json.load(f) |
| if isinstance(parsed, dict): |
| data = parsed |
| except (FileNotFoundError, json.JSONDecodeError): |
| data = {} |
|
|
| |
| for pid, prov in (data or {}).items(): |
| if not isinstance(prov, dict): |
| continue |
| prov.setdefault("prefix", "") |
| if "model_aliases" not in prov or not isinstance(prov.get("model_aliases"), dict): |
| legacy = prov.get("models") or [] |
| if isinstance(legacy, list) and legacy and isinstance(legacy[0], str): |
| prov["model_aliases"] = {m: m for m in legacy} |
| prov.setdefault("available_models", list(legacy)) |
| else: |
| prov["model_aliases"] = {} |
| prov.setdefault("available_models", list((prov.get("model_aliases") or {}).keys())) |
|
|
| _dynamic_providers = data |
|
|
| |
| |
| for prov in _dynamic_providers.values(): |
| try: |
| _register_provider_in_upstreams(prov) |
| except Exception as e: |
| logger.warning("Failed to register provider on load: %s", e) |
|
|
| logger.info("Loaded %d dynamic provider(s) from storage", len(_dynamic_providers)) |
|
|
|
|
| def _save_providers(): |
| """Persist dynamic providers to the configured GitHub Gist as |
| ``models_config.json`` and write a local backup copy. |
| |
| Silently no-ops when no Gist credentials are configured, but the local |
| backup write is always attempted so development setups keep state across |
| restarts. |
| """ |
| payload = json.dumps(_dynamic_providers, indent=2, sort_keys=True) |
|
|
| |
| try: |
| os.makedirs(os.path.dirname(PROVIDERS_FILE), exist_ok=True) |
| with open(PROVIDERS_FILE, "w") as f: |
| f.write(payload) |
| except Exception as e: |
| logger.warning("Failed to write local providers backup: %s", e) |
|
|
| |
| if not (_github_token and _models_gist_id): |
| return |
| try: |
| resp = httpx.patch( |
| f"https://api.github.com/gists/{_models_gist_id}", |
| headers={ |
| "Authorization": f"Bearer {_github_token}", |
| "Content-Type": "application/json", |
| }, |
| json={"files": {MODELS_CONFIG_FILENAME: {"content": payload}}}, |
| timeout=10, |
| ) |
| if resp.status_code == 200: |
| logger.info( |
| "Providers saved to gist %s (%s)", |
| _models_gist_id, MODELS_CONFIG_FILENAME, |
| ) |
| else: |
| logger.warning( |
| "Models gist save failed: %s %s", resp.status_code, resp.text[:200] |
| ) |
| except Exception as e: |
| logger.warning("Models gist save error: %s", e) |
|
|
|
|
| def _register_provider_in_upstreams(prov: dict): |
| """Register provider's API key for its chat endpoint. |
| |
| With prefix-based routing we no longer pollute the global UPSTREAMS dict |
| with raw model ids β routing happens via :func:`_get_upstream_for_model`. |
| We still keep the api-key mapping so that downstream proxy code that looks |
| up the key by URL keeps working. |
| """ |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
| keys = prov.get("keys", []) |
| if keys: |
| UPSTREAM_API_KEYS.setdefault(chat_url, keys[0]) |
|
|
|
|
| @app.get("/admin/providers") |
| async def admin_list_providers(request: Request): |
| """List all dynamically imported providers.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| out = [] |
| for pid, prov in _dynamic_providers.items(): |
| aliases = prov.get("model_aliases") or {} |
| prefix = (prov.get("prefix") or "").strip() |
| exposed_models = [] |
| for upstream_id, exposed in aliases.items(): |
| if not exposed: |
| continue |
| exposed_models.append({ |
| "upstream": upstream_id, |
| "alias": exposed, |
| "full_id": f"{prefix}/{exposed}" if prefix else exposed, |
| }) |
| out.append({ |
| "id": pid, |
| "name": prov.get("name", ""), |
| "prefix": prefix, |
| "base_url": prov.get("base_url", ""), |
| "key_count": len(prov.get("keys", [])), |
| "available_models": prov.get("available_models", []) or list(aliases.keys()), |
| "models": exposed_models, |
| "model_aliases": aliases, |
| "added_at": prov.get("added_at", 0), |
| }) |
| return {"providers": out} |
|
|
|
|
| @app.post("/admin/providers/preview") |
| async def admin_preview_provider(request: Request, body: ProviderImportRequest): |
| """Preview models available from an OpenAI-compatible provider (no persistence).""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| base_url = body.base_url.rstrip("/") |
| models_url = f"{base_url}/models" |
|
|
| headers = {"Content-Type": "application/json"} |
| if body.keys: |
| headers["Authorization"] = f"Bearer {body.keys[0]}" |
|
|
| try: |
| async with httpx.AsyncClient(timeout=15) as client: |
| resp = await client.get(models_url, headers=headers) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Failed to reach provider: {e}"}}) |
|
|
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Provider /models returned {resp.status_code}"}}) |
|
|
| try: |
| data = resp.json() |
| except Exception: |
| raise HTTPException(status_code=502, detail={"error": {"message": "Invalid JSON from provider"}}) |
|
|
| |
| if isinstance(data, dict) and "data" in data: |
| models = [m.get("id") for m in data["data"] if m.get("id")] |
| elif isinstance(data, list): |
| models = [m.get("id") if isinstance(m, dict) else str(m) for m in data] |
| else: |
| raise HTTPException(status_code=502, detail={"error": {"message": "Unexpected response shape from /models"}}) |
|
|
| models = [m for m in models if m] |
| return { |
| "name": body.name, |
| "base_url": base_url, |
| "model_count": len(models), |
| "models": models, |
| } |
|
|
|
|
| @app.post("/admin/providers/import") |
| async def admin_import_provider(request: Request, body: ProviderImportRequest): |
| """Import an OpenAI-compatible provider: fetch models, register routing, persist.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| if not body.keys: |
| raise HTTPException(status_code=400, detail={"error": {"message": "At least one API key required"}}) |
|
|
| base_url = body.base_url.rstrip("/") |
| models_url = f"{base_url}/models" |
|
|
| headers = {"Content-Type": "application/json", "Authorization": f"Bearer {body.keys[0]}"} |
|
|
| try: |
| async with httpx.AsyncClient(timeout=15) as client: |
| resp = await client.get(models_url, headers=headers) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Failed to reach provider: {e}"}}) |
|
|
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Provider /models returned {resp.status_code}"}}) |
|
|
| data = resp.json() |
| if isinstance(data, dict) and "data" in data: |
| models = [m.get("id") for m in data["data"] if m.get("id")] |
| elif isinstance(data, list): |
| models = [m.get("id") if isinstance(m, dict) else str(m) for m in data] |
| else: |
| models = [] |
|
|
| available = [m for m in models if m] |
|
|
| |
| |
| incoming_aliases = body.model_aliases or {} |
| cleaned_aliases: dict = {} |
| available_set = set(available) |
| for upstream_id, exposed in incoming_aliases.items(): |
| if not upstream_id or not exposed: |
| continue |
| if upstream_id not in available_set: |
| |
| continue |
| cleaned_aliases[upstream_id] = str(exposed).strip() |
|
|
| if not cleaned_aliases: |
| cleaned_aliases = {m: m for m in available} |
|
|
| |
| prefix = (body.prefix or "").strip().strip("/") |
|
|
| pid = body.name.lower().replace(" ", "-") + "-" + uuid.uuid4().hex[:6] |
| prov = { |
| "name": body.name, |
| "prefix": prefix, |
| "base_url": base_url, |
| "keys": list(body.keys), |
| "available_models": available, |
| "model_aliases": cleaned_aliases, |
| "added_at": int(time.time()), |
| } |
| _dynamic_providers[pid] = prov |
| _register_provider_in_upstreams(prov) |
| _save_providers() |
|
|
| |
| |
| global _available_models |
| _available_models = list(UPSTREAMS.keys()) + ["gpt-image-1", "gpt-image-2"] + sorted(GEMINI_HUB_MODELS) |
| logger.info( |
| "Imported provider %s (prefix=%r) with %d selected models (%d available)", |
| pid, prefix, len(cleaned_aliases), len(available), |
| ) |
|
|
| return { |
| "id": pid, |
| "name": prov["name"], |
| "prefix": prefix, |
| "base_url": base_url, |
| "key_count": len(body.keys), |
| "model_count": len(cleaned_aliases), |
| "available_count": len(available), |
| "models": [ |
| {"upstream": k, "alias": v, "full_id": f"{prefix}/{v}" if prefix else v} |
| for k, v in cleaned_aliases.items() |
| ], |
| } |
|
|
|
|
| @app.patch("/admin/providers/{provider_id}") |
| async def admin_update_provider(request: Request, provider_id: str, body: ProviderUpdateRequest): |
| """Update an existing dynamic provider: rename, change prefix, edit model selection/aliases, or replace keys.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| if provider_id not in _dynamic_providers: |
| raise HTTPException(status_code=404, detail={"error": {"message": "Provider not found"}}) |
|
|
| prov = _dynamic_providers[provider_id] |
|
|
| if body.name is not None: |
| new_name = body.name.strip() |
| if new_name: |
| prov["name"] = new_name |
|
|
| if body.prefix is not None: |
| prov["prefix"] = body.prefix.strip().strip("/") |
|
|
| if body.keys is not None: |
| new_keys = [k for k in body.keys if k] |
| prov["keys"] = new_keys |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
| if new_keys: |
| UPSTREAM_API_KEYS[chat_url] = new_keys[0] |
| elif chat_url in UPSTREAM_API_KEYS: |
| del UPSTREAM_API_KEYS[chat_url] |
|
|
| if body.model_aliases is not None: |
| available = set(prov.get("available_models") or []) |
| cleaned: dict = {} |
| for upstream_id, exposed in body.model_aliases.items(): |
| if not upstream_id or not exposed: |
| continue |
| |
| if available and upstream_id not in available: |
| continue |
| cleaned[upstream_id] = str(exposed).strip() |
| prov["model_aliases"] = cleaned |
|
|
| _save_providers() |
|
|
| aliases = prov.get("model_aliases") or {} |
| return { |
| "success": True, |
| "id": provider_id, |
| "name": prov.get("name", ""), |
| "prefix": prov.get("prefix", ""), |
| "key_count": len(prov.get("keys", [])), |
| "model_count": len(aliases), |
| "models": [ |
| {"upstream": k, "alias": v, "full_id": f"{prov.get('prefix','')}/{v}" if prov.get('prefix') else v} |
| for k, v in aliases.items() |
| ], |
| } |
|
|
|
|
| @app.delete("/admin/providers/{provider_id}") |
| async def admin_delete_provider(request: Request, provider_id: str): |
| """Remove a dynamic provider and unregister its routes.""" |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
|
|
| if provider_id not in _dynamic_providers: |
| raise HTTPException(status_code=404, detail={"error": {"message": "Provider not found"}}) |
|
|
| prov = _dynamic_providers.pop(provider_id) |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
|
|
| |
| legacy_models = list((prov.get("model_aliases") or {}).keys()) + (prov.get("available_models") or []) |
| for model_id in legacy_models: |
| if model_id in UPSTREAMS: |
| UPSTREAMS[model_id] = [u for u in UPSTREAMS[model_id] if u.get("url") != chat_url] |
| if not UPSTREAMS[model_id]: |
| del UPSTREAMS[model_id] |
|
|
| |
| sharing = any( |
| p.get("base_url", "").rstrip("/") == base_url |
| for p in _dynamic_providers.values() |
| ) |
| if not sharing and chat_url in UPSTREAM_API_KEYS: |
| del UPSTREAM_API_KEYS[chat_url] |
|
|
| _save_providers() |
| global _available_models |
| _available_models = list(UPSTREAMS.keys()) + ["gpt-image-1", "gpt-image-2"] + sorted(GEMINI_HUB_MODELS) |
| return {"success": True} |
|
|
|
|
| |
|
|
| def _normalize_balance(data: dict, provider_name: str) -> dict: |
| """Normalize different balance API response formats into a common shape.""" |
| result = {"provider": provider_name, "currency": "USD", "balance": None, "used": None, "total": None, "raw": data} |
| |
| |
| if "balance_infos" in data and isinstance(data["balance_infos"], list) and data["balance_infos"]: |
| info = data["balance_infos"][0] |
| result["currency"] = info.get("currency", "USD") |
| result["total"] = float(info.get("total_balance", 0)) + float(info.get("granted_balance", 0)) |
| result["used"] = float(info.get("used_balance", 0)) if "used_balance" in info else None |
| if result["total"] is not None and result["used"] is not None: |
| result["balance"] = result["total"] - result["used"] |
| else: |
| result["balance"] = result["total"] |
| return result |
| |
| |
| if "data" in data and isinstance(data["data"], dict): |
| inner = data["data"] |
| result["total"] = float(inner.get("total_credits", 0)) |
| result["used"] = float(inner.get("total_usage", 0)) |
| if result["total"] is not None and result["used"] is not None: |
| result["balance"] = result["total"] - result["used"] |
| return result |
| |
| |
| for k in ("balance", "available", "credit", "remaining"): |
| if k in data and data[k] is not None: |
| result["balance"] = float(data[k]) |
| break |
| for k in ("total", "total_credits", "limit"): |
| if k in data and data[k] is not None: |
| result["total"] = float(data[k]) |
| break |
| for k in ("used", "usage", "spent"): |
| if k in data and data[k] is not None: |
| result["used"] = float(data[k]) |
| break |
| |
| return result |
|
|
|
|
| async def _check_balance_for_provider(provider_id: str) -> dict: |
| """Check balance for a single dynamic provider.""" |
| prov = _dynamic_providers.get(provider_id) |
| if not prov: |
| return {"provider_id": provider_id, "error": "Provider not found", "available": False} |
| |
| prov_name = (prov.get("name") or "").lower() |
| keys = prov.get("keys", []) |
| if not keys: |
| return {"provider_id": provider_id, "provider": prov_name, "error": "No API keys configured", "available": False} |
| |
| |
| balance_url = prov.get("balance_endpoint_url") or BALANCE_ENDPOINTS.get(prov_name) |
| if not balance_url: |
| return {"provider_id": provider_id, "provider": prov_name, "error": "Balance endpoint not configured for this provider", "available": False} |
| |
| try: |
| async with make_client() as client: |
| resp = await client.get( |
| balance_url, |
| headers={"Authorization": f"Bearer {keys[0]}"}, |
| timeout=10.0, |
| ) |
| if resp.status_code != 200: |
| return { |
| "provider_id": provider_id, |
| "provider": prov_name, |
| "error": f"HTTP {resp.status_code}", |
| "available": False, |
| } |
| data = resp.json() |
| normalized = _normalize_balance(data, prov_name) |
| normalized["provider_id"] = provider_id |
| normalized["available"] = True |
| normalized["checked_at"] = int(time.time()) |
| return normalized |
| except Exception as e: |
| return {"provider_id": provider_id, "provider": prov_name, "error": str(e), "available": False} |
|
|
|
|
| @app.get("/admin/balance/check/{provider_id}") |
| async def admin_balance_check(request: Request, provider_id: str): |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| result = await _check_balance_for_provider(provider_id) |
| return result |
|
|
|
|
| @app.get("/admin/balance/check-all") |
| async def admin_balance_check_all(request: Request): |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| results = [] |
| for pid in list(_dynamic_providers.keys()): |
| r = await _check_balance_for_provider(pid) |
| results.append(r) |
| return {"results": results, "checked_at": int(time.time())} |
|
|
|
|
| |
|
|
| class AppendKeysRequest(BaseModel): |
| keys: list = [] |
|
|
|
|
| @app.post("/admin/providers/{provider_id}/keys") |
| async def admin_append_provider_keys(request: Request, provider_id: str, body: AppendKeysRequest): |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| |
| if provider_id not in _dynamic_providers: |
| raise HTTPException(status_code=404, detail={"error": {"message": "Provider not found"}}) |
| |
| prov = _dynamic_providers[provider_id] |
| existing = prov.get("keys", []) |
| new_keys = [k for k in body.keys if k and k not in existing] |
| if not new_keys: |
| return {"success": True, "added": 0, "total": len(existing)} |
| |
| prov["keys"] = existing + new_keys |
| |
| |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
| UPSTREAM_API_KEYS[chat_url] = prov["keys"][0] |
| |
| _save_providers() |
| return {"success": True, "added": len(new_keys), "total": len(prov["keys"])} |
|
|
|
|
| @app.delete("/admin/providers/{provider_id}/keys/{idx}") |
| async def admin_remove_provider_key(request: Request, provider_id: str, idx: int): |
| auth = request.headers.get("authorization", "") |
| if not auth.startswith("Bearer ") or auth.replace("Bearer ", "") != MASTER_KEY: |
| raise HTTPException(status_code=401, detail={"error": {"message": "Invalid master key"}}) |
| |
| if provider_id not in _dynamic_providers: |
| raise HTTPException(status_code=404, detail={"error": {"message": "Provider not found"}}) |
| |
| prov = _dynamic_providers[provider_id] |
| keys = prov.get("keys", []) |
| if idx < 0 or idx >= len(keys): |
| raise HTTPException(status_code=400, detail={"error": {"message": "Invalid key index"}}) |
| |
| keys.pop(idx) |
| prov["keys"] = keys |
| |
| base_url = prov["base_url"].rstrip("/") |
| chat_url = f"{base_url}/chat/completions" |
| if keys: |
| UPSTREAM_API_KEYS[chat_url] = keys[0] |
| elif chat_url in UPSTREAM_API_KEYS: |
| del UPSTREAM_API_KEYS[chat_url] |
| |
| _save_providers() |
| return {"success": True, "remaining": len(keys)} |
|
|
|
|
| |
|
|
| class TurnstileVerifyRequest(BaseModel): |
| token: str |
|
|
|
|
| @app.post("/admin/verify-turnstile") |
| async def admin_verify_turnstile(body: TurnstileVerifyRequest): |
| if not TURNSTILE_SECRET_KEY: |
| return {"success": True, "note": "Turnstile not configured, bypassed"} |
| try: |
| async with make_client() as client: |
| resp = await client.post( |
| "https://challenges.cloudflare.com/turnstile/v0/siteverify", |
| data={"response": body.token, "secret": TURNSTILE_SECRET_KEY}, |
| timeout=10.0, |
| ) |
| result = resp.json() |
| return {"success": bool(result.get("success")), "raw": result} |
| except Exception as e: |
| return {"success": False, "error": str(e)} |
|
|
|
|
| |
|
|
| _ADMIN_NAV = [ |
| ("/admin/keys", "API Keys", "π"), |
| ("/admin/providers", "Providers", "π"), |
| ("/admin/balance", "Balance", "π°"), |
| ("/admin/models", "Models", "π€"), |
| ("/admin/settings", "Settings", "βοΈ"), |
| ] |
|
|
|
|
| def _admin_layout(title: str, active: str, body_html: str, turnstile_site_key: str = "") -> str: |
| """Render shared admin layout with sidebar navigation.""" |
| nav_items = "" |
| for href, label, icon in _ADMIN_NAV: |
| active_cls = "bg-amber-500/10 text-amber-400 border-amber-500/30" if href == active else "text-slate-400 border-transparent hover:text-slate-200 hover:border-slate-700" |
| nav_items += f'<a href="{href}" class="flex items-center gap-3 px-4 py-2.5 rounded-lg border text-sm font-medium transition {active_cls}"><span class="text-base">{icon}</span>{label}</a>\n' |
| |
| return f"""<!DOCTYPE html> |
| <html lang="en" class="antialiased"> |
| <head> |
| <meta charset="UTF-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| <title>APIarium | {title}</title> |
| <link rel="preconnect" href="https://fonts.googleapis.com"> |
| <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin> |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500;600&display=swap" rel="stylesheet"> |
| <script src="https://cdn.tailwindcss.com"></script> |
| <script src="https://challenges.cloudflare.com/turnstile/v0/api.js" async defer></script> |
| <script> |
| tailwind.config = {{ |
| theme: {{ |
| extend: {{ |
| fontFamily: {{ |
| sans: ['Inter', 'sans-serif'], |
| mono: ['JetBrains Mono', 'monospace'], |
| }}, |
| }} |
| }} |
| }} |
| const MASTER_KEY = localStorage.getItem('apiarium_master_key') || ''; |
| const TURNSTILE_SITE_KEY = '{turnstile_site_key}'; |
| function authHeaders() {{ return {{ 'Authorization': 'Bearer ' + MASTER_KEY, 'Content-Type': 'application/json' }}; }} |
| function checkAuth() {{ |
| if (!MASTER_KEY) {{ window.location.href = '/admin/login'; return false; }} |
| return true; |
| }} |
| </script> |
| </head> |
| <body class="bg-slate-950 text-slate-100 min-h-screen"> |
| <div class="flex min-h-screen"> |
| <aside class="w-64 bg-slate-900/50 border-r border-white/5 p-6 flex flex-col gap-2"> |
| <div class="mb-8"> |
| <h1 class="text-xl font-bold flex items-center gap-2">π APIarium</h1> |
| <p class="text-xs text-slate-500 mt-1">Admin Console</p> |
| </div> |
| <nav class="flex flex-col gap-1"> |
| {nav_items} |
| </nav> |
| <div class="mt-auto pt-6 border-t border-white/5"> |
| <a href="/" class="text-sm text-slate-500 hover:text-slate-300">β Back to site</a> |
| <button onclick="localStorage.removeItem('apiarium_master_key'); window.location.href='/admin/login'" class="block mt-2 text-sm text-red-400/70 hover:text-red-400">Logout</button> |
| </div> |
| </aside> |
| <main class="flex-1 p-8 overflow-auto"> |
| <h2 class="text-2xl font-bold mb-6">{title}</h2> |
| {body_html} |
| </main> |
| </div> |
| </body> |
| </html>""" |
|
|
|
|
| @app.get("/admin") |
| async def admin_root(): |
| from fastapi.responses import RedirectResponse |
| return RedirectResponse(url="/admin/keys") |
|
|
|
|
| @app.get("/admin/login", response_class=HTMLResponse) |
| async def admin_login_page(): |
| html = """<!DOCTYPE html> |
| <html lang="en" class="antialiased"> |
| <head> |
| <meta charset="UTF-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| <title>APIarium | Admin Login</title> |
| <meta name="description" content="APIarium admin console β master key sign-in."> |
| <link rel="preconnect" href="https://fonts.googleapis.com"> |
| <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin> |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500;600&display=swap" rel="stylesheet"> |
| <script src="https://cdn.tailwindcss.com"></script> |
| <script src="https://challenges.cloudflare.com/turnstile/v0/api.js" async defer></script> |
| <style> |
| body { background-color: #000; color: #ededed; font-family: 'Inter', sans-serif; } |
| .bg-grid-pattern { |
| background-image: linear-gradient(to right, #ffffff05 1px, transparent 1px), |
| linear-gradient(to bottom, #ffffff05 1px, transparent 1px); |
| background-size: 40px 40px; |
| mask-image: radial-gradient(circle at center, black, transparent 80%); |
| -webkit-mask-image: radial-gradient(circle at center, black, transparent 80%); |
| } |
| </style> |
| </head> |
| <body class="min-h-screen flex items-center justify-center p-4 relative"> |
| <div class="absolute inset-0 bg-grid-pattern pointer-events-none z-0"></div> |
| |
| <div class="w-full max-w-md relative z-10"> |
| <div class="text-center mb-8"> |
| <span class="text-5xl" aria-label="APIarium">π</span> |
| <h1 class="text-2xl font-semibold tracking-tight mt-3">APIarium Admin</h1> |
| <p class="text-xs font-mono text-gray-500 mt-2 px-2.5 py-0.5 rounded-full border border-gray-800 inline-block">v0.4.0 Β· Admin Console</p> |
| </div> |
| |
| <div class="bg-gray-900 border border-gray-800 rounded-2xl p-8 shadow-2xl"> |
| <label class="block text-[10px] font-mono text-gray-500 uppercase tracking-wider mb-2">Master Key</label> |
| <input type="password" id="master-key" placeholder="apiarium_master_..." class="w-full bg-black border border-gray-800 rounded-lg px-4 py-3 text-sm font-mono text-gray-100 focus:outline-none focus:border-amber-500/60 mb-4 placeholder:text-gray-700 transition-colors"> |
| <div id="turnstile-widget" class="flex justify-center mb-4"></div> |
| <div id="login-error" class="text-red-400 text-sm bg-red-500/10 border border-red-500/30 rounded-lg px-3 py-2 mb-4 hidden font-mono"></div> |
| <button id="login-btn" class="w-full bg-white hover:bg-gray-100 active:scale-[0.99] text-black font-semibold px-5 py-3 rounded-lg transition-all text-sm flex items-center justify-center gap-2"> |
| <span>Sign in</span> |
| <svg class="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M13 7l5 5m0 0l-5 5m5-5H6"/></svg> |
| </button> |
| <div class="mt-6 pt-5 border-t border-gray-800 text-center"> |
| <a href="/" class="text-[10px] font-mono text-gray-500 hover:text-gray-300 transition-colors">β Back to homepage</a> |
| </div> |
| </div> |
| </div> |
| <script> |
| let turnstileToken = null; |
| const TURNSTILE_SITE_KEY = '__TURNSTILE_SITE_KEY__'; |
| |
| if (TURNSTILE_SITE_KEY) { |
| turnstile.render('#turnstile-widget', { |
| sitekey: TURNSTILE_SITE_KEY, |
| callback: (token) => { turnstileToken = token; } |
| }); |
| } |
| |
| document.getElementById('login-btn').addEventListener('click', async () => { |
| const errEl = document.getElementById('login-error'); |
| errEl.classList.add('hidden'); |
| const key = document.getElementById('master-key').value.trim(); |
| if (!key) { |
| errEl.textContent = 'Master key required'; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| if (TURNSTILE_SITE_KEY && !turnstileToken) { |
| errEl.textContent = 'Please complete the Turnstile challenge'; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| |
| // Verify turnstile if configured |
| if (TURNSTILE_SITE_KEY) { |
| try { |
| const v = await fetch('/admin/verify-turnstile', { |
| method: 'POST', |
| headers: {'Content-Type': 'application/json'}, |
| body: JSON.stringify({token: turnstileToken}) |
| }).then(r => r.json()); |
| if (!v.success) { |
| errEl.textContent = 'Turnstile verification failed'; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| } catch (e) { |
| errEl.textContent = 'Turnstile verification error: ' + e.message; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| } |
| |
| // Test master key against /admin/keys |
| try { |
| const r = await fetch('/admin/keys', {headers: {'Authorization': 'Bearer ' + key}}); |
| if (r.status === 401) { |
| errEl.textContent = 'Invalid master key'; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| if (!r.ok) throw new Error('HTTP ' + r.status); |
| } catch (e) { |
| errEl.textContent = 'Connection error: ' + e.message; |
| errEl.classList.remove('hidden'); |
| return; |
| } |
| |
| localStorage.setItem('apiarium_master_key', key); |
| window.location.href = '/admin/keys'; |
| }); |
| </script> |
| </body> |
| </html>""" |
| html = html.replace("__TURNSTILE_SITE_KEY__", TURNSTILE_SITE_KEY or "") |
| return HTMLResponse(content=html) |
|
|
|
|
| @app.get("/admin/keys", response_class=HTMLResponse) |
| async def admin_keys_page(): |
| body = """ |
| <div id="keys-list" class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <div class="text-slate-500 text-sm">Loading...</div> |
| </div> |
| <script> |
| if (!checkAuth()) throw new Error('no auth'); |
| (async () => { |
| const r = await fetch('/admin/keys', {headers: authHeaders()}).then(r => r.json()); |
| const el = document.getElementById('keys-list'); |
| if (!r.keys || r.keys.length === 0) { |
| el.innerHTML = '<div class="text-slate-500 text-sm">No admin-created keys yet.</div>'; |
| return; |
| } |
| let html = '<table class="w-full text-sm"><thead><tr class="text-left text-slate-400 text-xs uppercase"><th class="py-2">Key</th><th class="py-2">IP</th><th class="py-2">Label</th><th class="py-2">Created</th></tr></thead><tbody>'; |
| r.keys.forEach(k => { |
| html += `<tr class="border-t border-white/5"><td class="py-2 font-mono text-xs">${k.key.slice(0,16)}...</td><td class="py-2">${k.ip||''}</td><td class="py-2">${k.label||''}</td><td class="py-2 text-slate-500">${k.created_at||''}</td></tr>`; |
| }); |
| html += '</tbody></table>'; |
| el.innerHTML = html; |
| })(); |
| </script> |
| """ |
| return HTMLResponse(content=_admin_layout("API Keys", "/admin/keys", body, TURNSTILE_SITE_KEY)) |
|
|
|
|
| @app.get("/admin/providers", response_class=HTMLResponse) |
| async def admin_providers_page(): |
| body = """ |
| <div class="grid md:grid-cols-2 gap-6"> |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <h3 class="text-lg font-semibold mb-4">Add New Provider</h3> |
| <form id="add-provider-form" class="flex flex-col gap-3"> |
| <input type="text" name="name" placeholder="Provider name (e.g. deepseek)" required class="bg-black/30 border border-white/10 rounded-lg px-4 py-2 text-sm focus:outline-none focus:border-amber-500/50"> |
| <input type="text" name="base_url" placeholder="Base URL (e.g. https://api.deepseek.com/v1)" required class="bg-black/30 border border-white/10 rounded-lg px-4 py-2 text-sm focus:outline-none focus:border-amber-500/50"> |
| <input type="text" name="balance_endpoint_url" placeholder="Balance endpoint URL (optional override)" class="bg-black/30 border border-white/10 rounded-lg px-4 py-2 text-sm focus:outline-none focus:border-amber-500/50"> |
| <textarea name="models" placeholder="Models (one per line)" rows="2" class="bg-black/30 border border-white/10 rounded-lg px-4 py-2 text-sm font-mono focus:outline-none focus:border-amber-500/50 resize-none"></textarea> |
| <textarea name="keys" placeholder="API keys (one per line)" rows="2" class="bg-black/30 border border-white/10 rounded-lg px-4 py-2 text-sm font-mono focus:outline-none focus:border-amber-500/50 resize-none"></textarea> |
| <button type="submit" class="bg-amber-500 hover:bg-amber-400 text-black font-semibold py-2 rounded-lg text-sm">Add Provider</button> |
| <div id="add-result" class="text-sm"></div> |
| </form> |
| </div> |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <h3 class="text-lg font-semibold mb-4">Registered Providers</h3> |
| <div id="providers-list">Loading...</div> |
| </div> |
| </div> |
| <script> |
| if (!checkAuth()) throw new Error('no auth'); |
| async function loadProviders() { |
| const r = await fetch('/admin/providers', {headers: authHeaders()}).then(r => r.json()); |
| const el = document.getElementById('providers-list'); |
| if (!r.providers || r.providers.length === 0) { |
| el.innerHTML = '<div class="text-slate-500 text-sm">No dynamic providers yet.</div>'; |
| return; |
| } |
| let html = '<div class="flex flex-col gap-3">'; |
| r.providers.forEach(p => { |
| const keysCount = (p.keys || []).length; |
| const modelsCount = (p.models || []).length; |
| html += `<div class="bg-black/30 border border-white/5 rounded-lg p-4"> |
| <div class="flex items-center justify-between mb-2"> |
| <div class="font-semibold">${p.name}</div> |
| <button onclick="deleteProvider('${p.id}')" class="text-red-400/70 hover:text-red-400 text-xs">Delete</button> |
| </div> |
| <div class="text-xs text-slate-500 font-mono truncate mb-2">${p.base_url}</div> |
| <div class="flex gap-3 text-xs text-slate-400"> |
| <span>π ${keysCount} keys</span> |
| <span>π€ ${modelsCount} models</span> |
| </div> |
| <div class="mt-2 flex gap-2"> |
| <button onclick="appendKeyPrompt('${p.id}')" class="text-xs bg-amber-500/20 hover:bg-amber-500/30 text-amber-300 px-2 py-1 rounded">+ Add Key</button> |
| <button onclick="checkBalance('${p.id}')" class="text-xs bg-blue-500/20 hover:bg-blue-500/30 text-blue-300 px-2 py-1 rounded">π° Balance</button> |
| </div> |
| <div id="balance-${p.id}" class="mt-2 text-xs hidden"></div> |
| </div>`; |
| }); |
| html += '</div>'; |
| el.innerHTML = html; |
| } |
| async function deleteProvider(id) { |
| if (!confirm('Delete this provider?')) return; |
| await fetch('/admin/providers/' + id, {method: 'DELETE', headers: authHeaders()}); |
| loadProviders(); |
| } |
| async function appendKeyPrompt(id) { |
| const k = prompt('Enter new API key:'); |
| if (!k) return; |
| const r = await fetch('/admin/providers/' + id + '/keys', { |
| method: 'POST', headers: authHeaders(), |
| body: JSON.stringify({keys: [k]}) |
| }).then(r => r.json()); |
| alert(r.success ? `Added ${r.added} key(s)` : 'Failed: ' + JSON.stringify(r)); |
| loadProviders(); |
| } |
| async function checkBalance(id) { |
| const el = document.getElementById('balance-' + id); |
| el.classList.remove('hidden'); |
| el.innerHTML = '<span class="text-slate-400">Checking...</span>'; |
| const r = await fetch('/admin/balance/check/' + id, {headers: authHeaders()}).then(r => r.json()); |
| if (!r.available) { |
| el.innerHTML = `<span class="text-red-400">${r.error || 'Unavailable'}</span>`; |
| return; |
| } |
| el.innerHTML = `<span class="text-green-400">Balance: ${r.balance ?? '?'} ${r.currency || ''} (total: ${r.total ?? '?'}, used: ${r.used ?? '?'})</span>`; |
| } |
| document.getElementById('add-provider-form').addEventListener('submit', async (e) => { |
| e.preventDefault(); |
| const f = e.target; |
| const payload = { |
| name: f.name.value.trim(), |
| base_url: f.base_url.value.trim(), |
| balance_endpoint_url: f.balance_endpoint_url.value.trim() || null, |
| models: f.models.value.split('\n').map(s => s.trim()).filter(Boolean), |
| keys: f.keys.value.split('\n').map(s => s.trim()).filter(Boolean), |
| }; |
| const r = await fetch('/admin/providers/import', { |
| method: 'POST', headers: authHeaders(), |
| body: JSON.stringify(payload) |
| }); |
| const resEl = document.getElementById('add-result'); |
| if (r.ok) { |
| resEl.innerHTML = '<span class="text-green-400">β Added</span>'; |
| f.reset(); |
| loadProviders(); |
| } else { |
| const j = await r.json().catch(() => ({})); |
| resEl.innerHTML = '<span class="text-red-400">Error: ' + (j.detail?.error?.message || r.status) + '</span>'; |
| } |
| }); |
| loadProviders(); |
| </script> |
| """ |
| return HTMLResponse(content=_admin_layout("Providers", "/admin/providers", body, TURNSTILE_SITE_KEY)) |
|
|
|
|
| @app.get("/admin/balance", response_class=HTMLResponse) |
| async def admin_balance_page(): |
| body = """ |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <div class="flex items-center justify-between mb-4"> |
| <h3 class="text-lg font-semibold">Balance Overview</h3> |
| <button id="refresh-all" class="bg-amber-500 hover:bg-amber-400 text-black font-semibold px-4 py-2 rounded-lg text-sm">Refresh All</button> |
| </div> |
| <div id="balance-results"><div class="text-slate-500 text-sm">Click "Refresh All" to check balances.</div></div> |
| </div> |
| <script> |
| if (!checkAuth()) throw new Error('no auth'); |
| document.getElementById('refresh-all').addEventListener('click', async () => { |
| const el = document.getElementById('balance-results'); |
| el.innerHTML = '<div class="text-slate-400 text-sm">Checking all providers...</div>'; |
| const r = await fetch('/admin/balance/check-all', {headers: authHeaders()}).then(r => r.json()); |
| if (!r.results || r.results.length === 0) { |
| el.innerHTML = '<div class="text-slate-500 text-sm">No dynamic providers configured.</div>'; |
| return; |
| } |
| let html = '<table class="w-full text-sm"><thead><tr class="text-left text-slate-400 text-xs uppercase"><th class="py-2">Provider</th><th class="py-2">Balance</th><th class="py-2">Total</th><th class="py-2">Used</th><th class="py-2">Status</th></tr></thead><tbody>'; |
| r.results.forEach(x => { |
| const status = x.available |
| ? '<span class="text-green-400">β OK</span>' |
| : '<span class="text-red-400">' + (x.error || 'Error') + '</span>'; |
| html += `<tr class="border-t border-white/5"> |
| <td class="py-2">${x.provider || x.provider_id}</td> |
| <td class="py-2 font-mono">${x.balance ?? 'β'} ${x.currency || ''}</td> |
| <td class="py-2 text-slate-400 font-mono">${x.total ?? 'β'}</td> |
| <td class="py-2 text-slate-400 font-mono">${x.used ?? 'β'}</td> |
| <td class="py-2">${status}</td> |
| </tr>`; |
| }); |
| html += '</tbody></table>'; |
| el.innerHTML = html; |
| }); |
| </script> |
| """ |
| return HTMLResponse(content=_admin_layout("Balance", "/admin/balance", body, TURNSTILE_SITE_KEY)) |
|
|
|
|
| @app.get("/admin/models", response_class=HTMLResponse) |
| async def admin_models_page(): |
| body = """ |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <div class="flex items-center justify-between mb-4"> |
| <h3 class="text-lg font-semibold">Available Models</h3> |
| <span class="text-xs text-slate-500">Models are listed with provider prefixes where applicable</span> |
| </div> |
| <div id="models-list">Loading...</div> |
| </div> |
| <script> |
| if (!checkAuth()) throw new Error('no auth'); |
| (async () => { |
| const r = await fetch('/admin/models', {headers: authHeaders()}).then(r => r.json()); |
| const el = document.getElementById('models-list'); |
| const models = r.models || []; |
| if (models.length === 0) { |
| el.innerHTML = '<div class="text-slate-500 text-sm">No models registered.</div>'; |
| return; |
| } |
| let html = '<div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-2">'; |
| models.forEach(m => { |
| html += `<div class="bg-black/30 border border-white/5 rounded-lg px-3 py-2 text-xs font-mono">${m}</div>`; |
| }); |
| html += '</div>'; |
| el.innerHTML = html; |
| })(); |
| </script> |
| """ |
| return HTMLResponse(content=_admin_layout("Models", "/admin/models", body, TURNSTILE_SITE_KEY)) |
|
|
|
|
| @app.get("/admin/settings", response_class=HTMLResponse) |
| async def admin_settings_page(): |
| body = """ |
| <div class="grid md:grid-cols-2 gap-6"> |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <h3 class="text-lg font-semibold mb-4">System Info</h3> |
| <div class="space-y-2 text-sm"> |
| <div class="flex justify-between"><span class="text-slate-400">Turnstile:</span><span id="turnstile-status" class="font-mono">...</span></div> |
| <div class="flex justify-between"><span class="text-slate-400">GitHub Gist:</span><span id="gist-status" class="font-mono">...</span></div> |
| <div class="flex justify-between"><span class="text-slate-400">HF Dataset:</span><span class="font-mono">__HF_DATASET__</span></div> |
| </div> |
| </div> |
| <div class="bg-slate-900/30 border border-white/5 rounded-xl p-6"> |
| <h3 class="text-lg font-semibold mb-4">Blocked IPs</h3> |
| <div id="blocked-ips" class="space-y-1 text-sm font-mono text-slate-400"> |
| __BLOCKED_IPS_HTML__ |
| </div> |
| </div> |
| </div> |
| <script> |
| if (!checkAuth()) throw new Error('no auth'); |
| document.getElementById('turnstile-status').textContent = '__TURNSTILE_STATUS__'; |
| document.getElementById('gist-status').textContent = '__GIST_STATUS__'; |
| </script> |
| """ |
| body = body.replace("__HF_DATASET__", HF_DATASET or "not configured") |
| body = body.replace("__BLOCKED_IPS_HTML__", "<br>".join(BLOCKED_IPS) if BLOCKED_IPS else "None") |
| body = body.replace("__TURNSTILE_STATUS__", "β Configured" if TURNSTILE_SECRET_KEY else "β Disabled") |
| body = body.replace("__GIST_STATUS__", "β Configured" if (GITHUB_TOKEN and GIST_ID) else "β Disabled") |
| return HTMLResponse(content=_admin_layout("Settings", "/admin/settings", body, TURNSTILE_SITE_KEY)) |
|
|
|
|
| |
|
|
| def _gemini_hub_extract_prompt(messages: list) -> str: |
| """Flatten OpenAI chat messages into a single prompt string for Gemini Hub. |
| |
| We concatenate every user/system turn (skipping tool noise) and return |
| the joined plain text. Multi-modal ``content`` arrays are reduced to their |
| textual parts; image parts are dropped because the dashboard exposes only |
| text-to-* generation today. |
| """ |
| parts: list[str] = [] |
| for msg in messages: |
| if not isinstance(msg, dict): |
| continue |
| role = msg.get("role") |
| if role not in ("system", "user"): |
| continue |
| content = msg.get("content") |
| if isinstance(content, str): |
| parts.append(content.strip()) |
| elif isinstance(content, list): |
| for piece in content: |
| if isinstance(piece, dict) and piece.get("type") in ("text", "input_text"): |
| txt = piece.get("text") or piece.get("input_text") or "" |
| if txt: |
| parts.append(str(txt).strip()) |
| return "\n\n".join(p for p in parts if p) |
|
|
|
|
| def _gemini_hub_chat_envelope(model: str, content: str, *, usage_prompt: int = 0) -> dict: |
| """Wrap ``content`` in a non-streaming OpenAI chat.completion response.""" |
| return { |
| "id": f"chatcmpl-{uuid.uuid4().hex}", |
| "object": "chat.completion", |
| "created": int(time.time()), |
| "model": model, |
| "choices": [ |
| { |
| "index": 0, |
| "message": {"role": "assistant", "content": content}, |
| "finish_reason": "stop", |
| } |
| ], |
| "usage": { |
| "prompt_tokens": usage_prompt, |
| "completion_tokens": len(content) // 4, |
| "total_tokens": usage_prompt + len(content) // 4, |
| }, |
| } |
|
|
|
|
| async def _gemini_hub_chat(model: str, messages: list, body: dict) -> dict: |
| """Route a chat completion to the Gemini Hub backend based on ``model``. |
| |
| Returns an OpenAI-compatible ``chat.completion`` dict. Raises HTTPException |
| on upstream failures so the outer ``chat_completions`` handler reports them |
| consistently with the rest of the stack. |
| """ |
| prompt = _gemini_hub_extract_prompt(messages) |
| if not prompt: |
| raise HTTPException( |
| status_code=400, |
| detail={"error": {"message": "No textual prompt found in messages"}}, |
| ) |
|
|
| |
| aspect_ratio = body.get("aspect_ratio") or body.get("size") |
| duration = body.get("duration_seconds") or body.get("duration") |
| resolution = body.get("resolution") |
|
|
| if model in GEMINI_HUB_VIDEO_MODELS: |
| payload: dict = {"prompt": prompt} |
| if aspect_ratio in ("16:9", "9:16"): |
| payload["aspect_ratio"] = aspect_ratio |
| if duration and str(duration) in ("4", "6", "8"): |
| payload["duration_seconds"] = str(duration) |
| if resolution in ("720p", "1080p"): |
| payload["resolution"] = resolution |
| |
| if isinstance(body.get("image_url"), str): |
| payload["image_url"] = body["image_url"] |
|
|
| url = f"{GEMINI_HUB_URL}/call/generar_video" |
| logger.info("GeminiHub video request model=%s prompt=%s", model, prompt[:80]) |
| async with httpx.AsyncClient(timeout=600.0) as client: |
| try: |
| resp = await client.post(url, json=payload) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub unreachable: {e}"}}) |
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub {resp.status_code}: {resp.text[:300]}"}}) |
| data = resp.json() or {} |
| if data.get("error"): |
| raise HTTPException(status_code=400, detail={"error": {"message": str(data["error"])}}) |
|
|
| video_url = data.get("video_url") or data.get("url") |
| if not video_url: |
| raise HTTPException(status_code=502, detail={"error": {"message": "No video URL returned"}}) |
| |
| |
| return _gemini_hub_chat_envelope(model, video_url, usage_prompt=len(prompt) // 4) |
|
|
| if model in GEMINI_HUB_IMAGE_MODELS: |
| payload = {"prompt": prompt, "model": "pro" if model.endswith("-pro") else "flash"} |
| if aspect_ratio in ("1:1", "16:9", "9:16", "4:3", "3:4"): |
| payload["aspect_ratio"] = aspect_ratio |
| url = f"{GEMINI_HUB_URL}/call/generar_imagen" |
| logger.info("GeminiHub image request model=%s prompt=%s", model, prompt[:80]) |
| async with httpx.AsyncClient(timeout=180.0) as client: |
| try: |
| resp = await client.post(url, json=payload) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub unreachable: {e}"}}) |
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub {resp.status_code}: {resp.text[:300]}"}}) |
| data = resp.json() or {} |
| if data.get("error"): |
| raise HTTPException(status_code=400, detail={"error": {"message": str(data["error"])}}) |
| image_url = data.get("image_url") or data.get("url") |
| if not image_url: |
| raise HTTPException(status_code=502, detail={"error": {"message": "No image URL returned"}}) |
| |
| |
| return _gemini_hub_chat_envelope( |
| model, |
| f"\n\n{image_url}", |
| usage_prompt=len(prompt) // 4, |
| ) |
|
|
| |
| |
| payload = {"query": prompt, "top_k": int(body.get("top_k") or 5)} |
| url = f"{GEMINI_HUB_URL}/query" |
| logger.info("GeminiHub text request model=%s prompt=%s", model, prompt[:80]) |
| async with httpx.AsyncClient(timeout=120.0) as client: |
| try: |
| resp = await client.post(url, json=payload) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub unreachable: {e}"}}) |
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub {resp.status_code}: {resp.text[:300]}"}}) |
| data = resp.json() or {} |
| answer = data.get("response") or data.get("text") or "" |
| if not answer: |
| raise HTTPException(status_code=502, detail={"error": {"message": "Gemini Hub returned empty response"}}) |
| return _gemini_hub_chat_envelope(model, answer, usage_prompt=len(prompt) // 4) |
|
|
|
|
| |
|
|
| _OPENAI_SIZE_TO_ASPECT = { |
| "1024x1024": "1:1", |
| "1024x1792": "9:16", |
| "1792x1024": "16:9", |
| "768x1024": "3:4", |
| "1024x768": "4:3", |
| } |
|
|
|
|
| def _openai_size_to_aspect_ratio(size: str | None, allowed: set[str], default: str) -> str: |
| """Map OpenAI-style ``"WIDTHxHEIGHT"`` sizes to the aspect ratios the |
| Gemini Hub tools understand. Falls back to ``default`` for unknown values.""" |
| if not size: |
| return default |
| if size in allowed: |
| return size |
| mapped = _OPENAI_SIZE_TO_ASPECT.get(size) |
| if mapped and mapped in allowed: |
| return mapped |
| return default |
|
|
|
|
| async def _gemini_hub_generate_image(prompt: str, *, model_alias: str, size: str | None) -> dict: |
| """Call Gemini Hub's ``generar_imagen`` tool and return the raw payload. |
| |
| ``model_alias`` is one of GEMINI_HUB_IMAGE_MODELS. We pick the ``pro`` or |
| ``flash`` upstream variant based on its suffix. |
| """ |
| aspect = _openai_size_to_aspect_ratio(size, {"1:1", "16:9", "9:16", "4:3", "3:4"}, "1:1") |
| payload = { |
| "prompt": prompt, |
| "model": "pro" if model_alias.endswith("-pro") else "flash", |
| "aspect_ratio": aspect, |
| } |
| url = f"{GEMINI_HUB_URL}/call/generar_imagen" |
| async with httpx.AsyncClient(timeout=180.0) as client: |
| resp = await client.post(url, json=payload) |
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Image upstream HTTP {resp.status_code}"}}) |
| data = resp.json() or {} |
| if data.get("error"): |
| raise HTTPException(status_code=400, detail={"error": {"message": str(data["error"])}}) |
| if not (data.get("image_url") or data.get("url")): |
| raise HTTPException(status_code=502, detail={"error": {"message": "Image upstream returned no URL"}}) |
| return data |
|
|
|
|
| async def _gemini_hub_generate_video(prompt: str, *, size: str | None, duration_seconds, resolution, image_url) -> dict: |
| """Call Gemini Hub's ``generar_video`` tool (Veo 3) and return the raw payload.""" |
| payload: dict = {"prompt": prompt} |
| aspect = _openai_size_to_aspect_ratio(size, {"16:9", "9:16"}, "16:9") |
| payload["aspect_ratio"] = aspect |
| if duration_seconds and str(duration_seconds) in ("4", "6", "8"): |
| payload["duration_seconds"] = str(duration_seconds) |
| if resolution in ("720p", "1080p"): |
| payload["resolution"] = resolution |
| if isinstance(image_url, str) and image_url.strip(): |
| payload["image_url"] = image_url.strip() |
|
|
| url = f"{GEMINI_HUB_URL}/call/generar_video" |
| async with httpx.AsyncClient(timeout=600.0) as client: |
| resp = await client.post(url, json=payload) |
| if resp.status_code != 200: |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Video upstream HTTP {resp.status_code}"}}) |
| data = resp.json() or {} |
| if data.get("error"): |
| raise HTTPException(status_code=400, detail={"error": {"message": str(data["error"])}}) |
| if not (data.get("video_url") or data.get("url")): |
| raise HTTPException(status_code=502, detail={"error": {"message": "Video upstream returned no URL"}}) |
| return data |
|
|
|
|
| def _enforce_auth_and_limits(api_key: str, client_ip: str) -> None: |
| """Shared 401 / 429 gate used by image and video endpoints.""" |
| valid, msg = verify_key_ip(api_key, client_ip) |
| if not valid: |
| raise HTTPException(status_code=401, detail={"error": {"message": msg}}) |
| rl_ok, rl_detail, cooldown = check_rate_limit(api_key) |
| if not rl_ok: |
| raise HTTPException( |
| status_code=429, |
| detail={ |
| "error": {"message": rl_detail, "type": "rate_limit_error", "code": "rate_limit_exceeded"}, |
| "content": f"Rate limited, wait {int(cooldown)}s", |
| }, |
| ) |
|
|
|
|
| |
|
|
| @app.post("/v1/images/generations") |
| async def generate_images(request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)): |
| """OpenAI-compatible image generation. |
| |
| Routes by model id: |
| β’ ``gpt-image-1`` / ``gpt-image-2`` β internal image backend |
| β’ ``gemini-3-image`` / ``gemini-3-flash-image`` / ``gemini-3.1-flash-image`` |
| (optionally suffixed ``-pro``) β Gemini Hub generar_imagen |
| |
| Response is always OpenAI-shaped: ``{"created": β¦, "data": [{"url"|"b64_json", "revised_prompt"}]}``. |
| """ |
| api_key = credentials.credentials |
| client_ip = get_client_ip(request) |
| _enforce_auth_and_limits(api_key, client_ip) |
|
|
| body = await request.json() |
| model = body.get("model", "gpt-image-1") |
| prompt = body.get("prompt") |
| n = int(body.get("n", 1) or 1) |
| size = body.get("size", "1024x1024") |
| response_format = body.get("response_format", "url") |
|
|
| if not prompt: |
| raise HTTPException(status_code=400, detail={"error": {"message": "prompt is required"}}) |
|
|
| logger.info("IMG_REQ model=%s prompt=%s n=%d size=%s", model, str(prompt)[:60], n, size) |
|
|
| |
| if model in GEMINI_HUB_IMAGE_MODELS: |
| data = await _gemini_hub_generate_image(prompt=prompt, model_alias=model, size=size) |
| image_url = data.get("image_url") or data.get("url") or "" |
| |
| |
| images = [{"url": image_url, "revised_prompt": prompt} for _ in range(max(1, n))] |
| return {"created": int(time.time()), "data": images} |
|
|
| |
| if model in ("gpt-image-1", "gpt-image-2"): |
| try: |
| async with httpx.AsyncClient(timeout=60.0) as client: |
| resp = await client.post( |
| f"{CHATGPT_IMAGE_API_URL}/api/generate", |
| json={"prompt": prompt, "n": n, "size": size}, |
| headers={"Content-Type": "application/json"}, |
| ) |
| resp.raise_for_status() |
| upstream_data = resp.json() |
| except Exception as e: |
| logger.error("Image API error: %s", str(e)) |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Image upstream error: {e}"}}) |
|
|
| images = [] |
| if "images" in upstream_data: |
| for img in upstream_data["images"][:n]: |
| if response_format == "url": |
| images.append({"url": img.get("url", ""), "revised_prompt": img.get("revised_prompt", prompt)}) |
| else: |
| images.append({"b64_json": img.get("b64_json", ""), "revised_prompt": img.get("revised_prompt", prompt)}) |
| elif "url" in upstream_data: |
| images.append({"url": upstream_data["url"], "revised_prompt": upstream_data.get("revised_prompt", prompt)}) |
| return {"created": int(time.time()), "data": images} |
|
|
| raise HTTPException(status_code=400, detail={"error": {"message": f"Model {model!r} is not an image model"}}) |
|
|
|
|
| |
|
|
| @app.post("/v1/videos/generations") |
| async def generate_videos(request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)): |
| """Video generation endpoint shaped like ``/v1/images/generations``. |
| |
| Accepts any of the Veo 3 model ids (``veo-3``, ``veo-3-fast``, |
| ``veo-3.1-fast``). Body fields: |
| |
| * ``prompt`` *(required)* β text description |
| * ``size`` β ``"16:9"`` / ``"9:16"`` (or an OpenAI WxH string) |
| * ``duration_seconds`` β ``"4"`` | ``"6"`` | ``"8"`` (default ``"8"``) |
| * ``resolution`` β ``"720p"`` | ``"1080p"`` (default ``"720p"``) |
| * ``image_url`` β seed frame for image-to-video (optional) |
| |
| Response: ``{"created": β¦, "data": [{"url": "β¦"}]}``. |
| """ |
| api_key = credentials.credentials |
| client_ip = get_client_ip(request) |
| _enforce_auth_and_limits(api_key, client_ip) |
|
|
| body = await request.json() |
| model = body.get("model", "veo-3") |
| prompt = body.get("prompt") |
|
|
| if not prompt: |
| raise HTTPException(status_code=400, detail={"error": {"message": "prompt is required"}}) |
| if model not in GEMINI_HUB_VIDEO_MODELS: |
| raise HTTPException(status_code=400, detail={"error": {"message": f"Model {model!r} is not a video model"}}) |
|
|
| n = int(body.get("n", 1) or 1) |
| logger.info("VID_REQ model=%s prompt=%s", model, str(prompt)[:60]) |
|
|
| data = await _gemini_hub_generate_video( |
| prompt=prompt, |
| size=body.get("size"), |
| duration_seconds=body.get("duration_seconds") or body.get("duration"), |
| resolution=body.get("resolution"), |
| image_url=body.get("image_url"), |
| ) |
| video_url = data.get("video_url") or data.get("url") or "" |
| videos = [{"url": video_url, "revised_prompt": prompt} for _ in range(max(1, n))] |
| return {"created": int(time.time()), "data": videos} |
|
|
| |
|
|
| @app.get("/v1/keys/me") |
| async def get_my_key(request: Request, cf_token: str = None): |
| """Return caller's IP and auto-generate a key if they don't have one.""" |
| client_ip = get_client_ip(request) |
| existing_key = find_key_by_ip(client_ip) |
|
|
| if existing_key: |
| return { |
| "ip": client_ip, |
| "key": existing_key, |
| "rpm": RPM_LIMIT, |
| "status": "existing", |
| } |
|
|
| |
| vpn_result = await check_vpn(client_ip) |
| if vpn_result["vpn"]: |
| raise HTTPException( |
| status_code=403, |
| detail={"error": {"message": "VPN/Proxy detected. Please turn off VPN and try again.", "type": "verification_error", "code": "vpn_blocked", "provider": vpn_result["provider"]}} |
| ) |
|
|
| |
|
|
| |
| new_key = generate_api_key() |
| keys = _load_keys() |
| keys[new_key] = { |
| "ip": client_ip, |
| "created": int(time.time()), |
| "rpm": RPM_LIMIT, |
| } |
| _save_keys(keys) |
| logger.info("Auto-generated key: %s for IP: %s", new_key, client_ip) |
|
|
| return { |
| "ip": client_ip, |
| "key": new_key, |
| "rpm": RPM_LIMIT, |
| "status": "new", |
| } |
|
|
|
|
|
|
|
|
|
|
| @app.post("/v1/chat/completions") |
| async def chat_completions(request: Request, _auth=Depends(verify_request)): |
| try: |
| body = await request.json() |
| except Exception: |
| raise HTTPException(status_code=400, detail="Invalid JSON body") |
|
|
| messages = body.get("messages", []) |
| if not messages: |
| raise HTTPException(status_code=400, detail="messages is required") |
|
|
| model_raw = body.get("model", "gpt-4o") |
| model = normalize_model(model_raw) |
| stream = body.get("stream", False) |
| temperature = body.get("temperature", 0.7) |
| max_tokens = body.get("max_tokens") or body.get("max_completion_tokens") or 4096 |
|
|
| |
| safe_body = {k: v for k, v in body.items() if k != "messages"} |
| safe_body["msg_count"] = len(messages) |
| safe_body["has_tools"] = bool(body.get("tools")) |
| logger.info("CHAT_REQ ip=%s model=%s->%s stream=%s body_keys=%s", |
| get_client_ip(request), model_raw, model, stream, list(safe_body.keys())) |
|
|
| |
| |
| |
| |
| if model in GEMINI_HUB_MODELS: |
| try: |
| completion = await _gemini_hub_chat(model, messages, body) |
| except HTTPException: |
| raise |
| except Exception as e: |
| logger.exception("Gemini Hub adapter failed: %s", e) |
| raise HTTPException(status_code=502, detail={"error": {"message": f"Gemini Hub error: {e}"}}) |
|
|
| if not stream: |
| return completion |
|
|
| |
| async def _gemini_hub_stream(): |
| choice = completion["choices"][0] |
| chunk = { |
| "id": completion["id"], |
| "object": "chat.completion.chunk", |
| "created": completion["created"], |
| "model": completion["model"], |
| "choices": [{ |
| "index": 0, |
| "delta": {"role": "assistant", "content": choice["message"]["content"]}, |
| "finish_reason": None, |
| }], |
| } |
| yield f"data: {json.dumps(chunk)}\n\n" |
| final = { |
| "id": completion["id"], |
| "object": "chat.completion.chunk", |
| "created": completion["created"], |
| "model": completion["model"], |
| "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], |
| } |
| yield f"data: {json.dumps(final)}\n\n" |
| yield "data: [DONE]\n\n" |
| return StreamingResponse(_gemini_hub_stream(), media_type="text/event-stream") |
| tools = body.get("tools") |
| tool_choice = body.get("tool_choice") |
| req_id = uuid.uuid4().hex |
| created = int(time.time()) |
|
|
| |
| client_ip = get_client_ip(request) |
|
|
| |
| auth_header = request.headers.get("authorization", "") |
| api_key = auth_header.replace("Bearer ", "") if auth_header.startswith("Bearer ") else "" |
|
|
| |
| if stream: |
| logger.info("CHAT_STREAM_START req_id=%s model=%s", req_id, model) |
| async def generate(): |
| all_content = "" |
| all_reasoning = "" |
| all_tool_calls = [] |
| has_tool_calls = False |
| current_messages = list(messages) |
|
|
| for attempt in range(MAX_CONTINUATIONS + 1): |
| logger.info("Attempt %d/%d req_id=%s", attempt + 1, MAX_CONTINUATIONS + 1, req_id) |
|
|
| try: |
| |
| req_body = {**body, "messages": current_messages} if attempt > 0 else body |
| upstream, client = await send_upstream_request(req_body, model) |
| except Exception as e: |
| logger.exception("Upstream request failed req_id=%s: %s", req_id, e) |
| break |
|
|
| if upstream.status_code != 200: |
| try: |
| body_text = await upstream.aread() |
| logger.error("Upstream status=%s req_id=%s body=%s", upstream.status_code, req_id, body_text.decode(errors="replace")) |
| finally: |
| await upstream.aclose() |
| await client.aclose() |
| break |
|
|
| attempt_content = "" |
| attempt_reasoning = "" |
| got_done = False |
| timed_out = False |
| gateway_tool_mode = getattr(upstream, "upstream_url", "") == AI_GATEWAY_URL and bool(tools) |
|
|
| try: |
| async for raw_line in upstream.aiter_lines(): |
| line = raw_line.strip() |
| if not line: |
| continue |
|
|
| clean = line.removeprefix("data:").strip() |
| if clean == "[DONE]": |
| got_done = True |
| break |
|
|
| parsed = parse_upstream_delta(clean) |
| content = parsed.get("content") |
| tool_calls = parsed.get("tool_calls") |
| reasoning = parsed.get("reasoning_content") |
| finish_reason = parsed.get("finish_reason") |
|
|
| if content: |
| attempt_content += content |
| all_content += content |
| if reasoning: |
| attempt_reasoning += reasoning |
| all_reasoning += reasoning |
| if tool_calls: |
| merge_tool_calls(all_tool_calls, tool_calls) |
| has_tool_calls = True |
|
|
| if gateway_tool_mode: |
| continue |
|
|
| if reasoning: |
| reasoning_chunk = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": {"reasoning_content": reasoning}, "finish_reason": None}], |
| } |
| yield f"data: {json.dumps(reasoning_chunk)}\n\n" |
|
|
| if tool_calls: |
| tc_chunk = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": {"tool_calls": tool_calls}, "finish_reason": None}], |
| } |
| yield f"data: {json.dumps(tc_chunk)}\n\n" |
|
|
| if not content: |
| continue |
|
|
| delta = {"content": content} |
|
|
| chunk = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": delta, "finish_reason": None}], |
| } |
| yield f"data: {json.dumps(chunk)}\n\n" |
|
|
| except httpx.ReadTimeout: |
| timed_out = True |
| logger.warning("Read timeout attempt=%d req_id=%s", attempt + 1, req_id) |
| except httpx.RemoteProtocolError: |
| logger.warning("Protocol closed attempt=%d req_id=%s", attempt + 1, req_id) |
| except Exception as e: |
| logger.exception("Stream error attempt=%d req_id=%s error=%s", attempt + 1, req_id, e) |
| finally: |
| await upstream.aclose() |
| await client.aclose() |
|
|
| if gateway_tool_mode and not has_tool_calls: |
| parsed_gateway_tool_calls = _parse_gateway_tool_calls(attempt_content) |
| if parsed_gateway_tool_calls: |
| all_tool_calls = parsed_gateway_tool_calls |
| has_tool_calls = True |
| attempt_content = "" |
| all_content = "" |
| for i, tc in enumerate(all_tool_calls): |
| delta = { |
| "tool_calls": [{ |
| "index": i, |
| "id": tc.get("id"), |
| "type": tc.get("type", "function"), |
| "function": { |
| "name": (tc.get("function") or {}).get("name", ""), |
| "arguments": (tc.get("function") or {}).get("arguments", ""), |
| }, |
| }] |
| } |
| chunk = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": delta, "finish_reason": None}], |
| } |
| yield f"data: {json.dumps(chunk)}\n\n" |
| elif attempt_content: |
| delta = {"content": attempt_content} |
| if attempt_reasoning: |
| delta["reasoning_content"] = attempt_reasoning |
| chunk = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": delta, "finish_reason": None}], |
| } |
| yield f"data: {json.dumps(chunk)}\n\n" |
|
|
| if not attempt_content.strip() and not attempt_reasoning.strip() and not has_tool_calls: |
| logger.warning("No content on attempt=%d req_id=%s", attempt + 1, req_id) |
| break |
|
|
| if has_tool_calls: |
| logger.info("Tool calls present, skipping auto-continue req_id=%s", req_id) |
| break |
|
|
| if got_done and not looks_incomplete(all_content): |
| logger.info("Completed naturally req_id=%s", req_id) |
| break |
|
|
| if attempt >= MAX_CONTINUATIONS: |
| logger.warning("Max continuations reached req_id=%s", req_id) |
| break |
|
|
| continuation = ( |
| "The response was cut off. " |
| "Continue only from the last sentence you left off. " |
| "Never start over. Never repeat previous text. " |
| "Continue directly without a brief transition. " |
| "Write more compactly if possible and complete the response." |
| ) |
| if timed_out: |
| continuation += " First complete by summarizing the remaining missing sections." |
|
|
| current_messages = list(messages) |
| current_messages.append({"role": "assistant", "content": all_content}) |
| current_messages.append({"role": "user", "content": continuation}) |
|
|
| final_reason = "tool_calls" if has_tool_calls else "stop" |
|
|
| |
| prompt_text = json.dumps(messages) |
| completion_text = all_content or "" |
| prompt_tokens = max(1, len(prompt_text) // 4) |
| completion_tokens = max(1, len(completion_text) // 4) |
|
|
| stop = { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "delta": {}, "finish_reason": final_reason}], |
| "usage": { |
| "prompt_tokens": prompt_tokens, |
| "completion_tokens": completion_tokens, |
| "total_tokens": prompt_tokens + completion_tokens, |
| }, |
| } |
| yield f"data: {json.dumps(stop)}\n\n" |
|
|
| |
| log_chat(client_ip, model, messages, all_content, api_key) |
|
|
| yield "data: [DONE]\n\n" |
|
|
| return StreamingResponse( |
| generate(), |
| media_type="text/event-stream", |
| headers={ |
| "Cache-Control": "no-cache, no-transform", |
| "X-Accel-Buffering": "no", |
| "Connection": "keep-alive", |
| }, |
| ) |
|
|
| |
| logger.info("CHAT_NONSTREAM_START req_id=%s model=%s", req_id, model) |
| full_content = "" |
| full_reasoning = "" |
| all_tool_calls = [] |
| has_tool_calls = False |
| current_messages = list(messages) |
|
|
| for attempt in range(MAX_CONTINUATIONS + 1): |
| logger.info("Non-stream attempt %d/%d req_id=%s", attempt + 1, MAX_CONTINUATIONS + 1, req_id) |
|
|
| try: |
| req_body = {**body, "messages": current_messages} if attempt > 0 else body |
| upstream, client = await send_upstream_request(req_body, model) |
| except Exception as e: |
| raise HTTPException(status_code=502, detail=f"Upstream error: {e}") |
|
|
| if upstream.status_code != 200: |
| error_body = await upstream.aread() |
| await upstream.aclose() |
| await client.aclose() |
| raise HTTPException(status_code=upstream.status_code, detail=error_body.decode(errors="replace")) |
|
|
| attempt_content = "" |
| attempt_reasoning = "" |
| got_done = False |
| timed_out = False |
|
|
| try: |
| async for raw_line in upstream.aiter_lines(): |
| line = raw_line.strip() |
| if not line: |
| continue |
| clean = line.removeprefix("data:").strip() |
| if clean == "[DONE]": |
| got_done = True |
| break |
|
|
| parsed = parse_upstream_delta(clean) |
| content = parsed.get("content") |
| tool_calls = parsed.get("tool_calls") |
| reasoning = parsed.get("reasoning_content") |
|
|
| if content: |
| attempt_content += content |
| full_content += content |
| if reasoning: |
| attempt_reasoning += reasoning |
| full_reasoning += reasoning |
| if tool_calls: |
| merge_tool_calls(all_tool_calls, tool_calls) |
| has_tool_calls = True |
|
|
| except httpx.ReadTimeout: |
| timed_out = True |
| except Exception as e: |
| logger.warning("Non-stream interrupted req_id=%s error=%s", req_id, e) |
| finally: |
| await upstream.aclose() |
| await client.aclose() |
|
|
| if not attempt_content.strip() and not attempt_reasoning.strip() and not has_tool_calls: |
| break |
| if has_tool_calls: |
| break |
| if got_done and not looks_incomplete(full_content): |
| break |
| if attempt >= MAX_CONTINUATIONS: |
| break |
|
|
| continuation = ( |
| "The response was cut off. " |
| "Continue only from where you left off. " |
| "Do not repeat previous text. " |
| "Complete the remaining sections in a more compact format." |
| ) |
| if timed_out: |
| continuation += " Summarize to finish if necessary." |
|
|
| current_messages = list(messages) |
| current_messages.append({"role": "assistant", "content": full_content}) |
| current_messages.append({"role": "user", "content": continuation}) |
|
|
| if tools and not has_tool_calls: |
| parsed_gateway_tool_calls = _parse_gateway_tool_calls(full_content) |
| if parsed_gateway_tool_calls: |
| all_tool_calls = parsed_gateway_tool_calls |
| has_tool_calls = True |
| full_content = "" |
|
|
| |
| |
| |
| |
| normalized_content = full_content if full_content else "" |
| if not normalized_content and full_reasoning and not has_tool_calls: |
| normalized_content = full_reasoning |
| if not normalized_content and not has_tool_calls: |
| normalized_content = "Model returned no visible answer" |
|
|
| message = {"role": "assistant", "content": normalized_content} |
| if full_reasoning: |
| message["reasoning_content"] = full_reasoning |
| if has_tool_calls: |
| message["tool_calls"] = all_tool_calls |
| |
| if not full_content: |
| message["content"] = "" |
|
|
| finish_reason = "tool_calls" if has_tool_calls else "stop" |
|
|
| |
| log_chat(client_ip, model, messages, full_content, api_key) |
|
|
| |
| prompt_text = json.dumps(messages) |
| completion_text = normalized_content or "" |
| prompt_tokens = max(1, len(prompt_text) // 4) |
| completion_tokens = max(1, len(completion_text) // 4) |
|
|
| return { |
| "id": f"chatcmpl-{req_id}", |
| "object": "chat.completion", |
| "created": created, |
| "model": model, |
| "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], |
| "usage": { |
| "prompt_tokens": prompt_tokens, |
| "completion_tokens": completion_tokens, |
| "total_tokens": prompt_tokens + completion_tokens, |
| }, |
| } |
|
|
|
|