From 7ef590df84d0a6ad452fa6dfe88c5f42394b4bfb Mon Sep 17 00:00:00 2001
From: Zayd <zaydsimjee@gmail.com>
Date: Tue, 1 Jul 2025 18:07:38 -0700
Subject: [PATCH] Passes through extra_ properties on "custom" llm provider
 (#12185)

* Passes through headers on "custom" llm provider

* add test

* adds extra_body support for custom llm providers
---
 litellm/main.py                 |  2 ++
 tests/test_litellm/test_main.py | 60 +++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index 1f45894d76..8dff82a9b8 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -3221,6 +3221,7 @@ def completion(  # type: ignore # noqa: PLR0915
             prompt = " ".join([message["content"] for message in messages])  # type: ignore
             resp = litellm.module_level_client.post(
                 url,
+                headers=headers,
                 json={
                     "model": model,
                     "params": {
@@ -3230,6 +3231,7 @@ def completion(  # type: ignore # noqa: PLR0915
                         "top_p": top_p,
                         "top_k": kwargs.get("top_k"),
                     },
+                    **kwargs.get("extra_body", {}),
                 },
             )
             response_json = resp.json()
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index e69292f569..201313aec7 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -268,6 +268,66 @@ def test_bedrock_latency_optimized_inference():
         assert json_data["performanceConfig"]["latency"] == "optimized"
 
 
+def test_custom_provider_with_extra_headers():
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post:
+        response = litellm.completion(
+            model="custom/custom",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            headers={"X-Custom-Header": "custom-value"},
+            api_base="https://example.com/api/v1",
+        )
+
+        mock_post.assert_called_once()
+        assert mock_post.call_args[1]["headers"]["X-Custom-Header"] == "custom-value"
+
+def test_custom_provider_with_extra_body():
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+    
+    with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post:
+        response = litellm.completion(
+            model="custom/custom",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            extra_body={"X-Custom-BodyValue": "custom-value", "X-Custom-BodyValue2": "custom-value2"},
+            api_base="https://example.com/api/v1",
+        )
+        mock_post.assert_called_once()
+
+        assert mock_post.call_args[1]["json"]["X-Custom-BodyValue"] == "custom-value"
+        assert mock_post.call_args[1]["json"] == {
+            'model': 'custom',
+            'params': {
+                'prompt': ['Hello, how are you?'],
+                'max_tokens': None,
+                'temperature': None,
+                'top_p': None,
+                'top_k': None
+            },
+            'X-Custom-BodyValue': 'custom-value',
+            'X-Custom-BodyValue2': 'custom-value2'
+        }
+
+    # test that extra_body is not passed if not provided
+    with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post:
+        response = litellm.completion(
+            model="custom/custom",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            api_base="https://example.com/api/v1",
+        )
+        mock_post.assert_called_once()
+        assert mock_post.call_args[1]["json"] == {
+            'model': 'custom',
+            'params': {
+                'prompt': ['Hello, how are you?'],
+                'max_tokens': None,
+                'temperature': None,
+                'top_p': None,
+                'top_k': None
+            }
+        }
+
+
 @pytest.fixture(autouse=True)
 def set_openrouter_api_key():
     original_api_key = os.environ.get("OPENROUTER_API_KEY")