phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(vertex_httpx.py): allow passing extra headers

Browse source 
Closes https://github.com/BerriAI/litellm/pull/4327
This commit is contained in:
Krrish Dholakia 2024-06-20 21:14:22 -07:00
parent 3a510582c2
commit dfd5882c31
3 changed files with 5 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/vertex_httpx.py
View file

@ -775,6 +775,8 @@ class VertexLLM(BaseLLM):
} }
if auth_header is not None: if auth_header is not None:
headers["Authorization"] = f"Bearer {auth_header}" headers["Authorization"] = f"Bearer {auth_header}"
if extra_headers is not None:
headers.update(extra_headers)
## LOGGING ## LOGGING
logging_obj.pre_call( logging_obj.pre_call(

1
litellm/main.py
View file

@ -1931,6 +1931,7 @@ def completion(
custom_llm_provider=custom_llm_provider, custom_llm_provider=custom_llm_provider,
client=client, client=client,
api_base=api_base, api_base=api_base,
extra_headers=extra_headers,
) )
elif custom_llm_provider == "vertex_ai": elif custom_llm_provider == "vertex_ai":

2
litellm/tests/test_amazing_vertex_completion.py
View file

@ -877,6 +877,7 @@ async def test_gemini_pro_httpx_custom_api_base(provider):
response_format={"type": "json_object"}, response_format={"type": "json_object"},
client=client, client=client,
api_base="my-custom-api-base", api_base="my-custom-api-base",
extra_headers={"hello": "world"},
) )
except Exception as e: except Exception as e:
pass pass
@ -884,6 +885,7 @@ async def test_gemini_pro_httpx_custom_api_base(provider):
mock_call.assert_called_once() mock_call.assert_called_once()
assert "my-custom-api-base:generateContent" == mock_call.call_args.kwargs["url"] assert "my-custom-api-base:generateContent" == mock_call.call_args.kwargs["url"]
assert "hello" in mock_call.call_args.kwargs["headers"]
@pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call") @pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call")