refactor upstream files and classes

Author: shroominic

routstr/algorithm.py

@@ -6,7 +6,7 @@
 
 if TYPE_CHECKING:
     from .payment.models import Model
-    from .upstream import UpstreamProvider
+    from .upstream import BaseUpstreamProvider
 
 logger = get_logger(__name__)
 
@@ -59,7 +59,7 @@ def calculate_model_cost_score(model: "Model") -> float:
     return total_cost
 
 
-def get_provider_penalty(provider: "UpstreamProvider") -> float:
+def get_provider_penalty(provider: "BaseUpstreamProvider") -> float:
     """Calculate a penalty multiplier for certain providers.
 
     This allows applying policy-based adjustments beyond pure cost.
@@ -86,9 +86,9 @@ def get_provider_penalty(provider: "UpstreamProvider") -> float:
 
 def should_prefer_model(
     candidate_model: "Model",
-    candidate_provider: "UpstreamProvider",
+    candidate_provider: "BaseUpstreamProvider",
     current_model: "Model",
-    current_provider: "UpstreamProvider",
+    current_provider: "BaseUpstreamProvider",
     alias: str,
 ) -> bool:
     """Determine if candidate model should replace current model for an alias.
@@ -166,10 +166,10 @@ def alias_priority(model: "Model") -> int:
 
 
 def create_model_mappings(
-    upstreams: list["UpstreamProvider"],
+    upstreams: list["BaseUpstreamProvider"],
     overrides_by_id: dict[str, tuple],
     disabled_model_ids: set[str],
-) -> tuple[dict[str, "Model"], dict[str, "UpstreamProvider"], dict[str, "Model"]]:
+) -> tuple[dict[str, "Model"], dict[str, "BaseUpstreamProvider"], dict[str, "Model"]]:
     """Create optimal model mappings based on cost and provider preferences.
 
     This is the main entry point for the algorithm. It processes all upstream providers
@@ -196,12 +196,12 @@ def create_model_mappings(
     from .upstream import resolve_model_alias
 
     model_instances: dict[str, "Model"] = {}
-    provider_map: dict[str, "UpstreamProvider"] = {}
+    provider_map: dict[str, "BaseUpstreamProvider"] = {}
     unique_models: dict[str, "Model"] = {}
 
     # Separate OpenRouter from other providers
-    openrouter: "UpstreamProvider" | None = None
-    other_upstreams: list["UpstreamProvider"] = []
+    openrouter: "BaseUpstreamProvider" | None = None
+    other_upstreams: list["BaseUpstreamProvider"] = []
 
     for upstream in upstreams:
         base_url = getattr(upstream, "base_url", "")
@@ -215,7 +215,7 @@ def get_base_model_id(model_id: str) -> str:
         return model_id.split("/", 1)[1] if "/" in model_id else model_id
 
     def _maybe_set_alias(
-        alias: str, model: "Model", provider: "UpstreamProvider"
+        alias: str, model: "Model", provider: "BaseUpstreamProvider"
     ) -> None:
         """Set alias to model/provider if not set or if new model is preferred."""
         existing_model = model_instances.get(alias)
@@ -233,7 +233,7 @@ def _maybe_set_alias(
                 provider_map[alias] = provider
 
     def process_provider_models(
-        upstream: "UpstreamProvider", is_openrouter: bool = False
+        upstream: "BaseUpstreamProvider", is_openrouter: bool = False
     ) -> None:
         """Process all models from a given provider."""
         upstream_prefix = getattr(upstream, "upstream_name", None)

routstr/proxy.py

@@ -23,14 +23,14 @@
     get_max_cost_for_model,
 )
 from .payment.models import Model
-from .upstream import UpstreamProvider, init_upstreams
+from .upstream import BaseUpstreamProvider, init_upstreams
 
 logger = get_logger(__name__)
 proxy_router = APIRouter()
 
-_upstreams: list[UpstreamProvider] = []
+_upstreams: list[BaseUpstreamProvider] = []
 _model_instances: dict[str, Model] = {}  # All aliases -> Model
-_provider_map: dict[str, UpstreamProvider] = {}  # All aliases -> Provider
+_provider_map: dict[str, BaseUpstreamProvider] = {}  # All aliases -> Provider
 _unique_models: dict[str, Model] = {}  # Unique model.id -> Model (no duplicates)
 
 
@@ -53,7 +53,7 @@ async def reinitialize_upstreams() -> None:
     await refresh_model_maps()
 
 
-def get_upstreams() -> list[UpstreamProvider]:
+def get_upstreams() -> list[BaseUpstreamProvider]:
     """Get the initialized upstream providers.
 
     Returns:
