fix: keep existing header values while merging tracking headers for llm_request.config.http_options in Gemini.generate_content_async

PiperOrigin-RevId: 789013693
This commit is contained in:
Xuan Yang
2025-07-30 13:09:56 -07:00
committed by Copybara-Service
parent 3be1bb37d9
commit 6191412b07
2 changed files with 21 additions and 4 deletions
+20 -3
View File
@@ -122,9 +122,9 @@ class Gemini(BaseLlm):
if llm_request.config:
if not llm_request.config.http_options:
llm_request.config.http_options = types.HttpOptions()
if not llm_request.config.http_options.headers:
llm_request.config.http_options.headers = {}
llm_request.config.http_options.headers.update(self._tracking_headers)
llm_request.config.http_options.headers = self._merge_tracking_headers(
llm_request.config.http_options.headers
)
if stream:
responses = await self.api_client.aio.models.generate_content_stream(
@@ -336,6 +336,23 @@ class Gemini(BaseLlm):
llm_request.config.system_instruction = None
await self._adapt_computer_use_tool(llm_request)
def _merge_tracking_headers(self, headers: dict[str, str]) -> dict[str, str]:
"""Merge tracking headers to the given headers."""
headers = headers or {}
for key, tracking_header_value in self._tracking_headers.items():
custom_value = headers.get(key, None)
if not custom_value:
headers[key] = tracking_header_value
continue
# Merge tracking headers with existing headers and avoid duplicates.
value_parts = tracking_header_value.split(' ')
for custom_value_part in custom_value.split(' '):
if custom_value_part not in value_parts:
value_parts.append(custom_value_part)
headers[key] = ' '.join(value_parts)
return headers
def _build_function_declaration_log(
func_decl: types.FunctionDeclaration,