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

google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py

@@ -0,0 +1,268 @@
+# 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.
+
+
+import logging
+from typing import Literal
+from typing import Optional
+
+from fastapi.encoders import jsonable_encoder
+from pydantic import BaseModel
+
+from ....auth.auth_credential import AuthCredential
+from ....auth.auth_credential import AuthCredentialTypes
+from ....auth.auth_schemes import AuthScheme
+from ....auth.auth_schemes import AuthSchemeType
+from ....auth.auth_tool import AuthConfig
+from ...tool_context import ToolContext
+from ..auth.credential_exchangers.auto_auth_credential_exchanger import AutoAuthCredentialExchanger
+from ..auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
+from ..auth.credential_exchangers.base_credential_exchanger import BaseAuthCredentialExchanger
+
+logger = logging.getLogger(__name__)
+
+AuthPreparationState = Literal["pending", "done"]
+
+
+class AuthPreparationResult(BaseModel):
+  """Result of the credential preparation process."""
+
+  state: AuthPreparationState
+  auth_scheme: Optional[AuthScheme] = None
+  auth_credential: Optional[AuthCredential] = None
+
+
+class ToolContextCredentialStore:
+  """Handles storage and retrieval of credentials within a ToolContext."""
+
+  def __init__(self, tool_context: ToolContext):
+    self.tool_context = tool_context
+
+  def get_credential_key(
+      self,
+      auth_scheme: Optional[AuthScheme],
+      auth_credential: Optional[AuthCredential],
+  ) -> str:
+    """Generates a unique key for the given auth scheme and credential."""
+    scheme_name = (
+        f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}"
+        if auth_scheme
+        else ""
+    )
+    credential_name = (
+        f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}"
+        if auth_credential
+        else ""
+    )
+    # no need to prepend temp: namespace, session state is a copy, changes to
+    # it won't be persisted , only changes in event_action.state_delta will be
+    # persisted. temp: namespace will be cleared after current run. but tool
+    # want access token to be there stored across runs
+
+    return f"{scheme_name}_{credential_name}_existing_exchanged_credential"
+
+  def get_credential(
+      self,
+      auth_scheme: Optional[AuthScheme],
+      auth_credential: Optional[AuthCredential],
+  ) -> Optional[AuthCredential]:
+    if not self.tool_context:
+      return None
+
+    token_key = self.get_credential_key(auth_scheme, auth_credential)
+    # TODO try not to use session state, this looks a hacky way, depend on
+    # session implementation, we don't want session to persist the token,
+    # meanwhile we want the token shared across runs.
+    serialized_credential = self.tool_context.state.get(token_key)
+    if not serialized_credential:
+      return None
+    return AuthCredential.model_validate(serialized_credential)
+
+  def store_credential(
+      self,
+      key: str,
+      auth_credential: Optional[AuthCredential],
+  ):
+    if self.tool_context:
+      serializable_credential = jsonable_encoder(
+          auth_credential, exclude_none=True
+      )
+      self.tool_context.state[key] = serializable_credential
+
+  def remove_credential(self, key: str):
+    del self.tool_context.state[key]
+
+
+class ToolAuthHandler:
+  """Handles the preparation and exchange of authentication credentials for tools."""
+
+  def __init__(
+      self,
+      tool_context: ToolContext,
+      auth_scheme: Optional[AuthScheme],
+      auth_credential: Optional[AuthCredential],
+      credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
+      credential_store: Optional["ToolContextCredentialStore"] = None,
+  ):
+    self.tool_context = tool_context
+    self.auth_scheme = (
+        auth_scheme.model_copy(deep=True) if auth_scheme else None
+    )
+    self.auth_credential = (
+        auth_credential.model_copy(deep=True) if auth_credential else None
+    )
+    self.credential_exchanger = (
+        credential_exchanger or AutoAuthCredentialExchanger()
+    )
+    self.credential_store = credential_store
+    self.should_store_credential = True
+
+  @classmethod
+  def from_tool_context(
+      cls,
+      tool_context: ToolContext,
+      auth_scheme: Optional[AuthScheme],
+      auth_credential: Optional[AuthCredential],
+      credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
+  ) -> "ToolAuthHandler":
+    """Creates a ToolAuthHandler instance from a ToolContext."""
+    credential_store = ToolContextCredentialStore(tool_context)
+    return cls(
+        tool_context,
+        auth_scheme,
+        auth_credential,
+        credential_exchanger,
+        credential_store,
+    )
+
+  def _handle_existing_credential(
+      self,
+  ) -> Optional[AuthPreparationResult]:
+    """Checks for and returns an existing, exchanged credential."""
+    if self.credential_store:
+      existing_credential = self.credential_store.get_credential(
+          self.auth_scheme, self.auth_credential
+      )
+      if existing_credential:
+        return AuthPreparationResult(
+            state="done",
+            auth_scheme=self.auth_scheme,
+            auth_credential=existing_credential,
+        )
+    return None
+
+  def _exchange_credential(
+      self, auth_credential: AuthCredential
+  ) -> Optional[AuthPreparationResult]:
+    """Handles an OpenID Connect authorization response."""
+
+    exchanged_credential = None
+    try:
+      exchanged_credential = self.credential_exchanger.exchange_credential(
+          self.auth_scheme, auth_credential
+      )
+    except Exception as e:
+      logger.error("Failed to exchange credential: %s", e)
+    return exchanged_credential
+
+  def _store_credential(self, auth_credential: AuthCredential) -> None:
+    """stores the auth_credential."""
+
+    if self.credential_store:
+      key = self.credential_store.get_credential_key(
+          self.auth_scheme, self.auth_credential
+      )
+      self.credential_store.store_credential(key, auth_credential)
+
+  def _reqeust_credential(self) -> None:
+    """Handles the case where an OpenID Connect or OAuth2 authentication request is needed."""
+    if self.auth_scheme.type_ in (
+        AuthSchemeType.openIdConnect,
+        AuthSchemeType.oauth2,
+    ):
+      if not self.auth_credential or not self.auth_credential.oauth2:
+        raise ValueError(
+            f"auth_credential is empty for scheme {self.auth_scheme.type_}."
+            "Please create AuthCredential using OAuth2Auth."
+        )
+
+      if not self.auth_credential.oauth2.client_id:
+        raise AuthCredentialMissingError(
+            "OAuth2 credentials client_id is missing."
+        )
+
+      if not self.auth_credential.oauth2.client_secret:
+        raise AuthCredentialMissingError(
+            "OAuth2 credentials client_secret is missing."
+        )
+
+    self.tool_context.request_credential(
+        AuthConfig(
+            auth_scheme=self.auth_scheme,
+            raw_auth_credential=self.auth_credential,
+        )
+    )
+    return None
+
+  def _get_auth_response(self) -> AuthCredential:
+    return self.tool_context.get_auth_response(
+        AuthConfig(
+            auth_scheme=self.auth_scheme,
+            raw_auth_credential=self.auth_credential,
+        )
+    )
+
+  def _request_credential(self, auth_config: AuthConfig):
+    if not self.tool_context:
+      return
+    self.tool_context.request_credential(auth_config)
+
+  def prepare_auth_credentials(
+      self,
+  ) -> AuthPreparationResult:
+    """Prepares authentication credentials, handling exchange and user interaction."""
+
+    # no auth is needed
+    if not self.auth_scheme:
+      return AuthPreparationResult(state="done")
+
+    # Check for existing credential.
+    existing_result = self._handle_existing_credential()
+    if existing_result:
+      return existing_result
+
+    # fetch credential from adk framework
+    # Some auth scheme like OAuth2 AuthCode & OpenIDConnect may require
+    # multi-step exchange:
+    # client_id , client_secret -> auth_uri -> auth_code -> access_token
+    # -> bearer token
+    # adk framework supports exchange access_token already
+    fetched_credential = self._get_auth_response() or self.auth_credential
+
+    exchanged_credential = self._exchange_credential(fetched_credential)
+
+    if exchanged_credential:
+      self._store_credential(exchanged_credential)
+      return AuthPreparationResult(
+          state="done",
+          auth_scheme=self.auth_scheme,
+          auth_credential=exchanged_credential,
+      )
+    else:
+      self._reqeust_credential()
+      return AuthPreparationResult(
+          state="pending",
+          auth_scheme=self.auth_scheme,
+          auth_credential=self.auth_credential,
+      )

