From 83879d2a3d77e52785586537cd90c87a664bb010 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 30 Dec 2024 10:12:56 -0800
Subject: [PATCH] test_rerank_response_assertions (#7476)

---
 litellm/types/rerank.py              | 14 ++++++-------
 tests/llm_translation/test_rerank.py | 31 ++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py
index 019d90a0ae65..7ac9d850d984 100644
--- a/litellm/types/rerank.py
+++ b/litellm/types/rerank.py
@@ -30,19 +30,19 @@ class OptionalRerankParams(TypedDict, total=False):
 
 
 class RerankBilledUnits(TypedDict, total=False):
-    search_units: int
-    total_tokens: int
+    search_units: Optional[int]
+    total_tokens: Optional[int]
 
 
 class RerankTokens(TypedDict, total=False):
-    input_tokens: int
-    output_tokens: int
+    input_tokens: Optional[int]
+    output_tokens: Optional[int]
 
 
 class RerankResponseMeta(TypedDict, total=False):
-    api_version: dict
-    billed_units: RerankBilledUnits
-    tokens: RerankTokens
+    api_version: Optional[dict]
+    billed_units: Optional[RerankBilledUnits]
+    tokens: Optional[RerankTokens]
 
 
 class RerankResponseResult(TypedDict):
diff --git a/tests/llm_translation/test_rerank.py b/tests/llm_translation/test_rerank.py
index 48a6fe0ca260..82efa92dfda3 100644
--- a/tests/llm_translation/test_rerank.py
+++ b/tests/llm_translation/test_rerank.py
@@ -20,6 +20,7 @@
 import pytest
 
 import litellm
+from litellm.types.rerank import RerankResponse
 from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@@ -354,3 +355,33 @@ async def test_basic_rerank_caching(sync_mode, top_n_1, top_n_2, expect_cache_hi
     assert response.results is not None
 
     assert_response_shape(response, custom_llm_provider="cohere")
+
+
+def test_rerank_response_assertions():
+    r = RerankResponse(
+        **{
+            "id": "ab0fcca0-b617-11ef-b292-0242ac110002",
+            "results": [
+                {"index": 2, "relevance_score": 0.9958819150924683, "document": None},
+                {"index": 0, "relevance_score": 0.001293411129154265, "document": None},
+                {
+                    "index": 1,
+                    "relevance_score": 7.641685078851879e-05,
+                    "document": None,
+                },
+                {
+                    "index": 3,
+                    "relevance_score": 7.621097756782547e-05,
+                    "document": None,
+                },
+            ],
+            "meta": {
+                "api_version": None,
+                "billed_units": None,
+                "tokens": None,
+                "warnings": None,
+            },
+        }
+    )
+
+    assert_response_shape(r, custom_llm_provider="custom")