NotDiamond Auto Router for Open WebUI

README.md

NotDiamond Auto Router

This script is a pipe that uses the NotDiamond API to automatically route prompts to the most optimal language model based on a specified tradeoff (cost, performance, or latency).

It analyzes the user's prompt and selects the best-fit model from a predefined list, then seamlessly passes the request to the chosen model.

Installation

1. Place the Script: Copy the notdiamond_auto_router.py file into the /data/pipes directory inside the installation volume.
2. Install Dependencies: The system will automatically detect the requirements: httpx line in the script's docstring and install the necessary Python package. A restart may be required.
3. Restart: Restart the instance for the new pipe to be recognized.

Configuration

To use the router, you need to create a "Custom Model" that uses this pipe.

1. Go to Settings > Models.
2. Under "Model," pull a model, or select a model you want to use.
3. Once the model is pulled, click the edit button next to the model name.
4. In the edit screen, set the "Pipe" to NotDiamond Auto Router.
5. Configure the following settings:

   • NOTDIAMOND_API_KEY: (Required) Your API key from NotDiamond.

   • NOTDIAMOND_BASE_URL: The base URL for the NotDiamond API. Defaults to https://api.notdiamond.ai/v2.

   • MODEL_MAPPING: (Required) A list mapping NotDiamond model identifiers to your model tags. Each line should be in the format: notdiamond/model-id, model-tag.

     Example:

     google/gemini-1.5-pro, openrouter/google/gemini-1.5-pro
     google/gemini-1.5-flash, openrouter/google/gemini-1.5-flash
     

   • TRADEOFF: The routing preference. Can be cost, performance, or latency. Defaults to cost.

6. Save the custom model.

Now, when you select and use this custom model in a chat, the pipe will automatically route your prompt to the best model as determined by NotDiamond.

notdiamond_auto_router.py

"""
title: NotDiamond Auto Router
author: karilaa-dev
author_url: https://github.com/karilaa-dev
gist_url: https://gist.github.com/karilaa-dev/ce8a3032b103d5425a482bfde1da5240
version: 1.0.0
description: Auto-routes prompts to optimal model using NotDiamond API with model-specific configurations
requirements: httpx
"""

from pydantic import BaseModel, Field
import httpx
from typing import Optional, Callable, Awaitable, Dict, Any, List

from open_webui.main import generate_chat_completions
from open_webui.models.users import Users
from open_webui.models.models import Models


