diff --git a/README.md b/README.md index 5ce0ca5e6..759398d52 100644 --- a/README.md +++ b/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 diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 73a512c7a..1f9560755 100644 --- a/pageindex/config.yaml +++ b/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" \ No newline at end of file +if_add_node_text: "no" diff --git a/pageindex/utils.py b/pageindex/utils.py index 92fc46d85..41f3fa966 100644 --- a/pageindex/utils.py +++ b/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)) - diff --git a/tests/test_atlascloud_litellm.py b/tests/test_atlascloud_litellm.py new file mode 100644 index 000000000..fea063db5 --- /dev/null +++ b/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, + }]