google-adk 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

google/adk/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py

@@ -0,0 +1,201 @@
+# Copyright 2025 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.
+
+from typing import Optional
+from unittest.mock import MagicMock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import HttpAuth
+from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.sessions.session import Session
+from google.adk.tools.openapi_tool.auth.auth_helpers import openid_dict_to_scheme_credential
+from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_credential
+from google.adk.tools.openapi_tool.auth.credential_exchangers.auto_auth_credential_exchanger import OAuth2CredentialExchanger
+from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler
+from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolContextCredentialStore
+from google.adk.tools.tool_context import ToolContext
+import pytest
+
+
+# Helper function to create a mock ToolContext
+def create_mock_tool_context():
+  return ToolContext(
+      function_call_id='test-fc-id',
+      invocation_context=InvocationContext(
+          agent=LlmAgent(name='test'),
+          session=Session(app_name='test', user_id='123', id='123'),
+          invocation_id='123',
+          session_service=InMemorySessionService(),
+      ),
+  )
+
+
+# Test cases for OpenID Connect
+class MockOpenIdConnectCredentialExchanger(OAuth2CredentialExchanger):
+
+  def __init__(
+      self, expected_scheme, expected_credential, expected_access_token
+  ):
+    self.expected_scheme = expected_scheme
+    self.expected_credential = expected_credential
+    self.expected_access_token = expected_access_token
+
+  def exchange_credential(
+      self,
+      auth_scheme: AuthScheme,
+      auth_credential: Optional[AuthCredential] = None,
+  ) -> AuthCredential:
+    if auth_credential.oauth2 and (
+        auth_credential.oauth2.auth_response_uri
+        or auth_credential.oauth2.auth_code
+    ):
+      auth_code = (
+          auth_credential.oauth2.auth_response_uri
+          if auth_credential.oauth2.auth_response_uri
+          else auth_credential.oauth2.auth_code
+      )
+      # Simulate the token exchange
+      updated_credential = AuthCredential(
+          auth_type=AuthCredentialTypes.HTTP,  # Store as a bearer token
+          http=HttpAuth(
+              scheme='bearer',
+              credentials=HttpCredentials(
+                  token=auth_code + self.expected_access_token
+              ),
+          ),
+      )
+      return updated_credential
+
+    # simulate the case of getting auth_uri
+    return None
+
+
+def get_mock_openid_scheme_credential():
+  config_dict = {
+      'authorization_endpoint': 'test.com',
+      'token_endpoint': 'test.com',
+  }
+  scopes = ['test_scope']
+  credential_dict = {
+      'client_id': '123',
+      'client_secret': '456',
+      'redirect_uri': 'test.com',
+  }
+  return openid_dict_to_scheme_credential(config_dict, scopes, credential_dict)
+
+
+# Fixture for the OpenID Connect security scheme
+@pytest.fixture
+def openid_connect_scheme():
+  scheme, _ = get_mock_openid_scheme_credential()
+  return scheme
+
+
+# Fixture for a base OpenID Connect credential
+@pytest.fixture
+def openid_connect_credential():
+  _, credential = get_mock_openid_scheme_credential()
+  return credential
+
+
+def test_openid_connect_no_auth_response(
+    openid_connect_scheme, openid_connect_credential
+):
+  # Setup Mock exchanger
+  mock_exchanger = MockOpenIdConnectCredentialExchanger(
+      openid_connect_scheme, openid_connect_credential, None
+  )
+  tool_context = create_mock_tool_context()
+  credential_store = ToolContextCredentialStore(tool_context=tool_context)
+  handler = ToolAuthHandler(
+      tool_context,
+      openid_connect_scheme,
+      openid_connect_credential,
+      credential_exchanger=mock_exchanger,
+      credential_store=credential_store,
+  )
+  result = handler.prepare_auth_credentials()
+  assert result.state == 'pending'
+  assert result.auth_credential == openid_connect_credential
+
+
+def test_openid_connect_with_auth_response(
+    openid_connect_scheme, openid_connect_credential, monkeypatch
+):
+  mock_exchanger = MockOpenIdConnectCredentialExchanger(
+      openid_connect_scheme,
+      openid_connect_credential,
+      'test_access_token',
+  )
+  tool_context = create_mock_tool_context()
+
+  mock_auth_handler = MagicMock()
+  mock_auth_handler.get_auth_response.return_value = AuthCredential(
+      auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+      oauth2=OAuth2Auth(auth_response_uri='test_auth_response_uri'),
+  )
+  mock_auth_handler_path = 'google.adk.tools.tool_context.AuthHandler'
+  monkeypatch.setattr(
+      mock_auth_handler_path, lambda *args, **kwargs: mock_auth_handler
+  )
+
+  credential_store = ToolContextCredentialStore(tool_context=tool_context)
+  handler = ToolAuthHandler(
+      tool_context,
+      openid_connect_scheme,
+      openid_connect_credential,
+      credential_exchanger=mock_exchanger,
+      credential_store=credential_store,
+  )
+  result = handler.prepare_auth_credentials()
+  assert result.state == 'done'
+  assert result.auth_credential.auth_type == AuthCredentialTypes.HTTP
+  assert 'test_access_token' in result.auth_credential.http.credentials.token
+  # Verify that the credential was stored:
+  stored_credential = credential_store.get_credential(
+      openid_connect_scheme, openid_connect_credential
+  )
+  assert stored_credential == result.auth_credential
+  mock_auth_handler.get_auth_response.assert_called_once()
+
+
+def test_openid_connect_existing_token(
+    openid_connect_scheme, openid_connect_credential
+):
+  _, existing_credential = token_to_scheme_credential(
+      'oauth2Token', 'header', 'bearer', '123123123'
+  )
+  tool_context = create_mock_tool_context()
+  # Store the credential to simulate existing credential
+  credential_store = ToolContextCredentialStore(tool_context=tool_context)
+  key = credential_store.get_credential_key(
+      openid_connect_scheme, openid_connect_credential
+  )
+  credential_store.store_credential(key, existing_credential)
+
+  handler = ToolAuthHandler(
+      tool_context,
+      openid_connect_scheme,
+      openid_connect_credential,
+      credential_store=credential_store,
+  )
+  result = handler.prepare_auth_credentials()
+  assert result.state == 'done'
+  assert result.auth_credential == existing_credential

