diff --git a/router.py b/router.py index 7cf3ada..ab54cc7 100644 --- a/router.py +++ b/router.py @@ -331,6 +331,7 @@ app.add_middleware( ) default_headers={ "HTTP-Referer": "https://nomyo.ai", + "Referer": "https://nomyo.ai", "X-Title": "NOMYO Router", } @@ -797,18 +798,19 @@ class fetch: Internal function that performs the actual HTTP request to fetch available models. This is called by available_models() after checking caches and in-flight requests. """ - headers = None + headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} if api_key is not None: - headers = {"Authorization": "Bearer " + api_key} + headers["Authorization"] = "Bearer " + api_key + ep_base = endpoint.rstrip("/") if endpoint in config.llama_server_endpoints and "/v1" not in endpoint: - endpoint_url = f"{endpoint}/v1/models" + endpoint_url = f"{ep_base}/v1/models" key = "data" elif "/v1" in endpoint or endpoint in config.llama_server_endpoints: - endpoint_url = f"{endpoint}/models" + endpoint_url = f"{ep_base}/models" key = "data" else: - endpoint_url = f"{endpoint}/api/tags" + endpoint_url = f"{ep_base}/api/tags" key = "models" client: aiohttp.ClientSession = get_session(endpoint) @@ -817,13 +819,12 @@ class fetch: await _ensure_success(resp) data = await resp.json() - items = data.get(key, []) - models = {item.get("id") or item.get("name") for item in items if item.get("id") or item.get("name")} + items = data.get(key, []) + models = {item.get("id") or item.get("name") for item in items if item.get("id") or item.get("name")} - # Update cache with lock protection - async with _models_cache_lock: - _models_cache[endpoint] = (models, time.time()) - return models + async with _models_cache_lock: + _models_cache[endpoint] = (models, time.time()) + return models except Exception as e: # Treat any error as if the endpoint offers no models message = _format_connection_issue(endpoint_url, e) @@ -1077,12 +1078,12 @@ class fetch: if _is_fresh(_available_error_cache[endpoint], 300): return [] - client: aiohttp.ClientSession = get_session(endpoint) - headers = None + headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} if api_key is not None: - headers = {"Authorization": "Bearer " + api_key} + headers["Authorization"] = "Bearer " + api_key - request_url = f"{endpoint}{route}" + request_url = f"{endpoint.rstrip('/')}/{route.lstrip('/')}" + client: aiohttp.ClientSession = get_session(endpoint) req_kwargs = {} if timeout is not None: req_kwargs["timeout"] = aiohttp.ClientTimeout(total=timeout) @@ -3984,11 +3985,21 @@ async def startup_event() -> None: ssl_context = ssl.create_default_context() connector = aiohttp.TCPConnector(limit=0, limit_per_host=512, ssl=ssl_context) timeout = aiohttp.ClientTimeout(total=60, connect=15, sock_read=120, sock_connect=15) - session = aiohttp.ClientSession(connector=connector, timeout=timeout) + session = aiohttp.ClientSession( + connector=connector, + timeout=timeout, + headers={"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")}, + ) app_state["connector"] = connector app_state["session"] = session + # Create httpx clients for external OpenAI endpoints (Google, etc.) + # aiohttp strips Referer headers for cross-origin requests, so we use httpx + for ep in config.endpoints: + if is_ext_openai_endpoint(ep): + app_state["httpx_clients"][ep] = httpx.AsyncClient(timeout=30.0) + # Create per-endpoint Unix socket sessions for .sock endpoints for ep in config.llama_server_endpoints: if _is_unix_socket_endpoint(ep):