From 4d5f2f3f306ea4635b448f091383a0c1fd2bc5f6 Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Mon, 3 Aug 2026 05:29:21 +0800 Subject: [PATCH] feat: add Atlas Cloud model provider Signed-off-by: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> --- agent/core/direct_models.py | 35 +++++++++++++++++ agent/core/llm_params.py | 45 ++++++++++++++++++++- agent/core/model_switcher.py | 27 +++++++++---- agent/main.py | 13 ++++-- tests/unit/test_cli_local_models.py | 61 +++++++++++++++++++++++++++++ tests/unit/test_cli_rendering.py | 33 ++++++++++++++++ tests/unit/test_llm_params.py | 41 +++++++++++++++++++ 7 files changed, 242 insertions(+), 13 deletions(-) create mode 100644 agent/core/direct_models.py diff --git a/agent/core/direct_models.py b/agent/core/direct_models.py new file mode 100644 index 00000000..c34b55d8 --- /dev/null +++ b/agent/core/direct_models.py @@ -0,0 +1,35 @@ +"""Helpers for direct OpenAI-compatible cloud model ids.""" + +DIRECT_MODEL_PROVIDERS: dict[str, dict[str, str]] = { + "atlas/": { + "base_url_env": "ATLASCLOUD_BASE_URL", + "base_url_default": "https://api.atlascloud.ai/v1", + "api_key_env": "ATLASCLOUD_API_KEY", + }, +} + +DIRECT_MODEL_PREFIXES = tuple(DIRECT_MODEL_PROVIDERS) + + +def direct_model_provider(model_id: str) -> dict[str, str] | None: + """Return provider config for a direct model id, if supported.""" + for prefix, config in DIRECT_MODEL_PROVIDERS.items(): + if model_id.startswith(prefix): + return config + return None + + +def direct_model_name(model_id: str) -> str | None: + """Return the upstream model name with the direct provider prefix removed.""" + for prefix in DIRECT_MODEL_PREFIXES: + if model_id.startswith(prefix): + name = model_id[len(prefix) :] + return name or None + return None + + +def is_direct_model_id(model_id: str) -> bool: + """Return True for non-empty, whitespace-free direct model ids.""" + if not model_id or any(char.isspace() for char in model_id): + return False + return direct_model_name(model_id) is not None diff --git a/agent/core/llm_params.py b/agent/core/llm_params.py index d2f821c2..3d06ea7a 100644 --- a/agent/core/llm_params.py +++ b/agent/core/llm_params.py @@ -7,6 +7,10 @@ import os +from agent.core.direct_models import ( + direct_model_name, + direct_model_provider, +) from agent.core.hf_tokens import resolve_hf_router_token from agent.core.local_models import ( LOCAL_MODEL_API_KEY_DEFAULT, @@ -47,7 +51,7 @@ class UnsupportedEffortError(ValueError): """ -def _local_api_base(base_url: str) -> str: +def _openai_api_base(base_url: str) -> str: base = base_url.strip().rstrip("/") if base.endswith("/v1"): return base @@ -82,7 +86,36 @@ def _resolve_local_model_params( ) return { "model": f"openai/{local_name}", - "api_base": _local_api_base(raw_base), + "api_base": _openai_api_base(raw_base), + "api_key": api_key, + } + + +def _resolve_direct_model_params( + model_name: str, + reasoning_effort: str | None = None, + strict: bool = False, +) -> dict: + if reasoning_effort and strict: + raise UnsupportedEffortError( + "Direct OpenAI-compatible endpoints don't accept reasoning_effort" + ) + + upstream_name = direct_model_name(model_name) + provider = direct_model_provider(model_name) + if upstream_name is None or provider is None: + raise ValueError(f"Unsupported direct model id: {model_name}") + + raw_base = os.environ.get(provider["base_url_env"]) or provider[ + "base_url_default" + ] + api_key = os.environ.get(provider["api_key_env"], "").strip() + if not api_key: + raise ValueError(f"{provider['api_key_env']} is required for {model_name}") + + return { + "model": f"openai/{upstream_name}", + "api_base": _openai_api_base(raw_base), "api_key": api_key, } @@ -96,6 +129,9 @@ def _resolve_llm_params( """ Build LiteLLM kwargs for a given model id. + • ``atlas/`` — Atlas Cloud's OpenAI-compatible endpoint, using + ``ATLASCLOUD_API_KEY`` and optional ``ATLASCLOUD_BASE_URL``. + • ``ollama/``, ``vllm/``, ``lm_studio/``, and ``llamacpp/`` — local OpenAI-compatible endpoints. The id prefix selects a configurable localhost base URL, and the model suffix is sent @@ -126,6 +162,11 @@ def _resolve_llm_params( if is_reserved_local_model_id(normalized_model): raise ValueError(f"Unsupported local model id: {normalized_model}") + if direct_model_provider(normalized_model) is not None: + return _resolve_direct_model_params( + normalized_model, reasoning_effort, strict + ) + if local_model_provider(normalized_model) is not None: return _resolve_local_model_params(normalized_model, reasoning_effort, strict) diff --git a/agent/core/model_switcher.py b/agent/core/model_switcher.py index 5ece764d..eeada78b 100644 --- a/agent/core/model_switcher.py +++ b/agent/core/model_switcher.py @@ -19,6 +19,10 @@ from litellm import acompletion +from agent.core.direct_models import ( + DIRECT_MODEL_PREFIXES, + is_direct_model_id, +) from agent.core.effort_probe import ProbeInconclusive, probe_effort from agent.core.llm_params import _resolve_llm_params from agent.core.local_models import ( @@ -59,6 +63,7 @@ def is_valid_model_id(model_id: str) -> bool: """Loose format check — lets users pick any model id. Accepts: + • atlas/ (Atlas Cloud) • ollama/, vllm/, lm_studio/, llamacpp/ • /[:] (HF router; tag = provider or policy) • huggingface//[:] (same, optional LiteLLM prefix) @@ -69,6 +74,10 @@ def is_valid_model_id(model_id: str) -> bool: if not model_id: return False normalized_model_id = strip_huggingface_model_prefix(model_id) or model_id + if is_direct_model_id(normalized_model_id): + return True + if any(normalized_model_id.startswith(prefix) for prefix in DIRECT_MODEL_PREFIXES): + return False if is_local_model_id(normalized_model_id): return True if is_reserved_local_model_id(normalized_model_id): @@ -92,7 +101,7 @@ def _print_hf_routing_info(model_id: str, console) -> bool: against the router catalog when possible; the probe below covers provider availability for uncataloged ids. """ - if is_local_model_id(model_id): + if is_direct_model_id(model_id) or is_local_model_id(model_id): return True from agent.core import hf_router_catalog as cat @@ -163,6 +172,7 @@ def print_model_listing(config, console) -> None: console.print( "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M3:novita').\n" "Add ':fastest', ':cheapest', ':preferred', or ':' to override routing.\n" + "Use 'atlas/' for Atlas Cloud.\n" "Use 'ollama/', 'vllm/', 'lm_studio/', or " "'llamacpp/' for local OpenAI-compatible endpoints.[/dim]" ) @@ -173,11 +183,12 @@ def print_invalid_id(arg: str, console) -> None: console.print( "[dim]Expected:\n" " • /[:tag] (HF router — paste from huggingface.co)\n" + " • atlas/ (Atlas Cloud)\n" " • ollama/ | vllm/ | lm_studio/ | llamacpp/[/dim]" ) -async def _probe_local_model(model_id: str) -> None: +async def _probe_direct_model(model_id: str) -> None: params = _resolve_llm_params(model_id) await asyncio.wait_for( acompletion( @@ -208,15 +219,15 @@ async def probe_and_switch_model( * ✗ hard error (auth, model-not-found, quota) — we reject the switch and keep the current model so the user isn't stranded - For non-local models, transient errors (5xx, timeout) complete the switch + For HF Router models, transient errors (5xx, timeout) complete the switch with a yellow warning; the next real call re-surfaces the error if it's - persistent. Local models reject every probe error, including timeouts, and - keep the current model. + persistent. Direct and local models reject every probe error, including + timeouts, and keep the current model. """ - if is_local_model_id(model_id): - console.print(f"[dim]checking local model {model_id}...[/dim]") + if is_direct_model_id(model_id) or is_local_model_id(model_id): + console.print(f"[dim]checking model {model_id}...[/dim]") try: - await _probe_local_model(model_id) + await _probe_direct_model(model_id) except Exception as e: console.print(f"[bold red]Switch failed:[/bold red] {e}") console.print(f"[dim]Keeping current model: {config.model_name}[/dim]") diff --git a/agent/main.py b/agent/main.py index 6f29c06b..280d5004 100644 --- a/agent/main.py +++ b/agent/main.py @@ -28,6 +28,7 @@ from agent.core import model_switcher from agent.core.hf_access import fetch_whoami_v2, normalize_hf_user_plan from agent.core.hf_tokens import resolve_hf_token +from agent.core.direct_models import is_direct_model_id from agent.core.local_models import is_local_model_id from agent.core.model_ids import strip_huggingface_model_prefix from agent.core.session import OpType @@ -1193,9 +1194,12 @@ async def main(model: str | None = None, sandbox_tools: bool = False): local_mode = _is_local_tool_runtime(config) # HF token — required for Hub-backed models/tools and sandbox tools, but - # not for local LLMs using only local filesystem tools. + # not for direct/local LLMs using only local filesystem tools. hf_token = resolve_hf_token() - if not hf_token and (not is_local_model_id(config.model_name) or not local_mode): + non_router_model = is_direct_model_id(config.model_name) or is_local_model_id( + config.model_name + ) + if not hf_token and (not non_router_model or not local_mode): hf_token = await _prompt_and_save_hf_token(prompt_session) # Resolve username and plan from one whoami-v2 request for banner and CTAs. @@ -1447,7 +1451,10 @@ async def headless_main( local_mode = _is_local_tool_runtime(config) hf_token = resolve_hf_token() - if not hf_token and (not is_local_model_id(config.model_name) or not local_mode): + non_router_model = is_direct_model_id(config.model_name) or is_local_model_id( + config.model_name + ) + if not hf_token and (not non_router_model or not local_mode): print( "ERROR: No HF token found. Set HF_TOKEN or run `hf auth login`.", file=sys.stderr, diff --git a/tests/unit/test_cli_local_models.py b/tests/unit/test_cli_local_models.py index 7c675c59..f208e925 100644 --- a/tests/unit/test_cli_local_models.py +++ b/tests/unit/test_cli_local_models.py @@ -2,6 +2,7 @@ from agent.config import load_config from agent.core import model_switcher +from agent.core.direct_models import is_direct_model_id from agent.core.local_models import is_local_model_id from agent.main import CLI_CONFIG_PATH @@ -13,11 +14,18 @@ def test_local_model_helper_accepts_supported_prefixes(): assert is_local_model_id("llamacpp/unsloth/Qwen3.5-2B") +def test_direct_model_helper_accepts_atlas_prefix(): + assert is_direct_model_id("atlas/deepseek-ai/deepseek-v4-pro") + assert not is_direct_model_id("atlas/") + assert not is_direct_model_id("atlas/model with spaces") + + def test_model_switcher_accepts_supported_local_prefixes(): assert model_switcher.is_valid_model_id("ollama/llama3.1:8b") assert model_switcher.is_valid_model_id("vllm/meta-llama/Llama-3.1-8B") assert model_switcher.is_valid_model_id("lm_studio/google/gemma-3-4b") assert model_switcher.is_valid_model_id("llamacpp/llama-3.1-8b") + assert model_switcher.is_valid_model_id("atlas/deepseek-ai/deepseek-v4-pro") def test_model_switcher_rejects_empty_or_whitespace_local_ids(): @@ -26,6 +34,8 @@ def test_model_switcher_rejects_empty_or_whitespace_local_ids(): assert not model_switcher.is_valid_model_id("lm_studio/") assert not model_switcher.is_valid_model_id("llamacpp/") assert not model_switcher.is_valid_model_id("ollama/llama 3.1") + assert not model_switcher.is_valid_model_id("atlas/") + assert not model_switcher.is_valid_model_id("atlas/model with spaces") def test_openai_compat_prefix_is_not_supported(): @@ -104,6 +114,57 @@ def print(self, *args, **kwargs): assert "extra_body" not in calls[0] +@pytest.mark.asyncio +async def test_probe_and_switch_atlas_model_uses_direct_endpoint(monkeypatch): + calls = [] + + async def fake_acompletion(**kwargs): + calls.append(kwargs) + return object() + + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-secret") + monkeypatch.setattr(model_switcher, "acompletion", fake_acompletion) + + class Config: + model_name = "anthropic/claude-opus-4.8:fal-ai" + reasoning_effort = "high" + + class Session: + def __init__(self): + self.model_id = None + self.model_effective_effort = {} + + def update_model(self, model_id): + self.model_id = model_id + + class Console: + def print(self, *args, **kwargs): + pass + + session = Session() + model_id = "atlas/deepseek-ai/deepseek-v4-pro" + await model_switcher.probe_and_switch_model( + model_id, + Config(), + session, + Console(), + hf_token=None, + ) + + assert session.model_id == model_id + assert session.model_effective_effort[model_id] is None + assert calls == [ + { + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 1, + "stream": False, + "model": "openai/deepseek-ai/deepseek-v4-pro", + "api_base": "https://api.atlascloud.ai/v1", + "api_key": "atlas-secret", + } + ] + + @pytest.mark.asyncio async def test_probe_and_switch_local_model_rejects_probe_errors(monkeypatch): async def failing_acompletion(**kwargs): diff --git a/tests/unit/test_cli_rendering.py b/tests/unit/test_cli_rendering.py index 285e943e..5a364fb3 100644 --- a/tests/unit/test_cli_rendering.py +++ b/tests/unit/test_cli_rendering.py @@ -260,6 +260,39 @@ def fake_banner(*, model=None, hf_user=None, tool_runtime=None): await main_mod.main() +@pytest.mark.asyncio +async def test_atlas_model_local_runtime_skips_hf_token_prompt(monkeypatch): + class StopAfterBanner(Exception): + pass + + async def fail_prompt(_prompt_session): + raise AssertionError("Atlas direct model should not prompt for an HF token") + + def fake_banner(*, model=None, hf_user=None, tool_runtime=None): + assert model == "atlas/deepseek-ai/deepseek-v4-pro" + assert hf_user is None + assert tool_runtime == "local filesystem" + raise StopAfterBanner + + monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0) + monkeypatch.setattr(main_mod, "PromptSession", lambda: object()) + monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: None) + monkeypatch.setattr(main_mod, "_prompt_and_save_hf_token", fail_prompt) + monkeypatch.setattr( + main_mod, + "load_config", + lambda _path, **_kwargs: SimpleNamespace( + model_name="atlas/deepseek-ai/deepseek-v4-pro", + mcpServers={}, + tool_runtime="local", + ), + ) + monkeypatch.setattr(main_mod, "print_banner", fake_banner) + + with pytest.raises(StopAfterBanner): + await main_mod.main() + + @pytest.mark.asyncio async def test_local_model_sandbox_runtime_prompts_for_hf_token(monkeypatch): class StopAfterBanner(Exception): diff --git a/tests/unit/test_llm_params.py b/tests/unit/test_llm_params.py index a985025a..3a0151d4 100644 --- a/tests/unit/test_llm_params.py +++ b/tests/unit/test_llm_params.py @@ -90,6 +90,47 @@ def test_resolve_ollama_params_adds_v1_and_uses_default_key(monkeypatch): } +def test_resolve_atlas_params_uses_default_endpoint(monkeypatch): + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-secret") + monkeypatch.delenv("ATLASCLOUD_BASE_URL", raising=False) + + params = _resolve_llm_params("atlas/deepseek-ai/deepseek-v4-pro") + + assert params == { + "model": "openai/deepseek-ai/deepseek-v4-pro", + "api_base": "https://api.atlascloud.ai/v1", + "api_key": "atlas-secret", + } + + +def test_resolve_atlas_params_supports_endpoint_override(monkeypatch): + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-secret") + monkeypatch.setenv("ATLASCLOUD_BASE_URL", "https://example.com/openai/v1/") + + params = _resolve_llm_params("atlas/custom/model") + + assert params["model"] == "openai/custom/model" + assert params["api_base"] == "https://example.com/openai/v1" + + +def test_resolve_atlas_params_requires_api_key(monkeypatch): + monkeypatch.delenv("ATLASCLOUD_API_KEY", raising=False) + + with pytest.raises(ValueError, match="ATLASCLOUD_API_KEY"): + _resolve_llm_params("atlas/deepseek-ai/deepseek-v4-pro") + + +def test_atlas_params_reject_reasoning_effort_in_strict_mode(monkeypatch): + monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-secret") + + with pytest.raises(UnsupportedEffortError, match="reasoning_effort"): + _resolve_llm_params( + "atlas/deepseek-ai/deepseek-v4-pro", + reasoning_effort="high", + strict=True, + ) + + def test_resolve_vllm_params_keeps_existing_v1_and_trims_slash(monkeypatch): monkeypatch.delenv("VLLM_API_KEY", raising=False) monkeypatch.setenv("VLLM_BASE_URL", "http://localhost:8000/v1/")