Add Atlas Cloud LiteLLM provider alias

pageindex/config.yaml

@@ -6,10 +6,13 @@
 # providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6").
 # index_model: "gpt-5.6-luna"
 # chat_model: "gpt-5.6-sol"
+# Atlas Cloud example (requires ATLASCLOUD_API_KEY):
+# index_model: "atlascloud/qwen/qwen3.5-flash"
+# chat_model: "atlascloud/qwen/qwen3.5-flash"
 toc_check_page_num: 20
 max_page_num_each_node: 10
 max_token_num_each_node: 20000
 if_add_node_id: "yes"
 if_add_node_summary: "yes"
 if_add_doc_description: "no"
-if_add_node_text: "no"
+if_add_node_text: "no"

pageindex/utils.py

@@ -105,6 +105,27 @@ def run_off_loop(func, *args):
         return pool.submit(func, *args).result()
 
 
+ATLASCLOUD_MODEL_PREFIX = "atlascloud/"
+ATLASCLOUD_API_BASE = "https://api.atlascloud.ai/v1"
+
+
+def _atlascloud_kwargs(model):
+    """Credential kwargs for an ``atlascloud/`` model, or ``{}`` for anything else.
+
+    Atlas Cloud is reached over LiteLLM's ``openai/`` wire form, so the endpoint
+    and key have to travel as completion kwargs rather than being inferred from
+    the provider name."""
+    if not model or not model.removeprefix("litellm/").startswith(ATLASCLOUD_MODEL_PREFIX):
+        return {}
+    api_key = os.getenv("ATLASCLOUD_API_KEY")
+    if not api_key:
+        raise ValueError("ATLASCLOUD_API_KEY is required when using Atlas Cloud models.")
+    return {
+        "api_base": os.getenv("ATLASCLOUD_API_BASE", ATLASCLOUD_API_BASE),
+        "api_key": api_key,
+    }
+
+
 def _litellm_model(model):
     """Normalize to LiteLLM's grammar (``litellm/`` strips, bare names get
     the ``openai/`` wire form — same as the chat lane) and refuse an
@@ -113,6 +134,13 @@ def _litellm_model(model):
     if not model:
         return model
     model = _strip_prefix(model, "litellm/")
+    if model.startswith(ATLASCLOUD_MODEL_PREFIX):
+        atlas_model = model[len(ATLASCLOUD_MODEL_PREFIX):]
+        if not atlas_model:
+            raise ValueError("Atlas Cloud model must be provided after 'atlascloud/'.")
+        # Atlas Cloud speaks the OpenAI wire format; the endpoint arrives as a
+        # completion kwarg from _atlascloud_kwargs().
+        return f"openai/{atlas_model}"
     if "/" not in model:
         model = f"openai/{model}"
     import litellm
@@ -163,6 +191,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
     max_retries = 10
     messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
     backend = _llm_backend.get()
+    atlas = _atlascloud_kwargs(model)
     model = _litellm_model(model)
     _repair_litellm_types()
     _quiet_litellm()
@@ -174,6 +203,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
                 "drop_params": True,
                 # the loop is the retry policy; the merge lets a backend override win
                 "max_retries": 0,
+                **atlas,
                 **(backend or {}),
             })
             content = response.choices[0].message.content
@@ -200,6 +230,7 @@ async def llm_acompletion(model, prompt):
     max_retries = 10
     messages = [{"role": "user", "content": prompt}]
     backend = _llm_backend.get()
+    atlas = _atlascloud_kwargs(model)
     model = _litellm_model(model)
     _repair_litellm_types()
     _quiet_litellm()
@@ -210,6 +241,7 @@ async def llm_acompletion(model, prompt):
                 "messages": messages,
                 "drop_params": True,
                 "max_retries": 0,
+                **atlas,
                 **(backend or {}),
             })
             return response.choices[0].message.content

tests/test_atlascloud_litellm.py

@@ -0,0 +1,147 @@
+import asyncio
+import os
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from pageindex.utils import (
+    ATLASCLOUD_API_BASE,
+    _atlascloud_kwargs,
+    _litellm_model,
+    _llm_backend,
+    llm_acompletion,
+    llm_completion,
+)
+
+
+def prepare_litellm_call(model):
+    """Test shim mirroring the pre-refactor helper: upstream now splits this
+    into `_litellm_model()` (name normalization) and `_atlascloud_kwargs()`
+    (credentials), and `llm_completion` merges the two."""
+    return _litellm_model(model), _atlascloud_kwargs(model)
+
+
+def completion_response(content):
+    return SimpleNamespace(
+        choices=[
+            SimpleNamespace(
+                finish_reason="stop",
+                message=SimpleNamespace(content=content),
+            )
+        ]
+    )
+
+
+def test_prepare_litellm_call_keeps_regular_models():
+    # Upstream's _litellm_model() now gives bare names the openai/ wire form;
+    # what matters here is that a non-Atlas model carries no Atlas credentials.
+    model, kwargs = prepare_litellm_call("gpt-4o")
+    assert model == "openai/gpt-4o"
+    assert kwargs == {}
+
+
+def test_prepare_litellm_call_strips_litellm_prefix():
+    model, kwargs = prepare_litellm_call("litellm/anthropic/claude-sonnet-4")
+    assert model == "anthropic/claude-sonnet-4"
+    assert kwargs == {}
+
+
+def test_prepare_litellm_call_maps_atlascloud_models(monkeypatch):
+    monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
+    model, kwargs = prepare_litellm_call("atlascloud/qwen/qwen3.5-flash")
+    assert model == "openai/qwen/qwen3.5-flash"
+    assert kwargs == {
+        "api_base": ATLASCLOUD_API_BASE,
+        "api_key": "test-key",
+    }
+
+
+def test_prepare_litellm_call_respects_custom_atlascloud_base(monkeypatch):
+    monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
+    monkeypatch.setenv("ATLASCLOUD_API_BASE", "https://atlas.example/v1")
+    model, kwargs = prepare_litellm_call("litellm/atlascloud/deepseek-ai/deepseek-v4-pro")
+    assert model == "openai/deepseek-ai/deepseek-v4-pro"
+    assert kwargs["api_base"] == "https://atlas.example/v1"
+    assert kwargs["api_key"] == "test-key"
+
+
+def test_prepare_litellm_call_requires_atlascloud_api_key(monkeypatch):
+    monkeypatch.delenv("ATLASCLOUD_API_KEY", raising=False)
+    with pytest.raises(ValueError, match="ATLASCLOUD_API_KEY"):
+        prepare_litellm_call("atlascloud/qwen/qwen3.5-flash")
+
+
+def test_llm_completion_routes_atlascloud_through_litellm(monkeypatch):
+    calls = []
+
+    def completion(**kwargs):
+        calls.append(kwargs)
+        return completion_response("sync response")
+
+    monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
+    monkeypatch.setitem(sys.modules, "litellm", SimpleNamespace(completion=completion))
+
+    result = llm_completion("atlascloud/qwen/qwen3.5-flash", "hello")
+
+    assert result == "sync response"
+    assert calls == [{
+        "api_base": ATLASCLOUD_API_BASE,
+        "api_key": "test-key",
+        "drop_params": True,
+        "messages": [{"role": "user", "content": "hello"}],
+        "model": "openai/qwen/qwen3.5-flash",
+        "max_retries": 0,
+    }]
+
+
+def test_index_backend_overrides_atlascloud_defaults(monkeypatch):
+    calls = []
+
+    def completion(**kwargs):
+        calls.append(kwargs)
+        return completion_response("sync response")
+
+    monkeypatch.setenv("ATLASCLOUD_API_KEY", "environment-key")
+    monkeypatch.setitem(sys.modules, "litellm", SimpleNamespace(completion=completion))
+    token = _llm_backend.set({
+        "api_base": "https://backend.example/v1",
+        "api_key": "backend-key",
+        "timeout": 30,
+    })
+    try:
+        result = llm_completion("atlascloud/qwen/qwen3.5-flash", "hello")
+    finally:
+        _llm_backend.reset(token)
+
+    assert result == "sync response"
+    assert calls[0]["api_base"] == "https://backend.example/v1"
+    assert calls[0]["api_key"] == "backend-key"
+    assert calls[0]["timeout"] == 30
+
+
+def test_llm_acompletion_routes_atlascloud_through_litellm(monkeypatch):
+    calls = []
+
+    async def acompletion(**kwargs):
+        calls.append(kwargs)
+        return completion_response("async response")
+
+    monkeypatch.setenv("ATLASCLOUD_API_KEY", "test-key")
+    monkeypatch.setitem(sys.modules, "litellm", SimpleNamespace(acompletion=acompletion))
+
+    result = asyncio.run(
+        llm_acompletion("atlascloud/qwen/qwen3.5-flash", "hello")
+    )
+
+    assert result == "async response"
+    assert calls == [{
+        "api_base": ATLASCLOUD_API_BASE,
+        "api_key": "test-key",
+        "drop_params": True,
+        "messages": [{"role": "user", "content": "hello"}],
+        "model": "openai/qwen/qwen3.5-flash",
+        "max_retries": 0,
+    }]