From 1b31e2ed0ab712a6d5dd1e518bc5a8042a5315bc Mon Sep 17 00:00:00 2001
From: Ritika shrestha <87307821+ritsth@users.noreply.github.com>
Date: Mon, 29 Jun 2026 16:56:38 -0700
Subject: [PATCH] fix(deepinfra): strip only trailing /openai from rerank
 api_base

get_complete_url used api_base.replace("openai", ""), which removes the
substring "openai" from anywhere in the URL. Any api_base that contains
"openai" outside the intended trailing path segment (for example a gateway
host like https://openai-gw.example.com/v1/openai) was silently corrupted into
a malformed endpoint.

Strip only a trailing "/openai" segment instead, leaving the rest of the host
and path untouched. The DeepInfra default base and existing happy paths are
unchanged.
---
 .../llms/deepinfra/rerank/transformation.py   |  8 +++--
 .../test_deepinfra_rerank_transformation.py   | 29 +++++++++++++------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py
index 82069e4e195..51c6eec127c 100644
--- a/litellm/llms/deepinfra/rerank/transformation.py
+++ b/litellm/llms/deepinfra/rerank/transformation.py
@@ -52,8 +52,12 @@ class DeepinfraRerankConfig(BaseRerankConfig):
                 "Deepinfra API Base is required. api_base=None. Set in call or via `DEEPINFRA_API_BASE` env var."
             )
 
-        # Remove 'openai' from the base if present
-        api_base_clean = api_base.replace("openai", "") if "openai" in api_base else api_base
+        # Strip only a trailing '/openai' segment (DeepInfra's OpenAI-compatible base
+        # ends in '/openai'); a plain str.replace would also corrupt an api_base that
+        # contains 'openai' elsewhere, e.g. in the host or a proxy path.
+        api_base_clean = api_base.rstrip("/")
+        if api_base_clean.endswith("/openai"):
+            api_base_clean = api_base_clean[: -len("/openai")]
 
         # Remove any trailing slashes for consistency, then add one
         api_base_clean = api_base_clean.rstrip("/") + "/"
diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
index a5411078cf7..8223aa70456 100644
--- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
+++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
@@ -43,6 +43,23 @@ class TestDeepinfraRerankTransform:
         with pytest.raises(ValueError, match="Deepinfra API Base is required"):
             self.config.get_complete_url(None, model)
 
+    def test_get_complete_url_preserves_openai_in_host(self):
+        """Regression: only a trailing '/openai' segment should be stripped.
+
+        A plain str.replace("openai", "") corrupted any api_base that contained
+        'openai' elsewhere (e.g. a gateway host or proxy path), producing a
+        malformed URL. The host/path must be preserved.
+        """
+        model = "Qwen/Qwen3-Reranker-0.6B"
+
+        # 'openai' in the host, with a trailing '/openai' segment to strip
+        url = self.config.get_complete_url("https://openai-gw.mycorp.com/v1/openai", model)
+        assert url == "https://openai-gw.mycorp.com/v1/inference/Qwen/Qwen3-Reranker-0.6B"
+
+        # 'openai' in the host, no trailing '/openai' segment to strip
+        url = self.config.get_complete_url("https://my-openai-proxy.example.com/v1", model)
+        assert url == "https://my-openai-proxy.example.com/v1/inference/Qwen/Qwen3-Reranker-0.6B"
+
     def test_map_cohere_rerank_params_basic(self):
         """Test basic parameter mapping for DeepInfra rerank."""
         params = self.config.map_cohere_rerank_params(
@@ -166,9 +183,7 @@ class TestDeepinfraRerankTransform:
         assert result._hidden_params["model"] == self.model
 
         # Verify logging was called
-        mock_logging.post_call.assert_called_once_with(
-            original_response=mock_response.text
-        )
+        mock_logging.post_call.assert_called_once_with(original_response=mock_response.text)
 
     def test_transform_rerank_response_minimal(self):
         """Test response transformation with minimal data."""
@@ -245,12 +260,8 @@ class TestDeepinfraRerankTransform:
                 query="query1",
                 documents=documents,
             )
-            assert (
-                params["queries"] == expected_queries
-            ), f"Failed for {len(documents)} documents"
-            assert len(params["queries"]) == len(
-                documents
-            ), "Queries length must match documents length"
+            assert params["queries"] == expected_queries, f"Failed for {len(documents)} documents"
+            assert len(params["queries"]) == len(documents), "Queries length must match documents length"
 
     def test_get_error_class_basic(self):
         """Test error class generation for basic error."""