diff --git a/pageindex/config.yaml b/pageindex/config.yaml index d3786ca12..557ec8948 100644 --- a/pageindex/config.yaml +++ b/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" \ No newline at end of file +if_add_node_text: "no" diff --git a/pageindex/utils.py b/pageindex/utils.py index f23995057..8f5ff4405 100644 --- a/pageindex/utils.py +++ b/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 diff --git a/tests/test_atlascloud_litellm.py b/tests/test_atlascloud_litellm.py new file mode 100644 index 000000000..398571a28 --- /dev/null +++ b/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, + }]