flexgate 0.1.0__py3-none-any.whl

flexgate/proxy.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import json
+import logging
+
+import httpx
+from starlette.responses import JSONResponse, StreamingResponse
+
+from flexgate.config import ProviderConfig
+
+logger = logging.getLogger("flexgate.proxy")
+
+MULTIMODAL_MODELS = {"MiniMax-M3"}
+_IMAGE_NOT_SUPPORTED_NOTE = "[flexgate] image message is not handled because multimodal is not supported by the current model"
+
+
+def _model_supports_multimodal(model: str) -> bool:
+    return model in MULTIMODAL_MODELS
+
+
+def _has_image_content(body_json: dict) -> bool:
+    for msg in body_json.get("messages", []):
+        content = msg.get("content")
+        if isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict) and block.get("type") in ("image", "image_url"):
+                    return True
+    return False
+
+
+def _strip_images_from_request(body_json: dict) -> None:
+    for msg in body_json.get("messages", []):
+        content = msg.get("content")
+        if not isinstance(content, list):
+            continue
+        remaining = [b for b in content if not (isinstance(b, dict) and b.get("type") in ("image", "image_url"))]
+        if not remaining:
+            remaining = [{"type": "text", "text": _IMAGE_NOT_SUPPORTED_NOTE}]
+        else:
+            remaining.append({"type": "text", "text": _IMAGE_NOT_SUPPORTED_NOTE})
+        msg["content"] = remaining
+
+
+def _add_note_to_response(data: dict) -> None:
+    for block in data.get("content", []):
+        if isinstance(block, dict) and block.get("type") == "text":
+            block["text"] += "\n" + _IMAGE_NOT_SUPPORTED_NOTE
+            return
+    data.setdefault("content", []).append({"type": "text", "text": _IMAGE_NOT_SUPPORTED_NOTE})
+
+
+def _error_json(error_type: str, message: str, status: int) -> JSONResponse:
+    return JSONResponse(
+        {"type": "error", "error": {"type": error_type, "message": message}},
+        status_code=status,
+    )
+
+
+def _build_upstream_headers(
+    incoming_headers: dict[str, str], provider: ProviderConfig
+) -> dict[str, str]:
+    return {
+        "content-type": "application/json",
+        "x-api-key": provider.api_key,
+        "anthropic-version": incoming_headers.get("anthropic-version", "2023-06-01"),
+        "anthropic-beta": incoming_headers.get("anthropic-beta", ""),
+    }
+
+
+async def handle_request(
+    client: httpx.AsyncClient,
+    body: bytes,
+    body_json: dict,
+    incoming_headers: dict[str, str],
+    provider: ProviderConfig,
+    model_override: str | None,
+) -> JSONResponse | StreamingResponse:
+    if model_override:
+        body_json["model"] = model_override
+        body = json.dumps(body_json, ensure_ascii=False).encode()
+
+    url = f"{provider.base_url}/v1/messages"
+    headers = _build_upstream_headers(incoming_headers, provider)
+
+    is_stream = body_json.get("stream", False)
+    model = body_json.get("model", "")
+    strip_images = not _model_supports_multimodal(model) and _has_image_content(body_json)
+
+    if strip_images:
+        _strip_images_from_request(body_json)
+        body = json.dumps(body_json, ensure_ascii=False).encode()
+        logger.info("Stripped image content for non-multimodal model: %s", model)
+
+    if is_stream:
+        return await _stream_proxy(client, url, headers, body, provider.name, strip_images)
+    return await _regular_proxy(client, url, headers, body, provider.name, strip_images)
+
+
+async def _regular_proxy(
+    client: httpx.AsyncClient,
+    url: str,
+    headers: dict[str, str],
+    body: bytes,
+    provider_name: str,
+    images_stripped: bool = False,
+) -> JSONResponse:
+    try:
+        resp = await client.post(url, headers=headers, content=body)
+    except httpx.ConnectError as exc:
+        logger.error("Provider %s unreachable: %s", provider_name, exc)
+        return _error_json("api_error", f"Provider unreachable: {provider_name}", 502)
+
+    try:
+        data = resp.json()
+    except Exception:
+        return _error_json("api_error", resp.text[:500], resp.status_code)
+
+    if images_stripped and resp.status_code == 200:
+        _add_note_to_response(data)
+
+    return JSONResponse(data, status_code=resp.status_code)
+
+
+async def _stream_proxy(
+    client: httpx.AsyncClient,
+    url: str,
+    headers: dict[str, str],
+    body: bytes,
+    provider_name: str,
+    images_stripped: bool = False,
+) -> StreamingResponse:
+    try:
+        req = client.build_request("POST", url, headers=headers, content=body)
+        upstream = await client.send(req, stream=True)
+    except httpx.ConnectError as exc:
+        logger.error("Provider %s unreachable: %s", provider_name, exc)
+        return _error_json("api_error", f"Provider unreachable: {provider_name}", 502)
+
+    if upstream.status_code != 200:
+        error_body = await upstream.aread()
+        await upstream.aclose()
+        try:
+            data = json.loads(error_body)
+        except Exception:
+            data = {"type": "error", "error": {"type": "api_error", "message": error_body.decode(errors="replace")[:500]}}
+        return JSONResponse(data, status_code=upstream.status_code)
+
+    async def generate():
+        try:
+            async for chunk in upstream.aiter_bytes():
+                yield chunk
+        except httpx.ReadError:
+            logger.error("Upstream %s connection lost during streaming", provider_name)
+        finally:
+            await upstream.aclose()
+
+    return StreamingResponse(
+        generate(),
+        status_code=200,
+        headers={
+            "content-type": "text/event-stream",
+            "cache-control": "no-cache",
+            "connection": "keep-alive",
+        },
+    )

