Skip to content

Commit f82ceb0

Browse files
google-genai-botcopybara-github
authored andcommitted
chore: Replace proxy methods with utils implementation
PiperOrigin-RevId: 864773118
1 parent 33012e6 commit f82ceb0

2 files changed

Lines changed: 6 additions & 12 deletions

File tree

src/google/adk/models/apigee_llm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from google.genai import types
2626
from typing_extensions import override
2727

28+
from ..utils._google_client_headers import merge_tracking_headers
2829
from ..utils.env_utils import is_env_enabled
2930
from .google_llm import Gemini
3031

@@ -145,7 +146,7 @@ def api_client(self) -> Client:
145146
kwargs_for_http_options['api_version'] = self._api_version
146147
http_options = types.HttpOptions(
147148
base_url=self._proxy_url,
148-
headers=self._merge_tracking_headers(self._custom_headers),
149+
headers=merge_tracking_headers(self._custom_headers),
149150
retry_options=self.retry_options,
150151
**kwargs_for_http_options,
151152
)

src/google/adk/models/google_llm.py

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ async def generate_content_async(
192192
if llm_request.config:
193193
if not llm_request.config.http_options:
194194
llm_request.config.http_options = types.HttpOptions()
195-
llm_request.config.http_options.headers = self._merge_tracking_headers(
195+
llm_request.config.http_options.headers = merge_tracking_headers(
196196
llm_request.config.http_options.headers
197197
)
198198

@@ -303,7 +303,7 @@ def api_client(self) -> Client:
303303

304304
return Client(
305305
http_options=types.HttpOptions(
306-
headers=self._tracking_headers(),
306+
headers=get_tracking_headers(),
307307
retry_options=self.retry_options,
308308
)
309309
)
@@ -316,9 +316,6 @@ def _api_backend(self) -> GoogleLLMVariant:
316316
else GoogleLLMVariant.GEMINI_API
317317
)
318318

319-
def _tracking_headers(self) -> dict[str, str]:
320-
return get_tracking_headers()
321-
322319
@cached_property
323320
def _live_api_version(self) -> str:
324321
if self._api_backend == GoogleLLMVariant.VERTEX_AI:
@@ -334,7 +331,7 @@ def _live_api_client(self) -> Client:
334331

335332
return Client(
336333
http_options=types.HttpOptions(
337-
headers=self._tracking_headers(), api_version=self._live_api_version
334+
headers=get_tracking_headers(), api_version=self._live_api_version
338335
)
339336
)
340337

@@ -358,7 +355,7 @@ async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
358355
if not llm_request.live_connect_config.http_options.headers:
359356
llm_request.live_connect_config.http_options.headers = {}
360357
llm_request.live_connect_config.http_options.headers = (
361-
self._merge_tracking_headers(
358+
merge_tracking_headers(
362359
llm_request.live_connect_config.http_options.headers
363360
)
364361
)
@@ -451,10 +448,6 @@ async def _preprocess_request(self, llm_request: LlmRequest) -> None:
451448
llm_request.config.system_instruction = None
452449
await self._adapt_computer_use_tool(llm_request)
453450

454-
def _merge_tracking_headers(self, headers: dict[str, str]) -> dict[str, str]:
455-
"""Merge tracking headers to the given headers."""
456-
return merge_tracking_headers(headers)
457-
458451

459452
def _build_function_declaration_log(
460453
func_decl: types.FunctionDeclaration,

0 commit comments

Comments
 (0)