fix(rerank): stamp a fresh response id when Voyage, watsonx, or Fireworks omit one
Some checks failed
ai-gateway image / ai-gateway release image (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled

This commit is contained in:
mateo-berri 2026-09-13 01:34:22 -07:00
parent 4d2352d0b5
commit 31bd4d34ed
6 changed files with 76 additions and 31 deletions

View file

@ -250,8 +250,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
rerank_results.append(rerank_result)
# Use model name as id if no id is provided
response_id: Final = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4())
response_id: Final = raw_response_json.get("id") or str(uuid.uuid4())
return RerankResponse(
id=response_id,

View file

@ -9,6 +9,7 @@ from typing import Any, Final
import httpx
from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
@ -127,7 +128,7 @@ class VoyageRerankConfig(BaseRerankConfig):
rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
return RerankResponse(
id=_json_response.get("id", f"voyage-rerank-{model}"),
id=_json_response.get("id") or str(uuid.uuid4()),
results=transformed_results,
meta=rerank_meta,
)

View file

@ -191,7 +191,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
transformed_results.append(transformed_result)
response_id: Final = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4())
response_id: Final = raw_response_json.get("id") or str(uuid.uuid4())
# Extract usage information
_tokens: Final = RerankTokens(

View file

@ -3,6 +3,7 @@ Tests for Fireworks AI rerank transformation functionality.
"""
import json
import uuid
from unittest.mock import MagicMock
import httpx
@ -181,8 +182,7 @@ class TestFireworksAIRerankTransform:
)
# Verify response structure
# Fireworks AI doesn't return "id", so it uses "model" as the id
assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
assert uuid.UUID(result.id).version == 4
assert len(result.results) == 2
assert result.results[0]["index"] == 0
assert result.results[0]["relevance_score"] == 0.95
@ -229,16 +229,14 @@ class TestFireworksAIRerankTransform:
logging_obj=mock_logging,
)
# Fireworks AI doesn't return "id", so it uses "model" as the id
assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
assert uuid.UUID(result.id).version == 4
assert len(result.results) == 2
assert result.results[0]["index"] == 0
assert result.results[0]["relevance_score"] == 0.95
# Document should not be present
assert "document" not in result.results[0]
def test_transform_rerank_response_missing_id(self):
"""Test response transformation when id is missing (should use model name or generate UUID)."""
def test_transform_rerank_response_missing_id_stamps_a_fresh_id_per_call(self):
response_data = {
"object": "list",
"model": "accounts/fireworks/models/qwen3-reranker-8b",
@ -248,23 +246,22 @@ class TestFireworksAIRerankTransform:
"usage": {"total_tokens": 10},
}
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
def transform() -> str:
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
return self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=RerankResponse(),
logging_obj=MagicMock(),
).id
mock_logging = MagicMock()
model_response = RerankResponse()
first, second = transform(), transform()
result = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
)
# Should use model name when id is missing
assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
assert first != second
assert "accounts/fireworks/models/qwen3-reranker-8b" not in (first, second)
def test_transform_rerank_response_missing_results(self):
"""Test that missing results raises ValueError."""

View file

@ -3,6 +3,7 @@ Tests for Voyage AI rerank transformation functionality.
"""
import json
import uuid
from unittest.mock import MagicMock, patch
import httpx
@ -258,6 +259,33 @@ class TestVoyageRerankTransform:
assert "Failed to parse response" in str(exc_info.value)
def test_transform_rerank_response_without_id_stamps_a_fresh_id_per_call(self):
response_data = {
"object": "list",
"data": [{"relevance_score": 0.5, "index": 0}],
"model": "rerank-2.5",
"usage": {"total_tokens": 10},
}
def transform() -> str:
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.text = json.dumps(response_data)
mock_response.headers = {}
return self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=RerankResponse(),
logging_obj=MagicMock(),
).id
first, second = transform(), transform()
assert uuid.UUID(first).version == 4
assert first != second
assert f"voyage-rerank-{self.model}" not in (first, second)
def test_get_supported_cohere_rerank_params(self):
"""Test getting supported parameters for Voyage AI rerank."""
supported_params = self.config.get_supported_cohere_rerank_params(self.model)

View file

@ -120,9 +120,7 @@ class TestIBMWatsonXRerankTransform:
logging_obj=mock_logging,
)
# Verify response structure
# IBM watsonx.ai doesn't return "id", so it uses "model" as the id
assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
assert uuid.UUID(result.id).version == 4
assert len(result.results) == 2
assert result.results[0]["index"] == 0
assert result.results[0]["relevance_score"] == 6.53515625
@ -172,9 +170,7 @@ class TestIBMWatsonXRerankTransform:
logging_obj=mock_logging,
)
# Verify response structure
# IBM watsonx.ai doesn't return "id", so it uses "model" as the id
assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2"
assert uuid.UUID(result.id).version == 4
assert len(result.results) == 2
assert result.results[0]["index"] == 0
@ -231,6 +227,30 @@ class TestIBMWatsonXRerankTransform:
logging_obj=mock_logging,
)
def test_transform_rerank_response_without_id_stamps_a_fresh_id_per_call(self):
response_data = {
"model_id": self.model,
"results": [{"index": 0, "score": 1.5}],
"input_token_count": 12,
}
def transform() -> str:
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
return self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=RerankResponse(),
logging_obj=MagicMock(),
).id
first, second = transform(), transform()
assert first != second
assert self.model not in (first, second)
def test_get_supported_cohere_rerank_params(self):
"""Test getting supported parameters for IBM watsonx.ai rerank."""
supported_params = self.config.get_supported_cohere_rerank_params(self.model)