Add Atlas Cloud LiteLLM provider alias

README.md

@@ -163,6 +163,18 @@ Create a `.env` file in the root directory with your LLM API key. Multi-LLM is s
 OPENAI_API_KEY=your_openai_key_here
 ```
 
+For Atlas Cloud's OpenAI-compatible endpoint, set `ATLASCLOUD_API_KEY` and use the
+`atlascloud/` model prefix:
+
+```bash
+ATLASCLOUD_API_KEY=your_atlascloud_key_here
+```
+
+```yaml
+model: "atlascloud/qwen/qwen3.5-flash"
+retrieve_model: "atlascloud/qwen/qwen3.5-flash"
+```
+
 ### 3. Generate PageIndex structure for your PDF
 
 ```bash

pageindex/config.yaml

@@ -1,10 +1,11 @@
 model: "gpt-4o-2024-11-20"
 # model: "anthropic/claude-sonnet-4-6"
+# model: "atlascloud/qwen/qwen3.5-flash"
 retrieve_model: "gpt-5.4"  # defaults to `model` if not set
 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

@@ -24,15 +24,40 @@
 
 litellm.drop_params = True
 
+ATLASCLOUD_API_BASE = "https://api.atlascloud.ai/v1"
+ATLASCLOUD_MODEL_PREFIX = "atlascloud/"
+
+
+def prepare_litellm_call(model):
+    """Normalize PageIndex model aliases into LiteLLM completion kwargs."""
+    if not model:
+        return model, {}
+
+    model = model.removeprefix("litellm/")
+    if not model.startswith(ATLASCLOUD_MODEL_PREFIX):
+        return model, {}
+
+    atlas_model = model[len(ATLASCLOUD_MODEL_PREFIX):]
+    if not atlas_model:
+        raise ValueError("Atlas Cloud model must be provided after 'atlascloud/'.")
+
+    api_key = os.getenv("ATLASCLOUD_API_KEY")
+    if not api_key:
+        raise ValueError("ATLASCLOUD_API_KEY is required when using Atlas Cloud models.")
+
+    return f"openai/{atlas_model}", {
+        "api_base": os.getenv("ATLASCLOUD_API_BASE", ATLASCLOUD_API_BASE),
+        "api_key": api_key,
+    }
+
 def count_tokens(text, model=None):
     if not text:
         return 0
     return litellm.token_counter(model=model, text=text)
 
 
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
-    if model:
-        model = model.removeprefix("litellm/")
+    model, provider_kwargs = prepare_litellm_call(model)
     max_retries = 10
     messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
     for i in range(max_retries):
@@ -41,6 +66,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
                 model=model,
                 messages=messages,
                 temperature=0,
+                **provider_kwargs,
             )
             content = response.choices[0].message.content
             if return_finish_reason:
@@ -61,8 +87,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
 
 
 async def llm_acompletion(model, prompt):
-    if model:
-        model = model.removeprefix("litellm/")
+    model, provider_kwargs = prepare_litellm_call(model)
     max_retries = 10
     messages = [{"role": "user", "content": prompt}]
     for i in range(max_retries):
@@ -71,6 +96,7 @@ async def llm_acompletion(model, prompt):
                 model=model,
                 messages=messages,
                 temperature=0,
+                **provider_kwargs,
             )
             return response.choices[0].message.content
         except Exception as e:
@@ -708,4 +734,3 @@ def print_tree(tree, indent=0):
 def print_wrapped(text, width=100):
     for line in text.splitlines():
         print(textwrap.fill(line, width=width))
-

tests/test_atlascloud_litellm.py

@@ -0,0 +1,45 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from pageindex.utils import ATLASCLOUD_API_BASE, prepare_litellm_call
+
+
+def test_prepare_litellm_call_keeps_regular_models():
+    model, kwargs = prepare_litellm_call("gpt-4o")
+    assert model == "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")