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

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

@@ -1,201 +0,0 @@
-# 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

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

@@ -1,147 +0,0 @@
-# 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

@@ -1,167 +0,0 @@
-# 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

@@ -1,141 +0,0 @@
-# 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