scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

test(vertex_ai): expand test coverage for return_documents parameter edge cases

Browse source
This commit is contained in:
somuai 2026-08-26 17:49:09 +05:30
parent 3df9c6c258
commit e441d38d17
1 changed files with 34 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
View file

@ -391,6 +391,40 @@ class TestVertexAIRerankTransform:
assert "document" not in result_request_data.results[0]
assert "document" not in result_request_data.results[1]
# Test with litellm_params={"return_documents": True}
result_litellm_params = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
litellm_params={"return_documents": True},
)
assert result_litellm_params.results[0]["document"]["text"] == "Content of document 1"
# Test with request_data={"return_documents": True}
result_req_data_true = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=mock_logging,
request_data={"return_documents": True},
)
assert result_req_data_true.results[0]["document"]["text"] == "Content of document 1"
# Test with records missing content
no_content_response_data = {"records": [{"id": "0", "score": 0.9}]}
mock_no_content = MagicMock(spec=httpx.Response)
mock_no_content.json.return_value = no_content_response_data
mock_no_content.text = json.dumps(no_content_response_data)
result_no_content = self.config.transform_rerank_response(
model=self.model,
raw_response=mock_no_content,
model_response=model_response,
logging_obj=mock_logging,
optional_params={"return_documents": True},
)
assert "document" not in result_no_content.results[0]
def test_transform_rerank_response_with_ignore_record_details(self):
"""Test response transformation when ignoreRecordDetailsInResponse=true."""
# Mock response with only IDs (when ignoreRecordDetailsInResponse=true)