from urllib.parse import urlparse, urlunparse, parse_qs, urlencode
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
_original_check_proxy = VertexBase._check_custom_proxy
def fixed_check_custom_proxy(
self, api_base: str | None, custom_llm_provider: str,
gemini_api_key: str | None, endpoint: str, stream: bool | None,
auth_header: str | None, url: str
) -> tuple[str | None, str]:
is_helicone_gemini_proxy = (
api_base is not None
and "gateway.helicone.ai" in api_base
and custom_llm_provider == "gemini"
)
if is_helicone_gemini_proxy:
final_url = api_base
auth = None
if stream:
parsed = urlparse(final_url)
qs = parse_qs(parsed.query)
qs['alt'] = ['sse']
final_url = urlunparse((parsed.scheme, parsed.netloc, parsed.path, parsed.params, urlencode(qs, doseq=True), parsed.fragment))
return (auth, final_url)
elif _original_check_proxy:
return _original_check_proxy(
self, api_base, custom_llm_provider, gemini_api_key,
endpoint, stream, auth_header, url
)
else:
return (auth_header, url)
VertexBase._check_custom_proxy = fixed_check_custom_proxy