""" title: OpenRouter API author: ross1996 version: 0.1.0 license: MIT """ import os import json import requests import time from typing import List, Union, Generator, Dict from pydantic import BaseModel, Field DEBUG = False class Pipe: class Valves(BaseModel): """Configuration for OpenRouter API.""" OPENROUTER_API_BASE_URL: str = Field(default="https://openrouter.ai/api/v1") OPENROUTER_API_KEY: str = Field(default="") FREE_ONLY: bool = Field(default=False) def __init__(self): self.type = "manifold" self.id = "openrouter" self.name = "openrouter/" self.valves = self.Valves( **{ "OPENROUTER_API_KEY": os.getenv("OPENROUTER_API_KEY", ""), "FREE_ONLY": os.getenv("FREE_ONLY", "false").lower() == "true", } ) def _debug(self, message: str): if DEBUG: print(message) def _get_headers(self) -> Dict[str, str]: if not self.valves.OPENROUTER_API_KEY: raise ValueError("OPENROUTER_API_KEY is not set") return { "Authorization": f"Bearer {self.valves.OPENROUTER_API_KEY}", "Content-Type": "application/json", } def _handle_response(self, response: requests.Response) -> dict: try: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: self._debug(f"HTTPError: {e.response.text}") raise except ValueError as e: self._debug(f"Invalid JSON response: {response.text}") raise def _format_model_id(self, model_id: str) -> str: """Formats the model ID to be compatible with OpenRouter API.""" # Remove both 'openrouter.' and 'openroutermodels.' prefixes if present if model_id.startswith("openrouter."): model_id = model_id[len("openrouter.") :] elif model_id.startswith("openroutermodels."): model_id = model_id[len("openroutermodels.") :] return model_id def get_openrouter_models(self) -> List[Dict[str, str]]: url = f"{self.valves.OPENROUTER_API_BASE_URL}/models" try: self._debug(f"Fetching models from {url}") response = requests.get(url, headers=self._get_headers()) models_data = self._handle_response(response).get("data", []) if self.valves.FREE_ONLY: models_data = [ model for model in models_data if "free" in model.get("id", "").lower() ] return [ { "id": model.get("id", "unknown"), "name": model.get("name", "Unknown Model"), } for model in models_data ] except Exception as e: self._debug(f"Failed to fetch models: {e}") return [{"id": "openrouter", "name": str(e)}] def pipes(self) -> List[dict]: return self.get_openrouter_models() def pipe(self, body: dict) -> Union[str, Generator[str, None, None]]: try: model = self._format_model_id(body["model"]) messages = body["messages"] stream = body.get("stream", False) if DEBUG: self._debug("Incoming body:") self._debug(json.dumps(body, indent=2)) if stream: return self.stream_response(model, messages) return self.get_completion(model, messages) except KeyError as e: error_msg = f"Missing required key in body: {e}" self._debug(error_msg) return f"Error: {error_msg}" except Exception as e: self._debug(f"Error in pipe method: {e}") return f"Error: {e}" def stream_response( self, model: str, messages: List[dict], retries: int = 5 ) -> Generator[str, None, None]: url = f"{self.valves.OPENROUTER_API_BASE_URL}/chat/completions" payload = {"model": model, "messages": messages, "stream": True} self._debug(f"Streaming response from {url}") self._debug(f"Payload: {json.dumps(payload, indent=2)}") for attempt in range(retries): try: response = requests.post( url, json=payload, headers=self._get_headers(), stream=True ) response.raise_for_status() for line in response.iter_lines(): if line: try: line_data = line.decode("utf-8").lstrip("data: ") if line_data == "[DONE]": break event = json.loads(line_data) delta_content = ( event.get("choices", [{}])[0] .get("delta", {}) .get("content") ) if delta_content: yield delta_content except json.JSONDecodeError: self._debug(f"Failed to decode stream line: {line}") continue return except requests.RequestException as e: if response.status_code == 429 and attempt < retries - 1: wait_time = 2**attempt self._debug( f"Rate limited (429). Retrying after {wait_time} seconds..." ) time.sleep(wait_time) else: self._debug(f"Stream request failed: {e}") yield f"Error: {str(e)}" def get_completion(self, model: str, messages: List[dict], retries: int = 3) -> str: url = f"{self.valves.OPENROUTER_API_BASE_URL}/chat/completions" payload = {"model": model, "messages": messages} for attempt in range(retries): try: self._debug( f"Attempt {attempt + 1}: Sending completion request to {url}" ) response = requests.post(url, json=payload, headers=self._get_headers()) data = self._handle_response(response) return data["choices"][0]["message"]["content"] except requests.RequestException as e: if response.status_code == 429 and attempt < retries - 1: wait_time = 2**attempt self._debug( f"Rate limited (429). Retrying after {wait_time} seconds..." ) time.sleep(wait_time) else: self._debug(f"Completion request failed: {e}") return f"Error: {str(e)}"