google-adk 0.4.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

google/adk/models/anthropic_llm.py

@@ -19,12 +19,14 @@ from __future__ import annotations
 from functools import cached_property
 import logging
 import os
+from typing import Any
 from typing import AsyncGenerator
 from typing import Generator
 from typing import Iterable
 from typing import Literal
-from typing import Optional, Union
+from typing import Optional
 from typing import TYPE_CHECKING
+from typing import Union
 
 from anthropic import AnthropicVertex
 from anthropic import NOT_GIVEN
@@ -41,7 +43,7 @@ if TYPE_CHECKING:
 
 __all__ = ["Claude"]
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 MAX_TOKEN = 1024
 
@@ -139,18 +141,36 @@ def message_to_generate_content_response(
           role="model",
           parts=[content_block_to_part(cb) for cb in message.content],
       ),
+      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
+          ),
+      ),
       # 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 _update_type_string(value_dict: dict[str, Any]):
+  """Updates 'type' field to expected JSON schema format."""
+  if "type" in value_dict:
+    value_dict["type"] = value_dict["type"].lower()
+
+  if "items" in value_dict:
+    # 'type' field could exist for items as well, this would be the case if
+    # items represent primitive types.
+    _update_type_string(value_dict["items"])
+
+    if "properties" in value_dict["items"]:
+      # There could be properties as well on the items, especially if the items
+      # are complex object themselves. We recursively traverse each individual
+      # property as well and fix the "type" value.
+      for _, value in value_dict["items"]["properties"].items():
+        _update_type_string(value)
+
+
 def function_declaration_to_tool_param(
     function_declaration: types.FunctionDeclaration,
 ) -> anthropic_types.ToolParam:
@@ -163,8 +183,7 @@ def function_declaration_to_tool_param(
   ):
     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()
