From 06f6deff6cfaede099e48000d48eb69adbcde76b Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:20:59 +0800 Subject: [PATCH] Add Atlas Cloud LiteLLM support --- README.md | 12 +++++++++ pageindex/config.yaml | 3 ++- pageindex/utils.py | 35 +++++++++++++++++++++---- tests/test_atlascloud_litellm.py | 45 ++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 6 deletions(-) create mode 100644 tests/test_atlascloud_litellm.py diff --git a/README.md b/README.md index 594b70e8e..5c6be292d 100644 --- a/README.md +++ b/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 diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 591fe9331..011ffb457 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -1,5 +1,6 @@ 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 @@ -7,4 +8,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" \ No newline at end of file +if_add_node_text: "no" diff --git a/pageindex/utils.py b/pageindex/utils.py index 235dd09cc..956ccfbcb 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -24,6 +24,32 @@ 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 @@ -31,8 +57,7 @@ def count_tokens(text, model=None): 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)) - diff --git a/tests/test_atlascloud_litellm.py b/tests/test_atlascloud_litellm.py new file mode 100644 index 000000000..92c8563ed --- /dev/null +++ b/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")