flexgate/registry.py

@@ -0,0 +1,22 @@
+"""Registry of known upstream LLM services.
+
+`KNOWN_BASE_URLS` maps a service name prefix to its Anthropic-compatible
+base_url. `flexgate sync` uses it to auto-import providers that exist in
+Infisical but not yet in the config: the Infisical folder name is matched
+against these prefixes on a dash boundary, longest prefix first
+(e.g. 'minimax-tmy' -> minimax's base_url).
+
+To onboard a new service, add one entry here — no other code changes needed.
+Providers already present in the config are merged on top of this table at
+sync time, so their (possibly newer) base_urls take precedence.
+"""
+
+from __future__ import annotations
+
+# prefix -> base_url
+KNOWN_BASE_URLS: dict[str, str] = {
+    "minimax": "https://api.minimaxi.com/anthropic",
+    "zai": "https://api.z.ai/api/anthropic",
+    "ustc": "https://api.llm.ustc.edu.cn",
+    "kimi": "https://api.kimi.com/coding",
+}

flexgate/router.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from flexgate.config import GatewayConfig, ProviderConfig, RouteConfig
+
+logger = logging.getLogger("flexgate")
+
+
+class NoRouteMatchError(Exception):
+    def __init__(self, model: str) -> None:
+        self.model = model
+        super().__init__(f"No route matches model: {model}")
+
+
+def _current_minutes() -> int:
+    now = datetime.now()
+    return now.hour * 60 + now.minute
+
+
+def _in_window(minutes: int, start: int, end: int) -> bool:
+    if start == end:
+        return True
+    if start < end:
+        return start <= minutes < end
+    # Overnight wrap: e.g. 22:00-06:00
+    return minutes >= start or minutes < end
+
+
+def _match_route(routes: list[RouteConfig], model: str) -> tuple[ProviderConfig, str | None] | None:
+    for route in routes:
+        if route.pattern.search(model):
+            return route.provider_name, route.model
+    return None
+
+
+def resolve(config: GatewayConfig, model: str) -> tuple[ProviderConfig, str | None, str]:
+    """Return (provider, model_override, schedule_name) for the matching route.
+
+    model_override is always a concrete model name: if the matched route omits
+    'model', it falls back to the provider's first available_models entry.
+    """
+    now = _current_minutes()
+
+    for entry in config.schedule:
+        if _in_window(now, entry.start_minutes, entry.end_minutes):
+            match = _match_route(entry.routes, model)
+            if match:
+                prov_name, model_override = match
+                provider = config.providers[prov_name]
+                model_override = _resolve_model(provider, model_override)
+                label = entry.name or f"{entry.start_minutes//60:02d}:{entry.start_minutes%60:02d}-{entry.end_minutes//60:02d}:{entry.end_minutes%60:02d}"
+                logger.debug("schedule [%s] matched for model %s", label, model)
+                return provider, model_override, label
+
+    match = _match_route(config.routes, model)
+    if match:
+        prov_name, model_override = match
+        provider = config.providers[prov_name]
+        model_override = _resolve_model(provider, model_override)
+        return provider, model_override, "default"
+
+    raise NoRouteMatchError(model)
+
+
+def _resolve_model(provider: ProviderConfig, model_override: str | None) -> str:
+    """Ensure a concrete model: fall back to the provider's first available model."""
+    if model_override:
+        return model_override
+    if provider.available_models:
+        return provider.available_models[0]
+    raise ValueError(
+        f"Provider '{provider.name}' has no model override and no available_models to fall back to"
+    )

