mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-27 05:56:14 +00:00
Yash Raj PandeyandClaude Opus 5.5
5fa30ae404
`CohereReranker` and `VoyageAIReranker` read the environment variable before the `api_key` argument. When `COHERE_API_KEY` or `VOYAGE_API_KEY` is set, the reranker ignores an explicit `api_key` and uses the account from the environment. The docs say the opposite. `docs/web/integrations/reranking/cohere.mdx` and `voyageai.mdx` use the environment variable only when `api_key` is not provided. `JinaReranker` and `WatsonxReranker` already put the explicit key first, and #4515 makes the same change for `OpenaiReranker`. ## Fix `CohereReranker._client` and `VoyageAIReranker._client` now use `self.api_key` first and the environment variable as the fallback. ## Tests `python/python/tests/test_reranker_credentials.py` replaces the SDK module with a stub whose `Client` records the API key. It covers both keys set, each key alone, and no key. With the stub, the test runs without `cohere` or `voyageai` installed. - On `main`, the 2 cases with both keys set fail and the other 6 pass. With the fix, all 8 pass. - `test_rerankers.py` with the new file: 24 passed, 13 skipped (the skipped tests need real API keys). - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed files copied over it. In that wheel, both reranker files are byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
75 lines
2.0 KiB
Python
75 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from lancedb.rerankers import CohereReranker, VoyageAIReranker
|
|
|
|
RERANKERS = [
|
|
pytest.param("cohere", "COHERE_API_KEY", CohereReranker, {}, id="cohere"),
|
|
pytest.param(
|
|
"voyageai",
|
|
"VOYAGE_API_KEY",
|
|
VoyageAIReranker,
|
|
{"model_name": "rerank-2.5"},
|
|
id="voyageai",
|
|
),
|
|
]
|
|
|
|
|
|
def fake_sdk(monkeypatch, module_name):
|
|
"""Replace the SDK module with one whose Client records the API key."""
|
|
keys = []
|
|
|
|
def client(api_key=None):
|
|
keys.append(api_key)
|
|
return object()
|
|
|
|
module = types.ModuleType(module_name)
|
|
module.Client = client
|
|
monkeypatch.setitem(sys.modules, module_name, module)
|
|
return keys
|
|
|
|
|
|
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
|
|
@pytest.mark.parametrize(
|
|
"explicit_key, environment_key, expected_key",
|
|
[
|
|
("explicit-test-key", "environment-test-key", "explicit-test-key"),
|
|
("explicit-test-key", None, "explicit-test-key"),
|
|
(None, "environment-test-key", "environment-test-key"),
|
|
],
|
|
)
|
|
def test_reranker_api_key_precedence(
|
|
monkeypatch,
|
|
module_name,
|
|
env_var,
|
|
reranker_cls,
|
|
kwargs,
|
|
explicit_key,
|
|
environment_key,
|
|
expected_key,
|
|
):
|
|
keys = fake_sdk(monkeypatch, module_name)
|
|
if environment_key is None:
|
|
monkeypatch.delenv(env_var, raising=False)
|
|
else:
|
|
monkeypatch.setenv(env_var, environment_key)
|
|
|
|
reranker_cls(api_key=explicit_key, **kwargs)._client
|
|
assert keys == [expected_key]
|
|
|
|
|
|
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
|
|
def test_reranker_requires_api_key(
|
|
monkeypatch, module_name, env_var, reranker_cls, kwargs
|
|
):
|
|
keys = fake_sdk(monkeypatch, module_name)
|
|
monkeypatch.delenv(env_var, raising=False)
|
|
with pytest.raises(ValueError, match=f"{env_var} not set"):
|
|
reranker_cls(**kwargs)._client
|
|
assert keys == []
|