orcarouter.py

"""
title: OrcaRouter
authors: OrcaRouter team
author_url: https://www.orcarouter.ai
funding_url: https://www.orcarouter.ai
version: 0.1.0
required_open_webui_version: 0.5.0
license: MIT
"""

import os
import re
import time
from typing import Any, AsyncGenerator, Dict, List, Optional, Union

import httpx
from pydantic import BaseModel, Field


# OrcaRouter exposes models across multiple modalities (chat, embedding, TTS,
# image, video). The Open WebUI model selector is for chat completions only,
# so we filter the /v1/models response to chat-capable models.
_NON_CHAT_NAME_KEYWORDS = (
    "embedding",
    "tts",
    "whisper",
    "transcrib",
    "rerank",
    "imagen",
    "dall-e",
    "gpt-image",
    "grok-imagine",
)
# Model IDs that match these endpoint-type sets are routed through endpoints
# other than /v1/chat/completions (e.g. /v1/responses, /v1/completions). They
# would return 404 if surfaced as chat models.
_RESPONSES_ONLY_PATTERN = re.compile(r"^openai/gpt-5(\.\d+)?-pro")


def _is_chat_model(model: Dict[str, Any]) -> bool:
    model_id = (model.get("id") or "").lower()
    if not model_id:
        return False

    endpoint_types = set(model.get("supported_endpoint_types") or [])
    if "image-generation" in endpoint_types or "openai-video" in endpoint_types:
        return False

    output_modalities = set(model.get("output_modalities") or [])
    if "image" in output_modalities or "video" in output_modalities:
        return False

    if any(kw in model_id for kw in _NON_CHAT_NAME_KEYWORDS):
        return False
    if model_id.endswith("-speech"):
        return False

    # Models that only expose /v1/responses or /v1/completions are not usable
    # through /v1/chat/completions.
    if "openai-response" in endpoint_types and "openai" not in endpoint_types:
        return False
    if "codex" in model_id:
        return False
    if _RESPONSES_ONLY_PATTERN.match(model_id):
        return False

    return True


class Pipe:
    class Valves(BaseModel):
        ORCAROUTER_API_BASE_URL: str = Field(
            default="https://api.orcarouter.ai/v1",
            description="Base URL for OrcaRouter API endpoints.",
        )
        ORCAROUTER_API_KEY: str = Field(
            default="",
            description="Global API key used to list models and (if no per-user key is set) to send requests. OrcaRouter keys start with 'sk-orca-'.",
        )
        NAME_PREFIX: str = Field(
            default="",
            description="Optional prefix prepended to each model name in the selector. Empty by default since OrcaRouter model IDs already carry a namespace like 'openai/gpt-5.5'.",
        )
        MODEL_CACHE_TTL_SECONDS: int = Field(
            default=600,
            description="How long to cache the model list before refetching from /v1/models.",
        )
        REQUEST_TIMEOUT_SECONDS: int = Field(
            default=300,
            description="HTTP timeout for chat completion requests. Reasoning models like openai/gpt-5-pro can take minutes to cold-start.",
        )

    class UserValves(BaseModel):
        ORCAROUTER_API_KEY: str = Field(
            default="",
            description="Per-user API key. If set, overrides the global key for this user's requests.",
        )

    def __init__(self) -> None:
        self.type = "manifold"
        self.id = "orcarouter"
        self.name = "orcarouter/"
        self.valves = self.Valves(
            **{"ORCAROUTER_API_KEY": os.getenv("ORCAROUTER_API_KEY", "")}
        )
        self._model_cache: Optional[List[Dict[str, str]]] = None
        self._model_cache_time: float = 0.0

    def _headers(self, api_key: str) -> Dict[str, str]:
        # Attribution headers mirror what Open WebUI sends for OpenRouter so
        # the OrcaRouter dashboard can group Open WebUI traffic.
        return {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://openwebui.com/",
            "X-Title": "Open WebUI",
        }

    async def pipes(self) -> List[Dict[str, str]]:
        if not self.valves.ORCAROUTER_API_KEY:
            return [
                {
                    "id": "error",
                    "name": "ORCAROUTER_API_KEY is not set. Configure it under Functions → OrcaRouter → Valves.",
                }
            ]

        now = time.time()
        if (
            self._model_cache is not None
            and (now - self._model_cache_time) < self.valves.MODEL_CACHE_TTL_SECONDS
        ):
            return self._model_cache

        url = f"{self.valves.ORCAROUTER_API_BASE_URL.rstrip('/')}/models"
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                r = await client.get(url, headers=self._headers(self.valves.ORCAROUTER_API_KEY))
                r.raise_for_status()
                payload = r.json()
        except Exception as exc:
            return [
                {
                    "id": "error",
                    "name": f"Could not fetch models from OrcaRouter: {exc}",
                }
            ]

        models = [
            {
                "id": model["id"],
                "name": f'{self.valves.NAME_PREFIX}{model.get("name") or model["id"]}',
            }
            for model in payload.get("data", [])
            if _is_chat_model(model)
        ]

        self._model_cache = models
        self._model_cache_time = now
        return models

    def _resolve_api_key(self, __user__: Optional[Dict[str, Any]]) -> str:
        if __user__:
            user_valves = __user__.get("valves")
            if user_valves is not None:
                key = getattr(user_valves, "ORCAROUTER_API_KEY", "")
                if key:
                    return key
        return self.valves.ORCAROUTER_API_KEY

    def _strip_manifold_prefix(self, model: str) -> str:
        # Open WebUI prepends the manifold id and a dot, e.g.
        # "orcarouter.openai/gpt-5.5". Strip the leading "<manifold>." once.
        prefix = f"{self.id}."
        if model.startswith(prefix):
            return model[len(prefix) :]
        return model

    async def pipe(
        self,
        body: Dict[str, Any],
        __user__: Optional[Dict[str, Any]] = None,
    ) -> Union[str, AsyncGenerator[str, None], Dict[str, Any]]:
        api_key = self._resolve_api_key(__user__)
        if not api_key:
            return "Error: no OrcaRouter API key configured. Set it under the Function's Valves (admin) or your User Valves."

        model = body.get("model") or ""
        payload = {**body, "model": self._strip_manifold_prefix(model)}
        # 'user' in the body upstream-side is not always accepted; the wrapper
        # already authenticates via the API key, so we pass it through unchanged.

        url = f"{self.valves.ORCAROUTER_API_BASE_URL.rstrip('/')}/chat/completions"
        headers = self._headers(api_key)
        timeout = httpx.Timeout(self.valves.REQUEST_TIMEOUT_SECONDS)

        stream = bool(body.get("stream"))

        if stream:
            async def event_stream() -> AsyncGenerator[str, None]:
                async with httpx.AsyncClient(timeout=timeout) as client:
                    async with client.stream(
                        "POST", url, json=payload, headers=headers
                    ) as r:
                        r.raise_for_status()
                        async for line in r.aiter_lines():
                            yield line

            return event_stream()

        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                r = await client.post(url, json=payload, headers=headers)
                r.raise_for_status()
                return r.json()
        except httpx.HTTPStatusError as exc:
            return f"Error from OrcaRouter ({exc.response.status_code}): {exc.response.text}"
        except Exception as exc:
            return f"Error: {exc}"