mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
fix(python): honor explicit Cohere and VoyageAI reranker API keys (#4546)
`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>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
bc4c32b779
commit
5fa30ae404
@@ -62,7 +62,7 @@ class CohereReranker(Reranker):
|
||||
"COHERE_API_KEY not set. Either set it in your environment or \
|
||||
pass it as `api_key` argument to the CohereReranker."
|
||||
)
|
||||
return cohere.Client(os.environ.get("COHERE_API_KEY") or self.api_key)
|
||||
return cohere.Client(self.api_key or os.environ.get("COHERE_API_KEY"))
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str):
|
||||
result_set = self._handle_empty_results(result_set)
|
||||
|
||||
@@ -64,7 +64,7 @@ class VoyageAIReranker(Reranker):
|
||||
pass it as `api_key` argument to the VoyageAIReranker."
|
||||
)
|
||||
return voyageai.Client(
|
||||
api_key=os.environ.get("VOYAGE_API_KEY") or self.api_key,
|
||||
api_key=self.api_key or os.environ.get("VOYAGE_API_KEY"),
|
||||
)
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str):
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
# 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 == []
|
||||
Reference in New Issue
Block a user