Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions config.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,10 @@
import os


_ATLAS_CLOUD_PREFIX = "atlascloud/"
_ATLAS_CLOUD_BASE_URL = "https://api.atlascloud.ai/v1"


def get_default_model(fallback: str = "gpt-5.4"):
"""Return the configured default model for standard agents."""
model = os.getenv("DEFAULT_MODEL", fallback)
Expand All @@ -12,18 +16,28 @@ def is_openai_provider() -> bool:
"""Return True when the configured provider is OpenAI (not LiteLLM).

OpenAI model IDs never contain a slash (e.g. 'gpt-5.4', 'o3').
Any 'provider/model' string (e.g. 'anthropic/claude-sonnet-4-6',
'litellm/gemini/gemini-3-flash') is treated as a LiteLLM-routed model.
Any provider-prefixed model, including Atlas Cloud and LiteLLM routes,
must not receive Responses API-only reasoning or web-search settings.
"""
return "/" not in os.getenv("DEFAULT_MODEL", "")


def _resolve(model: str):
"""Route 'provider/model' strings through LitellmModel.
"""Resolve Atlas Cloud and LiteLLM provider-prefixed model strings."""
if model.startswith(_ATLAS_CLOUD_PREFIX):
model_id = model[len(_ATLAS_CLOUD_PREFIX) :]
if not model_id:
raise ValueError("Atlas Cloud model ID is required after 'atlascloud/'")
api_key = os.getenv("ATLASCLOUD_API_KEY")
if not api_key:
raise RuntimeError("ATLASCLOUD_API_KEY environment variable is required")

from agents import OpenAIChatCompletionsModel
from openai import AsyncOpenAI

client = AsyncOpenAI(api_key=api_key, base_url=_ATLAS_CLOUD_BASE_URL)
return OpenAIChatCompletionsModel(model=model_id, openai_client=client)

Handles both explicit 'litellm/<model>' and bare 'provider/model' forms.
OpenAI model IDs contain no slash, so they pass through unchanged.
"""
if "/" not in model:
return model
bare = model[len("litellm/"):] if model.startswith("litellm/") else model
Expand Down
78 changes: 78 additions & 0 deletions tests/test_atlascloud_model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
from __future__ import annotations

import importlib
import sys
import types

import pytest


class FakeAsyncOpenAI:
def __init__(self, *, api_key, base_url):
self.api_key = api_key
self.base_url = base_url


class FakeChatCompletionsModel:
def __init__(self, *, model, openai_client):
self.model = model
self.openai_client = openai_client


class FakeLitellmModel:
def __init__(self, *, model):
self.model = model


@pytest.fixture
def config_module(monkeypatch):
agents = types.ModuleType("agents")
agents.OpenAIChatCompletionsModel = FakeChatCompletionsModel
openai = types.ModuleType("openai")
openai.AsyncOpenAI = FakeAsyncOpenAI
agency_swarm = types.ModuleType("agency_swarm")
agency_swarm.LitellmModel = FakeLitellmModel
monkeypatch.setitem(sys.modules, "agents", agents)
monkeypatch.setitem(sys.modules, "openai", openai)
monkeypatch.setitem(sys.modules, "agency_swarm", agency_swarm)
sys.modules.pop("config", None)
yield importlib.import_module("config")
sys.modules.pop("config", None)


def test_atlascloud_model_uses_chat_completions_client(monkeypatch, config_module):
monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
monkeypatch.setenv("DEFAULT_MODEL", "atlascloud/deepseek-ai/deepseek-v4-pro")

model = config_module.get_default_model()

assert isinstance(model, FakeChatCompletionsModel)
assert model.model == "deepseek-ai/deepseek-v4-pro"
assert model.openai_client.api_key == "test-key"
assert model.openai_client.base_url == "https://api.atlascloud.ai/v1"
assert config_module.is_openai_provider() is False


def test_atlascloud_model_requires_api_key(monkeypatch, config_module):
monkeypatch.delenv("ATLASCLOUD_API_KEY", raising=False)
monkeypatch.setenv("DEFAULT_MODEL", "atlascloud/deepseek-ai/deepseek-v4-pro")

with pytest.raises(RuntimeError, match="ATLASCLOUD_API_KEY"):
config_module.get_default_model()


def test_atlascloud_model_requires_model_id(monkeypatch, config_module):
monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
monkeypatch.setenv("DEFAULT_MODEL", "atlascloud/")

with pytest.raises(ValueError, match="model ID"):
config_module.get_default_model()


def test_existing_litellm_route_is_unchanged(monkeypatch, config_module):
monkeypatch.setenv("DEFAULT_MODEL", "anthropic/claude-sonnet")

model = config_module.get_default_model()

assert isinstance(model, FakeLitellmModel)
assert model.model == "anthropic/claude-sonnet"