+      _update_type_string(value_dict)
       properties[key] = value_dict
 
   return anthropic_types.ToolParam(
@@ -178,6 +197,12 @@ def function_declaration_to_tool_param(
 
 
 class Claude(BaseLlm):
+  """ "Integration with Claude models served from Vertex AI.
+
+  Attributes:
+    model: The name of the Claude model.
+  """
+
   model: str = "claude-3-5-sonnet-v2@20241022"
 
   @staticmethod

google/adk/models/base_llm.py

@@ -14,9 +14,9 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import AsyncGenerator
-from typing import TYPE_CHECKING
+from typing import AsyncGenerator, TYPE_CHECKING
 
+from google.genai import types
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
@@ -32,14 +32,13 @@ class BaseLlm(BaseModel):
 
   Attributes:
     model: The name of the LLM, e.g. gemini-1.5-flash or gemini-1.5-flash-001.
-    model_config: The model config
   """
 
   model_config = ConfigDict(
       # This allows us to use arbitrary types in the model. E.g. PIL.Image.
       arbitrary_types_allowed=True,
   )
-  """The model config."""
+  """The pydantic model config."""
 
   model: str
   """The name of the LLM, e.g. gemini-1.5-flash or gemini-1.5-flash-001."""
@@ -73,6 +72,48 @@ class BaseLlm(BaseModel):
     )
     yield  # AsyncGenerator requires a yield statement in function body.
 
+  def _maybe_append_user_content(self, llm_request: LlmRequest):
+    """Appends a user content, so that model can continue to output.
+
+    Args:
+      llm_request: LlmRequest, the request to send to the Gemini model.
+    """
+    # If no content is provided, append a user content to hint model response
+    # using system instruction.
+    if not llm_request.contents:
+      llm_request.contents.append(
+          types.Content(
+              role='user',
+              parts=[
+                  types.Part(
+                      text=(
+                          'Handle the requests as specified in the System'
+                          ' Instruction.'
+                      )
+                  )
+              ],
+          )
+      )
+      return
+
+    # Insert a user content to preserve user intent and to avoid empty
+    # model response.
+    if llm_request.contents[-1].role != 'user':
+      llm_request.contents.append(
+          types.Content(
+              role='user',
+              parts=[
+                  types.Part(
+                      text=(
+                          'Continue processing previous requests as instructed.'
+                          ' Exit or provide a summary if no more outputs are'
+                          ' needed.'
+                      )
+                  )
+              ],
+          )
+      )
+
   def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
     """Creates a live connection to the LLM.
 

google/adk/models/gemini_llm_connection.py

@@ -21,7 +21,7 @@ from google.genai import types
 from .base_llm_connection import BaseLlmConnection
 from .llm_response import LlmResponse
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class GeminiLlmConnection(BaseLlmConnection):
@@ -145,7 +145,20 @@ class GeminiLlmConnection(BaseLlmConnection):
             yield self.__build_full_text_response(text)
             text = ''
           yield llm_response
-
+        if (
+            message.server_content.input_transcription
+            and message.server_content.input_transcription.text
+        ):
+          user_text = message.server_content.input_transcription.text
+          parts = [
+              types.Part.from_text(
+                  text=user_text,
+              )
+          ]
+          llm_response = LlmResponse(
+              content=types.Content(role='user', parts=parts)
+          )
+          yield llm_response
         if (
             message.server_content.output_transcription
             and message.server_content.output_transcription.text

google/adk/models/google_llm.py

@@ -11,6 +11,8 @@
 # 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
 
 import contextlib
@@ -34,7 +36,7 @@ from .llm_response import LlmResponse
 if TYPE_CHECKING:
   from .llm_request import LlmRequest
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 _NEW_LINE = '\n'
 _EXCLUDED_PART_FIELD = {'inline_data': {'data'}}
@@ -121,6 +123,7 @@ class Gemini(BaseLlm):
               content=types.ModelContent(
                   parts=[types.Part.from_text(text=text)],
               ),
+              usage_metadata=llm_response.usage_metadata,
           )
           text = ''
         yield llm_response
@@ -174,9 +177,13 @@ class Gemini(BaseLlm):
   @cached_property
   def _live_api_client(self) -> Client:
     if self._api_backend == 'vertex':
+      # use beta version for vertex api
+      api_version = 'v1beta1'
       # use default api version for vertex
       return Client(
-          http_options=types.HttpOptions(headers=self._tracking_headers)
+          http_options=types.HttpOptions(
+              headers=self._tracking_headers, api_version=api_version
+          )
       )
     else:
       # use v1alpha for ml_dev
@@ -210,48 +217,6 @@ class Gemini(BaseLlm):
     ) as live_session:
       yield GeminiLlmConnection(live_session)
 
-  def _maybe_append_user_content(self, llm_request: LlmRequest):
-    """Appends a user content, so that model can continue to output.
-
-    Args:
-      llm_request: LlmRequest, the request to send to the Gemini model.
-    """
-    # If no content is provided, append a user content to hint model response
-    # using system instruction.
-    if not llm_request.contents:
-      llm_request.contents.append(
-          types.Content(
-              role='user',
-              parts=[
-                  types.Part(
-                      text=(
-                          'Handle the requests as specified in the System'
-                          ' Instruction.'
-                      )
-                  )
-              ],
-          )
-      )
-      return
-
-    # Insert a user content to preserve user intent and to avoid empty
-    # model response.
-    if llm_request.contents[-1].role != 'user':
-      llm_request.contents.append(
-          types.Content(
-              role='user',
-              parts=[
-                  types.Part(
-                      text=(
-                          'Continue processing previous requests as instructed.'
-                          ' Exit or provide a summary if no more outputs are'
-                          ' needed.'
-                      )
-                  )
-              ],
-          )
-      )
-
 
 def _build_function_declaration_log(
     func_decl: types.FunctionDeclaration,

google/adk/models/google_llm.py.orig

@@ -0,0 +1,305 @@
+# 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
+
+import contextlib
+from functools import cached_property
+import logging
+import sys
+from typing import AsyncGenerator
+from typing import cast
+from typing import TYPE_CHECKING
+
+from google.genai import Client
+from google.genai import types
+from typing_extensions import override
+
+from .. import version
+from .base_llm import BaseLlm
+from .base_llm_connection import BaseLlmConnection
+from .gemini_llm_connection import GeminiLlmConnection
+from .llm_response import LlmResponse
+
+if TYPE_CHECKING:
+  from .llm_request import LlmRequest
+
+logger = None
+
+_NEW_LINE = '\n'
+_EXCLUDED_PART_FIELD = {'inline_data': {'data'}}
+
+
+class Gemini(BaseLlm):
+  """Integration for Gemini models.
+
+  Attributes:
+    model: The name of the Gemini model.
+  """
+
+  model: str = 'gemini-1.5-flash'
+
+  @staticmethod
+  @override
+  def supported_models() -> list[str]:
+    """Provides the list of supported models.
+
+    Returns:
+      A list of supported models.
+    """
+
+    return [
+        r'gemini-.*',
+        # fine-tuned vertex endpoint pattern
+        r'projects\/.+\/locations\/.+\/endpoints\/.+',
+        # vertex gemini long name
+        r'projects\/.+\/locations\/.+\/publishers\/google\/models\/gemini.+',
+    ]
+
+  async def generate_content_async(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> AsyncGenerator[LlmResponse, None]:
+    """Sends a request to the Gemini model.
+
+    Args:
+      llm_request: LlmRequest, the request to send to the Gemini model.
+      stream: bool = False, whether to do streaming call.
+
+    Yields:
+      LlmResponse: The model response.
+    """
+
+    self._maybe_append_user_content(llm_request)
+
+    global logger
+    if not logger:
+      logger = logging.getLogger(__name__)
+
+    logger.info(
+        'Sending out request, model: %s, backend: %s, stream: %s',
+        llm_request.model,
+        self._api_backend,
+        stream,
+    )
+    logger.info(_build_request_log(llm_request))
+
+    print('********* Jack --> ')
+    for hh in logging.root.handlers:
+      print(hh, hh.level)
+    for hh in logger.handlers:
+      print(hh, hh.level)
+    print('********* Jack <-- ')
+
+    if stream:
+      responses = await self.api_client.aio.models.generate_content_stream(
+          model=llm_request.model,
+          contents=llm_request.contents,
+          config=llm_request.config,
+      )
+      response = None
+      text = ''
+      # for sse, similar as bidi (see receive method in gemini_llm_connecton.py),
+      # we need to mark those text content as partial and after all partial
+      # contents are sent, we send an accumulated event which contains all the
+      # previous partial content. The only difference is bidi rely on
+      # complete_turn flag to detect end while sse depends on finish_reason.
+      async for response in responses:
+        logger.info(_build_response_log(response))
+        llm_response = LlmResponse.create(response)
+        if (
+            llm_response.content
+            and llm_response.content.parts
+            and llm_response.content.parts[0].text
+        ):
+          text += llm_response.content.parts[0].text
+          llm_response.partial = True
+        elif text and (
+            not llm_response.content
+            or not llm_response.content.parts
+            # don't yield the merged text event when receiving audio data
+            or not llm_response.content.parts[0].inline_data
+        ):
+          yield LlmResponse(
+              content=types.ModelContent(
+                  parts=[types.Part.from_text(text=text)],
+              ),
+              usage_metadata=llm_response.usage_metadata,
+          )
+          text = ''
+        yield llm_response
+      if (
+          text
+          and response
+          and response.candidates
+          and response.candidates[0].finish_reason == types.FinishReason.STOP
+      ):
+        yield LlmResponse(
+            content=types.ModelContent(
+                parts=[types.Part.from_text(text=text)],
+            ),
+        )
+
+    else:
+      response = await self.api_client.aio.models.generate_content(
+          model=llm_request.model,
+          contents=llm_request.contents,
+          config=llm_request.config,
+      )
+      logger.info(_build_response_log(response))
+      yield LlmResponse.create(response)
+
+  @cached_property
+  def api_client(self) -> Client:
+    """Provides the api client.
+
+    Returns:
+      The api client.
+    """
+    return Client(
+        http_options=types.HttpOptions(headers=self._tracking_headers)
+    )
+
+  @cached_property
+  def _api_backend(self) -> str:
+    return 'vertex' if self.api_client.vertexai else 'ml_dev'
+
+  @cached_property
+  def _tracking_headers(self) -> dict[str, str]:
+    framework_label = f'google-adk/{version.__version__}'
+    language_label = 'gl-python/' + sys.version.split()[0]
+    version_header_value = f'{framework_label} {language_label}'
+    tracking_headers = {
+        'x-goog-api-client': version_header_value,
+        'user-agent': version_header_value,
+    }
+    return tracking_headers
+
+  @cached_property
+  def _live_api_client(self) -> Client:
+    if self._api_backend == 'vertex':
+      # use beta version for vertex api
+      api_version = 'v1beta1'
+      # use default api version for vertex
+      return Client(
+          http_options=types.HttpOptions(
+              headers=self._tracking_headers, api_version=api_version
+          )
+      )
+    else:
+      # use v1alpha for ml_dev
+      api_version = 'v1alpha'
+      return Client(
+          http_options=types.HttpOptions(
+              headers=self._tracking_headers, api_version=api_version
+          )
+      )
+
+  @contextlib.asynccontextmanager
+  async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
+    """Connects to the Gemini model and returns an llm connection.
+
+    Args:
+      llm_request: LlmRequest, the request to send to the Gemini model.
+
+    Yields:
+      BaseLlmConnection, the connection to the Gemini model.
+    """
+
+    llm_request.live_connect_config.system_instruction = types.Content(
+        role='system',
+        parts=[
+            types.Part.from_text(text=llm_request.config.system_instruction)
+        ],
+    )
+    llm_request.live_connect_config.tools = llm_request.config.tools
+    async with self._live_api_client.aio.live.connect(
+        model=llm_request.model, config=llm_request.live_connect_config
+    ) as live_session:
+      yield GeminiLlmConnection(live_session)
+
+
+def _build_function_declaration_log(
+    func_decl: types.FunctionDeclaration,
+) -> str:
+  param_str = '{}'
+  if func_decl.parameters and func_decl.parameters.properties:
+    param_str = str({
+        k: v.model_dump(exclude_none=True)
+        for k, v in func_decl.parameters.properties.items()
+    })
+  return_str = 'None'
+  if func_decl.response:
+    return_str = str(func_decl.response.model_dump(exclude_none=True))
+  return f'{func_decl.name}: {param_str} -> {return_str}'
+
+
+def _build_request_log(req: LlmRequest) -> str:
+  function_decls: list[types.FunctionDeclaration] = cast(
+      list[types.FunctionDeclaration],
+      req.config.tools[0].function_declarations if req.config.tools else [],
+  )
+  function_logs = (
+      [
+          _build_function_declaration_log(func_decl)
+          for func_decl in function_decls
+      ]
+      if function_decls
+      else []
+  )
+  contents_logs = [
+      content.model_dump_json(
+          exclude_none=True,
+          exclude={
+              'parts': {
+                  i: _EXCLUDED_PART_FIELD for i in range(len(content.parts))
+              }
+          },
+      )
+      for content in req.contents
+  ]
+
+  return f"""
+LLM Request:
+-----------------------------------------------------------
+System Instruction:
+{req.config.system_instruction}
+-----------------------------------------------------------
+Contents:
+{_NEW_LINE.join(contents_logs)}
+-----------------------------------------------------------
+Functions:
+{_NEW_LINE.join(function_logs)}
+-----------------------------------------------------------
+"""
+
+
+def _build_response_log(resp: types.GenerateContentResponse) -> str:
+  function_calls_text = []
+  if function_calls := resp.function_calls:
+    for func_call in function_calls:
+      function_calls_text.append(
+          f'name: {func_call.name}, args: {func_call.args}'
+      )
+  return f"""
+LLM Response:
+-----------------------------------------------------------
+Text:
+{resp.text}
+-----------------------------------------------------------
+Function calls:
+{_NEW_LINE.join(function_calls_text)}
+-----------------------------------------------------------
+Raw response:
+{resp.model_dump_json(exclude_none=True)}
+-----------------------------------------------------------
+"""