From f860cbca34f061e09704c245a29690e028e3eb6b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 29 Jun 2024 10:49:10 -0700 Subject: [PATCH] test test_pass_through_endpoint_rerank --- litellm/tests/test_pass_through_endpoints.py | 34 ++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/litellm/tests/test_pass_through_endpoints.py b/litellm/tests/test_pass_through_endpoints.py index 0287f6f421b..0f234dfa8bd 100644 --- a/litellm/tests/test_pass_through_endpoints.py +++ b/litellm/tests/test_pass_through_endpoints.py @@ -49,3 +49,37 @@ async def test_pass_through_endpoint(client, monkeypatch): # Assert the response assert response.status_code == 200 assert response.json() == {"message": "Mocked response"} + + +@pytest.mark.asyncio +async def test_pass_through_endpoint_rerank(client): + _cohere_api_key = os.environ.get("COHERE_API_KEY") + + # Define a pass-through endpoint + pass_through_endpoints = [ + { + "path": "/v1/rerank", + "target": "https://api.cohere.com/v1/rerank", + "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + } + ] + + # Initialize the pass-through endpoint + await initialize_pass_through_endpoints(pass_through_endpoints) + + _json_data = { + "model": "rerank-english-v3.0", + "query": "What is the capital of the United States?", + "top_n": 3, + "documents": [ + "Carson City is the capital city of the American state of Nevada." + ], + } + + # Make a request to the pass-through endpoint + response = client.post("/v1/rerank", json=_json_data) + + print("JSON response: ", _json_data) + + # Assert the response + assert response.status_code == 200