diff --git a/pageindex/client.py b/pageindex/client.py index 1f8ea5b..362bb9a 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -35,7 +35,7 @@ class PageIndexClient: api_key: PageIndex cloud API key. When provided, cloud mode is used and local-only params (model, storage_path, index_config, …) are ignored. model: LLM model for indexing (local mode only, default: gpt-4o-2024-11-20). - retrieve_model: LLM model for agent QA (local mode only, default: same as model). + retrieve_model: LLM model for agent QA (local mode only, default: gpt-5.4). storage_path: Directory for SQLite DB and files (local mode only, default: ./.pageindex). storage: Custom StorageEngine instance (local mode only). index_config: Advanced indexing parameters (local mode only, optional). @@ -312,7 +312,7 @@ class LocalClient(PageIndexClient): Args: model: LLM model for indexing (default: gpt-4o-2024-11-20) - retrieve_model: LLM model for agent QA (default: same as model) + retrieve_model: LLM model for agent QA (default: gpt-5.4) storage_path: Directory for SQLite DB and files (default: ./.pageindex) storage: Custom StorageEngine instance (default: SQLiteStorage) index_config: Advanced indexing parameters. Pass an IndexConfig instance diff --git a/pageindex/config.py b/pageindex/config.py index 4fe08b8..9b8d21d 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -18,7 +18,7 @@ class IndexConfig(BaseModel): model_config = {"extra": "forbid"} model: str = "gpt-4o-2024-11-20" - retrieve_model: str | None = None + retrieve_model: str | None = "gpt-5.4" # None = follow `model` toc_check_page_num: int = 20 max_page_num_each_node: int = 10 max_token_num_each_node: int = 20000 diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 591fe93..1073ff1 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -1,6 +1,6 @@ model: "gpt-4o-2024-11-20" # model: "anthropic/claude-sonnet-4-6" -retrieve_model: "gpt-5.4" # defaults to `model` if not set +retrieve_model: "gpt-5.4" # set to null to follow `model` toc_check_page_num: 20 max_page_num_each_node: 10 max_token_num_each_node: 20000 diff --git a/tests/test_client.py b/tests/test_client.py index de179a4..99a903e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -69,6 +69,25 @@ def test_delete_collection(tmp_path): assert "papers" not in client.list_collections() +def test_retrieve_model_defaults_to_strong_reasoner(tmp_path): + """Retrieval must not silently follow the (cheaper) indexing model.""" + client = LocalClient(model="gpt-4o", storage_path=str(tmp_path / "pi")) + assert client._backend.get_retrieve_model() == "gpt-5.4" + + +def test_explicit_retrieve_model_wins(tmp_path): + client = LocalClient(model="gpt-4o", retrieve_model="ollama/llama3", + storage_path=str(tmp_path / "pi")) + assert client._backend.get_retrieve_model() == "litellm/ollama/llama3" + + +def test_retrieve_model_none_follows_model(tmp_path): + from pageindex.config import IndexConfig + client = LocalClient(model="gpt-4o", storage_path=str(tmp_path / "pi"), + index_config=IndexConfig(retrieve_model=None)) + assert client._backend.get_retrieve_model() == "gpt-4o" + + def test_register_parser(tmp_path): client = LocalClient(model="gpt-4o", storage_path=str(tmp_path / "pi")) class FakeParser: diff --git a/tests/test_config.py b/tests/test_config.py index ee92307..84b032b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -6,7 +6,7 @@ from pageindex.config import IndexConfig def test_defaults(): config = IndexConfig() assert config.model == "gpt-4o-2024-11-20" - assert config.retrieve_model is None + assert config.retrieve_model == "gpt-5.4" assert config.toc_check_page_num == 20