""" title: OpenRouter Cost/Balance Tracking author: TomBayne version: 1.0.0 """ import logging from functools import lru_cache from typing import Dict, List, Optional, Callable import requests from pydantic import BaseModel, Field # Optional token counting via tiktoken; fall back to estimate if unavailable try: import tiktoken # type: ignore _has_tiktoken = True except ImportError: _has_tiktoken = False logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # ---------- Token Counting Helpers ---------- def _count_tokens(text: str, model: str) -> int: if _has_tiktoken: try: enc = tiktoken.encoding_for_model(model) return len(enc.encode(text)) except Exception: pass # Fallback: 1 token ≈ 4 characters return max(1, round(len(text) / 4)) def _message_tokens(messages: List[Dict[str, str]], model: str) -> int: return sum(_count_tokens(m.get("content", ""), model) for m in messages) # ---------- Model Pricing Lookup (cached) ---------- @lru_cache(maxsize=64) def _model_prices(model_id: str) -> tuple[float, float]: try: resp = requests.get("https://openrouter.ai/api/v1/models", timeout=5) resp.raise_for_status() for m in resp.json().get("data", []): if m.get("id") == model_id: p = m.get("pricing", {}) return float(p.get("prompt", 0)), float(p.get("completion", 0)) except Exception as exc: logger.warning("Price lookup failed for %s – %s", model_id, exc) return 0.0, 0.0 # ---------- Filter Definition ---------- class Filter: class Valves(BaseModel): openrouter_api_key: str = Field( "ENTER_KEY", description="Your OpenRouter API key" ) def __init__(self, valves: Optional[dict] = None): self.valves = self.Valves(**(valves or {})) self._messages: List[Dict[str, str]] = [] self._model_id: str = "" def inlet(self, body: Dict, __user__: Optional[dict] = None) -> Dict: # Store messages and model for token counting self._messages = body.get("messages", []) self._model_id = body.get("model", "") return body async def outlet( self, body: Dict, __metadata__: Optional[dict] = None, __user__: Optional[dict] = None, __event_emitter__: Optional[Callable] = None, ) -> Dict: # Count prompt and completion tokens prompt_tokens = _message_tokens(self._messages, self._model_id) completions = [ c.get("message", {}).get("content", "") for c in body.get("choices", []) ] completion_tokens = sum( _count_tokens(text, self._model_id) for text in completions ) # Compute cost prompt_price, completion_price = _model_prices(self._model_id) cost = prompt_tokens * prompt_price + completion_tokens * completion_price # Fetch remaining credits if API key provided remaining: Optional[float] = None key = self.valves.openrouter_api_key if key and key != "ENTER_KEY": try: r = requests.get( "https://openrouter.ai/api/v1/credits", headers={"Authorization": f"Bearer {key}"}, timeout=5, ) r.raise_for_status() data = r.json().get("data", {}) remaining = float(data.get("total_credits", 0)) - float( data.get("total_usage", 0) ) except Exception as exc: logger.warning("Credits lookup failed – %s", exc) # Build banner banner = f"💸 Cost: ${cost:.7f}" if remaining is not None: banner += f" | 💰 Remaining: ${remaining:.7f}" # Emit status bar update if __event_emitter__: await __event_emitter__({"type": "status", "data": {"description": banner}}) prefix = banner + "\n\n" for choice in body.get("choices", []): msg = choice.get("message", {}) content = msg.get("content", "") if not content.startswith("💸 Cost:"): msg["content"] = prefix + content return body