@@ -67,7 +67,7 @@ def get_model_instance(model_id: str) -> Model | None:
     return _model_instances.get(model_id)
 
 
-def get_provider_for_model(model_id: str) -> UpstreamProvider | None:
+def get_provider_for_model(model_id: str) -> BaseUpstreamProvider | None:
     """Get UpstreamProvider for model ID from global cache."""
     return _provider_map.get(model_id)
 

routstr/upstream/__init__.py

@@ -0,0 +1,35 @@
+from .anthropic import AnthropicUpstreamProvider
+from .azure import AzureUpstreamProvider
+from .base import BaseUpstreamProvider
+from .generic import GenericUpstreamProvider
+from .helpers import (
+    _instantiate_provider,
+    _seed_providers_from_settings,
+    get_all_models_with_overrides,
+    get_model_with_override,
+    init_upstreams,
+    refresh_upstreams_models_periodically,
+    resolve_model_alias,
+)
+from .ollama import OllamaUpstreamProvider
+from .openai import OpenAIUpstreamProvider
+from .openrouter import OpenRouterUpstreamProvider
+
+__all__ = [
+    # upstreams
+    "AnthropicUpstreamProvider",
+    "AzureUpstreamProvider",
+    "BaseUpstreamProvider",
+    "GenericUpstreamProvider",
+    "OllamaUpstreamProvider",
+    "OpenAIUpstreamProvider",
+    "OpenRouterUpstreamProvider",
+    # helpers
+    "resolve_model_alias",
+    "get_all_models_with_overrides",
+    "get_model_with_override",
+    "refresh_upstreams_models_periodically",
+    "init_upstreams",
+    "_seed_providers_from_settings",
+    "_instantiate_provider",
+]

routstr/upstream/anthropic.py

@@ -0,0 +1,23 @@
+from ..payment.models import Model, async_fetch_openrouter_models
+from .base import BaseUpstreamProvider
+
+
+class AnthropicUpstreamProvider(BaseUpstreamProvider):
+    """Upstream provider specifically configured for Anthropic API."""
+
+    def __init__(self, api_key: str, provider_fee: float = 1.01):
+        self.upstream_name = "anthropic"
+        super().__init__(
+            base_url="https://api.anthropic.com/v1",
+            api_key=api_key,
+            provider_fee=provider_fee,
+        )
+
+    def transform_model_name(self, model_id: str) -> str:
+        """Strip 'anthropic/' prefix for Anthropic API compatibility."""
+        return model_id.removeprefix("anthropic/")
+
+    async def fetch_models(self) -> list[Model]:
+        """Fetch Anthropic models from OpenRouter API filtered by anthropic source."""
+        models_data = await async_fetch_openrouter_models(source_filter="anthropic")
+        return [Model(**model) for model in models_data]  # type: ignore

routstr/upstream/azure.py

@@ -0,0 +1,46 @@
+from typing import Mapping
+
+from .base import BaseUpstreamProvider
+
+
+class AzureUpstreamProvider(BaseUpstreamProvider):
+    """Upstream provider specifically configured for Azure OpenAI Service."""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        api_version: str,
+        provider_fee: float = 1.01,
+    ):
+        """Initialize Azure provider with API key and version.
+
+        Args:
+            base_url: Azure OpenAI endpoint base URL
+            api_key: Azure OpenAI API key for authentication
+            api_version: Azure OpenAI API version (e.g., "2024-02-15-preview")
+            provider_fee: Provider fee multiplier (default 1.01 for 1% fee)
+        """
+        super().__init__(
+            base_url=base_url,
+            api_key=api_key,
+            provider_fee=provider_fee,
+        )
+        self.api_version = api_version
+
+    def prepare_params(
+        self, path: str, query_params: Mapping[str, str] | None
+    ) -> Mapping[str, str]:
+        """Prepare query parameters for Azure OpenAI, adding API version.
+
+        Args:
+            path: Request path
+            query_params: Original query parameters from the client
+
+        Returns:
+            Query parameters dict with Azure API version added for chat completions
+        """
+        params = dict(query_params or {})
+        if path.endswith("chat/completions"):
+            params["api-version"] = self.api_version
+        return params

routstr/upstreams/upstream.py renamed to routstr/upstream/base.py

@@ -25,15 +25,14 @@
     Pricing,
     _calculate_usd_max_costs,
     _update_model_sats_pricing,
-    async_fetch_openrouter_models,
 )
 from ..payment.price import sats_usd_price
 from ..wallet import recieve_token, send_token
 
 logger = get_logger(__name__)
 
 
-class UpstreamProvider:
+class BaseUpstreamProvider:
     """Provider for forwarding requests to an upstream AI service API."""
 
     base_url: str
