feat: support injecting custom LLM clients into Gemini and AnthropicLlm

Merge https://github.com/google/adk-python/pull/5035

Enable injecting pre-configured LLM clients into Gemini and AnthropicLlm models to support multi-agent systems with distinct configurations.

Fixes #5027

PiperOrigin-RevId: 968151995

src/google/adk/models/anthropic_llm.py

@@ -824,6 +824,11 @@ class AnthropicLlm(BaseLlm):
   model: str = "claude-sonnet-4-20250514"
   max_tokens: int = 8192
 
+  client: Optional[Union[AsyncAnthropic, AsyncAnthropicVertex]] = Field(
+      default=None, exclude=True
+  )
+  """An optional pre-configured Anthropic client."""
+
   @classmethod
   @override
   def supported_models(cls) -> list[str]:
@@ -1151,6 +1156,8 @@ class AnthropicLlm(BaseLlm):
 
   @cached_property
   def _anthropic_client(self) -> AsyncAnthropic | AsyncAnthropicVertex:
+    if self.client:
+      return self.client
     client = AsyncAnthropic()
     # Let the SDK run its own credential resolution first, then ask the client
     # what it found. Enumerating credential sources here would reject setups
@@ -1190,6 +1197,10 @@ class Claude(AnthropicLlm):
   @cached_property
   @override
   def _anthropic_client(self) -> AsyncAnthropicVertex:
+    if self.client is not None:
+      if not isinstance(self.client, AsyncAnthropicVertex):
+        raise ValueError("Claude requires an AsyncAnthropicVertex client.")
+      return self.client
     project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
     location = os.environ.get("GOOGLE_CLOUD_LOCATION")
 

src/google/adk/models/apigee_llm.py

@@ -28,6 +28,7 @@ from typing import AsyncGenerator
 from typing import Generator
 from typing import Optional
 from typing import TYPE_CHECKING
+import warnings
 
 from google.adk import version as adk_version
 from google.genai import types
@@ -116,6 +117,7 @@ class ApigeeLlm(Gemini):
       retry_options: Optional[types.HttpRetryOptions] = None,
       api_type: ApiType | str = ApiType.UNKNOWN,
       credentials: Credentials | None = None,
+      client: Client | None = None,
   ) -> None:
     """Initializes the Apigee LLM backend.
 
@@ -152,9 +154,10 @@ class ApigeeLlm(Gemini):
         additional OAuth scopes (e.g., `userinfo.email` for tokeninfo-based
         caller identification). When omitted, the default `genai.Client`
         authentication flow is used.
+      client: An optional pre-configured google-genai Client.
     """  # fmt: skip
 
-    super().__init__(model=model, retry_options=retry_options)
+    super().__init__(model=model, retry_options=retry_options, client=client)
     # Validate the model string. Create a helper method to validate the model
     # string.
     if not _validate_model_string(model):
@@ -200,6 +203,22 @@ class ApigeeLlm(Gemini):
     self._user_agent = f'google-adk/{adk_version.__version__}'
     self._credentials = credentials
 
+    if client:
+      if self._proxy_url or self._custom_headers:
+        warnings.warn(
+            'Both client and proxy_url/custom_headers were provided. The'
+            ' injected client will be used as-is for GENAI calls, and'
+            ' proxy_url/custom_headers will be ignored. Ensure the injected'
+            ' client is pre-configured with the correct proxy and headers.',
+            UserWarning,
+        )
+      if self._api_type == ApigeeLlm.ApiType.CHAT_COMPLETIONS:
+        warnings.warn(
+            'An injected client was provided but ApiType is CHAT_COMPLETIONS. '
+            'The injected client will be ignored for CHAT_COMPLETIONS calls.',
+            UserWarning,
+        )
+
   @classmethod
   @override
   def supported_models(cls) -> list[str]:
@@ -264,6 +283,9 @@ class ApigeeLlm(Gemini):
     Returns:
       The api client.
     """
+    if self.client:
+      return self.client
+
     from google.genai import Client
 
     http_options = types.HttpOptions(

src/google/adk/models/google_llm.py

@@ -37,6 +37,8 @@ from google.genai.errors import ClientError
 from pydantic import Field
 from typing_extensions import override
 
+from google import genai
+
 from ..utils._google_client_headers import get_tracking_headers
 from ..utils._google_client_headers import merge_tracking_headers
 from ..utils.context_utils import Aclosing
@@ -121,6 +123,13 @@ class Gemini(BaseLlm):
 
   model: str = 'gemini-2.5-flash'
 
+  client: Optional[genai.Client] = Field(default=None, exclude=True)
+  """An optional pre-configured google-genai Client.
+
+  When provided, this client will be used for all API calls instead of
+  constructing a new one from environment variables or other attributes.
+  """
+
   client_kwargs: Optional[dict[str, Any]] = Field(
       default=None, exclude=True, repr=False
   )
@@ -381,6 +390,9 @@ class Gemini(BaseLlm):
     Returns:
       The api client.
     """