google/adk/tests/unittests/tools/retrieval/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 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.
+

google/adk/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py

@@ -0,0 +1,147 @@
+# Copyright 2025 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.
+
+from google.adk.agents import Agent
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
+from google.genai import types
+
+from ... import utils
+
+
+def noop_tool(x: str) -> str:
+  return x
+
+
+def test_vertex_rag_retrieval_for_gemini_1_x():
+  responses = [
+      'response1',
+  ]
+  mockModel = utils.MockModel.create(responses=responses)
+  mockModel.model = 'gemini-1.5-pro'
+
+  # Calls the first time.
+  agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      tools=[
+          VertexAiRagRetrieval(
+              name='rag_retrieval',
+              description='rag_retrieval',
+              rag_corpora=[
+                  'projects/123456789/locations/us-central1/ragCorpora/1234567890'
+              ],
+          )
+      ],
+  )
+  runner = utils.InMemoryRunner(agent)
+  events = runner.run('test1')
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 1
+  assert utils.simplify_contents(mockModel.requests[0].contents) == [
+      ('user', 'test1'),
+  ]
+  assert len(mockModel.requests[0].config.tools) == 1
+  assert (
+      mockModel.requests[0].config.tools[0].function_declarations[0].name
+      == 'rag_retrieval'
+  )
+  assert mockModel.requests[0].tools_dict['rag_retrieval'] is not None
+
+
+def test_vertex_rag_retrieval_for_gemini_1_x_with_another_function_tool():
+  responses = [
+      'response1',
+  ]
+  mockModel = utils.MockModel.create(responses=responses)
+  mockModel.model = 'gemini-1.5-pro'
+
+  # Calls the first time.
+  agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      tools=[
+          VertexAiRagRetrieval(
+              name='rag_retrieval',
+              description='rag_retrieval',
+              rag_corpora=[
+                  'projects/123456789/locations/us-central1/ragCorpora/1234567890'
+              ],
+          ),
+          FunctionTool(func=noop_tool),
+      ],
+  )
+  runner = utils.InMemoryRunner(agent)
+  events = runner.run('test1')
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 1
+  assert utils.simplify_contents(mockModel.requests[0].contents) == [
+      ('user', 'test1'),
+  ]
+  assert len(mockModel.requests[0].config.tools[0].function_declarations) == 2
+  assert (
+      mockModel.requests[0].config.tools[0].function_declarations[0].name
+      == 'rag_retrieval'
+  )
+  assert (
+      mockModel.requests[0].config.tools[0].function_declarations[1].name
+      == 'noop_tool'
+  )
+  assert mockModel.requests[0].tools_dict['rag_retrieval'] is not None
+
+
+def test_vertex_rag_retrieval_for_gemini_2_x():
+  responses = [
+      'response1',
+  ]
+  mockModel = utils.MockModel.create(responses=responses)
+  mockModel.model = 'gemini-2.0-flash'
+
+  # Calls the first time.
+  agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      tools=[
+          VertexAiRagRetrieval(
+              name='rag_retrieval',
+              description='rag_retrieval',
+              rag_corpora=[
+                  'projects/123456789/locations/us-central1/ragCorpora/1234567890'
+              ],
+          )
+      ],
+  )
+  runner = utils.InMemoryRunner(agent)
+  events = runner.run('test1')
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 1
+  assert utils.simplify_contents(mockModel.requests[0].contents) == [
+      ('user', 'test1'),
+  ]
+  assert len(mockModel.requests[0].config.tools) == 1
+  assert mockModel.requests[0].config.tools == [
+      types.Tool(
+          retrieval=types.Retrieval(
+              vertex_rag_store=types.VertexRagStore(
+                  rag_corpora=[
+                      'projects/123456789/locations/us-central1/ragCorpora/1234567890'
+                  ]
+              )
+          )
+      )
+  ]
+  assert 'rag_retrieval' not in mockModel.requests[0].tools_dict