@@ -1702,111 +1701,3 @@ def get_cached_model_by_id(self, model_id: str) -> Model | None:
             Model object or None if not found
         """
         return self._models_by_id.get(model_id)
-
-
-class OpenAIUpstreamProvider(UpstreamProvider):
-    """Upstream provider specifically configured for OpenAI API."""
-
-    def __init__(self, api_key: str, provider_fee: float = 1.01):
-        self.upstream_name = "openai"
-        super().__init__(
-            base_url="https://api.openai.com/v1",
-            api_key=api_key,
-            provider_fee=provider_fee,
-        )
-
-    def transform_model_name(self, model_id: str) -> str:
-        """Strip 'openai/' prefix for OpenAI API compatibility."""
-        return model_id.removeprefix("openai/")
-
-    async def fetch_models(self) -> list[Model]:
-        """Fetch OpenAI models from OpenRouter API filtered by openai source."""
-        models_data = await async_fetch_openrouter_models(source_filter="openai")
-        return [Model(**model) for model in models_data]  # type: ignore
-
-
-class AnthropicUpstreamProvider(UpstreamProvider):
-    """Upstream provider specifically configured for Anthropic API."""
-
-    def __init__(self, api_key: str, provider_fee: float = 1.01):
-        self.upstream_name = "anthropic"
-        super().__init__(
-            base_url="https://api.anthropic.com/v1",
-            api_key=api_key,
-            provider_fee=provider_fee,
-        )
-
-    def transform_model_name(self, model_id: str) -> str:
-        """Strip 'anthropic/' prefix for Anthropic API compatibility."""
-        return model_id.removeprefix("anthropic/")
-
-    async def fetch_models(self) -> list[Model]:
-        """Fetch Anthropic models from OpenRouter API filtered by anthropic source."""
-        models_data = await async_fetch_openrouter_models(source_filter="anthropic")
-        return [Model(**model) for model in models_data]  # type: ignore
-
-
-class AzureUpstreamProvider(UpstreamProvider):
-    """Upstream provider specifically configured for Azure OpenAI Service."""
-
-    def __init__(
-        self,
-        base_url: str,
-        api_key: str,
-        api_version: str,
-        provider_fee: float = 1.01,
-    ):
-        """Initialize Azure provider with API key and version.
-
-        Args:
-            base_url: Azure OpenAI endpoint base URL
-            api_key: Azure OpenAI API key for authentication
-            api_version: Azure OpenAI API version (e.g., "2024-02-15-preview")
-            provider_fee: Provider fee multiplier (default 1.01 for 1% fee)
-        """
-        super().__init__(
-            base_url=base_url,
-            api_key=api_key,
-            provider_fee=provider_fee,
-        )
-        self.api_version = api_version
-
-    def prepare_params(
-        self, path: str, query_params: Mapping[str, str] | None
-    ) -> Mapping[str, str]:
-        """Prepare query parameters for Azure OpenAI, adding API version.
-
-        Args:
-            path: Request path
-            query_params: Original query parameters from the client
-
-        Returns:
-            Query parameters dict with Azure API version added for chat completions
-        """
-        params = dict(query_params or {})
-        if path.endswith("chat/completions"):
-            params["api-version"] = self.api_version
-        return params
-
-
-class OpenRouterUpstreamProvider(UpstreamProvider):
-    """Upstream provider specifically configured for OpenRouter API."""
-
-    def __init__(self, api_key: str, provider_fee: float = 1.06):
-        """Initialize OpenRouter provider with API key.
-
-        Args:
-            api_key: OpenRouter API key for authentication
-            provider_fee: Provider fee multiplier (default 1.06 for 6% fee)
-        """
-        self.upstream_name = "openrouter"
-        super().__init__(
-            base_url="https://openrouter.ai/api/v1",
-            api_key=api_key,
-            provider_fee=provider_fee,
-        )
-
-    async def fetch_models(self) -> list[Model]:
-        """Fetch all OpenRouter models."""
-        models_data = await async_fetch_openrouter_models()
-        return [Model(**model) for model in models_data]  # type: ignore

routstr/upstreams/generic.py renamed to routstr/upstream/generic.py

@@ -4,7 +4,7 @@
 
 import httpx
 
-from .upstream import UpstreamProvider
+from .base import BaseUpstreamProvider
 
 if TYPE_CHECKING:
     from ..payment.models import Model
@@ -14,7 +14,7 @@
 logger = get_logger(__name__)
 
 
-class GenericUpstreamProvider(UpstreamProvider):
+class GenericUpstreamProvider(BaseUpstreamProvider):
     """Generic upstream provider that can fetch models from any OpenAI-compatible API."""
 
     def __init__(