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

fix gemini streaming test

Browse source
This commit is contained in:
Ishaan Jaff 2024-09-10 17:50:24 -07:00
parent 26aa20df73
commit 8d773e66af
1 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -1070,6 +1070,7 @@ class VertexLLM(VertexBase):
vertex_project: Optional[str] = None, vertex_project: Optional[str] = None,
vertex_location: Optional[str] = None, vertex_location: Optional[str] = None,
vertex_credentials: Optional[str] = None, vertex_credentials: Optional[str] = None,
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None, extra_headers: Optional[dict] = None,
) -> CustomStreamWrapper: ) -> CustomStreamWrapper:
request_body = await async_transform_request_body(**data) # type: ignore 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( auth_header, api_base = self._get_token_and_url(
model=model, model=model,
gemini_api_key=None, gemini_api_key=gemini_api_key,
auth_header=_auth_header, auth_header=_auth_header,
vertex_project=vertex_project, vertex_project=vertex_project,
vertex_location=vertex_location, vertex_location=vertex_location,
@ -1151,6 +1152,7 @@ class VertexLLM(VertexBase):
vertex_project: Optional[str] = None, vertex_project: Optional[str] = None,
vertex_location: Optional[str] = None, vertex_location: Optional[str] = None,
vertex_credentials: Optional[str] = None, vertex_credentials: Optional[str] = None,
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None, extra_headers: Optional[dict] = None,
) -> Union[ModelResponse, CustomStreamWrapper]: ) -> Union[ModelResponse, CustomStreamWrapper]:
@ -1166,7 +1168,7 @@ class VertexLLM(VertexBase):
auth_header, api_base = self._get_token_and_url( auth_header, api_base = self._get_token_and_url(
model=model, model=model,
gemini_api_key=None, gemini_api_key=gemini_api_key,
auth_header=_auth_header, auth_header=_auth_header,
vertex_project=vertex_project, vertex_project=vertex_project,
vertex_location=vertex_location, vertex_location=vertex_location,
@ -1283,6 +1285,7 @@ class VertexLLM(VertexBase):
vertex_project=vertex_project, vertex_project=vertex_project,
vertex_location=vertex_location, vertex_location=vertex_location,
vertex_credentials=vertex_credentials, vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
custom_llm_provider=custom_llm_provider, custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers, extra_headers=extra_headers,
) )
@ -1305,6 +1308,7 @@ class VertexLLM(VertexBase):
vertex_project=vertex_project, vertex_project=vertex_project,
vertex_location=vertex_location, vertex_location=vertex_location,
vertex_credentials=vertex_credentials, vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
custom_llm_provider=custom_llm_provider, custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers, extra_headers=extra_headers,
) )