google/adk/tools/preload_memory_tool.py

@@ -0,0 +1,72 @@
+# 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 __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from .base_tool import BaseTool
+from .tool_context import ToolContext
+
+if TYPE_CHECKING:
+  from ..models import LlmRequest
+
+
+class PreloadMemoryTool(BaseTool):
+  """A tool that preloads the memory for the current user."""
+
+  def __init__(self):
+    # Name and description are not used because this tool only
+    # changes llm_request.
+    super().__init__(name='preload_memory', description='preload_memory')
+
+  @override
+  async def process_llm_request(
+      self,
+      *,
+      tool_context: ToolContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    parts = tool_context.user_content.parts
+    if not parts or not parts[0].text:
+      return
+    query = parts[0].text
+    response = tool_context.search_memory(query)
+    if not response.memories:
+      return
+    memory_text = ''
+    for memory in response.memories:
+      time_str = datetime.fromtimestamp(memory.events[0].timestamp).isoformat()
+      memory_text += f'Time: {time_str}\n'
+      for event in memory.events:
+        # TODO: support multi-part content.
+        if (
+            event.content
+            and event.content.parts
+            and event.content.parts[0].text
+        ):
+          memory_text += f'{event.author}: {event.content.parts[0].text}\n'
+    si = f"""The following content is from your previous conversations with the user.
+They may be useful for answering the user's current query.
+<PAST_CONVERSATIONS>
+{memory_text}
+</PAST_CONVERSATIONS>
+"""
+    llm_request.append_instructions([si])
+
+
+preload_memory_tool = PreloadMemoryTool()

google/adk/tools/retrieval/__init__.py

@@ -0,0 +1,36 @@
+# 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 .base_retrieval_tool import BaseRetrievalTool
+from .files_retrieval import FilesRetrieval
+from .llama_index_retrieval import LlamaIndexRetrieval
+
+__all__ = [
+    'BaseRetrievalTool',
+    'FilesRetrieval',
+    'LlamaIndexRetrieval',
+]
+
+try:
+  from .vertex_ai_rag_retrieval import VertexAiRagRetrieval
+
+  __all__.append('VertexAiRagRetrieval')
+except ImportError:
+  import logging
+
+  logger = logging.getLogger(__name__)
+  logger.debug(
+      'The Vertex sdk is not installed. If you want to use the Vertex RAG with'
+      ' agents, please install it. If not, you can ignore this warning.'
+  )

google/adk/tools/retrieval/base_retrieval_tool.py

@@ -0,0 +1,37 @@
+# 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.genai import types
+from typing_extensions import override
+
+from ..base_tool import BaseTool
+
+
+class BaseRetrievalTool(BaseTool):
+
+  @override
+  def _get_declaration(self) -> types.FunctionDeclaration:
+    return types.FunctionDeclaration(
+        name=self.name,
+        description=self.description,
+        parameters=types.Schema(
+            type=types.Type.OBJECT,
+            properties={
+                'query': types.Schema(
+                    type=types.Type.STRING,
+                    description='The query to retrieve.',
+                ),
+            },
+        ),
+    )

google/adk/tools/retrieval/files_retrieval.py

@@ -0,0 +1,33 @@
+# 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.
+
+"""Provides data for the agent."""
+
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core import VectorStoreIndex
+
+from .llama_index_retrieval import LlamaIndexRetrieval
+
+
+class FilesRetrieval(LlamaIndexRetrieval):
+
+  def __init__(self, *, name: str, description: str, input_dir: str):
+
+    self.input_dir = input_dir
+
+    print(f'Loading data from {input_dir}')
+    retriever = VectorStoreIndex.from_documents(
+        SimpleDirectoryReader(input_dir).load_data()
+    ).as_retriever()
+    super().__init__(name=name, description=description, retriever=retriever)

google/adk/tools/retrieval/llama_index_retrieval.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""Provides data for the agent."""
+
+from __future__ import annotations
+
+from typing import Any
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from ..tool_context import ToolContext
+from .base_retrieval_tool import BaseRetrievalTool
+
+if TYPE_CHECKING:
+  from llama_index.core.base.base_retriever import BaseRetriever
+
+
+class LlamaIndexRetrieval(BaseRetrievalTool):
+
+  def __init__(self, *, name: str, description: str, retriever: BaseRetriever):
+    super().__init__(name=name, description=description)
+    self.retriever = retriever
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    return self.retriever.retrieve(args['query'])[0].text

google/adk/tools/retrieval/vertex_ai_rag_retrieval.py

@@ -0,0 +1,107 @@
+# 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.
+
+"""A retrieval tool that uses Vertex AI RAG to retrieve data."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+from typing import TYPE_CHECKING
+
+from google.genai import types
+from typing_extensions import override
+from vertexai.preview import rag
+
+from ..tool_context import ToolContext
+from .base_retrieval_tool import BaseRetrievalTool
+
+if TYPE_CHECKING:
+  from ...models.llm_request import LlmRequest
+
+logger = logging.getLogger(__name__)
+
+
+class VertexAiRagRetrieval(BaseRetrievalTool):
+  """A retrieval tool that uses Vertex AI RAG (Retrieval-Augmented Generation) to retrieve data."""
+
+  def __init__(
+      self,
+      *,
+      name: str,
+      description: str,
+      rag_corpora: list[str] = None,
+      rag_resources: list[rag.RagResource] = None,
+      similarity_top_k: int = None,
+      vector_distance_threshold: float = None,
+  ):
+    super().__init__(name=name, description=description)
+    self.vertex_rag_store = types.VertexRagStore(
+        rag_corpora=rag_corpora,
+        rag_resources=rag_resources,
+        similarity_top_k=similarity_top_k,
+        vector_distance_threshold=vector_distance_threshold,
+    )
+
+  @override
+  async def process_llm_request(
+      self,
+      *,
+      tool_context: ToolContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    # Use Gemini built-in Vertex AI RAG tool for Gemini 2 models.
+    if llm_request.model and llm_request.model.startswith('gemini-2'):
+      llm_request.config = (
+          types.GenerateContentConfig()
+          if not llm_request.config
+          else llm_request.config
+      )
+      llm_request.config.tools = (
+          [] if not llm_request.config.tools else llm_request.config.tools
+      )
+      llm_request.config.tools.append(
+          types.Tool(
+              retrieval=types.Retrieval(vertex_rag_store=self.vertex_rag_store)
+          )
+      )
+    else:
+      # Add the function declaration to the tools
+      await super().process_llm_request(
+          tool_context=tool_context, llm_request=llm_request
+      )
+
+  @override
+  async def run_async(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Any:
+
+    response = rag.retrieval_query(
+        text=args['query'],
+        rag_resources=self.vertex_rag_store.rag_resources,
+        rag_corpora=self.vertex_rag_store.rag_corpora,
+        similarity_top_k=self.vertex_rag_store.similarity_top_k,
+        vector_distance_threshold=self.vertex_rag_store.vector_distance_threshold,
+    )
+
+    logging.debug('RAG raw response: %s', response)
+
+    return (
+        f'No matching result found with the config: {self.vertex_rag_store}'
+        if not response.contexts.contexts
+        else [context.text for context in response.contexts.contexts]
+    )

google/adk/tools/tool_context.py

@@ -0,0 +1,90 @@
+# 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 __future__ import annotations
+
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from ..agents.callback_context import CallbackContext
+from ..auth.auth_credential import AuthCredential
+from ..auth.auth_handler import AuthHandler
+from ..auth.auth_tool import AuthConfig
+
+if TYPE_CHECKING:
+  from ..agents.invocation_context import InvocationContext
+  from ..events.event_actions import EventActions
+  from ..memory.base_memory_service import SearchMemoryResponse
+
+
+class ToolContext(CallbackContext):
+  """The context of the tool.
+
+  This class provides the context for a tool invocation, including access to
+  the invocation context, function call ID, event actions, and authentication
+  response. It also provides methods for requesting credentials, retrieving
+  authentication responses, listing artifacts, and searching memory.
+
+  Attributes:
+    invocation_context: The invocation context of the tool.
+    function_call_id: The function call id of the current tool call. This id was
+      returned in the function call event from LLM to identify a function call.
+      If LLM didn't return this id, ADK will assign one to it. This id is used
+      to map function call response to the original function call.
+    event_actions: The event actions of the current tool call.
+  """
+
+  def __init__(
+      self,
+      invocation_context: InvocationContext,
+      *,
+      function_call_id: Optional[str] = None,
+      event_actions: Optional[EventActions] = None,
+  ):
+    super().__init__(invocation_context, event_actions=event_actions)
+    self.function_call_id = function_call_id
+
+  @property
+  def actions(self) -> EventActions:
+    return self._event_actions
+
+  def request_credential(self, auth_config: AuthConfig) -> None:
+    if not self.function_call_id:
+      raise ValueError('function_call_id is not set.')
+    self._event_actions.requested_auth_configs[self.function_call_id] = (
+        AuthHandler(auth_config).generate_auth_request()
+    )
+
+  def get_auth_response(self, auth_config: AuthConfig) -> AuthCredential:
+    return AuthHandler(auth_config).get_auth_response(self.state)
+
+  def list_artifacts(self) -> list[str]:
+    """Lists the filenames of the artifacts attached to the current session."""
+    if self._invocation_context.artifact_service is None:
+      raise ValueError('Artifact service is not initialized.')
+    return self._invocation_context.artifact_service.list_artifact_keys(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        session_id=self._invocation_context.session.id,
+    )
+
+  def search_memory(self, query: str) -> 'SearchMemoryResponse':
+    """Searches the memory of the current user."""
+    if self._invocation_context.memory_service is None:
+      raise ValueError('Memory service is not available.')
+    return self._invocation_context.memory_service.search_memory(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        query=query,
+    )