scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

fix(infinity): handle object rerank documents

Browse source
This commit is contained in:
nightcityblade 2026-10-07 15:38:26 +00:00
parent e1bdd42726
commit a58a7582b1
2 changed files with 12 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/llms/infinity/rerank/transformation.py
View file

@ -31,12 +31,16 @@ class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]):
"""The token counters Infinity reports in the ``usage`` block of a rerank response."""
class _InfinityRerankDocument(TypedDict):
text: ReadOnly[str]
class _InfinityRerankResult(TypedDict):
"""One scored document in an Infinity ``/v1/rerank`` response."""
index: ReadOnly[int]
relevance_score: ReadOnly[float]
document: ReadOnly[str]
document: ReadOnly[str | _InfinityRerankDocument]
class _InfinityRerankResponse(TypedDict):
@ -104,8 +108,6 @@ class InfinityRerankConfig(CohereRerankConfig):
) -> RerankResponse:
"""
Transform Infinity rerank response
No transformation required, Infinity follows Cohere API response format
"""
try:
raw_response_json: Final = _parse_rerank_response(raw_response)
@ -129,8 +131,10 @@ class InfinityRerankConfig(CohereRerankConfig):
index=result.get("index"),
relevance_score=result.get("relevance_score"),
)
if result.get("document"):
_rerank_response["document"] = RerankResponseDocument(text=result.get("document"))
if document := result.get("document"):
_rerank_response["document"] = RerankResponseDocument(
text=document.get("text") if isinstance(document, Mapping) else document
)
cohere_results.append(_rerank_response)
if cohere_results is None:
raise ValueError(f"No results found in the response={raw_response_json}")

5
tests/unit/llms/infinity/test_infinity.py
View file

@ -69,8 +69,9 @@ async def test_infinity_rerank():
assert_response_shape(response, custom_llm_provider="infinity")
@pytest.mark.parametrize("provider_document", ["hello", {"text": "hello"}])
@pytest.mark.asyncio()
async def test_infinity_rerank_with_return_documents():
async def test_infinity_rerank_with_return_documents(provider_document: str | dict[str, str]):
mock_response = AsyncMock()
mock_response = AsyncMock()
@ -78,7 +79,7 @@ async def test_infinity_rerank_with_return_documents():
def return_val():
return {
"id": "cmpl-mockid",
"results": [{"index": 0, "relevance_score": 0.95, "document": "hello"}],
"results": [{"index": 0, "relevance_score": 0.95, "document": provider_document}],
"usage": {"prompt_tokens": 100, "total_tokens": 150},
}