64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
"""Trusted upstream discovery. Never return provider URLs, keys or raw metadata."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import time
|
|
|
|
import httpx
|
|
|
|
|
|
class CatalogUnavailable(Exception):
|
|
pass
|
|
|
|
|
|
class ModelCatalog:
|
|
def __init__(self, *, base_url=None, api_key=None, client=None):
|
|
self.base_url = (base_url or os.getenv('HERMES_UPSTREAM_URL', 'http://192.168.86.5:2099/v1')).rstrip('/')
|
|
self.api_key = api_key if api_key is not None else os.getenv('HERMES_UPSTREAM_API_KEY', '')
|
|
self.client = client or httpx.AsyncClient(timeout=10, trust_env=False, follow_redirects=False)
|
|
self._owns_client = client is None
|
|
self._cache = None
|
|
self._expires = 0
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def aclose(self):
|
|
if self._owns_client:
|
|
await self.client.aclose()
|
|
|
|
async def models(self):
|
|
async with self._lock:
|
|
if self._cache is not None and time.monotonic() < self._expires:
|
|
return self._cache
|
|
try:
|
|
if not self.api_key:
|
|
raise CatalogUnavailable()
|
|
async with self.client.stream('GET', self.base_url + '/models', headers={'Authorization': 'Bearer ' + self.api_key}) as response:
|
|
response.raise_for_status()
|
|
body = bytearray()
|
|
async for chunk in response.aiter_bytes():
|
|
body.extend(chunk)
|
|
if len(body) > 2 * 1024 * 1024:
|
|
raise CatalogUnavailable()
|
|
import json
|
|
data = json.loads(body).get('data')
|
|
if not isinstance(data, list):
|
|
raise CatalogUnavailable()
|
|
models = []
|
|
seen = set()
|
|
for item in data[:1000]:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
model = item.get('id')
|
|
if not isinstance(model, str) or not 1 <= len(model) <= 200 or not all(33 <= ord(c) <= 126 for c in model) or model in seen:
|
|
continue
|
|
seen.add(model)
|
|
models.append({'id': model, 'name': model, 'thinking_levels': ['none', 'minimal', 'low', 'medium', 'high', 'xhigh']})
|
|
if not models:
|
|
raise CatalogUnavailable()
|
|
self._cache = models
|
|
self._expires = time.monotonic() + 60
|
|
return models
|
|
except (httpx.HTTPError, ValueError, TypeError, AttributeError, CatalogUnavailable):
|
|
raise CatalogUnavailable() from None
|