mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-27 22:16:14 +00:00
Yash Raj PandeyandClaude Opus 5.5
2af7d88ec7
`OpenaiReranker._rerank` asks the chat model for a relevance score per document. The model can return the documents in any order, for example sorted by score. The reranker then replaces the `text` column with the documents in that order and appends the scores in that order. The other columns keep the input order. So after a rerank, the `text` and `_relevance_score` of a row belong to a different row. The `id`, the vector and every other column still describe the original row. If the model leaves out a document, `append_column` raises `ArrowInvalid`, because the new column has fewer rows than the table. ## Fix `_rerank` keeps the `text` column unchanged, so it also keeps its position and type. It maps each score back to its row by the document content. If the model leaves out a document, `_rerank` raises a `ValueError` that gives the number of documents with no score. ## Tests `python/python/tests/test_openai_reranker.py` replaces the chat client with a stub that returns a fixed reply. The tests do not need `openai` or an API key. - One test returns the documents sorted by score. It checks that each `id` keeps its own `text` and score. On `main`, the row with id 1 gets the text `cats purr`, which belongs to id 2. - One test returns 2 of 3 documents. It checks for the `ValueError` with the count of missing documents. On `main`, the call raises `ArrowInvalid: Added column's length must match table's length`. - Both tests fail on `main` and pass with the fix. - `test_rerankers.py`: 16 passed, 13 skipped, on the lancedb 0.40.0 wheel with and without the fix. 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 file copied over it. In that wheel, `rerankers/openai.py` and `rerankers/base.py` are byte-identical to `main`. The Node SDK has no OpenAI reranker, so there is no parity change. This change does not conflict with #4515, which edits `_client` in the same file. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
from lancedb.rerankers import OpenaiReranker
|
|
|
|
|
|
def reranker_replying(documents):
|
|
"""Return an OpenaiReranker whose chat client replies with these documents."""
|
|
reply = json.dumps({"documents": documents})
|
|
|
|
def create(**kwargs):
|
|
message = SimpleNamespace(content=reply)
|
|
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
|
|
|
reranker = OpenaiReranker(api_key="test-key")
|
|
reranker.__dict__["_client"] = SimpleNamespace(
|
|
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
|
)
|
|
return reranker
|
|
|
|
|
|
def vector_results():
|
|
return pa.table(
|
|
{
|
|
"id": [1, 2, 3],
|
|
"text": ["dogs bark", "cats purr", "birds sing"],
|
|
"_distance": [0.1, 0.2, 0.3],
|
|
}
|
|
)
|
|
|
|
|
|
def test_scores_stay_with_their_rows_when_the_model_reorders():
|
|
reranker = reranker_replying(
|
|
[
|
|
{"content": "cats purr", "relevance_score": 0.9},
|
|
{"content": "birds sing", "relevance_score": 0.5},
|
|
{"content": "dogs bark", "relevance_score": 0.1},
|
|
]
|
|
)
|
|
|
|
result = reranker.rerank_vector("which animal purrs?", vector_results())
|
|
|
|
assert result.select(["id", "text"]).to_pylist() == [
|
|
{"id": 2, "text": "cats purr"},
|
|
{"id": 3, "text": "birds sing"},
|
|
{"id": 1, "text": "dogs bark"},
|
|
]
|
|
assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1])
|
|
|
|
|
|
def test_a_row_the_model_leaves_out_raises():
|
|
reranker = reranker_replying(
|
|
[
|
|
{"content": "cats purr", "relevance_score": 0.9},
|
|
{"content": "dogs bark", "relevance_score": 0.1},
|
|
]
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="1 of 3"):
|
|
reranker.rerank_vector("which animal purrs?", vector_results())
|