mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-08-10 11:13:08 +00:00
Merge pull request #104 from Routstr/auto-load-models
auto load models in memory
This commit is contained in:
+74
-16
@@ -2,6 +2,7 @@ import asyncio
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from urllib.request import urlopen
|
||||
|
||||
from fastapi import APIRouter
|
||||
from pydantic.v1 import BaseModel
|
||||
@@ -51,31 +52,88 @@ class Model(BaseModel):
|
||||
MODELS: list[Model] = []
|
||||
|
||||
|
||||
def fetch_openrouter_models(source_filter: str | None = None) -> list[dict]:
|
||||
"""Fetches model information from OpenRouter API."""
|
||||
base_url = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
|
||||
|
||||
try:
|
||||
with urlopen(f"{base_url}/models") as response:
|
||||
data = json.loads(response.read().decode("utf-8"))
|
||||
|
||||
models_data: list[dict] = []
|
||||
for model in data.get("data", []):
|
||||
model_id = model.get("id", "")
|
||||
|
||||
if source_filter:
|
||||
source_prefix = f"{source_filter}/"
|
||||
if not model_id.startswith(source_prefix):
|
||||
continue
|
||||
|
||||
model = dict(model)
|
||||
model["id"] = model_id[len(source_prefix) :]
|
||||
model_id = model["id"]
|
||||
|
||||
if (
|
||||
"(free)" in model.get("name", "")
|
||||
or model_id == "openrouter/auto"
|
||||
or model_id == "google/gemini-2.5-pro-exp-03-25"
|
||||
):
|
||||
continue
|
||||
|
||||
models_data.append(model)
|
||||
|
||||
return models_data
|
||||
except Exception as e:
|
||||
print(f"Error fetching models from OpenRouter API: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def load_models() -> list[Model]:
|
||||
"""Load model definitions from a JSON file.
|
||||
"""Load model definitions from a JSON file or auto-generate from OpenRouter API.
|
||||
|
||||
The file path can be specified via the ``MODELS_PATH`` environment variable.
|
||||
If ``models.json`` is not found, the bundled ``models.example.json`` is used
|
||||
as a fallback. If neither file exists or an error occurs while loading, an
|
||||
empty list is returned.
|
||||
If a user-provided models.json exists, it will be used. Otherwise, models are
|
||||
automatically fetched from OpenRouter API in memory. If the example file exists
|
||||
and no user file is provided, it will be used as a fallback.
|
||||
"""
|
||||
|
||||
models_path = Path(os.environ.get("MODELS_PATH", "models.json"))
|
||||
if not models_path.exists():
|
||||
example = Path(__file__).resolve().parent.parent / "models.example.json"
|
||||
if example.exists():
|
||||
models_path = example
|
||||
else:
|
||||
return []
|
||||
|
||||
try:
|
||||
with models_path.open("r") as f:
|
||||
data = json.load(f)
|
||||
except Exception as e: # pragma: no cover - log and continue
|
||||
print(f"Error loading models from {models_path}: {e}")
|
||||
# Check if user has actively provided a models.json file
|
||||
if models_path.exists():
|
||||
print(f"Loading models from user-provided file: {models_path}")
|
||||
try:
|
||||
with models_path.open("r") as f:
|
||||
data = json.load(f)
|
||||
return [Model(**model) for model in data.get("models", [])]
|
||||
except Exception as e:
|
||||
print(f"Error loading models from {models_path}: {e}")
|
||||
# Fall through to auto-generation
|
||||
|
||||
# Check for example file as fallback
|
||||
example = Path(__file__).resolve().parent.parent.parent / "models.example.json"
|
||||
if example.exists():
|
||||
print(f"Loading models from example file: {example}")
|
||||
try:
|
||||
with example.open("r") as f:
|
||||
data = json.load(f)
|
||||
return [Model(**model) for model in data.get("models", [])]
|
||||
except Exception as e:
|
||||
print(f"Error loading models from {example}: {e}")
|
||||
# Fall through to auto-generation
|
||||
|
||||
# Auto-generate models from OpenRouter API
|
||||
print("Auto-generating models from OpenRouter API")
|
||||
source_filter = os.getenv("SOURCE")
|
||||
source_filter = source_filter if source_filter and source_filter.strip() else None
|
||||
|
||||
models_data = fetch_openrouter_models(source_filter=source_filter)
|
||||
if not models_data:
|
||||
print("Failed to fetch models from OpenRouter API")
|
||||
return []
|
||||
|
||||
return [Model(**model) for model in data.get("models", [])]
|
||||
print(f"Successfully fetched {len(models_data)} models from OpenRouter API")
|
||||
return [Model(**model) for model in models_data]
|
||||
|
||||
|
||||
MODELS = load_models()
|
||||
|
||||
Reference in New Issue
Block a user