+    if self.client:
+      return self.client
+
     from google.genai import Client
 
     base_url, api_version = self._base_url_and_api_version
@@ -450,6 +462,9 @@ class Gemini(BaseLlm):
 
   @cached_property
   def _live_api_client(self) -> Client:
+    if self.client:
+      return self.client
+
     from google.genai import Client
 
     base_url, _ = self._base_url_and_api_version

tests/unittests/models/test_custom_client.py

@@ -0,0 +1,285 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for custom client injection in ADK models."""
+
+# pylint: disable=protected-access
+
+from unittest import mock
+
+from anthropic import AsyncAnthropic
+from anthropic import AsyncAnthropicVertex
+from anthropic import types as anthropic_types
+from google.adk.models.anthropic_llm import AnthropicLlm
+from google.adk.models.anthropic_llm import Claude
+from google.adk.models.apigee_llm import ApigeeLlm
+from google.adk.models.google_llm import Gemini
+from google.adk.models.llm_request import LlmRequest
+from google.genai import Client
+from google.genai import types
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
+
+def test_gemini_custom_client():
+  """Verify that Gemini uses the provided custom client."""
+  mock_client = mock.MagicMock(spec=Client)
+  gemini = Gemini(model="gemini-1.5-flash", client=mock_client)
+
+  assert gemini.api_client is mock_client
+  # Verify it persists (cached_property)
+  assert gemini.api_client is mock_client
+  assert gemini._live_api_client is mock_client
+
+
+@pytest.mark.asyncio
+async def test_gemini_uses_custom_client_in_connect():
+  """Verify that Gemini connect uses the provided custom client."""
+  mock_client = mock.MagicMock(spec=Client)
+  mock_live_session = mock.AsyncMock()
+
+  class MockLiveConnect:
+
+    async def __aenter__(self):
+      return mock_live_session
+
+    async def __aexit__(self, *args):
+      pass
+
+  mock_client.aio.live.connect.return_value = MockLiveConnect()
+
+  gemini = Gemini(model="gemini-1.5-flash", client=mock_client)
+  request = LlmRequest(
+      model="gemini-1.5-flash",
+  )
+
+  async with gemini.connect(request) as connection:
+    mock_client.aio.live.connect.assert_called_once()
+    assert connection._gemini_session is mock_live_session
+
+
+def test_anthropic_custom_client():
+  """Verify that AnthropicLlm uses the provided custom client."""
+  mock_client = mock.MagicMock(spec=AsyncAnthropic)
+  anthropic_llm = AnthropicLlm(
+      model="claude-3-5-sonnet-20241022", client=mock_client
+  )
+
+  assert anthropic_llm._anthropic_client is mock_client
+
+
+@pytest.mark.asyncio
+async def test_gemini_uses_custom_client_in_call():
+  """Verify that Gemini calls use the provided custom client's methods."""
+  mock_client = mock.MagicMock(spec=Client)
+  # Mock the nested aio.models.generate_content
+  mock_aio_models = mock_client.aio.models
+
+  gemini = Gemini(model="gemini-1.5-flash", client=mock_client)
+
+  request = LlmRequest(
+      model="gemini-1.5-flash",
+      contents=[Content(role="user", parts=[Part.from_text(text="Hi")])],
+  )
+
+  # Mock the response
+  mock_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=Content(
+                  role="model", parts=[Part.from_text(text="Hello")]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  async def mock_coro(*_, **__):
+    return mock_response
+
+  mock_aio_models.generate_content.return_value = mock_coro()
+
+  # We use stream=False to simplify the mock
+  responses = [
+      r async for r in gemini.generate_content_async(request, stream=False)
+  ]
+
+  assert len(responses) == 1
+  assert responses[0].content.parts[0].text == "Hello"
+  mock_aio_models.generate_content.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_uses_custom_client_in_call():
+  """Verify that AnthropicLlm calls use the provided custom client's methods."""
+  mock_client = mock.MagicMock(spec=AsyncAnthropic)
+  mock_messages = mock_client.messages
+
+  anthropic_llm = AnthropicLlm(
+      model="claude-3-5-sonnet-20241022", client=mock_client
+  )
+
+  request = LlmRequest(
+      model="claude-3-5-sonnet-20241022",
+      contents=[Content(role="user", parts=[Part.from_text(text="Hi")])],
+  )
+
+  mock_response = anthropic_types.Message(
+      id="msg_test",
+      content=[anthropic_types.TextBlock(text="Hello", type="text")],
+      model="claude-3-5-sonnet-20241022",
+      role="assistant",
+      stop_reason="end_turn",
+      type="message",
+      usage=anthropic_types.Usage(input_tokens=1, output_tokens=1),
+  )
+
+  async def mock_coro(*_, **__):
+    return mock_response
+
+  mock_messages.create.return_value = mock_coro()
+
+  responses = [
+      r
+      async for r in anthropic_llm.generate_content_async(request, stream=False)
+  ]
+
+  assert len(responses) == 1
+  assert responses[0].content.parts[0].text == "Hello"
+  mock_messages.create.assert_called()
+
+
+def test_apigee_custom_client():
+  """Verify that ApigeeLlm uses the provided custom client."""
+  mock_client = mock.MagicMock(spec=Client)
+  apigee_llm = ApigeeLlm(
+      model="apigee/gemini/gemini-1.5-flash", client=mock_client
+  )
+
+  assert apigee_llm.api_client is mock_client
+  # Verify it persists (cached_property)
+  assert apigee_llm.api_client is mock_client
+
+
+@pytest.mark.asyncio
+async def test_apigee_uses_custom_client_in_call():
+  """Verify that ApigeeLlm calls use the provided custom client's methods."""
+  mock_client = mock.MagicMock(spec=Client)
+  mock_aio_models = mock_client.aio.models
+
+  apigee_llm = ApigeeLlm(
+      model="apigee/gemini/gemini-1.5-flash", client=mock_client
+  )
+
+  request = LlmRequest(
+      model="apigee/gemini/gemini-1.5-flash",
+      contents=[Content(role="user", parts=[Part.from_text(text="Hi")])],
+  )
+
+  mock_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=Content(
+                  role="model", parts=[Part.from_text(text="Hello")]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  async def mock_coro(*_, **__):
+    return mock_response
+
+  mock_aio_models.generate_content.return_value = mock_coro()
+
+  responses = [
+      r async for r in apigee_llm.generate_content_async(request, stream=False)
+  ]
+
+  assert len(responses) == 1
+  assert responses[0].content.parts[0].text == "Hello"
+  mock_aio_models.generate_content.assert_called()
+
+
+def test_claude_custom_client():
+  """Verify that Claude uses the provided custom client."""
+  mock_client = mock.MagicMock(spec=AsyncAnthropicVertex)
+  claude = Claude(
+      model="projects/p/locations/l/publishers/google/models/claude-3-5-sonnet-v2@20241022",
+      client=mock_client,
+  )
+
+  assert claude._anthropic_client is mock_client
+
+
+def test_claude_rejects_non_vertex_client():
+  """Verify that Claude rejects an AsyncAnthropic client."""
+  mock_client = mock.MagicMock(spec=AsyncAnthropic)
+  claude = Claude(
+      model="projects/p/locations/l/publishers/google/models/claude-3-5-sonnet-v2@20241022",
+      client=mock_client,
+  )
+
+  with pytest.raises(
+      ValueError, match="Claude requires an AsyncAnthropicVertex client."
+  ):
+    _ = claude._anthropic_client
+
+
+def test_apigee_custom_client_warnings():
+  """Verify warnings when custom client is used with conflicting options in ApigeeLlm."""
+  mock_client = mock.MagicMock(spec=Client)
+
+  # Warning when proxy_url is also provided
+  with pytest.warns(
+      UserWarning, match="Both client and proxy_url/custom_headers"
+  ):
+    ApigeeLlm(
+        model="apigee/gemini/gemini-1.5-flash",
+        client=mock_client,
+        proxy_url="http://example.com",
+    )
+
+  # Warning when custom_headers are also provided
+  with pytest.warns(
+      UserWarning, match="Both client and proxy_url/custom_headers"
+  ):
+    ApigeeLlm(
+        model="apigee/gemini/gemini-1.5-flash",
+        client=mock_client,
+        custom_headers={"X-Test": "test"},
+    )
+
+  # Warning when api_type is CHAT_COMPLETIONS
+  with pytest.warns(
+      UserWarning, match="injected client will be ignored for CHAT_COMPLETIONS"
+  ):
+    ApigeeLlm(
+        model="apigee/openai/gpt-4",
+        client=mock_client,
+        api_type=ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+    )
+
+  # Warning when proxy_url is set via env var
+  with mock.patch.dict(
+      "os.environ", {"APIGEE_PROXY_URL": "http://example.com"}
+  ):
+    with pytest.warns(
+        UserWarning, match="Both client and proxy_url/custom_headers"
+    ):
+      ApigeeLlm(
+          model="apigee/gemini/gemini-1.5-flash",
+          client=mock_client,
+      )