google/adk/tests/unittests/tools/test_agent_tool.py

@@ -0,0 +1,167 @@
+# Copyright 2025 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.
+
+from google.adk.agents import Agent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.tools.agent_tool import AgentTool
+from google.genai.types import Part
+from pydantic import BaseModel
+import pytest
+from pytest import mark
+
+from .. import utils
+
+pytestmark = pytest.mark.skip(
+    reason='Skipping until tool.func evaluations are fixed (async)'
+)
+
+
+function_call_custom = Part.from_function_call(
+    name='tool_agent', args={'custom_input': 'test1'}
+)
+
+function_call_no_schema = Part.from_function_call(
+    name='tool_agent', args={'request': 'test1'}
+)
+
+function_response_custom = Part.from_function_response(
+    name='tool_agent', response={'custom_output': 'response1'}
+)
+
+function_response_no_schema = Part.from_function_response(
+    name='tool_agent', response={'result': 'response1'}
+)
+
+
+def change_state_callback(callback_context: CallbackContext):
+  callback_context.state['state_1'] = 'changed_value'
+  print('change_state_callback: ', callback_context.state)
+
+
+def test_no_schema():
+  mock_model = utils.MockModel.create(
+      responses=[
+          function_call_no_schema,
+          'response1',
+          'response2',
+      ]
+  )
+
+  tool_agent = Agent(
+      name='tool_agent',
+      model=mock_model,
+  )
+
+  root_agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[AgentTool(agent=tool_agent)],
+  )
+
+  runner = utils.InMemoryRunner(root_agent)
+
+  assert utils.simplify_events(runner.run('test1')) == [
+      ('root_agent', function_call_no_schema),
+      ('root_agent', function_response_no_schema),
+      ('root_agent', 'response2'),
+  ]
+
+
+def test_update_state():
+  """The agent tool can read and change parent state."""
+
+  mock_model = utils.MockModel.create(
+      responses=[
+          function_call_no_schema,
+          '{"custom_output": "response1"}',
+          'response2',
+      ]
+  )
+
+  tool_agent = Agent(
+      name='tool_agent',
+      model=mock_model,
+      instruction='input: {state_1}',
+      before_agent_callback=change_state_callback,
+  )
+
+  root_agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[AgentTool(agent=tool_agent)],
+  )
+
+  runner = utils.InMemoryRunner(root_agent)
+  runner.session.state['state_1'] = 'state1_value'
+
+  runner.run('test1')
+  assert (
+      'input: changed_value' in mock_model.requests[1].config.system_instruction
+  )
+  assert runner.session.state['state_1'] == 'changed_value'
+
+
+@mark.parametrize(
+    'env_variables',
+    [
+        'GOOGLE_AI',
+        # TODO(wanyif): re-enable after fix.
+        # 'VERTEX',
+    ],
+    indirect=True,
+)
+def test_custom_schema():
+  class CustomInput(BaseModel):
+    custom_input: str
+
+  class CustomOutput(BaseModel):
+    custom_output: str
+
+  mock_model = utils.MockModel.create(
+      responses=[
+          function_call_custom,
+          '{"custom_output": "response1"}',
+          'response2',
+      ]
+  )
+
+  tool_agent = Agent(
+      name='tool_agent',
+      model=mock_model,
+      input_schema=CustomInput,
+      output_schema=CustomOutput,
+      output_key='tool_output',
+  )
+
+  root_agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[AgentTool(agent=tool_agent)],
+  )
+
+  runner = utils.InMemoryRunner(root_agent)
+  runner.session.state['state_1'] = 'state1_value'
+
+  assert utils.simplify_events(runner.run('test1')) == [
+      ('root_agent', function_call_custom),
+      ('root_agent', function_response_custom),
+      ('root_agent', 'response2'),
+  ]
+
+  assert runner.session.state['tool_output'] == {'custom_output': 'response1'}
+
+  assert len(mock_model.requests) == 3
+  # The second request is the tool agent request.
+  assert mock_model.requests[1].config.response_schema == CustomOutput
+  assert mock_model.requests[1].config.response_mime_type == 'application/json'

