hermes-hub/src/antigravity_provider/router/connection_preflight.py

"""Validate supplied credentials before creating any profile or writing auth files."""
import json
import urllib.error
import urllib.request
from urllib.parse import urlsplit

DEFAULT_URLS = {
    "openrouter": "https://openrouter.ai/api/v1",
    "nvidia": "https://integrate.api.nvidia.com/v1",
    "local": "http://127.0.0.1:8081/v1", "vllm": "http://127.0.0.1:8081/v1",
    "ollama": "http://127.0.0.1:11434", "claude": "https://api.anthropic.com/v1",
    "opencode-go": "https://opencode.ai/zen/go/v1", "grok": "https://api.x.ai/v1",
    "openai-codex": "https://api.openai.com/v1",
}


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise ValueError("Перенаправление API запрещено: укажите конечный URL сервера")


def validate_connection(provider, token="", base_url="", preferred_model=""):
    provider = {"nvidia-nim": "nvidia", "local-llm": "local", "llama.cpp": "local"}.get(provider, provider)
    try:
        base_url = (base_url or DEFAULT_URLS.get(provider, "")).rstrip("/")
        parsed = urlsplit(base_url)
        if parsed.scheme not in ("http", "https") or not parsed.hostname or parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment:
            raise ValueError("Укажите HTTP(S) URL сервера без пароля, параметров и фрагмента")
        if provider not in DEFAULT_URLS:
            raise ValueError("Для этого провайдера используйте вход через авторизацию")
        if provider not in ("local", "vllm", "ollama") and not token:
            raise ValueError("Не указан API-ключ")

        # Подписочные аккаунты Codex (ChatGPT) и Claude работают по OAuth, а не
        # по платформенному ключу. Проверка ниже обращается к каталогу моделей,
        # где платформенному ключу нужны права api.model.read: подписочный
        # токен там получает 403 или 401, хотя аккаунт исправен и в Hermes
        # работает. Предупреждаем об этом заранее, чтобы отказ не выглядел
        # поломкой аккаунта.
        subscription_note = ""
        if provider in ("openai-codex", "codex", "claude", "anthropic"):
            subscription_note = (
                " Если это аккаунт по подписке (ChatGPT Plus/Pro или Claude Max), "
                "ключ здесь не подойдёт: такие аккаунты подключаются входом по ссылке "
                "на предыдущем шаге мастера, а не API-ключом."
            )
        headers = {"Content-Type": "application/json"}
        if provider == "claude":
            headers.update({"x-api-key": token, "anthropic-version": "2023-06-01"})
        elif token:
            headers["Authorization"] = "Bearer " + token
        opener = urllib.request.build_opener(_NoRedirect())

        def request(path, body=None):
            req = urllib.request.Request(base_url + path, headers=headers,
                data=json.dumps(body).encode() if body is not None else None)
            with opener.open(req, timeout=20) as response:
                payload = json.load(response)
            if not isinstance(payload, dict) or payload.get("error"):
                raise ValueError("Провайдер вернул ошибку или неверный JSON")
            return payload

        if provider == "openrouter":
            # The catalog is public: its HTTP 200 is not proof of a valid key.
            request("/key")
        if provider == "ollama":
            base_url = base_url.removesuffix("/v1")
            payload, field, key = request("/api/tags"), "models", "name"
        else:
            payload, field, key = request("/models"), "data", "id"
        entries = payload.get(field)
        if not isinstance(entries, list):
            raise ValueError("Провайдер не вернул список моделей")
        models = sorted({m[key] for m in entries if isinstance(m, dict) and isinstance(m.get(key), str) and m[key]})
        if provider == "nvidia":
            if not models:
                raise ValueError("Каталог NVIDIA пуст: проверить ключ тестовым запросом невозможно")
            if preferred_model and preferred_model not in models:
                raise ValueError("Выбранной модели нет в каталоге NVIDIA")
            # Каталог NVIDIA общий для всех, а доступ к конкретной модели даётся
            # по аккаунту. Успешный список уже доказывает, что ключ рабочий:
            # сервер опознал аккаунт и ответил. Пробный запрос — уточнение, а не
            # условие. Раньше его отказ («Function ... Not found for account»)
            # объявлялся провалом подключения, хотя ключ был верным.
            chat_models = [model for model in models if any(word in model.lower() for word in ('instruct', 'chat')) and not any(word in model.lower() for word in ('embed', 'guard', 'reward'))]
            probe_target = preferred_model or (chat_models[0] if chat_models else None)
            if probe_target:
                try:
                    result = request("/chat/completions", {"model": probe_target,
                        "messages": [{"role": "user", "content": "ping"}], "max_tokens": 1})
                    if not isinstance(result.get("choices"), list) or not result["choices"]:
                        probe_note = f"модель {probe_target} не вернула результат"
                    else:
                        probe_note = ""
                except Exception as probe_exc:
                    # 401 и 403 — ключ отвергнут, это отказ подключения.
                    # 404 и прочее — ключ принят, но модель аккаунту не выдана:
                    # NVIDIA отвечает «Function ... Not found for account <id>»,
                    # то есть аккаунт опознан. Валить подключение из-за этого нельзя.
                    code = getattr(probe_exc, "code", None)
                    if code in (401, 403):
                        raise
                    probe_note = f"модель {probe_target} недоступна вашему аккаунту ({str(probe_exc)[:120]})"
            else:
                probe_note = "чат-модель для пробы не найдена"
            if probe_note:
                return {"ok": True, "data": {"models": models, "base_url": base_url},
                        "message": f"Ключ принят, моделей в каталоге: {len(models)}. Проба: {probe_note}. Выберите доступную модель."}
        message = f"Подключено и проверено. Получено моделей: {len(models)}" if models else "Сервер отвечает; моделей пока нет"
        return {"ok": True, "message": message, "data": {"models": models, "base_url": base_url}}
    except Exception as exc:
        exc_str = str(exc)
        exc_lower = exc_str.lower()
        # У Ollama Cloud локального сервера нет вовсе: там только ключ. Если
        # ключ задан, а локальный адрес не отвечает, это не отказ подключения —
        # аккаунт облачный. Проверяем его по общедоступному каталогу и
        # принимаем, честно сказав, что локальные модели недоступны.
        if provider == "ollama" and token and any(
            k in exc_lower for k in ("connection refused", "winerror 10061", "errno 111", "refused", "failed to connect", "target machine actively refused")
        ):
            try:
                cloud = urllib.request.Request(
                    "https://ollama.com/api/tags",
                    headers={"Accept": "application/json", "Authorization": f"Bearer {token}"},
                )
                with urllib.request.urlopen(cloud, timeout=15) as response:
                    payload = json.load(response)
                cloud_models = sorted({
                    m.get("name") for m in (payload.get("models") or [])
                    if isinstance(m, dict) and isinstance(m.get("name"), str)
                })
                return {
                    "ok": True,
                    "message": (
                        f"Облачный аккаунт Ollama принят. Моделей в каталоге: {len(cloud_models)}. "
                        f"Локальный сервер по адресу {base_url} не отвечает — локальные модели недоступны."
                    ),
                    "data": {"models": cloud_models, "base_url": base_url, "cloud_only": True},
                }
            except Exception as cloud_exc:
                return {
                    "ok": False,
                    "message": (
                        f"Локальный сервер {base_url} не отвечает, и облачный каталог Ollama недоступен: "
                        f"{str(cloud_exc)[:160]}"
                    ),
                    "data": {"models": []},
                }
        if provider == "ollama" and any(k in exc_lower for k in ("connection refused", "winerror 10061", "errno 111", "111", "refused", "failed to connect", "target machine actively refused")):
            return {
                "ok": False,
                "message": f"Не удалось подключиться к {base_url} (Connection refused). Если Ollama работает на сервере или другой машине, укажите её сетевой адрес (например, http://192.168.1.81:11434).",
                "data": {"models": []},
            }
        if isinstance(exc, urllib.error.HTTPError):
            reason = (exc.reason or 'провайдер отклонил запрос')
            if exc.code in (401, 403) and subscription_note:
                reason = str(reason) + subscription_note
            try:
                body = json.loads(exc.read(4096).decode('utf-8', errors='replace'))
                detail = body.get('error') or body.get('detail')
                if isinstance(detail, dict):
                    detail = detail.get('message')
                if isinstance(detail, str) and detail.strip():
                    reason = detail[:500]
            except Exception:
                pass
            message = f"HTTP {exc.code}: {reason}"
        else:
            message = str(exc).strip() or type(exc).__name__
        if token:
            message = message.replace(token, "[скрыто]")
        return {"ok": False, "message": message, "data": {"models": []}}