Add Atlas Cloud LiteLLM provider alias

README.md

@@ -164,6 +164,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

@@ -2,6 +2,7 @@
 # For other providers, use "provider/model" format (e.g. "anthropic/claude-sonnet-4-6").
 model: "gpt-4o-2024-11-20"
 # model: "anthropic/claude-sonnet-4-6"
+# model: "atlascloud/qwen/qwen3.5-flash"
 summary_model: "gpt-5.6-luna"
 retrieve_model: "gpt-5.4"  # defaults to `model` if not set
 toc_check_page_num: 20
@@ -10,4 +11,4 @@ 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,6 +24,33 @@
 if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):
     os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY")
 
+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
@@ -56,10 +83,9 @@ def _is_unrecoverable(exc: Exception) -> bool:
 
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
     use_openai_sdk = _is_openai_model(model)
-    if model:
-        model = model.removeprefix("litellm/")
-        if use_openai_sdk:
-            model = model.removeprefix("openai/")
+    model, provider_kwargs = prepare_litellm_call(model)
+    if use_openai_sdk:
+        model = model.removeprefix("openai/")
     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):
@@ -80,6 +106,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
                     messages=messages,
                     temperature=0,
                     drop_params=True,
+                    **provider_kwargs,
                 )
             content = response.choices[0].message.content
             if return_finish_reason:
@@ -102,10 +129,9 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
 
 async def llm_acompletion(model, prompt):
     use_openai_sdk = _is_openai_model(model)
-    if model:
-        model = model.removeprefix("litellm/")
-        if use_openai_sdk:
-            model = model.removeprefix("openai/")
+    model, provider_kwargs = prepare_litellm_call(model)
+    if use_openai_sdk:
+        model = model.removeprefix("openai/")
     max_retries = 10
     messages = [{"role": "user", "content": prompt}]
     for i in range(max_retries):
@@ -126,6 +152,7 @@ async def llm_acompletion(model, prompt):
                     messages=messages,
                     temperature=0,
                     drop_params=True,
+                    **provider_kwargs,
                 )
             return response.choices[0].message.content
         except Exception as e:
@@ -974,4 +1001,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,111 @@
+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,
+    llm_acompletion,
+    llm_completion,
+    prepare_litellm_call,
+)
+
+
+def completion_response(content):
+    return SimpleNamespace(
+        choices=[
+            SimpleNamespace(
+                finish_reason="stop",
+                message=SimpleNamespace(content=content),
+            )
+        ]
+    )
+
+
+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")
+
+
+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",
+        "temperature": 0,
+    }]
+
+
+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",
+        "temperature": 0,
+    }]