mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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
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:
parent
4d2352d0b5
commit
31bd4d34ed
6 changed files with 76 additions and 31 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue