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

google/adk/flows/llm_flows/single_flow.py

@@ -0,0 +1,57 @@
+# 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.
+
+"""Implementation of single flow."""
+
+import logging
+
+from ...auth import auth_preprocessor
+from . import _code_execution
+from . import _nl_planning
+from . import basic
+from . import contents
+from . import identity
+from . import instructions
+from .base_llm_flow import BaseLlmFlow
+
+logger = logging.getLogger(__name__)
+
+
+class SingleFlow(BaseLlmFlow):
+  """SingleFlow is the LLM flows that handles tools calls.
+
+  A single flow only consider an agent itself and tools.
+  No sub-agents are allowed for single flow.
+  """
+
+  def __init__(self):
+    super().__init__()
+    self.request_processors += [
+        basic.request_processor,
+        auth_preprocessor.request_processor,
+        instructions.request_processor,
+        identity.request_processor,
+        contents.request_processor,
+        # Some implementations of NL Planning mark planning contents as thoughts
+        # in the post processor. Since these need to be unmarked, NL Planning
+        # should be after contents.
+        _nl_planning.request_processor,
+        # Code execution should be after the contents as it mutates the contents
+        # to optimize data files.
+        _code_execution.request_processor,
+    ]
+    self.response_processors += [
+        _nl_planning.response_processor,
+        _code_execution.response_processor,
+    ]

google/adk/memory/__init__.py

@@ -0,0 +1,35 @@
+# 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 .base_memory_service import BaseMemoryService
+from .in_memory_memory_service import InMemoryMemoryService
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    'BaseMemoryService',
+    'InMemoryMemoryService',
+]
+
+try:
+  from .vertex_ai_rag_memory_service import VertexAiRagMemoryService
+
+  __all__.append('VertexAiRagMemoryService')
+except ImportError:
+  logger.debug(
+      'The Vertex sdk is not installed. If you want to use the'
+      ' VertexAiRagMemoryService please install it. If not, you can ignore this'
+      ' warning.'
+  )

google/adk/memory/base_memory_service.py

@@ -0,0 +1,74 @@
+# 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 abc
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from ..events.event import Event
+from ..sessions.session import Session
+
+
+class MemoryResult(BaseModel):
+  """Represents a single memory retrieval result.
+
+  Attributes:
+      session_id: The session id associated with the memory.
+      events: A list of events in the session.
+  """
+  session_id: str
+  events: list[Event]
+
+
+class SearchMemoryResponse(BaseModel):
+  """Represents the response from a memory search.
+
+  Attributes:
+      memories: A list of memory results matching the search query.
+  """
+  memories: list[MemoryResult] = Field(default_factory=list)
+
+
+class BaseMemoryService(abc.ABC):
+  """Base class for memory services.
+
+  The service provides functionalities to ingest sessions into memory so that
+  the memory can be used for user queries.
+  """
+
+  @abc.abstractmethod
+  def add_session_to_memory(self, session: Session):
+    """Adds a session to the memory service.
+
+    A session may be added multiple times during its lifetime.
+
+    Args:
+        session: The session to add.
+    """
+
+  @abc.abstractmethod
+  def search_memory(
+      self, *, app_name: str, user_id: str, query: str
+  ) -> SearchMemoryResponse:
+    """Searches for sessions that match the query.
+
+    Args:
+        app_name: The name of the application.
+        user_id: The id of the user.
+        query: The query to search for.
+
+    Returns:
+        A SearchMemoryResponse containing the matching memories.
+    """

google/adk/memory/in_memory_memory_service.py

@@ -0,0 +1,62 @@
+# 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 ..events.event import Event
+from ..sessions.session import Session
+from .base_memory_service import BaseMemoryService
+from .base_memory_service import MemoryResult
+from .base_memory_service import SearchMemoryResponse
+
+
+class InMemoryMemoryService(BaseMemoryService):
+  """An in-memory memory service for prototyping purpose only.
+
+  Uses keyword matching instead of semantic search.
+  """
+
+  def __init__(self):
+    self.session_events: dict[str, list[Event]] = {}
+    """keys are app_name/user_id/session_id"""
+
+  def add_session_to_memory(self, session: Session):
+    key = f'{session.app_name}/{session.user_id}/{session.id}'
+    self.session_events[key] = [
+        event for event in session.events if event.content
+    ]
+
+  def search_memory(
+      self, *, app_name: str, user_id: str, query: str
+  ) -> SearchMemoryResponse:
+    """Prototyping purpose only."""
+    keywords = set(query.lower().split())
+    response = SearchMemoryResponse()
+    for key, events in self.session_events.items():
+      if not key.startswith(f'{app_name}/{user_id}/'):
+        continue
+      matched_events = []
+      for event in events:
+        if not event.content or not event.content.parts:
+          continue
+        parts = event.content.parts
+        text = '\n'.join([part.text for part in parts if part.text]).lower()
+        for keyword in keywords:
+          if keyword in text:
+            matched_events.append(event)
+            break
+      if matched_events:
+        session_id = key.split('/')[-1]
+        response.memories.append(
+            MemoryResult(session_id=session_id, events=matched_events)
+        )
+    return response

google/adk/memory/vertex_ai_rag_memory_service.py

@@ -0,0 +1,177 @@
+# 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 collections import OrderedDict
+import json
+import os
+import tempfile
+
+from google.genai import types
+from typing_extensions import override
+from vertexai.preview import rag
+
+from ..events.event import Event
+from ..sessions.session import Session
+from .base_memory_service import BaseMemoryService
+from .base_memory_service import MemoryResult
+from .base_memory_service import SearchMemoryResponse
+
+
+class VertexAiRagMemoryService(BaseMemoryService):
+  """A memory service that uses Vertex AI RAG for storage and retrieval."""
+
+  def __init__(
+      self,
+      rag_corpus: str = None,
+      similarity_top_k: int = None,
+      vector_distance_threshold: float = 10,
+  ):
+    """Initializes a VertexAiRagMemoryService.
+
+    Args:
+        rag_corpus: The name of the Vertex AI RAG corpus to use. Format:
+          ``projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}``
+          or ``{rag_corpus_id}``
+        similarity_top_k: The number of contexts to retrieve.
+        vector_distance_threshold: Only returns contexts with vector distance
+          smaller than the threshold..
+    """
+    self.vertex_rag_store = types.VertexRagStore(
+        rag_resources=[rag.RagResource(rag_corpus=rag_corpus)],
+        similarity_top_k=similarity_top_k,
+        vector_distance_threshold=vector_distance_threshold,
+    )
+
+  @override
+  def add_session_to_memory(self, session: Session):
+    with tempfile.NamedTemporaryFile(
+        mode="w", delete=False, suffix=".txt"
+    ) as temp_file:
+
+      output_lines = []
+      for event in session.events:
+        if not event.content or not event.content.parts:
+          continue
+        text_parts = [
+            part.text.replace("\n", " ")
+            for part in event.content.parts
+            if part.text
+        ]
+        if text_parts:
+          output_lines.append(
+              json.dumps({
+                  "author": event.author,
+                  "timestamp": event.timestamp,
+                  "text": ".".join(text_parts),
+              })
+          )
+      output_string = "\n".join(output_lines)
+      temp_file.write(output_string)
+      temp_file_path = temp_file.name
+    for rag_resource in self.vertex_rag_store.rag_resources:
+      rag.upload_file(
+          corpus_name=rag_resource.rag_corpus,
+          path=temp_file_path,
+          # this is the temp workaround as upload file does not support
+          # adding metadata, thus use display_name to store the session info.
+          display_name=f"{session.app_name}.{session.user_id}.{session.id}",
+      )
+
+    os.remove(temp_file_path)
+
+  @override
+  def search_memory(
+      self, *, app_name: str, user_id: str, query: str
+  ) -> SearchMemoryResponse:
+    """Searches for sessions that match the query using rag.retrieval_query."""
+    response = rag.retrieval_query(
+        text=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,
+    )
+
+    memory_results = []
+    session_events_map = OrderedDict()
+    for context in response.contexts.contexts:
+      # filter out context that is not related
+      # TODO: Add server side filtering by app_name and user_id.
+      # if not context.source_display_name.startswith(f"{app_name}.{user_id}."):
+      #   continue
+      session_id = context.source_display_name.split(".")[-1]
+      events = []
+      if context.text:
+        lines = context.text.split("\n")
+
+        for line in lines:
+          line = line.strip()
+          if not line:
+            continue
+
+          try:
+            # Try to parse as JSON
+            event_data = json.loads(line)
+
+            author = event_data.get("author", "")
+            timestamp = float(event_data.get("timestamp", 0))
+            text = event_data.get("text", "")
+
+            content = types.Content(parts=[types.Part(text=text)])
+            event = Event(author=author, timestamp=timestamp, content=content)
+            events.append(event)
+          except json.JSONDecodeError:
+            # Not valid JSON, skip this line
+            continue
+
+      if session_id in session_events_map:
+        session_events_map[session_id].append(events)
+      else:
+        session_events_map[session_id] = [events]
+
+    # Remove overlap and combine events from the same session.
+    for session_id, event_lists in session_events_map.items():
+      for events in _merge_event_lists(event_lists):
+        sorted_events = sorted(events, key=lambda e: e.timestamp)
+        memory_results.append(
+            MemoryResult(session_id=session_id, events=sorted_events)
+        )
+    return SearchMemoryResponse(memories=memory_results)
+
+
+def _merge_event_lists(event_lists: list[list[Event]]) -> list[list[Event]]:
+  """Merge event lists that have overlapping timestamps."""
+  merged = []
+  while event_lists:
+    current = event_lists.pop(0)
+    current_ts = {event.timestamp for event in current}
+    merge_found = True
+
+    # Keep merging until no new overlap is found.
+    while merge_found:
+      merge_found = False
+      remaining = []
+      for other in event_lists:
+        other_ts = {event.timestamp for event in other}
+        # Overlap exists, so we merge and use the merged list to check again
+        if current_ts & other_ts:
+          new_events = [e for e in other if e.timestamp not in current_ts]
+          current.extend(new_events)
+          current_ts.update(e.timestamp for e in new_events)
+          merge_found = True
+        else:
+          remaining.append(other)
+      event_lists = remaining
+    merged.append(current)
+  return merged

google/adk/models/__init__.py

@@ -0,0 +1,31 @@
+# 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.
+
+"""Defines the interface to support a model."""
+
+from .base_llm import BaseLlm
+from .google_llm import Gemini
+from .llm_request import LlmRequest
+from .llm_response import LlmResponse
+from .registry import LLMRegistry
+
+__all__ = [
+    'BaseLlm',
+    'Gemini',
+    'LLMRegistry',
+]
+
+
+for regex in Gemini.supported_models():
+  LLMRegistry.register(Gemini)

google/adk/models/anthropic_llm.py

@@ -0,0 +1,243 @@
+# 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.
+
+"""Anthropic integration for Claude models."""
+
+from __future__ import annotations
+
+from functools import cached_property
+import logging
+import os
+from typing import AsyncGenerator
+from typing import Generator
+from typing import Iterable
+from typing import Literal
+from typing import Optional, Union
+from typing import TYPE_CHECKING
+
+from anthropic import AnthropicVertex
+from anthropic import NOT_GIVEN
+from anthropic import types as anthropic_types
+from google.genai import types
+from pydantic import BaseModel
+from typing_extensions import override
+
+from .base_llm import BaseLlm
+from .llm_response import LlmResponse
+
+if TYPE_CHECKING:
+  from .llm_request import LlmRequest
+
+__all__ = ["Claude"]
+
+logger = logging.getLogger(__name__)
+
+MAX_TOKEN = 1024
+
+
+class ClaudeRequest(BaseModel):
+  system_instruction: str
+  messages: Iterable[anthropic_types.MessageParam]
+  tools: list[anthropic_types.ToolParam]
+
+
+def to_claude_role(role: Optional[str]) -> Literal["user", "assistant"]:
+  if role in ["model", "assistant"]:
+    return "assistant"
+  return "user"
+
+
+def to_google_genai_finish_reason(
+    anthropic_stop_reason: Optional[str],
+) -> types.FinishReason:
+  if anthropic_stop_reason in ["end_turn", "stop_sequence", "tool_use"]:
+    return "STOP"
+  if anthropic_stop_reason == "max_tokens":
+    return "MAX_TOKENS"
+  return "FINISH_REASON_UNSPECIFIED"
+
+
+def part_to_message_block(
+    part: types.Part,
+) -> Union[
+    anthropic_types.TextBlockParam,
+    anthropic_types.ImageBlockParam,
+    anthropic_types.ToolUseBlockParam,
+    anthropic_types.ToolResultBlockParam,
+]:
+  if part.text:
+    return anthropic_types.TextBlockParam(text=part.text, type="text")
+  if part.function_call:
+    assert part.function_call.name
+
+    return anthropic_types.ToolUseBlockParam(
+        id=part.function_call.id or "",
+        name=part.function_call.name,
+        input=part.function_call.args,
+        type="tool_use",
+    )
+  if part.function_response:
+    content = ""
+    if (
+        "result" in part.function_response.response
+        and part.function_response.response["result"]
+    ):
+      # Transformation is required because the content is a list of dict.
+      # ToolResultBlockParam content doesn't support list of dict. Converting
+      # to str to prevent anthropic.BadRequestError from being thrown.
+      content = str(part.function_response.response["result"])
+    return anthropic_types.ToolResultBlockParam(
+        tool_use_id=part.function_response.id or "",
+        type="tool_result",
+        content=content,
+        is_error=False,
+    )
+  raise NotImplementedError("Not supported yet.")
+
+
+def content_to_message_param(
+    content: types.Content,
+) -> anthropic_types.MessageParam:
+  return {
+      "role": to_claude_role(content.role),
+      "content": [part_to_message_block(part) for part in content.parts or []],
+  }
+
+
+def content_block_to_part(
+    content_block: anthropic_types.ContentBlock,
+) -> types.Part:
+  if isinstance(content_block, anthropic_types.TextBlock):
+    return types.Part.from_text(text=content_block.text)
+  if isinstance(content_block, anthropic_types.ToolUseBlock):
+    assert isinstance(content_block.input, dict)
+    part = types.Part.from_function_call(
+        name=content_block.name, args=content_block.input
+    )
+    part.function_call.id = content_block.id
+    return part
+  raise NotImplementedError("Not supported yet.")
+
+
+def message_to_generate_content_response(
+    message: anthropic_types.Message,
+) -> LlmResponse:
+
+  return LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[content_block_to_part(cb) for cb in message.content],
+      ),
+      # TODO: Deal with these later.
+      # finish_reason=to_google_genai_finish_reason(message.stop_reason),
+      # usage_metadata=types.GenerateContentResponseUsageMetadata(
+      #     prompt_token_count=message.usage.input_tokens,
+      #     candidates_token_count=message.usage.output_tokens,
+      #     total_token_count=(
+      #         message.usage.input_tokens + message.usage.output_tokens
+      #     ),
+      # ),
+  )
+
+
+def function_declaration_to_tool_param(
+    function_declaration: types.FunctionDeclaration,
+) -> anthropic_types.ToolParam:
+  assert function_declaration.name
+
+  properties = {}
+  if (
+      function_declaration.parameters
+      and function_declaration.parameters.properties
+  ):
+    for key, value in function_declaration.parameters.properties.items():
+      value_dict = value.model_dump(exclude_none=True)
+      if "type" in value_dict:
+        value_dict["type"] = value_dict["type"].lower()
+      properties[key] = value_dict
+
+  return anthropic_types.ToolParam(
+      name=function_declaration.name,
+      description=function_declaration.description or "",
+      input_schema={
+          "type": "object",
+          "properties": properties,
+      },
+  )
+
+
+class Claude(BaseLlm):
+  model: str = "claude-3-5-sonnet-v2@20241022"
+
+  @staticmethod
+  @override
+  def supported_models() -> list[str]:
+    return [r"claude-3-.*"]
+
+  @override
+  async def generate_content_async(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> AsyncGenerator[LlmResponse, None]:
+    messages = [
+        content_to_message_param(content)
+        for content in llm_request.contents or []
+    ]
+    tools = NOT_GIVEN
+    if (
+        llm_request.config
+        and llm_request.config.tools
+        and llm_request.config.tools[0].function_declarations
+    ):
+      tools = [
+          function_declaration_to_tool_param(tool)
+          for tool in llm_request.config.tools[0].function_declarations
+      ]
+    tool_choice = (
+        anthropic_types.ToolChoiceAutoParam(
+            type="auto",
+            # TODO: allow parallel tool use.
+            disable_parallel_tool_use=True,
+        )
+        if llm_request.tools_dict
+        else NOT_GIVEN
+    )
+    message = self._anthropic_client.messages.create(
+        model=llm_request.model,
+        system=llm_request.config.system_instruction,
+        messages=messages,
+        tools=tools,
+        tool_choice=tool_choice,
+        max_tokens=MAX_TOKEN,
+    )
+    logger.info(
+        "Claude response: %s",
+        message.model_dump_json(indent=2, exclude_none=True),
+    )
+    yield message_to_generate_content_response(message)
+
+  @cached_property
+  def _anthropic_client(self) -> AnthropicVertex:
+    if (
+        "GOOGLE_CLOUD_PROJECT" not in os.environ
+        or "GOOGLE_CLOUD_LOCATION" not in os.environ
+    ):
+      raise ValueError(
+          "GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION must be set for using"
+          " Anthropic on Vertex."
+      )
+
+    return AnthropicVertex(
+        project_id=os.environ["GOOGLE_CLOUD_PROJECT"],
+        region=os.environ["GOOGLE_CLOUD_LOCATION"],
+    )