google/adk/tests/unittests/tools/test_base_tool.py

@@ -0,0 +1,141 @@
+# Copyright 2025 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.
+
+from typing import Optional
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+
+class _TestingTool(BaseTool):
+
+  def __init__(
+      self,
+      declaration: Optional[types.FunctionDeclaration] = None,
+  ):
+    super().__init__(name='test_tool', description='test_description')
+    self.declaration = declaration
+
+  def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
+    return self.declaration
+
+
+def _create_tool_context() -> ToolContext:
+  session_service = InMemorySessionService()
+  session = session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  agent = SequentialAgent(name='test_agent')
+  invocation_context = InvocationContext(
+      invocation_id='invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  return ToolContext(invocation_context)
+
+
+@pytest.mark.asyncio
+async def test_process_llm_request_no_declaration():
+  tool = _TestingTool()
+  tool_context = _create_tool_context()
+  llm_request = LlmRequest()
+
+  await tool.process_llm_request(
+      tool_context=tool_context, llm_request=llm_request
+  )
+
+  assert llm_request.config is None
+
+
+@pytest.mark.asyncio
+async def test_process_llm_request_with_declaration():
+  declaration = types.FunctionDeclaration(
+      name='test_tool',
+      description='test_description',
+      parameters=types.Schema(
+          type=types.Type.STRING,
+          title='param_1',
+      ),
+  )
+  tool = _TestingTool(declaration)
+  llm_request = LlmRequest()
+  tool_context = _create_tool_context()
+
+  await tool.process_llm_request(
+      tool_context=tool_context, llm_request=llm_request
+  )
+
+  assert llm_request.config.tools[0].function_declarations == [declaration]
+
+
+@pytest.mark.asyncio
+async def test_process_llm_request_with_builtin_tool():
+  declaration = types.FunctionDeclaration(
+      name='test_tool',
+      description='test_description',
+      parameters=types.Schema(
+          type=types.Type.STRING,
+          title='param_1',
+      ),
+  )
+  tool = _TestingTool(declaration)
+  llm_request = LlmRequest(
+      config=types.GenerateContentConfig(
+          tools=[types.Tool(google_search=types.GoogleSearch())]
+      )
+  )
+  tool_context = _create_tool_context()
+
+  await tool.process_llm_request(
+      tool_context=tool_context, llm_request=llm_request
+  )
+
+  # function_declaration is added to another types.Tool without builtin tool.
+  assert llm_request.config.tools[1].function_declarations == [declaration]
+
+
+@pytest.mark.asyncio
+async def test_process_llm_request_with_builtin_tool_and_another_declaration():
+  declaration = types.FunctionDeclaration(
+      name='test_tool',
+      description='test_description',
+      parameters=types.Schema(
+          type=types.Type.STRING,
+          title='param_1',
+      ),
+  )
+  tool = _TestingTool(declaration)
+  llm_request = LlmRequest(
+      config=types.GenerateContentConfig(
+          tools=[
+              types.Tool(google_search=types.GoogleSearch()),
+              types.Tool(function_declarations=[types.FunctionDeclaration()]),
+          ]
+      )
+  )
+  tool_context = _create_tool_context()
+
+  await tool.process_llm_request(
+      tool_context=tool_context, llm_request=llm_request
+  )
+
+  # function_declaration is added to existing types.Tool with function_declaration.
+  assert llm_request.config.tools[1].function_declarations[1] == declaration