From 8d773e66af010f809f461f24adcad56bcc167698 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 10 Sep 2024 17:50:24 -0700 Subject: [PATCH] fix gemini streaming test --- .../gemini/vertex_and_google_ai_studio_gemini.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py index 231c39079d4..b4615ab3df9 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py @@ -1070,6 +1070,7 @@ class VertexLLM(VertexBase): vertex_project: Optional[str] = None, vertex_location: Optional[str] = None, vertex_credentials: Optional[str] = None, + gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, ) -> CustomStreamWrapper: request_body = await async_transform_request_body(**data) # type: ignore @@ -1086,7 +1087,7 @@ class VertexLLM(VertexBase): auth_header, api_base = self._get_token_and_url( model=model, - gemini_api_key=None, + gemini_api_key=gemini_api_key, auth_header=_auth_header, vertex_project=vertex_project, vertex_location=vertex_location, @@ -1151,6 +1152,7 @@ class VertexLLM(VertexBase): vertex_project: Optional[str] = None, vertex_location: Optional[str] = None, vertex_credentials: Optional[str] = None, + gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, ) -> Union[ModelResponse, CustomStreamWrapper]: @@ -1166,7 +1168,7 @@ class VertexLLM(VertexBase): auth_header, api_base = self._get_token_and_url( model=model, - gemini_api_key=None, + gemini_api_key=gemini_api_key, auth_header=_auth_header, vertex_project=vertex_project, vertex_location=vertex_location, @@ -1283,6 +1285,7 @@ class VertexLLM(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, + gemini_api_key=gemini_api_key, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, ) @@ -1305,6 +1308,7 @@ class VertexLLM(VertexBase): vertex_project=vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, + gemini_api_key=gemini_api_key, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, )