test_rerank_response_assertions (BerriAI#7476)

kp-forks · Dec 30, 2024 · 83879d2 · 83879d2
1 parent a003af6
commit 83879d2
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 7 deletions.
diff --git 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
@@ -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")