mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 10:11:46 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(providers): localize OAuth model discovery

Browse Source
This commit is contained in:
Xubin Ren
2026-08-29 21:22:20 +08:00
parent 7941450a5d
commit 1c6483147e
8 changed files with 848 additions and 853 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/providers/github_copilot_provider.py
+174 -4
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
import asyncio
import hashlib
import os
import time
import webbrowser
@@ -17,7 +18,12 @@ from oauth_cli_kit.models import OAuthToken
from oauth_cli_kit.storage import FileTokenStorage
from nanobot.providers.base import LLMResponse, ProviderCallContext
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
OAuthModelCatalogSnapshot,
)
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
from nanobot.providers.registry import ProviderModelSpec, find_by_name
DEFAULT_GITHUB_DEVICE_CODE_URL = "https://github.com/login/device/code"
DEFAULT_GITHUB_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
@@ -96,7 +102,9 @@ def login_github_copilot(
device_code = str(payload["device_code"])
user_code = str(payload["user_code"])
verify_url = str(payload.get("verification_uri") or payload.get("verification_uri_complete") or "")
verify_url = str(
payload.get("verification_uri") or payload.get("verification_uri_complete") or ""
)
verify_complete = str(payload.get("verification_uri_complete") or verify_url)
interval = max(1, int(payload.get("interval") or 5))
expires_in = int(payload.get("expires_in") or 900)
@@ -180,8 +188,6 @@ class GitHubCopilotProvider(OpenAICompatProvider):
*,
provider_name: str = "github_copilot",
):
from nanobot.providers.registry import find_by_name
self._copilot_access_token: str | None = None
self._copilot_expires_at: float = 0.0
self._copilot_token_lock: asyncio.Lock = asyncio.Lock()
@@ -217,7 +223,9 @@ class GitHubCopilotProvider(OpenAICompatProvider):
)
timeout = httpx.Timeout(20.0, connect=20.0)
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True, trust_env=True) as client:
async with httpx.AsyncClient(
timeout=timeout, follow_redirects=True, trust_env=True
) as client:
response = await client.get(
_resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL),
headers=_copilot_headers(github_token.access),
@@ -296,3 +304,165 @@ class GitHubCopilotProvider(OpenAICompatProvider):
on_tool_call_delta=on_tool_call_delta,
provider_context=provider_context,
)
def get_github_copilot_model_catalog(
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
storage = get_storage()
token = storage.load()
account_key = _catalog_account_key(getattr(token, "account_id", None))
cache_key = (
f"{storage.get_token_path()}\0{account_key}\0"
f"{_resolve('NANOBOT_COPILOT_BASE_URL', DEFAULT_COPILOT_BASE_URL)}"
)
return _GITHUB_COPILOT_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_github_copilot_model_catalog() -> None:
_GITHUB_COPILOT_MODEL_CATALOG.invalidate()
def _fetch_github_copilot_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
github_token = get_storage().load()
if not github_token or not github_token.access:
raise RuntimeError("GitHub Copilot is not logged in")
common_headers = {
"Accept": "application/json",
"User-Agent": USER_AGENT,
"Editor-Version": EDITOR_VERSION,
"Editor-Plugin-Version": EDITOR_PLUGIN_VERSION,
}
client_kwargs: dict[str, Any] = {"timeout": 20.0, "follow_redirects": True}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
exchange = client.get(
_resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL),
headers={**common_headers, "Authorization": f"token {github_token.access}"},
)
exchange.raise_for_status()
exchange_mapping = _catalog_mapping(exchange.json())
copilot_token = exchange_mapping.get("token")
if not isinstance(copilot_token, str) or not copilot_token:
raise RuntimeError("GitHub Copilot token exchange returned no token")
endpoint_base = _catalog_first_text(
_catalog_mapping(exchange_mapping.get("endpoints")),
"api",
)
base_url = endpoint_base or _resolve(
"NANOBOT_COPILOT_BASE_URL",
DEFAULT_COPILOT_BASE_URL,
)
models_url = (
base_url
if base_url.rstrip("/").endswith("/models")
else f"{base_url.rstrip('/')}/models"
)
response = client.get(
models_url,
headers={**common_headers, "Authorization": f"Bearer {copilot_token}"},
)
response.raise_for_status()
return _parse_github_copilot_models(response.json())
def _parse_github_copilot_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
rows = cast(dict[str, Any], payload).get("data") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("github_copilot")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
models: list[ProviderModelSpec] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _catalog_first_text(row, "id")
policy = _catalog_mapping(row.get("policy"))
endpoints = row.get("supported_endpoints")
if (
not wire_id
or wire_id in seen
or row.get("model_picker_enabled") is not True
or policy.get("state") == "disabled"
or (
isinstance(endpoints, list)
and "/chat/completions" not in cast(list[object], endpoints)
)
):
continue
seen.add(wire_id)
capabilities = _catalog_mapping(row.get("capabilities"))
supports = _catalog_mapping(capabilities.get("supports"))
limits = _catalog_mapping(capabilities.get("limits"))
fallback = fallback_by_id.get(wire_id)
models.append(
ProviderModelSpec(
id=f"github-copilot/{wire_id}",
label=(
_catalog_first_text(row, "name")
or (fallback.label if fallback is not None else wire_id)
),
description=(fallback.description if fallback is not None else ""),
owned_by="GitHub Copilot",
context_window=(
_catalog_positive_int(limits, "max_context_window_tokens")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_catalog_reasoning_efforts(supports.get("reasoning_effort")),
)
)
return tuple(models)
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_mapping(value: Any) -> dict[str, Any]:
return cast(dict[str, Any], value) if isinstance(value, dict) else {}
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
return tuple(
dict.fromkeys(
item.strip()
for item in cast(list[object], value)
if isinstance(item, str) and item.strip()
)
)
_GITHUB_COPILOT_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("github_copilot"),
fetch=_fetch_github_copilot_models,
)
nanobot/providers/oauth_model_catalog.py
+46 -610
View File
@@ -1,87 +1,51 @@
"""Online model discovery for OAuth providers with bounded local fallback."""
# oauth-cli-kit does not publish type stubs.
# pyright: reportMissingTypeStubs=false
"""Shared cache seam for OAuth provider model discovery."""
from __future__ import annotations
import base64
import hashlib
import json
import os
import threading
import time
from collections.abc import Callable, Sequence
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Literal, Protocol, cast
from typing import Literal
import httpx
from loguru import logger
from nanobot import __version__
DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.6"
DEFAULT_XAI_GROK_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models"
DEFAULT_OPENAI_CODEX_MODELS_URL = "https://chatgpt.com/backend-api/codex/models"
OPENAI_CODEX_CATALOG_CLIENT_VERSION = "0.144.0"
from nanobot.providers.registry import ProviderModelSpec
CatalogSource = Literal["remote", "cache", "stale", "fallback"]
class _XAIToken(Protocol):
@property
def access(self) -> str: ...
@property
def account_id(self) -> str | None: ...
@dataclass(frozen=True)
class OAuthModelInfo:
"""Normalized provider model metadata used by settings and runtimes."""
id: str
label: str
description: str = ""
owned_by: str = ""
context_window: int | None = None
reasoning_efforts: tuple[str, ...] = ()
supports_backend_search: bool = False
@property
def wire_id(self) -> str:
return self.id.split("/", 1)[-1]
@dataclass(frozen=True)
@dataclass(frozen=True, slots=True)
class OAuthModelCatalogSnapshot:
"""One usable catalog view, including where it came from."""
models: tuple[OAuthModelInfo, ...]
models: tuple[ProviderModelSpec, ...]
source: CatalogSource
fetched_at: float
message: str | None = None
def find(self, model: str) -> OAuthModelInfo | None:
def find(self, model: str) -> ProviderModelSpec | None:
wire_id = model.split("/", 1)[-1]
return next((item for item in self.models if item.wire_id == wire_id), None)
return next(
(item for item in self.models if item.id.split("/", 1)[-1] == wire_id),
None,
)
@dataclass(frozen=True)
@dataclass(frozen=True, slots=True)
class _CacheEntry:
snapshot: OAuthModelCatalogSnapshot
stored_at: float
class OAuthModelCatalog:
"""Cache remote discovery behind one thread-safe, failure-tolerant interface."""
"""Cache one provider's discovery behind a small failure-tolerant interface."""
def __init__(
self,
*,
fallback_models: Sequence[OAuthModelInfo],
fetch: Callable[[str | None], Sequence[OAuthModelInfo]],
fallback_models: Sequence[ProviderModelSpec],
fetch: Callable[[str | None], Sequence[ProviderModelSpec]],
fresh_ttl_s: float = 5 * 60,
stale_ttl_s: float = 24 * 60 * 60,
failure_ttl_s: float = 30,
@@ -108,7 +72,7 @@ class OAuthModelCatalog:
self._generation = 0
def get(self, *, cache_key: str, proxy: str | None = None) -> OAuthModelCatalogSnapshot:
"""Return a fresh catalog, sharing concurrent work and failing to a usable list."""
"""Return a fresh catalog, sharing concurrent work and retaining a fallback."""
while True:
with self._condition:
cached = self._cached_result(cache_key)
@@ -127,10 +91,7 @@ class OAuthModelCatalog:
if not models:
raise ValueError("provider returned an empty model catalog")
except Exception as exc:
logger.warning(
"OAuth model catalog refresh failed: type={}",
type(exc).__name__,
)
logger.warning("OAuth model catalog refresh failed: type={}", type(exc).__name__)
with self._condition:
invalidated = generation != self._generation
result = self._failure_result(cache_key) if not invalidated else None
@@ -157,7 +118,7 @@ class OAuthModelCatalog:
return result
def invalidate(self) -> None:
"""Drop cached work and prevent an older account refresh from being stored."""
"""Drop cached work and prevent an older identity refresh from being stored."""
with self._condition:
self._generation += 1
self._entries.clear()
@@ -168,8 +129,7 @@ class OAuthModelCatalog:
entry = self._entries.get(cache_key)
if entry is not None and now - entry.stored_at < self._fresh_ttl_s:
return replace(entry.snapshot, source="cache")
failure_until = self._failures.get(cache_key, 0)
if failure_until > now:
if self._failures.get(cache_key, 0) > now:
return self._stale_or_fallback(entry, now)
return None
@@ -183,9 +143,12 @@ class OAuthModelCatalog:
entry: _CacheEntry | None,
now: float,
) -> OAuthModelCatalogSnapshot:
message = "Could not refresh the online model list; showing cached models."
if entry is not None and now - entry.stored_at < self._stale_ttl_s:
return replace(entry.snapshot, source="stale", message=message)
return replace(
entry.snapshot,
source="stale",
message="Could not refresh the online model list; showing cached models.",
)
return OAuthModelCatalogSnapshot(
models=self._fallback_models,
source="fallback",
@@ -201,569 +164,42 @@ class OAuthModelCatalog:
self._entries[cache_key] = entry
_CURATED_XAI_GROK_MODELS = (
OAuthModelInfo(
id="xai-grok/grok-4.6",
label="Grok 4.6",
description="Grok via xAI subscription; X Search is enabled when supported.",
owned_by="xAI Grok",
context_window=500_000,
),
OAuthModelInfo(
id="xai-grok/grok-4.5",
label="Grok 4.5",
description="Grok via xAI subscription; X Search is enabled when supported.",
owned_by="xAI Grok",
context_window=500_000,
),
)
_CURATED_OPENAI_CODEX_MODELS = (
OAuthModelInfo(
id="openai-codex/gpt-5.6-sol",
label="GPT-5.6-Sol",
description="Latest frontier agentic coding model.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.6-terra",
label="GPT-5.6-Terra",
description="Balanced agentic coding model for everyday work.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.6-luna",
label="GPT-5.6-Luna",
description="Fast and affordable agentic coding model.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.5",
label="GPT-5.5",
description="Frontier model for complex coding, research, and real-world work.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.4",
label="GPT-5.4",
description="Strong model for everyday coding.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.4-mini",
label="GPT-5.4-Mini",
description="Small, fast, and cost-efficient model for simpler coding tasks.",
owned_by="OpenAI Codex",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
OAuthModelInfo(
id="openai-codex/gpt-5.3-codex-spark",
label="GPT-5.3-Codex-Spark",
description="Ultra-fast coding model.",
owned_by="OpenAI Codex",
context_window=128_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
)
_CURATED_GITHUB_COPILOT_MODELS = (
OAuthModelInfo(
id="github-copilot/gpt-4.1",
label="GPT-4.1",
description="GitHub Copilot chat model.",
owned_by="GitHub Copilot",
),
)
def curated_oauth_models(provider_name: str) -> tuple[OAuthModelInfo, ...]:
"""Return stable metadata used only to enrich or backstop online discovery."""
if provider_name == "xai_grok":
return _CURATED_XAI_GROK_MODELS
if provider_name == "openai_codex":
return _CURATED_OPENAI_CODEX_MODELS
if provider_name == "github_copilot":
return _CURATED_GITHUB_COPILOT_MODELS
return ()
def get_oauth_model_catalog(
provider_name: str,
*,
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
"""Discover models for a supported OAuth provider."""
if provider_name == "xai_grok":
cache_key = f"{_xai_oauth_storage_path()}\0{_xai_account_key()}\0{proxy or ''}"
return _XAI_GROK_CATALOG.get(cache_key=cache_key, proxy=proxy)
"""Discover models through the owning provider module."""
if provider_name == "openai_codex":
cache_key = (
f"{_openai_codex_storage_path()}\0{_openai_codex_account_key()}\0{proxy or ''}"
)
return _OPENAI_CODEX_CATALOG.get(cache_key=cache_key, proxy=proxy)
from nanobot.providers.openai_codex_provider import get_openai_codex_model_catalog
return get_openai_codex_model_catalog(proxy)
if provider_name == "xai_grok":
from nanobot.providers.xai_grok_provider import get_xai_grok_model_catalog
return get_xai_grok_model_catalog(proxy)
if provider_name == "github_copilot":
cache_key = (
f"{_github_copilot_storage_path()}\0{_github_copilot_account_key()}\0"
f"{_github_copilot_models_url()}"
)
return _GITHUB_COPILOT_CATALOG.get(cache_key=cache_key, proxy=proxy)
from nanobot.providers.github_copilot_provider import get_github_copilot_model_catalog
return get_github_copilot_model_catalog(proxy)
raise ValueError(f"OAuth model discovery is not available for {provider_name}")
def invalidate_oauth_model_catalog(provider_name: str) -> None:
"""Invalidate provider discovery after OAuth identity changes."""
catalog = _OAUTH_CATALOGS.get(provider_name)
if catalog is not None:
catalog.invalidate()
def _fetch_openai_codex_models(proxy: str | None) -> tuple[OAuthModelInfo, ...]:
from oauth_cli_kit import get_token as get_codex_token
token = get_codex_token(proxy=proxy)
account_id = getattr(token, "account_id", None)
if not isinstance(account_id, str) or not account_id:
raise RuntimeError("OpenAI Codex OAuth token has no account ID")
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_OPENAI_CODEX_MODELS_URL,
params={"client_version": OPENAI_CODEX_CATALOG_CLIENT_VERSION},
headers={
"Authorization": f"Bearer {token.access}",
"chatgpt-account-id": account_id,
"originator": "nanobot",
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
},
"""Invalidate provider discovery after its OAuth identity changes."""
if provider_name == "openai_codex":
from nanobot.providers.openai_codex_provider import (
invalidate_openai_codex_model_catalog,
)
response.raise_for_status()
return _parse_openai_codex_models(response.json())
invalidate_openai_codex_model_catalog()
elif provider_name == "xai_grok":
from nanobot.providers.xai_grok_provider import invalidate_xai_grok_model_catalog
def _parse_openai_codex_models(payload: Any) -> tuple[OAuthModelInfo, ...]:
rows = cast(dict[str, Any], payload).get("models") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
curated = {model.wire_id: model for model in _CURATED_OPENAI_CODEX_MODELS}
parsed: list[tuple[int, OAuthModelInfo]] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _first_text(row, "slug", "id")
if not wire_id or wire_id in seen or row.get("visibility") in {"hide", "none"}:
continue
seen.add(wire_id)
fallback = curated.get(wire_id)
label = _first_text(row, "display_name", "name")
description = _first_text(row, "description")
priority = row.get("priority")
parsed.append(
(
priority if isinstance(priority, int) and not isinstance(priority, bool) else 2**31,
OAuthModelInfo(
id=f"openai-codex/{wire_id}",
label=label or (fallback.label if fallback is not None else wire_id),
description=(
description
or (fallback.description if fallback is not None else "")
),
owned_by="OpenAI Codex",
context_window=(
_positive_int(row, "context_window")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=(
_reasoning_efforts(row.get("supported_reasoning_levels"))
or (fallback.reasoning_efforts if fallback is not None else ())
),
),
)
invalidate_xai_grok_model_catalog()
elif provider_name == "github_copilot":
from nanobot.providers.github_copilot_provider import (
invalidate_github_copilot_model_catalog,
)
parsed.sort(key=lambda item: item[0])
return tuple(model for _, model in parsed)
def _fetch_github_copilot_models(proxy: str | None) -> tuple[OAuthModelInfo, ...]:
from nanobot.providers.github_copilot_provider import (
DEFAULT_COPILOT_TOKEN_URL,
EDITOR_PLUGIN_VERSION,
EDITOR_VERSION,
USER_AGENT,
get_storage,
)
github_token = get_storage().load()
if not github_token or not github_token.access:
raise RuntimeError("GitHub Copilot is not logged in")
common_headers = {
"Accept": "application/json",
"User-Agent": USER_AGENT,
"Editor-Version": EDITOR_VERSION,
"Editor-Plugin-Version": EDITOR_PLUGIN_VERSION,
}
client_kwargs: dict[str, Any] = {"timeout": 20.0, "follow_redirects": True}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
exchange = client.get(
os.environ.get("NANOBOT_COPILOT_TOKEN_URL", "").strip()
or DEFAULT_COPILOT_TOKEN_URL,
headers={**common_headers, "Authorization": f"token {github_token.access}"},
)
exchange.raise_for_status()
exchange_payload = exchange.json()
exchange_mapping = _mapping(exchange_payload)
copilot_token = (
exchange_mapping.get("token") if exchange_mapping else None
)
if not isinstance(copilot_token, str) or not copilot_token:
raise RuntimeError("GitHub Copilot token exchange returned no token")
endpoint_base = _first_text(_mapping(exchange_mapping.get("endpoints")), "api")
if endpoint_base:
models_url = (
endpoint_base
if endpoint_base.rstrip("/").endswith("/models")
else f"{endpoint_base.rstrip('/')}/models"
)
else:
models_url = _github_copilot_models_url()
response = client.get(
models_url,
headers={**common_headers, "Authorization": f"Bearer {copilot_token}"},
)
response.raise_for_status()
return _parse_github_copilot_models(response.json())
def _parse_github_copilot_models(payload: Any) -> tuple[OAuthModelInfo, ...]:
rows = cast(dict[str, Any], payload).get("data") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
curated = {model.wire_id: model for model in _CURATED_GITHUB_COPILOT_MODELS}
models: list[OAuthModelInfo] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _first_text(row, "id")
policy = _mapping(row.get("policy"))
endpoints = row.get("supported_endpoints")
if (
not wire_id
or wire_id in seen
or row.get("model_picker_enabled") is not True
or policy.get("state") == "disabled"
or (
isinstance(endpoints, list)
and "/chat/completions" not in cast(list[object], endpoints)
)
):
continue
seen.add(wire_id)
capabilities = _mapping(row.get("capabilities"))
supports = _mapping(capabilities.get("supports"))
limits = _mapping(capabilities.get("limits"))
fallback = curated.get(wire_id)
models.append(
OAuthModelInfo(
id=f"github-copilot/{wire_id}",
label=(
_first_text(row, "name")
or (fallback.label if fallback is not None else wire_id)
),
description=(fallback.description if fallback is not None else ""),
owned_by="GitHub Copilot",
context_window=(
_positive_int(limits, "max_context_window_tokens")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_reasoning_efforts(supports.get("reasoning_effort")),
)
)
return tuple(models)
def _fetch_xai_grok_models(proxy: str | None) -> tuple[OAuthModelInfo, ...]:
token = _xai_oauth_token(proxy)
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_XAI_GROK_MODELS_URL,
headers=_build_xai_model_headers(token),
)
response.raise_for_status()
return _parse_xai_grok_models(response.json())
def _parse_xai_grok_models(payload: Any) -> tuple[OAuthModelInfo, ...]:
if isinstance(payload, dict):
payload_mapping = cast(dict[str, Any], payload)
rows: object = payload_mapping.get("data")
if not isinstance(rows, list):
rows = payload_mapping.get("models")
else:
rows = payload
if not isinstance(rows, list):
return ()
curated = {model.wire_id: model for model in _CURATED_XAI_GROK_MODELS}
models: list[OAuthModelInfo] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
meta_value = row.get("_meta")
meta = cast(dict[str, Any], meta_value) if isinstance(meta_value, dict) else {}
raw_id = next(
(
candidate.strip()
for candidate in (
row.get("id"),
row.get("model"),
row.get("modelId"),
row.get("name"),
meta.get("id"),
meta.get("model"),
meta.get("modelId"),
)
if isinstance(candidate, str) and candidate.strip()
),
None,
)
if raw_id is None:
continue
wire_id = raw_id.split("/", 1)[-1]
if wire_id in seen:
continue
seen.add(wire_id)
fallback = curated.get(wire_id)
model_id = f"xai-grok/{wire_id}"
label = _first_text(row, "display_name", "label", "name") or _first_text(
meta,
"display_name",
"label",
"name",
)
if not label or label == raw_id:
label = fallback.label if fallback is not None else wire_id
description = _first_text(row, "description") or _first_text(meta, "description")
owner = _first_text(row, "owned_by", "owner", "organization") or _first_text(
meta,
"owned_by",
"owner",
"organization",
)
models.append(
OAuthModelInfo(
id=model_id,
label=label,
description=(
description
or (fallback.description if fallback is not None else "")
),
owned_by=owner or (fallback.owned_by if fallback is not None else "xAI"),
context_window=(
_positive_int(row, "context_window", "context_length")
or _positive_int(meta, "context_window", "context_length")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_reasoning_efforts(
row.get("reasoning_efforts", meta.get("reasoning_efforts"))
),
supports_backend_search=_bool_field(
row,
"supports_backend_search",
"supportsBackendSearch",
),
)
)
return tuple(models)
def _first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _mapping(value: Any) -> dict[str, Any]:
return cast(dict[str, Any], value) if isinstance(value, dict) else {}
def _positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _bool_field(row: dict[str, Any], *keys: str) -> bool:
for key in keys:
value = row.get(key)
if isinstance(value, bool):
return value
meta = row.get("_meta")
if isinstance(meta, dict):
return _bool_field(cast(dict[str, Any], meta), *keys)
return False
def _reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
efforts: list[str] = []
for item in cast(list[object], value):
if isinstance(item, str):
effort = item.strip()
elif isinstance(item, dict):
effort = _first_text(cast(dict[str, Any], item), "effort", "value", "id")
else:
effort = ""
if effort and effort not in efforts:
efforts.append(effort)
return tuple(efforts)
def _xai_oauth_storage_path() -> Path:
from nanobot.providers.xai_oauth import get_xai_oauth_storage_path
return get_xai_oauth_storage_path()
def _account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _xai_account_key() -> str:
from nanobot.providers.xai_oauth import get_xai_oauth_login_status
token = get_xai_oauth_login_status()
return _account_key(getattr(token, "account_id", None))
def _openai_codex_storage_path() -> Path:
from oauth_cli_kit.providers import OPENAI_CODEX_PROVIDER
from oauth_cli_kit.storage import FileTokenStorage
return FileTokenStorage(token_filename=OPENAI_CODEX_PROVIDER.token_filename).get_token_path()
def _openai_codex_account_key() -> str:
from oauth_cli_kit.providers import OPENAI_CODEX_PROVIDER
from oauth_cli_kit.storage import FileTokenStorage
token = FileTokenStorage(token_filename=OPENAI_CODEX_PROVIDER.token_filename).load()
return _account_key(getattr(token, "account_id", None))
def _github_copilot_storage_path() -> Path:
from nanobot.providers.github_copilot_provider import get_storage
return get_storage().get_token_path()
def _github_copilot_account_key() -> str:
from nanobot.providers.github_copilot_provider import get_storage
token = get_storage().load()
return _account_key(getattr(token, "account_id", None))
def _github_copilot_models_url() -> str:
from nanobot.providers.github_copilot_provider import DEFAULT_COPILOT_BASE_URL
base_url = (
os.environ.get("NANOBOT_COPILOT_BASE_URL", "").strip()
or DEFAULT_COPILOT_BASE_URL
)
return f"{base_url.rstrip('/')}/models"
def _xai_oauth_token(proxy: str | None) -> _XAIToken:
from nanobot.providers.xai_oauth import get_xai_oauth_token
return get_xai_oauth_token(proxy=proxy)
def _build_xai_model_headers(token: _XAIToken) -> dict[str, str]:
from nanobot.providers.xai_oauth import XAI_CLIENT_VERSION
headers = {
"Authorization": f"Bearer {token.access}",
"X-XAI-Token-Auth": "xai-grok-cli",
"x-grok-client-version": XAI_CLIENT_VERSION,
"x-grok-client-identifier": "nanobot",
"x-grok-client-mode": "headless",
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
}
claims = _decode_access_token_claims(token.access)
user_id = claims.get("sub")
if claims.get("principal_type") == "Team":
user_id = claims.get("principal_id") or user_id
if isinstance(user_id, str) and user_id:
headers["x-userid"] = user_id
email = claims.get("email")
if not isinstance(email, str) or "@" not in email:
email = token.account_id if token.account_id and "@" in token.account_id else None
if email:
headers["x-email"] = email
return headers
def _decode_access_token_claims(token: str) -> dict[str, Any]:
parts = token.split(".")
if len(parts) < 2 or not parts[1]:
return {}
try:
decoded = base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))
claims = json.loads(decoded)
except (ValueError, TypeError):
return {}
return cast(dict[str, Any], claims) if isinstance(claims, dict) else {}
_XAI_GROK_CATALOG = OAuthModelCatalog(
fallback_models=_CURATED_XAI_GROK_MODELS,
fetch=_fetch_xai_grok_models,
)
_OPENAI_CODEX_CATALOG = OAuthModelCatalog(
fallback_models=_CURATED_OPENAI_CODEX_MODELS,
fetch=_fetch_openai_codex_models,
)
_GITHUB_COPILOT_CATALOG = OAuthModelCatalog(
fallback_models=_CURATED_GITHUB_COPILOT_MODELS,
fetch=_fetch_github_copilot_models,
)
_OAUTH_CATALOGS = {
"xai_grok": _XAI_GROK_CATALOG,
"openai_codex": _OPENAI_CODEX_CATALOG,
"github_copilot": _GITHUB_COPILOT_CATALOG,
}
invalidate_github_copilot_model_catalog()
nanobot/providers/openai_codex_provider.py
+173 -28
View File
@@ -14,7 +14,10 @@ from typing import Any, cast
import httpx
from loguru import logger
from oauth_cli_kit import get_token as get_codex_token
from oauth_cli_kit.providers import OPENAI_CODEX_PROVIDER
from oauth_cli_kit.storage import FileTokenStorage
from nanobot import __version__
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
@@ -22,6 +25,10 @@ from nanobot.providers.base import (
ProviderConversationState,
resolve_stream_idle_timeout_s,
)
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
OAuthModelCatalogSnapshot,
)
from nanobot.providers.openai_responses import (
ResponsesStreamCapture,
build_responses_state,
@@ -35,8 +42,11 @@ from nanobot.providers.openai_responses import (
responses_state_items,
responses_state_matches,
)
from nanobot.providers.registry import ProviderModelSpec, find_by_name
DEFAULT_CODEX_URL = "https://chatgpt.com/backend-api/codex/responses"
DEFAULT_OPENAI_CODEX_MODELS_URL = "https://chatgpt.com/backend-api/codex/models"
OPENAI_CODEX_CATALOG_CLIENT_VERSION = "0.144.0"
DEFAULT_ORIGINATOR = "nanobot"
_COMPACTION_RETAINED_CHAR_BUDGET = 256_000
@@ -87,9 +97,7 @@ class OpenAICodexProvider(LLMProvider):
model = model or self.default_model
sanitized_messages = self._sanitize_empty_content(messages)
sanitized_state = (
provider_context.conversation_state
if provider_context is not None
else None
provider_context.conversation_state if provider_context is not None else None
)
if sanitized_state is not None:
sanitized_state = sanitized_state.with_pending_messages(
@@ -168,11 +176,7 @@ class OpenAICodexProvider(LLMProvider):
)
compact_threshold = resolve_compact_threshold(
(
provider_context.context_window_tokens
if provider_context is not None
else None
),
(provider_context.context_window_tokens if provider_context is not None else None),
max_tokens,
)
if (
@@ -236,8 +240,12 @@ class OpenAICodexProvider(LLMProvider):
return response
async def chat(
self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None,
provider_context: ProviderCallContext | None = None,
@@ -264,8 +272,12 @@ class OpenAICodexProvider(LLMProvider):
)
async def chat_stream(
self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None,
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
@@ -344,11 +356,7 @@ def _without_response_item_ids(
sanitized_input.append(raw_item)
continue
item = cast(dict[str, Any], raw_item)
sanitized_input.append({
key: value
for key, value in item.items()
if key != "id"
})
sanitized_input.append({key: value for key, value in item.items() if key != "id"})
body = dict(request_body)
body["input"] = sanitized_input
@@ -444,15 +452,12 @@ async def _request_codex(
raw = text.decode("utf-8", "ignore")
retry_after = LLMProvider._extract_retry_after_from_headers(response.headers)
error_type, error_code = LLMProvider._extract_error_type_code(raw)
compaction_unsupported = (
response.status_code in {400, 404, 422}
and any(
marker in raw.lower()
for marker in (
"context_management",
"compact_threshold",
"compaction_trigger",
)
compaction_unsupported = response.status_code in {400, 404, 422} and any(
marker in raw.lower()
for marker in (
"context_management",
"compact_threshold",
"compaction_trigger",
)
)
raise _CodexHTTPError(
@@ -461,7 +466,9 @@ async def _request_codex(
retry_after=retry_after,
error_type=error_type,
error_code=error_code,
should_retry=_should_retry_status(response.status_code, error_type, error_code, raw),
should_retry=_should_retry_status(
response.status_code, error_type, error_code, raw
),
compaction_unsupported=compaction_unsupported,
)
capture = ResponsesStreamCapture()
@@ -534,7 +541,9 @@ def _codex_error_response(exc: Exception) -> LLMResponse:
default_detail = "HTTP request failed"
if status_code is not None and should_retry is None:
retry_content = None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
retry_content = (
None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
)
should_retry = _should_retry_status(
int(status_code),
getattr(exc, "error_type", None),
@@ -592,3 +601,139 @@ def _should_retry_status(
)
)
return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500
def get_openai_codex_model_catalog(
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
storage = FileTokenStorage(token_filename=OPENAI_CODEX_PROVIDER.token_filename)
token = storage.load()
account_id = getattr(token, "account_id", None)
account_key = _catalog_account_key(account_id)
cache_key = f"{storage.get_token_path()}\0{account_key}\0{proxy or ''}"
return _OPENAI_CODEX_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_openai_codex_model_catalog() -> None:
_OPENAI_CODEX_MODEL_CATALOG.invalidate()
def _fetch_openai_codex_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
token = get_codex_token(proxy=proxy)
account_id = getattr(token, "account_id", None)
if not isinstance(account_id, str) or not account_id:
raise RuntimeError("OpenAI Codex OAuth token has no account ID")
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_OPENAI_CODEX_MODELS_URL,
params={"client_version": OPENAI_CODEX_CATALOG_CLIENT_VERSION},
headers={
"Authorization": f"Bearer {token.access}",
"chatgpt-account-id": account_id,
"originator": DEFAULT_ORIGINATOR,
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
},
)
response.raise_for_status()
return _parse_openai_codex_models(response.json())
def _parse_openai_codex_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
rows = cast(dict[str, Any], payload).get("models") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("openai_codex")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
parsed: list[tuple[int, ProviderModelSpec]] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _catalog_first_text(row, "slug", "id")
if not wire_id or wire_id in seen or row.get("visibility") in {"hide", "none"}:
continue
seen.add(wire_id)
fallback = fallback_by_id.get(wire_id)
priority = row.get("priority")
parsed.append(
(
priority if isinstance(priority, int) and not isinstance(priority, bool) else 2**31,
ProviderModelSpec(
id=f"openai-codex/{wire_id}",
label=(
_catalog_first_text(row, "display_name", "name")
or (fallback.label if fallback is not None else wire_id)
),
description=(
_catalog_first_text(row, "description")
or (fallback.description if fallback is not None else "")
),
owned_by="OpenAI Codex",
context_window=(
_catalog_positive_int(row, "context_window")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=(
_catalog_reasoning_efforts(row.get("supported_reasoning_levels"))
or (fallback.reasoning_efforts if fallback is not None else ())
),
),
)
)
parsed.sort(key=lambda item: item[0])
return tuple(model for _, model in parsed)
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
efforts: list[str] = []
for item in cast(list[object], value):
if isinstance(item, str):
effort = item.strip()
elif isinstance(item, dict):
effort = _catalog_first_text(cast(dict[str, Any], item), "effort", "value", "id")
else:
effort = ""
if effort and effort not in efforts:
efforts.append(effort)
return tuple(efforts)
_OPENAI_CODEX_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("openai_codex"),
fetch=_fetch_openai_codex_models,
)
nanobot/providers/registry.py
+70 -32
View File
@@ -17,8 +17,6 @@ from typing import Any
from pydantic.alias_generators import to_snake
from nanobot.providers.oauth_model_catalog import curated_oauth_models
@dataclass(frozen=True)
class ProviderModelSpec:
@@ -27,7 +25,10 @@ class ProviderModelSpec:
id: str
label: str = ""
description: str = ""
owned_by: str = ""
context_window: int | None = None
reasoning_efforts: tuple[str, ...] = ()
supports_backend_search: bool = False
@dataclass(frozen=True)
@@ -152,7 +153,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
backend="openai_compat",
is_direct=True,
),
# === Azure OpenAI (direct API calls with API version 2024-10-21) =====
ProviderSpec(
name="azure_openai",
@@ -315,7 +315,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
detect_by_base_keyword="siliconflow",
default_api_base="https://api.siliconflow.cn/v1",
),
# Novita AI: OpenAI-compatible gateway for hosted model APIs.
ProviderSpec(
name="novita",
@@ -327,7 +326,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
detect_by_base_keyword="novita",
default_api_base="https://api.novita.ai/openai",
),
# VolcEngine (火山引擎): OpenAI-compatible gateway, pay-per-use models
ProviderSpec(
name="volcengine",
@@ -341,7 +339,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
thinking_style="thinking_type",
supports_max_completion_tokens=True,
),
# VolcEngine Coding Plan (火山引擎 Coding Plan): same key as volcengine
ProviderSpec(
name="volcengine_coding_plan",
@@ -355,7 +352,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
thinking_style="thinking_type",
supports_max_completion_tokens=True,
),
# BytePlus: VolcEngine international, pay-per-use models
ProviderSpec(
name="byteplus",
@@ -369,7 +365,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
strip_model_prefix=True,
thinking_style="thinking_type",
),
# BytePlus Coding Plan: same key as byteplus
ProviderSpec(
name="byteplus_coding_plan",
@@ -382,8 +377,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
strip_model_prefix=True,
thinking_style="thinking_type",
),
# === Standard providers (matched by model-name keywords) ===============
# Anthropic: native Anthropic SDK
ProviderSpec(
@@ -410,14 +403,56 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
env_key="",
display_name="OpenAI Codex",
model_catalog="hybrid",
builtin_models=tuple(
builtin_models=(
ProviderModelSpec(
id=model.id,
label=model.label,
description=model.description,
context_window=model.context_window,
)
for model in curated_oauth_models("openai_codex")
id="openai-codex/gpt-5.6-sol",
label="GPT-5.6-Sol",
description="Latest frontier agentic coding model.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.6-terra",
label="GPT-5.6-Terra",
description="Balanced agentic coding model for everyday work.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.6-luna",
label="GPT-5.6-Luna",
description="Fast and affordable agentic coding model.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.5",
label="GPT-5.5",
description="Frontier model for complex coding, research, and real-world work.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.4",
label="GPT-5.4",
description="Strong model for everyday coding.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.4-mini",
label="GPT-5.4-Mini",
description="Small, fast, and cost-efficient model for simpler coding tasks.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.3-codex-spark",
label="GPT-5.3-Codex-Spark",
description="Ultra-fast coding model.",
context_window=128_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
),
backend="openai_codex",
detect_by_base_keyword="codex",
@@ -431,14 +466,19 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
env_key="",
display_name="xAI Grok",
model_catalog="hybrid",
builtin_models=tuple(
builtin_models=(
ProviderModelSpec(
id=model.id,
label=model.label,
description=model.description,
context_window=model.context_window,
)
for model in curated_oauth_models("xai_grok")
id="xai-grok/grok-4.6",
label="Grok 4.6",
description="Grok via xAI subscription; X Search is enabled when supported.",
context_window=500_000,
),
ProviderModelSpec(
id="xai-grok/grok-4.5",
label="Grok 4.5",
description="Grok via xAI subscription; X Search is enabled when supported.",
context_window=500_000,
),
),
backend="xai_grok",
default_api_base="https://cli-chat-proxy.grok.com/v1",
@@ -451,14 +491,12 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
env_key="",
display_name="Github Copilot",
model_catalog="hybrid",
builtin_models=tuple(
builtin_models=(
ProviderModelSpec(
id=model.id,
label=model.label,
description=model.description,
context_window=model.context_window,
)
for model in curated_oauth_models("github_copilot")
id="github-copilot/gpt-4.1",
label="GPT-4.1",
description="GitHub Copilot chat model.",
),
),
backend="github_copilot",
default_api_base="https://api.githubcopilot.com",
@@ -730,7 +768,7 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
env_key="QIANFAN_API_KEY",
display_name="Qianfan",
backend="openai_compat",
default_api_base="https://qianfan.baidubce.com/v2"
default_api_base="https://qianfan.baidubce.com/v2",
),
)
nanobot/providers/xai_grok_provider.py
+237 -48
View File
@@ -3,6 +3,8 @@
from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import re
import uuid
@@ -20,21 +22,23 @@ from nanobot.providers.base import (
ToolCallRequest,
resolve_stream_idle_timeout_s,
)
from nanobot.providers.oauth_model_catalog import (
DEFAULT_XAI_GROK_MODEL,
get_oauth_model_catalog,
)
from nanobot.providers.oauth_model_catalog import OAuthModelCatalog, OAuthModelCatalogSnapshot
from nanobot.providers.openai_responses import (
consume_sse_with_reasoning,
convert_messages,
convert_tools,
)
from nanobot.providers.registry import ProviderModelSpec, find_by_name
from nanobot.providers.xai_oauth import (
XAI_CLIENT_VERSION,
get_xai_oauth_login_status,
get_xai_oauth_storage_path,
get_xai_oauth_token,
)
DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.6"
DEFAULT_XAI_GROK_URL = "https://cli-chat-proxy.grok.com/v1/responses"
DEFAULT_XAI_GROK_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models"
_HOSTED_SEARCH_MAX_TURNS = 5
_MAX_ERROR_BODY_CHARS = 1000
_SENSITIVE_ERROR_KEYS = {
@@ -81,9 +85,8 @@ class XAIGrokProvider(LLMProvider):
async def _supports_backend_search(self, model: str) -> bool:
catalog = await asyncio.to_thread(
get_oauth_model_catalog,
"xai_grok",
proxy=self.proxy,
get_xai_grok_model_catalog,
self.proxy,
)
if catalog.message:
logger.warning(
@@ -115,12 +118,8 @@ class XAIGrokProvider(LLMProvider):
token = await asyncio.to_thread(get_xai_oauth_token, proxy=self.proxy)
configured_tools = self._extra_body.get("tools")
tools_are_explicit = "tools" in self._extra_body
configured_hosted_search = (
isinstance(configured_tools, list)
and any(
_is_hosted_x_search_tool(tool)
for tool in cast(list[object], configured_tools)
)
configured_hosted_search = isinstance(configured_tools, list) and any(
_is_hosted_x_search_tool(tool) for tool in cast(list[object], configured_tools)
)
supports_backend_search = False
if not tools_are_explicit:
@@ -159,11 +158,9 @@ class XAIGrokProvider(LLMProvider):
# stopping after a single unsuccessful lookup.
body["max_turns"] = _HOSTED_SEARCH_MAX_TURNS
if self._extra_body:
body.update({
key: value
for key, value in self._extra_body.items()
if key != "tools"
})
body.update(
{key: value for key, value in self._extra_body.items() if key != "tools"}
)
if tools_are_explicit and not isinstance(configured_tools, list):
body["tools"] = configured_tools
@@ -198,18 +195,14 @@ class XAIGrokProvider(LLMProvider):
stage = "xai_request_after_oauth_refresh"
except _XAIIncompleteHostedToolError as exc:
retry_usage = _combine_usage(retry_usage, exc.usage)
cannot_recover_stream = (
exc.stream_output_emitted
and on_stream_recover is None
)
cannot_recover_stream = exc.stream_output_emitted and on_stream_recover is None
if hosted_tool_retried or cannot_recover_stream:
exc.usage = retry_usage
raise
hosted_tool_retried = True
stage = "hosted_tool_recovery"
logger.warning(
"xAI response ended with unfinished hosted tool(s): {}; "
"retrying once",
"xAI response ended with unfinished hosted tool(s): {}; retrying once",
", ".join(exc.tool_names),
)
if on_stream_recover is not None:
@@ -359,13 +352,9 @@ class _XAIIncompleteHostedToolError(RuntimeError):
usage: LLMUsage | None,
stream_output_emitted: bool = False,
) -> None:
names = [
str(event.get("name") or "hosted_tool")
for event in active_tools
]
names = [str(event.get("name") or "hosted_tool") for event in active_tools]
super().__init__(
"xAI ended the response before its hosted tool completed: "
+ ", ".join(names)
"xAI ended the response before its hosted tool completed: " + ", ".join(names)
)
self.tool_names = tuple(names)
self.usage = usage
@@ -427,9 +416,7 @@ async def _request_xai(
raise _build_xai_http_error(response.status_code, response.headers, raw)
result = await consume_sse_with_reasoning(
response,
on_content_delta=(
_forward_content_delta if on_content_delta is not None else None
),
on_content_delta=(_forward_content_delta if on_content_delta is not None else None),
# Always observe tool events so protocol validation also works for
# non-streaming callers that did not request UI progress callbacks.
on_tool_call_delta=_track_and_forward_tool_event,
@@ -441,12 +428,14 @@ async def _request_xai(
if result[2] != "error" and active_hosted_tools:
active = list(active_hosted_tools.values())
for event in active:
await _track_and_forward_tool_event({
**event,
"phase": "error",
"result": None,
"error": "xAI ended the response before this hosted tool completed.",
})
await _track_and_forward_tool_event(
{
**event,
"phase": "error",
"result": None,
"error": "xAI ended the response before this hosted tool completed.",
}
)
raise _XAIIncompleteHostedToolError(
active,
usage=result[3],
@@ -466,9 +455,7 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "start",
"call_id": str(call_id),
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(
event.get("input", event.get("arguments"))
),
"arguments": _xai_hosted_tool_arguments(event.get("input", event.get("arguments"))),
"result": None,
}
@@ -491,9 +478,7 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(item.get("action")),
"result": (
{"status": str(item.get("status") or "completed")}
if phase == "end"
else None
{"status": str(item.get("status") or "completed")} if phase == "end" else None
),
}
if event_type != "response.output_item.done" or item_type != "custom_tool_call":
@@ -509,9 +494,7 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "end",
"call_id": str(call_id),
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(
item.get("input", item.get("arguments"))
),
"arguments": _xai_hosted_tool_arguments(item.get("input", item.get("arguments"))),
# Keep the useful search subtype, but do not persist large hosted results
# in WebUI activity messages. The model answer already carries citations.
"result": {"name": tool_name},
@@ -662,3 +645,209 @@ def _should_retry_status(
)
)
return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500 # pyright: ignore[reportPrivateUsage]
def get_xai_grok_model_catalog(proxy: str | None = None) -> OAuthModelCatalogSnapshot:
token = get_xai_oauth_login_status()
account_key = _catalog_account_key(getattr(token, "account_id", None))
cache_key = f"{get_xai_oauth_storage_path()}\0{account_key}\0{proxy or ''}"
return _XAI_GROK_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_xai_grok_model_catalog() -> None:
_XAI_GROK_MODEL_CATALOG.invalidate()
def _fetch_xai_grok_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
token = get_xai_oauth_token(proxy=proxy)
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_XAI_GROK_MODELS_URL,
headers=_build_xai_model_headers(token.access, token.account_id),
)
response.raise_for_status()
return _parse_xai_grok_models(response.json())
def _parse_xai_grok_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
if isinstance(payload, dict):
payload_mapping = cast(dict[str, Any], payload)
rows: object = payload_mapping.get("data")
if not isinstance(rows, list):
rows = payload_mapping.get("models")
else:
rows = payload
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("xai_grok")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
models: list[ProviderModelSpec] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
meta = _catalog_mapping(row.get("_meta"))
raw_id = next(
(
candidate.strip()
for candidate in (
row.get("id"),
row.get("model"),
row.get("modelId"),
row.get("name"),
meta.get("id"),
meta.get("model"),
meta.get("modelId"),
)
if isinstance(candidate, str) and candidate.strip()
),
None,
)
if raw_id is None:
continue
wire_id = raw_id.split("/", 1)[-1]
if wire_id in seen:
continue
seen.add(wire_id)
fallback = fallback_by_id.get(wire_id)
label = _catalog_first_text(row, "display_name", "label", "name") or _catalog_first_text(
meta,
"display_name",
"label",
"name",
)
if not label or label == raw_id:
label = fallback.label if fallback is not None else wire_id
models.append(
ProviderModelSpec(
id=f"xai-grok/{wire_id}",
label=label,
description=(
_catalog_first_text(row, "description")
or _catalog_first_text(meta, "description")
or (fallback.description if fallback is not None else "")
),
owned_by=(
_catalog_first_text(row, "owned_by", "owner", "organization")
or _catalog_first_text(meta, "owned_by", "owner", "organization")
or (fallback.owned_by if fallback is not None else "xAI")
),
context_window=(
_catalog_positive_int(row, "context_window", "context_length")
or _catalog_positive_int(meta, "context_window", "context_length")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_catalog_reasoning_efforts(
row.get("reasoning_efforts", meta.get("reasoning_efforts"))
),
supports_backend_search=_catalog_bool_field(
row,
"supports_backend_search",
"supportsBackendSearch",
),
)
)
return tuple(models)
def _build_xai_model_headers(access_token: str, account_id: str | None) -> dict[str, str]:
headers = {
"Authorization": f"Bearer {access_token}",
"X-XAI-Token-Auth": "xai-grok-cli",
"x-grok-client-version": XAI_CLIENT_VERSION,
"x-grok-client-identifier": "nanobot",
"x-grok-client-mode": "headless",
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
}
claims = _decode_access_token_claims(access_token)
user_id = claims.get("sub")
if claims.get("principal_type") == "Team":
user_id = claims.get("principal_id") or user_id
if isinstance(user_id, str) and user_id:
headers["x-userid"] = user_id
email = claims.get("email")
if not isinstance(email, str) or "@" not in email:
email = account_id if account_id and "@" in account_id else None
if email:
headers["x-email"] = email
return headers
def _decode_access_token_claims(token: str) -> dict[str, Any]:
parts = token.split(".")
if len(parts) < 2 or not parts[1]:
return {}
try:
decoded = base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))
claims = json.loads(decoded)
except (ValueError, TypeError):
return {}
return cast(dict[str, Any], claims) if isinstance(claims, dict) else {}
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_mapping(value: Any) -> dict[str, Any]:
return cast(dict[str, Any], value) if isinstance(value, dict) else {}
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_bool_field(row: dict[str, Any], *keys: str) -> bool:
for key in keys:
value = row.get(key)
if isinstance(value, bool):
return value
meta = row.get("_meta")
return _catalog_bool_field(_catalog_mapping(meta), *keys) if isinstance(meta, dict) else False
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
efforts: list[str] = []
for item in cast(list[object], value):
if isinstance(item, str):
effort = item.strip()
elif isinstance(item, dict):
effort = _catalog_first_text(cast(dict[str, Any], item), "effort", "value", "id")
else:
effort = ""
if effort and effort not in efforts:
efforts.append(effort)
return tuple(efforts)
_XAI_GROK_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("xai_grok"),
fetch=_fetch_xai_grok_models,
)