flexgate/server.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import json
+import logging
+import signal
+import time
+from collections.abc import AsyncGenerator
+
+import httpx
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.routing import Route
+
+from flexgate.config import GatewayConfig, load_config
+from flexgate.proxy import handle_request
+from flexgate.router import NoRouteMatchError, resolve
+
+logger = logging.getLogger("flexgate")
+
+
+def _error_json(error_type: str, message: str, status: int) -> JSONResponse:
+    return JSONResponse(
+        {"type": "error", "error": {"type": error_type, "message": message}},
+        status_code=status,
+    )
+
+
+async def messages(request: Request) -> JSONResponse:
+    t0 = time.monotonic()
+    config: GatewayConfig = request.app.state.config
+    client: httpx.AsyncClient = request.app.state.client
+
+    body_bytes = await request.body()
+    try:
+        body_json = json.loads(body_bytes)
+    except json.JSONDecodeError:
+        return _error_json("invalid_request_error", "Invalid JSON in request body", 400)
+
+    model_name = body_json.get("model", "")
+
+    try:
+        provider, model_override, schedule_name = resolve(config, model_name)
+    except NoRouteMatchError as exc:
+        return _error_json("not_found_error", str(exc), 503)
+
+    display_model = model_override or model_name
+    logger.info("[%s] %s -> %s (%s)", schedule_name, model_name, provider.name, display_model)
+
+    incoming_headers = dict(request.headers)
+    resp = await handle_request(
+        client, body_bytes, body_json, incoming_headers, provider, model_override,
+    )
+
+    elapsed_ms = int((time.monotonic() - t0) * 1000)
+    logger.info(
+        "[%s] %s -> %s (%s) | %s | %dms",
+        schedule_name, model_name, provider.name, display_model, resp.status_code, elapsed_ms,
+    )
+    return resp
+
+
+def create_app(config: GatewayConfig, config_path: str | None = None) -> Starlette:
+    client = httpx.AsyncClient(
+        timeout=httpx.Timeout(connect=10, read=3600, write=10, pool=10),
+    )
+
+    routes = [Route("/v1/messages", messages, methods=["POST"])]
+
+    async def lifespan(app: Starlette) -> AsyncGenerator[None, None]:
+        app.state.config = config
+        app.state.client = client
+
+        def _handle_sigusr1(signum, frame):
+            try:
+                app.state.config = load_config(config_path)
+                logger.info("Config reloaded via SIGUSR1")
+            except Exception as e:
+                logger.error("Failed to reload config: %s", e)
+
+        signal.signal(signal.SIGUSR1, _handle_sigusr1)
+        yield
+        await client.aclose()
+
+    return Starlette(routes=routes, lifespan=lifespan)