diff --git a/src/octop/infra/agents/providers/probe.py b/src/octop/infra/agents/providers/probe.py index 7f4c5aa65..6bb3b9794 100644 --- a/src/octop/infra/agents/providers/probe.py +++ b/src/octop/infra/agents/providers/probe.py @@ -1,341 +1,352 @@ -"""Provider connectivity probes (shared by API and CLI).""" - -from __future__ import annotations - -import asyncio -import json -import logging -import time -from types import SimpleNamespace -from typing import Any - -import httpx - -from octop.infra.agents.providers import KIND_TO_PROTOCOL -from octop.infra.agents.providers.opencode_session import ensure_opencode_session_header - -logger = logging.getLogger(__name__) - -_DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1" -_FETCH_MODELS_TIMEOUT_S = 30.0 -_EMBEDDING_PROBE_TEXT = "ping" - - -def parse_provider_extra(row: Any) -> dict[str, Any]: - """Return the provider ``extra_json`` object, or ``{}`` when missing / invalid.""" - raw = getattr(row, "extra_json", None) - if not isinstance(raw, str) or not raw.strip(): - return {} - try: - extra = json.loads(raw) - except json.JSONDecodeError: - return {} - return extra if isinstance(extra, dict) else {} - - -def provider_headers(row: Any) -> dict[str, str]: - headers = parse_provider_extra(row).get("headers") - return dict(headers) if isinstance(headers, dict) else {} - - -def _is_codex_base_url(base_url: str | None) -> bool: - return bool(base_url and "chatgpt.com/backend-api/codex" in base_url) - - -def build_probe_chat_model(row: Any, *, model_id: str | None = None) -> Any: - """Construct a chat model from a provider row for probing.""" - from octop_harness.config import ModelConfig, ProviderConfig - from octop_harness.llm.factory import build_chat_model - - protocol = KIND_TO_PROTOCOL.get(row.kind, row.kind) - base_url = row.base_url or "https://api.openai.com/v1" - headers = ensure_opencode_session_header(row.name, provider_headers(row), base_url=base_url) - models = row.get_models() if hasattr(row, "get_models") else [] - mid = model_id or (models[0]["id"] if models else "gpt-4o-mini") - entry = next((m for m in models if m.get("id") == mid), None) - display_name = (entry or {}).get("name") or mid - model = ModelConfig(id=mid, name=display_name) - - if _is_codex_base_url(base_url): - from langchain_openai import ChatOpenAI - - kwargs: dict[str, Any] = { - "model": model.id, - "base_url": base_url, - "api_key": row.api_key or "", - "use_responses_api": True, - } - if headers: - kwargs["default_headers"] = dict(headers) - return ChatOpenAI(**kwargs) - - from octop.infra.agents.providers.store import resolve_provider_stream_usage - - provider = ProviderConfig( - id=row.name, - name=row.name, - protocol=protocol, - base_url=base_url, - api_key=row.api_key or "", - headers=headers, - stream_usage=resolve_provider_stream_usage( - provider_name=row.name, - extra=parse_provider_extra(row), - ), - ) - return build_chat_model(provider, model) - - -def make_probe_provider_row( - *, - name: str, - kind: str, - api_key: str | None, - base_url: str | None, - model_id: str, - extra_json: str | None = None, - embedding: bool = False, -) -> Any: - """Build a ProviderRow-like object for connectivity probes.""" - model: dict[str, Any] = {"id": model_id, "name": model_id} - if embedding: - model["embedding"] = True - model["task"] = "embedding" - return SimpleNamespace( - name=name, - kind=kind, - base_url=base_url, - api_key=api_key, - extra_json=extra_json, - get_models=lambda: [model], - ) - - -def _probe_model_id(row: Any, model_id: str | None) -> str: - models = row.get_models() if hasattr(row, "get_models") else [] - if model_id: - return model_id - if models: - return str(models[0].get("id") or "gpt-4o-mini") - return "gpt-4o-mini" - - -def _onnx_probe_model_id(row: Any, model_id: str | None) -> str: - """Explicit request wins, then the row's enabled entry, then the first one. - - ``_probe_model_id`` would take ``models[0]``, which for the ONNX row is - rarely the model the service is actually configured with. - """ - if model_id: - return model_id - models = row.get_models() if hasattr(row, "get_models") else [] - entries = [m for m in models if isinstance(m, dict) and str(m.get("id") or "").strip()] - chosen = next((m for m in entries if m.get("enabled")), None) or ( - entries[0] if entries else None - ) - return str(chosen["id"]) if chosen else "" - - -def _probe_model_entry(row: Any, model_id: str) -> dict[str, Any]: - models = row.get_models() if hasattr(row, "get_models") else [] - entry = next((m for m in models if isinstance(m, dict) and m.get("id") == model_id), None) - return entry if isinstance(entry, dict) else {} - - -def _should_probe_embedding(row: Any, *, model_id: str, embedding: bool | None) -> bool: - if embedding is True: - return True - from octop.infra.agents.providers.model_flags import is_embedding_model - - return is_embedding_model( - _probe_model_entry(row, model_id), - provider_name=getattr(row, "name", None), - provider_api_key=getattr(row, "api_key", None), - ) - - -def _embeddings_url(base_url: str | None) -> str: - root = (base_url or "").strip().rstrip("/") or _DEFAULT_OPENAI_BASE_URL - return f"{root}/embeddings" - - -def _friendly_probe_error(exc: BaseException | str, *, locale: str) -> str: - """Map raw provider exceptions / HTTP bodies to localized guidance when known.""" - from octop.i18n.domains.stream import exception_display_message, stream_error_message - - raw = exception_display_message(exc) - return stream_error_message(raw, locale) or raw - - -async def _probe_embedding_endpoint( - row: Any, *, model_id: str, locale: str = "en" -) -> dict[str, Any]: - """POST OpenAI-compatible ``{base}/embeddings`` and time the round-trip.""" - started = time.perf_counter() - url = _embeddings_url(getattr(row, "base_url", None)) - headers: dict[str, str] = {"Authorization": f"Bearer {getattr(row, 'api_key', None) or ''}"} - extra = ensure_opencode_session_header( - getattr(row, "name", None), - provider_headers(row), - base_url=getattr(row, "base_url", None), - ) - headers.update(extra) - try: - async with httpx.AsyncClient(timeout=_FETCH_MODELS_TIMEOUT_S) as client: - response = await client.post( - url, - headers=headers, - json={"model": model_id, "input": [_EMBEDDING_PROBE_TEXT]}, - ) - except Exception as exc: - logger.info( - "embedding probe failed for %s: %r (%s)", - getattr(row, "name", "?"), - exc, - type(exc).__name__, - ) - return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} - - if response.status_code >= 400: - detail = response.text.strip() - if len(detail) > 300: - detail = detail[:300] + "…" - error = f"HTTP {response.status_code} POST {url}" - if detail: - error = f"{error}: {detail}" - return {"ok": False, "error": _friendly_probe_error(error, locale=locale)} - - try: - payload = response.json() - except Exception: - return { - "ok": False, - "error": "response is not valid JSON (expected OpenAI-compatible /embeddings)", - } - - data = payload.get("data") if isinstance(payload, dict) else None - first = data[0] if isinstance(data, list) and data else None - vector = first.get("embedding") if isinstance(first, dict) else None - if not isinstance(vector, list) or not vector: - return { - "ok": False, - "error": "response is not OpenAI-compatible (expected {data: [{embedding, …}]})", - } - latency_ms = int((time.perf_counter() - started) * 1000) - return {"ok": True, "latency_ms": latency_ms} - - -async def probe_provider_row( - row: Any, - *, - model_id: str | None = None, - embedding: bool | None = None, - locale: str = "en", -) -> dict[str, Any]: - """Probe a provider: chat models get a one-token ping; embedding models POST /embeddings.""" - from octop.infra.agents.providers.model_flags import is_onnx_local_provider - - if is_onnx_local_provider( - getattr(row, "name", None), provider_api_key=getattr(row, "api_key", None) - ): - # Local runtime: probe on-device. It has no base URL, so the remote - # path below would post its placeholder API key to OpenAI. - from octop.infra.agents.providers.onnx_service import probe_local_model - - result = await probe_local_model(_onnx_probe_model_id(row, model_id)) - if result.get("latency_ms") is not None: - result["latency_ms"] = int(result["latency_ms"]) - if not result.get("ok") and result.get("error"): - result["error"] = _friendly_probe_error(str(result["error"]), locale=locale) - return result - - mid = _probe_model_id(row, model_id) - if _should_probe_embedding(row, model_id=mid, embedding=embedding): - return await _probe_embedding_endpoint(row, model_id=mid, locale=locale) - started = time.perf_counter() - try: - chat = build_probe_chat_model(row, model_id=mid) - result = await asyncio.wait_for(chat.ainvoke("ping"), timeout=30.0) - except Exception as exc: - logger.info( - "provider probe failed for %s: %r (%s)", - getattr(row, "name", "?"), - exc, - type(exc).__name__, - ) - return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} - latency_ms = int((time.perf_counter() - started) * 1000) - _ = getattr(result, "content", None) - return {"ok": True, "latency_ms": latency_ms} - - -def _models_list_url(base_url: str | None) -> str: - root = (base_url or "").strip().rstrip("/") or _DEFAULT_OPENAI_BASE_URL - return f"{root}/models" - - -async def fetch_openai_compatible_models( - *, - base_url: str | None, - api_key: str, - extra_headers: dict[str, str] | None = None, - locale: str = "en", - provider_name: str | None = None, -) -> dict[str, Any]: - """List models via OpenAI-compatible ``GET {base}/models``.""" - url = _models_list_url(base_url) - headers: dict[str, str] = {"Authorization": f"Bearer {api_key}"} - extra_headers = ensure_opencode_session_header(provider_name, extra_headers, base_url=base_url) - if extra_headers: - headers.update(extra_headers) - try: - async with httpx.AsyncClient(timeout=_FETCH_MODELS_TIMEOUT_S) as client: - response = await client.get(url, headers=headers) - except Exception as exc: - logger.info( - "provider fetch-models failed for %s: %r (%s)", - url, - exc, - type(exc).__name__, - ) - return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} - - if response.status_code >= 400: - detail = response.text.strip() - if len(detail) > 300: - detail = detail[:300] + "…" - error = f"HTTP {response.status_code}" - if detail: - error = f"{error}: {detail}" - return {"ok": False, "error": _friendly_probe_error(error, locale=locale)} - - try: - payload = response.json() - except Exception: - return { - "ok": False, - "error": "response is not valid JSON (expected OpenAI-compatible /models)", - } - - data = payload.get("data") if isinstance(payload, dict) else None - if not isinstance(data, list): - return { - "ok": False, - "error": "response is not OpenAI-compatible (expected {data: [{id, …}]})", - } - - models: list[dict[str, str]] = [] - seen: set[str] = set() - for item in data: - if not isinstance(item, dict): - continue - mid = item.get("id") - if not isinstance(mid, str) or not mid.strip(): - continue - mid = mid.strip() - if mid in seen: - continue - seen.add(mid) - models.append({"id": mid, "name": mid}) - models.sort(key=lambda m: m["id"]) - return {"ok": True, "models": models} +"""Provider connectivity probes (shared by API and CLI).""" + +from __future__ import annotations + +import asyncio +import json +import logging +import time +from types import SimpleNamespace +from typing import Any + +import httpx + +from octop.infra.agents.providers import KIND_TO_PROTOCOL +from octop.infra.agents.providers.opencode_session import ensure_opencode_session_header + +logger = logging.getLogger(__name__) + +_DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1" +_FETCH_MODELS_TIMEOUT_S = 30.0 +_EMBEDDING_PROBE_TEXT = "ping" + + +def parse_provider_extra(row: Any) -> dict[str, Any]: + """Return the provider ``extra_json`` object, or ``{}`` when missing / invalid.""" + raw = getattr(row, "extra_json", None) + if not isinstance(raw, str) or not raw.strip(): + return {} + try: + extra = json.loads(raw) + except json.JSONDecodeError: + return {} + return extra if isinstance(extra, dict) else {} + + +def provider_headers(row: Any) -> dict[str, str]: + headers = parse_provider_extra(row).get("headers") + return dict(headers) if isinstance(headers, dict) else {} + + +def _is_codex_base_url(base_url: str | None) -> bool: + return bool(base_url and "chatgpt.com/backend-api/codex" in base_url) + + +def build_probe_chat_model(row: Any, *, model_id: str | None = None) -> Any: + """Construct a chat model from a provider row for probing.""" + from octop_harness.config import ModelConfig, ProviderConfig + from octop_harness.llm.factory import build_chat_model + + protocol = KIND_TO_PROTOCOL.get(row.kind, row.kind) + base_url = row.base_url or "https://api.openai.com/v1" + headers = ensure_opencode_session_header(row.name, provider_headers(row), base_url=base_url) + models = row.get_models() if hasattr(row, "get_models") else [] + mid = model_id or (models[0]["id"] if models else "gpt-4o-mini") + entry = next((m for m in models if m.get("id") == mid), None) + display_name = (entry or {}).get("name") or mid + model = ModelConfig(id=mid, name=display_name) + + if _is_codex_base_url(base_url): + from langchain_openai import ChatOpenAI + + kwargs: dict[str, Any] = { + "model": model.id, + "base_url": base_url, + "api_key": row.api_key or "", + "use_responses_api": True, + } + if headers: + kwargs["default_headers"] = dict(headers) + return ChatOpenAI(**kwargs) + + from octop.infra.agents.providers.store import resolve_provider_stream_usage + + provider = ProviderConfig( + id=row.name, + name=row.name, + protocol=protocol, + base_url=base_url, + api_key=row.api_key or "", + headers=headers, + stream_usage=resolve_provider_stream_usage( + provider_name=row.name, + extra=parse_provider_extra(row), + ), + ) + return build_chat_model(provider, model) + + +def make_probe_provider_row( + *, + name: str, + kind: str, + api_key: str | None, + base_url: str | None, + model_id: str, + extra_json: str | None = None, + embedding: bool = False, +) -> Any: + """Build a ProviderRow-like object for connectivity probes.""" + model: dict[str, Any] = {"id": model_id, "name": model_id} + if embedding: + model["embedding"] = True + model["task"] = "embedding" + return SimpleNamespace( + name=name, + kind=kind, + base_url=base_url, + api_key=api_key, + extra_json=extra_json, + get_models=lambda: [model], + ) + + +def _probe_model_id(row: Any, model_id: str | None) -> str: + models = row.get_models() if hasattr(row, "get_models") else [] + if model_id: + return model_id + if models: + return str(models[0].get("id") or "gpt-4o-mini") + return "gpt-4o-mini" + + +def _onnx_probe_model_id(row: Any, model_id: str | None) -> str: + """Explicit request wins, then the row's enabled entry, then the first one. + + ``_probe_model_id`` would take ``models[0]``, which for the ONNX row is + rarely the model the service is actually configured with. + """ + if model_id: + return model_id + models = row.get_models() if hasattr(row, "get_models") else [] + entries = [m for m in models if isinstance(m, dict) and str(m.get("id") or "").strip()] + chosen = next((m for m in entries if m.get("enabled")), None) or ( + entries[0] if entries else None + ) + return str(chosen["id"]) if chosen else "" + + +def _probe_model_entry(row: Any, model_id: str) -> dict[str, Any]: + models = row.get_models() if hasattr(row, "get_models") else [] + entry = next((m for m in models if isinstance(m, dict) and m.get("id") == model_id), None) + return entry if isinstance(entry, dict) else {} + + +def _should_probe_embedding(row: Any, *, model_id: str, embedding: bool | None) -> bool: + if embedding is True: + return True + from octop.infra.agents.providers.model_flags import is_embedding_model + + return is_embedding_model( + _probe_model_entry(row, model_id), + provider_name=getattr(row, "name", None), + provider_api_key=getattr(row, "api_key", None), + ) + + +def _embeddings_url(base_url: str | None) -> str: + root = (base_url or "").strip().rstrip("/") or _DEFAULT_OPENAI_BASE_URL + return f"{root}/embeddings" + + +def _friendly_probe_error(exc: BaseException | str, *, locale: str) -> str: + """Map raw provider exceptions / HTTP bodies to localized guidance when known.""" + from octop.i18n.domains.stream import exception_display_message, stream_error_message + + raw = exception_display_message(exc) + return stream_error_message(raw, locale) or raw + + +async def _probe_embedding_endpoint( + row: Any, *, model_id: str, locale: str = "en" +) -> dict[str, Any]: + """POST OpenAI-compatible ``{base}/embeddings`` and time the round-trip.""" + started = time.perf_counter() + url = _embeddings_url(getattr(row, "base_url", None)) + headers: dict[str, str] = {"Authorization": f"Bearer {getattr(row, 'api_key', None) or ''}"} + extra = ensure_opencode_session_header( + getattr(row, "name", None), + provider_headers(row), + base_url=getattr(row, "base_url", None), + ) + headers.update(extra) + try: + async with httpx.AsyncClient(timeout=_FETCH_MODELS_TIMEOUT_S) as client: + response = await client.post( + url, + headers=headers, + json={"model": model_id, "input": [_EMBEDDING_PROBE_TEXT]}, + ) + except Exception as exc: + logger.info( + "embedding probe failed for %s: %r (%s)", + getattr(row, "name", "?"), + exc, + type(exc).__name__, + ) + return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} + + if response.status_code >= 400: + detail = response.text.strip() + if len(detail) > 300: + detail = detail[:300] + "…" + error = f"HTTP {response.status_code} POST {url}" + if detail: + error = f"{error}: {detail}" + return {"ok": False, "error": _friendly_probe_error(error, locale=locale)} + + try: + payload = response.json() + except Exception: + return { + "ok": False, + "error": "response is not valid JSON (expected OpenAI-compatible /embeddings)", + } + + data = payload.get("data") if isinstance(payload, dict) else None + first = data[0] if isinstance(data, list) and data else None + vector = first.get("embedding") if isinstance(first, dict) else None + if not isinstance(vector, list) or not vector: + return { + "ok": False, + "error": "response is not OpenAI-compatible (expected {data: [{embedding, …}]})", + } + latency_ms = int((time.perf_counter() - started) * 1000) + return {"ok": True, "latency_ms": latency_ms} + + +async def probe_provider_row( + row: Any, + *, + model_id: str | None = None, + embedding: bool | None = None, + locale: str = "en", +) -> dict[str, Any]: + """Probe a provider: chat models get a one-token ping; embedding models POST /embeddings.""" + from octop.infra.agents.providers.model_flags import is_onnx_local_provider + + if is_onnx_local_provider( + getattr(row, "name", None), provider_api_key=getattr(row, "api_key", None) + ): + # Local runtime: probe on-device. It has no base URL, so the remote + # path below would post its placeholder API key to OpenAI. + from octop.infra.agents.providers.onnx_service import probe_local_model + + result = await probe_local_model(_onnx_probe_model_id(row, model_id)) + if result.get("latency_ms") is not None: + result["latency_ms"] = int(result["latency_ms"]) + if not result.get("ok") and result.get("error"): + result["error"] = _friendly_probe_error(str(result["error"]), locale=locale) + return result + + mid = _probe_model_id(row, model_id) + if _should_probe_embedding(row, model_id=mid, embedding=embedding): + return await _probe_embedding_endpoint(row, model_id=mid, locale=locale) + started = time.perf_counter() + try: + chat = build_probe_chat_model(row, model_id=mid) + result = await asyncio.wait_for(chat.ainvoke("ping"), timeout=30.0) + except Exception as exc: + logger.info( + "provider probe failed for %s: %r (%s)", + getattr(row, "name", "?"), + exc, + type(exc).__name__, + ) + return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} + latency_ms = int((time.perf_counter() - started) * 1000) + _ = getattr(result, "content", None) + return {"ok": True, "latency_ms": latency_ms} + + +def _models_list_url(base_url: str | None) -> str: + root = (base_url or "").strip().rstrip("/") or _DEFAULT_OPENAI_BASE_URL + return f"{root}/models" + + +async def fetch_openai_compatible_models( + *, + base_url: str | None, + api_key: str, + extra_headers: dict[str, str] | None = None, + locale: str = "en", + provider_name: str | None = None, +) -> dict[str, Any]: + """List models via OpenAI-compatible ``GET {base}/models``.""" + url = _models_list_url(base_url) + headers: dict[str, str] = {"Authorization": f"Bearer {api_key}"} + extra_headers = ensure_opencode_session_header(provider_name, extra_headers, base_url=base_url) + if extra_headers: + headers.update(extra_headers) + params = {"client_version": "2026.6.9"} if _is_codex_base_url(base_url) else None + try: + async with httpx.AsyncClient(timeout=_FETCH_MODELS_TIMEOUT_S) as client: + response = await client.get(url, headers=headers, params=params) + except Exception as exc: + logger.info( + "provider fetch-models failed for %s: %r (%s)", + url, + exc, + type(exc).__name__, + ) + return {"ok": False, "error": _friendly_probe_error(exc, locale=locale)} + + if response.status_code >= 400: + detail = response.text.strip() + if len(detail) > 300: + detail = detail[:300] + "…" + error = f"HTTP {response.status_code}" + if detail: + error = f"{error}: {detail}" + return {"ok": False, "error": _friendly_probe_error(error, locale=locale)} + + try: + payload = response.json() + except Exception: + return { + "ok": False, + "error": "response is not valid JSON (expected OpenAI-compatible /models)", + } + + if _is_codex_base_url(base_url): + data = payload.get("models") if isinstance(payload, dict) else None + else: + data = payload.get("data") if isinstance(payload, dict) else None + if not isinstance(data, list): + return { + "ok": False, + "error": ( + "response is not Codex-compatible (expected {models: [{slug, display_name, …}]})" + if _is_codex_base_url(base_url) + else "response is not OpenAI-compatible (expected {data: [{id, …}]})" + ), + } + + models: list[dict[str, str]] = [] + seen: set[str] = set() + for item in data: + if not isinstance(item, dict): + continue + mid = item.get("slug") if _is_codex_base_url(base_url) else item.get("id") + if not isinstance(mid, str) or not mid.strip(): + continue + mid = mid.strip() + if mid in seen: + continue + seen.add(mid) + name = item.get("display_name") if _is_codex_base_url(base_url) else mid + models.append( + {"id": mid, "name": name.strip() if isinstance(name, str) and name.strip() else mid} + ) + models.sort(key=lambda m: m["id"]) + return {"ok": True, "models": models} diff --git a/tests/unit/test_provider_fetch_models.py b/tests/unit/test_provider_fetch_models.py index 949d0fee1..24ddb979c 100644 --- a/tests/unit/test_provider_fetch_models.py +++ b/tests/unit/test_provider_fetch_models.py @@ -1,176 +1,223 @@ -"""Unit tests for OpenAI-compatible remote model listing.""" - -from __future__ import annotations - -from typing import Any -from unittest.mock import AsyncMock, patch - -import httpx -import pytest - -from octop.infra.agents.providers.probe import fetch_openai_compatible_models - - -def _mock_response(status: int, payload: Any) -> httpx.Response: - request = httpx.Request("GET", "https://api.example.com/v1/models") - return httpx.Response(status, json=payload, request=request) - - -@pytest.mark.asyncio -async def test_fetch_models_parses_openai_list() -> None: - response = _mock_response( - 200, - {"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini", "owned_by": "system"}]}, - ) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="sk-test", - ) - - assert result["ok"] is True - assert result["models"] == [ - {"id": "gpt-4o", "name": "gpt-4o"}, - {"id": "gpt-4o-mini", "name": "gpt-4o-mini"}, - ] - mock_client.get.assert_awaited_once() - args, kwargs = mock_client.get.await_args - assert args[0] == "https://api.example.com/v1/models" - assert kwargs["headers"]["Authorization"] == "Bearer sk-test" - - -@pytest.mark.asyncio -async def test_fetch_models_defaults_openai_base_url() -> None: - response = _mock_response(200, {"data": []}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models(base_url=None, api_key="sk-test") - - assert result["ok"] is True - assert result["models"] == [] - assert mock_client.get.await_args.args[0] == "https://api.openai.com/v1/models" - - -@pytest.mark.asyncio -async def test_fetch_models_rejects_non_openai_payload() -> None: - response = _mock_response(200, {"models": ["gpt-4o"]}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="sk-test", - ) - - assert result["ok"] is False - assert "OpenAI" in result["error"] or "compatible" in result["error"].lower() - - -@pytest.mark.asyncio -async def test_fetch_models_auth_error() -> None: - response = _mock_response(401, {"error": {"message": "invalid key"}}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="bad", - ) - - assert result["ok"] is False - assert "API key" in result["error"] - assert "401" not in result["error"] - - -@pytest.mark.asyncio -async def test_fetch_models_insufficient_balance_zh() -> None: - response = _mock_response(402, {"error": {"message": "Insufficient Balance"}}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="sk-test", - locale="zh", - ) - - assert result["ok"] is False - assert "余额" in result["error"] or "额度" in result["error"] - assert "402" not in result["error"] - - -@pytest.mark.asyncio -async def test_fetch_models_merges_extra_headers() -> None: - response = _mock_response(200, {"data": [{"id": "m1"}]}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - await fetch_openai_compatible_models( - base_url="https://api.example.com/v1/", - api_key="sk-test", - extra_headers={"X-Custom": "1"}, - ) - - headers = mock_client.get.await_args.kwargs["headers"] - assert headers["Authorization"] == "Bearer sk-test" - assert headers["X-Custom"] == "1" - - -@pytest.mark.asyncio -async def test_fetch_models_connection_error() -> None: - mock_client = AsyncMock() - mock_client.get = AsyncMock(side_effect=httpx.ConnectError("boom")) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="sk-test", - ) - - assert result["ok"] is False - assert result["error"] - # ConnectError string often includes "connection" → timeout_network guidance - assert "401" not in result["error"] - - -@pytest.mark.asyncio -async def test_fetch_models_server_error_friendly() -> None: - response = _mock_response(503, {"error": {"message": "overloaded"}}) - mock_client = AsyncMock() - mock_client.get = AsyncMock(return_value=response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): - result = await fetch_openai_compatible_models( - base_url="https://api.example.com/v1", - api_key="sk-test", - locale="en", - ) - - assert result["ok"] is False - assert "temporarily unavailable" in result["error"].lower() - assert "503" not in result["error"] +"""Unit tests for OpenAI-compatible remote model listing.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from octop.infra.agents.providers.probe import fetch_openai_compatible_models + + +def _mock_response(status: int, payload: Any) -> httpx.Response: + request = httpx.Request("GET", "https://api.example.com/v1/models") + return httpx.Response(status, json=payload, request=request) + + +@pytest.mark.asyncio +async def test_fetch_models_parses_openai_list() -> None: + response = _mock_response( + 200, + {"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini", "owned_by": "system"}]}, + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + ) + + assert result["ok"] is True + assert result["models"] == [ + {"id": "gpt-4o", "name": "gpt-4o"}, + {"id": "gpt-4o-mini", "name": "gpt-4o-mini"}, + ] + mock_client.get.assert_awaited_once() + args, kwargs = mock_client.get.await_args + assert args[0] == "https://api.example.com/v1/models" + assert kwargs["headers"]["Authorization"] == "Bearer sk-test" + + +@pytest.mark.asyncio +async def test_fetch_models_defaults_openai_base_url() -> None: + response = _mock_response(200, {"data": []}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models(base_url=None, api_key="sk-test") + + assert result["ok"] is True + assert result["models"] == [] + assert mock_client.get.await_args.args[0] == "https://api.openai.com/v1/models" + + +@pytest.mark.asyncio +async def test_fetch_models_supports_codex_payload_and_client_version() -> None: + response = _mock_response( + 200, + { + "models": [ + {"slug": "gpt-6.1-sol", "display_name": "GPT-6.1 Sol"}, + {"slug": "gpt-6.1-sol", "display_name": "duplicate"}, + ] + }, + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://chatgpt.com/backend-api/codex", + api_key="codex-token", + ) + + assert result == { + "ok": True, + "models": [{"id": "gpt-6.1-sol", "name": "GPT-6.1 Sol"}], + } + kwargs = mock_client.get.await_args.kwargs + assert kwargs["params"] == {"client_version": "2026.6.9"} + + +@pytest.mark.asyncio +async def test_fetch_models_does_not_add_codex_params_for_openai() -> None: + response = _mock_response(200, {"data": [{"id": "gpt-4o"}]}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + ) + + assert mock_client.get.await_args.kwargs["params"] is None + + +@pytest.mark.asyncio +async def test_fetch_models_rejects_non_openai_payload() -> None: + response = _mock_response(200, {"models": ["gpt-4o"]}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + ) + + assert result["ok"] is False + assert "OpenAI" in result["error"] or "compatible" in result["error"].lower() + + +@pytest.mark.asyncio +async def test_fetch_models_auth_error() -> None: + response = _mock_response(401, {"error": {"message": "invalid key"}}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="bad", + ) + + assert result["ok"] is False + assert "API key" in result["error"] + assert "401" not in result["error"] + + +@pytest.mark.asyncio +async def test_fetch_models_insufficient_balance_zh() -> None: + response = _mock_response(402, {"error": {"message": "Insufficient Balance"}}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + locale="zh", + ) + + assert result["ok"] is False + assert "余额" in result["error"] or "额度" in result["error"] + assert "402" not in result["error"] + + +@pytest.mark.asyncio +async def test_fetch_models_merges_extra_headers() -> None: + response = _mock_response(200, {"data": [{"id": "m1"}]}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + await fetch_openai_compatible_models( + base_url="https://api.example.com/v1/", + api_key="sk-test", + extra_headers={"X-Custom": "1"}, + ) + + headers = mock_client.get.await_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer sk-test" + assert headers["X-Custom"] == "1" + + +@pytest.mark.asyncio +async def test_fetch_models_connection_error() -> None: + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + ) + + assert result["ok"] is False + assert result["error"] + # ConnectError string often includes "connection" → timeout_network guidance + assert "401" not in result["error"] + + +@pytest.mark.asyncio +async def test_fetch_models_server_error_friendly() -> None: + response = _mock_response(503, {"error": {"message": "overloaded"}}) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("octop.infra.agents.providers.probe.httpx.AsyncClient", return_value=mock_client): + result = await fetch_openai_compatible_models( + base_url="https://api.example.com/v1", + api_key="sk-test", + locale="en", + ) + + assert result["ok"] is False + assert "temporarily unavailable" in result["error"].lower() + assert "503" not in result["error"]