class Pipe:
    class Valves(BaseModel):
        NOTDIAMOND_API_KEY: str = Field(
            default="",
            description="NotDiamond API key for model routing (REQUIRED)",
        )
        NOTDIAMOND_BASE_URL: str = Field(
            default="https://api.notdiamond.ai/v2",
            description="NotDiamond API base URL (try 'https://proxy.notdiamond.ai/v1' if v2 doesn't work)",
        )
        MODEL_MAPPING: str = Field(
            default="google/gemini-2.5-pro, gemini-2.5-pro\ngoogle/gemini-2.5-flash, gemini-2.5-flash",
            description="Model mapping in format: 'notdiamond_model, openrouter_model_id' (one per line)",
        )
        TRADEOFF: str = Field(
            default="cost",
            description="Routing tradeoff preference ('cost', 'performance', or 'latency')",
        )

    def __init__(self):
        self.type = "pipe"
        self.id = "auto_router"
        self.valves = self.Valves()
        self._model_mapping_cache = None
        self._raw_model_mapping = ""

    def pipes(self):
        """Return available pipe models"""
        return [
            {
                "id": "auto_router",
                "name": "NotDiamond Auto Router"
            }
        ]

    async def _emit_status(self, event_emitter, level: str, message: str, done: bool = False):
        """Emit status event to Open WebUI"""
        if event_emitter:
            await event_emitter({
                "type": "status",
                "data": {
                    "status": "complete" if done else "in_progress",
                    "level": level,
                    "description": message,
                    "done": done,
                }
            })

    def _get_model_mapping(self) -> Dict[str, str]:
        """Get mapping from NotDiamond model IDs to OpenWebUI model IDs"""
        if (
            self.valves.MODEL_MAPPING == self._raw_model_mapping
            and self._model_mapping_cache is not None
        ):
            return self._model_mapping_cache

        mapping = {}
        
        if not self.valves.MODEL_MAPPING:
            return mapping
            
        lines = self.valves.MODEL_MAPPING.strip().split('\n')
        for line in lines:
            line = line.strip()
            if not line or ',' not in line:
                continue
                
            parts = [part.strip() for part in line.split(',', 1)]
            if len(parts) == 2:
                notdiamond_model, openwebui_model = parts
                mapping[notdiamond_model] = openwebui_model
        
        self._raw_model_mapping = self.valves.MODEL_MAPPING
        self._model_mapping_cache = mapping
        return mapping

    def _get_notdiamond_models(self) -> List[str]:
        """Get list of NotDiamond model IDs for API request"""
        return list(self._get_model_mapping().keys())

    async def _select_best_model(self, messages: List[Dict]) -> str:
        """Use NotDiamond API to select the best model and return OpenWebUI model ID"""
        if not self.valves.NOTDIAMOND_API_KEY:
            raise Exception("NotDiamond API key not provided")

        headers = {
            "Authorization": f"Bearer {self.valves.NOTDIAMOND_API_KEY}",
            "Content-Type": "application/json",
        }

        notdiamond_models = self._get_notdiamond_models()
        model_mapping = self._get_model_mapping()

        payload = {
            "messages": messages,
            "llm_providers": [
                {"provider": model.split("/")[0], "model": model.split("/")[1]}
                for model in notdiamond_models
            ],
            "tradeoff": self.valves.TRADEOFF,
        }

        endpoint = f"{self.valves.NOTDIAMOND_BASE_URL}/modelRouter/modelSelect"

        try:
            async with httpx.AsyncClient(timeout=15.0) as client:
                response = await client.post(endpoint, headers=headers, json=payload)
            response.raise_for_status()
            result = response.json()
        except httpx.HTTPStatusError as e:
            error_text = e.response.text[:300]
            raise Exception(f"NotDiamond API error: {e.response.status_code} - {error_text}") from e
        except httpx.RequestError as e:
            raise Exception(f"Request failed: {str(e)}") from e

        providers = result.get("providers", [])
        if not providers:
            raise Exception(f"NotDiamond API returned empty providers list: {result}")

        selected_provider = providers[0]
        provider_name = selected_provider.get("provider")
        model_name = selected_provider.get("model")

        if not provider_name or not model_name:
            raise Exception(f"NotDiamond API returned invalid provider information: {selected_provider}")

        selected_notdiamond_model = f"{provider_name}/{model_name}"

        selected_openwebui_model = model_mapping.get(selected_notdiamond_model)
        if not selected_openwebui_model:
            raise Exception(f"No OpenWebUI model mapping found for NotDiamond model: {selected_notdiamond_model}")

        return selected_openwebui_model

    async def _generate_with_selected_model(self, body: Dict, selected_model_id: str, __user__: Dict, __request__: Any) -> Any:
        """Generate response using the selected OpenWebUI model with its existing configuration"""
        try:
            new_body = {
                "messages": body["messages"],
                "model": selected_model_id,
                "stream": body.get("stream", True),
            }

            user = Users.get_user_by_id(__user__["id"]) if __user__ else None
            response = await generate_chat_completions(__request__, new_body, user)
            return response
        except Exception as e:
            raise

    async def pipe(
        self,
        body: Dict,
        __user__: Optional[Dict] = None,
        __event_emitter__: Optional[Callable[[Dict], Awaitable[None]]] = None,
        __request__: Optional[Any] = None,
    ) -> Any:
        """Main pipe function that handles the routing logic"""
        try:
            await self._emit_status(__event_emitter__, "info", "🧠 Starting NotDiamond model selection...")
            
            messages = body.get("messages", [])
            if not messages:
                raise Exception("No messages provided in request")
            
            try:
                selected_model_id = await self._select_best_model(messages)
            except Exception as e:
                await self._emit_status(__event_emitter__, "error", f"Model selection failed: {str(e)}", True)
                return {"error": f"NotDiamond routing failed: {str(e)}"}

            selected_model_display_name = selected_model_id
            try:
                model_info = Models.get_model_by_id(selected_model_id)
                if model_info and hasattr(model_info, 'name') and model_info.name:
                    selected_model_display_name = model_info.name
            except Exception:
                pass

            if not __request__:
                error_msg = "Cannot generate response: OpenWebUI request object not available."
                await self._emit_status(__event_emitter__, "error", error_msg, True)
                return {"error": error_msg}

            await self._emit_status(__event_emitter__, "info", f"🚀 Generating with {selected_model_display_name}")
            
            try:
                response = await self._generate_with_selected_model(body, selected_model_id, __user__, __request__)
                await self._emit_status(__event_emitter__, "info", f"Response generated using {selected_model_display_name}", True)
                return response
            except Exception as e:
                error_msg = f"Error generating response: {str(e)}"
                await self._emit_status(__event_emitter__, "error", error_msg, True)
                return {"error": error_msg}

        except Exception as e:
            error_msg = f"Error in NotDiamond router: {str(e)}"
            await self._emit_status(__event_emitter__, "error", error_msg, True)
            return {"error": error_msg}