diff --git a/routstr/proxy.py b/routstr/proxy.py index 55c7aad7..d153ad01 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -162,8 +162,6 @@ async def proxy( return create_error_response( "invalid_model", f"Model '{model_id}' not found", 400, request=request ) - print(model_id) - print(model_obj) upstream = get_provider_for_model(model_id) if not upstream: @@ -174,8 +172,6 @@ async def proxy( request=request, ) - print(upstream.upstream_name) - _max_cost_for_model = await get_max_cost_for_model( model=model_id, session=session, model_obj=model_obj ) diff --git a/routstr/upstream.py b/routstr/upstream.py index e6897915..c1744585 100644 --- a/routstr/upstream.py +++ b/routstr/upstream.py @@ -19,6 +19,7 @@ from .upstreams import ( OpenRouterUpstreamProvider, UpstreamProvider, ) +from .upstreams.generic import GenericUpstreamProvider logger = get_logger(__name__) @@ -453,6 +454,13 @@ def _instantiate_provider(provider_row: UpstreamProviderRow) -> UpstreamProvider return OllamaUpstreamProvider( provider_row.base_url, provider_row.api_key, provider_row.provider_fee ) + elif provider_row.provider_type == "generic": + return GenericUpstreamProvider( + provider_row.base_url, + provider_row.api_key, + provider_row.provider_fee, + provider_row.provider_type, + ) elif provider_row.provider_type == "custom": return UpstreamProvider( provider_row.base_url, provider_row.api_key, provider_row.provider_fee diff --git a/routstr/upstreams/generic.py b/routstr/upstreams/generic.py new file mode 100644 index 00000000..3fc72609 --- /dev/null +++ b/routstr/upstreams/generic.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import httpx + +from .upstream import UpstreamProvider + +if TYPE_CHECKING: + from ..payment.models import Model + +from ..core.logging import get_logger + +logger = get_logger(__name__) + + +class GenericUpstreamProvider(UpstreamProvider): + """Generic upstream provider that can fetch models from any OpenAI-compatible API.""" + + def __init__( + self, + base_url: str, + api_key: str = "", + provider_fee: float = 1.01, + upstream_name: str | None = None, + ): + """Initialize generic provider. + + Args: + base_url: Base URL of the upstream API endpoint + api_key: Optional API key for authentication + provider_fee: Provider fee multiplier (default 1.01 for 1% fee) + upstream_name: Optional name for the upstream provider + """ + self.upstream_name = upstream_name or "generic" + super().__init__( + base_url=base_url, + api_key=api_key, + provider_fee=provider_fee, + ) + + async def fetch_models(self) -> list[Model]: + """Fetch models from upstream API using /models endpoint.""" + from ..payment.models import Architecture, Model, Pricing, TopProvider + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + headers = {} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + response = await client.get(f"{self.base_url}/models", headers=headers) + response.raise_for_status() + data = response.json() + + models_list = [] + for model_data in data.get("data", []): + model_id = model_data.get("id", "") + if not model_id: + continue + + model_name = model_data.get("name", model_id) + created = model_data.get("created", 0) + owned_by = model_data.get("owned_by", "unknown") + model_spec = model_data.get("model_spec", {}) + + context_length = 4096 + if model_spec.get("availableContextTokens"): + context_length = model_spec["availableContextTokens"] + elif any( + pattern in model_id.lower() for pattern in ["32k", "32000"] + ): + context_length = 32768 + elif any( + pattern in model_id.lower() for pattern in ["16k", "16000"] + ): + context_length = 16384 + elif any(pattern in model_id.lower() for pattern in ["8k", "8000"]): + context_length = 8192 + elif "gpt-4" in model_id.lower(): + context_length = 8192 + elif "claude" in model_id.lower(): + context_length = 200000 + + pricing_info = model_spec.get("pricing", {}) + input_pricing = pricing_info.get("input", {}) + output_pricing = pricing_info.get("output", {}) + + prompt_price = input_pricing.get("usd", 0.001) / 1000000 + completion_price = output_pricing.get("usd", 0.001) / 1000000 + + capabilities = model_spec.get("capabilities", {}) + input_modalities = ["text"] + output_modalities = ["text"] + + if capabilities.get("supportsVision", False): + input_modalities.append("image") + + modality = "text" + if capabilities.get("supportsVision", False): + modality = "text->text" + + spec_name = model_spec.get("name", model_name) + description = f"{spec_name}" + if owned_by != "unknown": + description += f" via {owned_by}" + + models_list.append( + Model( + id=model_id, + name=spec_name, + created=created, + description=description, + context_length=context_length, + architecture=Architecture( + modality=modality, + input_modalities=input_modalities, + output_modalities=output_modalities, + tokenizer="unknown", + instruct_type=None, + ), + pricing=Pricing( + prompt=prompt_price, + completion=completion_price, + request=0.0, + image=0.0, + web_search=0.0, + internal_reasoning=0.0, + max_prompt_cost=0.001, + max_completion_cost=0.001, + max_cost=0.001, + ), + sats_pricing=None, + per_request_limits=None, + top_provider=TopProvider( + context_length=context_length, + max_completion_tokens=context_length // 2, + is_moderated=False, + ), + enabled=True, + upstream_provider_id=None, + canonical_slug=None, + ) + ) + + logger.info( + f"Fetched {len(models_list)} models from {self.upstream_name}", + extra={"model_count": len(models_list), "base_url": self.base_url}, + ) + return models_list + + except Exception as e: + logger.error( + f"Failed to fetch models from {self.upstream_name} API: {e}", + extra={ + "error": str(e), + "error_type": type(e).__name__, + "base_url": self.base_url, + }, + ) + return []