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

google/adk/models/lite_llm.py

@@ -51,7 +51,7 @@ from .base_llm import BaseLlm
 from .llm_request import LlmRequest
 from .llm_response import LlmResponse
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 _NEW_LINE = "\n"
 _EXCLUDED_PART_FIELD = {"inline_data": {"data"}}
@@ -67,6 +67,12 @@ class TextChunk(BaseModel):
   text: str
 
 
+class UsageMetadataChunk(BaseModel):
+  prompt_tokens: int
+  completion_tokens: int
+  total_tokens: int
+
+
 class LiteLLMClient:
   """Provides acompletion method (for better testability)."""
 
@@ -172,19 +178,19 @@ def _content_to_message_param(
     tool_calls = []
     content_present = False
     for part in content.parts:
-        if part.function_call:
-            tool_calls.append(
-                ChatCompletionMessageToolCall(
-                    type="function",
-                    id=part.function_call.id,
-                    function=Function(
-                        name=part.function_call.name,
-                        arguments=part.function_call.args,
-                    ),
-                )
+      if part.function_call:
+        tool_calls.append(
+            ChatCompletionMessageToolCall(
+                type="function",
+                id=part.function_call.id,
+                function=Function(
+                    name=part.function_call.name,
+                    arguments=part.function_call.args,
+                ),
             )
-        elif part.text or part.inline_data:
-            content_present = True
+        )
+      elif part.text or part.inline_data:
+        content_present = True
 
     final_content = message_content if content_present else None
 
@@ -344,15 +350,20 @@ def _function_declaration_to_tool_param(
 def _model_response_to_chunk(
     response: ModelResponse,
 ) -> Generator[
-    Tuple[Optional[Union[TextChunk, FunctionChunk]], Optional[str]], None, None
+    Tuple[
+        Optional[Union[TextChunk, FunctionChunk, UsageMetadataChunk]],
+        Optional[str],
+    ],
+    None,
+    None,
 ]:
-  """Converts a litellm message to text or function chunk.
+  """Converts a litellm message to text, function or usage metadata chunk.
 
   Args:
     response: The response from the model.
 
   Yields:
-    A tuple of text or function chunk and finish reason.
+    A tuple of text or function or usage metadata chunk and finish reason.
   """
 
   message = None
@@ -384,11 +395,21 @@ def _model_response_to_chunk(
   if not message:
     yield None, None
 
+  # Ideally usage would be expected with the last ModelResponseStream with a
+  # finish_reason set. But this is not the case we are observing from litellm.
+  # So we are sending it as a separate chunk to be set on the llm_response.
+  if response.get("usage", None):
+    yield UsageMetadataChunk(
+        prompt_tokens=response["usage"].get("prompt_tokens", 0),
+        completion_tokens=response["usage"].get("completion_tokens", 0),
+        total_tokens=response["usage"].get("total_tokens", 0),
+    ), None
+
 
 def _model_response_to_generate_content_response(
     response: ModelResponse,
 ) -> LlmResponse:
-  """Converts a litellm response to LlmResponse.
+  """Converts a litellm response to LlmResponse. Also adds usage metadata.
 
   Args:
     response: The model response.
@@ -403,7 +424,15 @@ def _model_response_to_generate_content_response(
 
   if not message:
     raise ValueError("No message in response")
-  return _message_to_generate_content_response(message)
+
+  llm_response = _message_to_generate_content_response(message)
+  if response.get("usage", None):
+    llm_response.usage_metadata = types.GenerateContentResponseUsageMetadata(
+        prompt_token_count=response["usage"].get("prompt_tokens", 0),
+        candidates_token_count=response["usage"].get("completion_tokens", 0),
+        total_token_count=response["usage"].get("total_tokens", 0),
+    )
+  return llm_response
 
 
 def _message_to_generate_content_response(
@@ -453,9 +482,9 @@ def _get_completion_inputs(
   for content in llm_request.contents or []:
     message_param_or_list = _content_to_message_param(content)
     if isinstance(message_param_or_list, list):
-        messages.extend(message_param_or_list)
-    elif message_param_or_list: # Ensure it's not None before appending
-        messages.append(message_param_or_list)
+      messages.extend(message_param_or_list)
+    elif message_param_or_list:  # Ensure it's not None before appending
+      messages.append(message_param_or_list)
 
   if llm_request.config.system_instruction:
     messages.insert(
@@ -573,7 +602,6 @@ class LiteLlm(BaseLlm):
   Attributes:
     model: The name of the LiteLlm model.
     llm_client: The LLM client to use for the model.
-    model_config: The model config.
   """
 
   llm_client: LiteLLMClient = Field(default_factory=LiteLLMClient)
@@ -611,7 +639,8 @@ class LiteLlm(BaseLlm):
       LlmResponse: The model response.
     """
 
-    logger.info(_build_request_log(llm_request))
+    self._maybe_append_user_content(llm_request)
+    logger.debug(_build_request_log(llm_request))
 
     messages, tools = _get_completion_inputs(llm_request)
 
@@ -628,6 +657,10 @@ class LiteLlm(BaseLlm):
       function_args = ""
       function_id = None
       completion_args["stream"] = True
+      aggregated_llm_response = None
+      aggregated_llm_response_with_tool_call = None
+      usage_metadata = None
+
       for part in self.llm_client.completion(**completion_args):
         for chunk, finish_reason in _model_response_to_chunk(part):
           if isinstance(chunk, FunctionChunk):
@@ -645,32 +678,55 @@ class LiteLlm(BaseLlm):
                 ),
                 is_partial=True,
             )
+          elif isinstance(chunk, UsageMetadataChunk):
+            usage_metadata = types.GenerateContentResponseUsageMetadata(
+                prompt_token_count=chunk.prompt_tokens,
+                candidates_token_count=chunk.completion_tokens,
+                total_token_count=chunk.total_tokens,
+            )
+
           if finish_reason == "tool_calls" and function_id:
-            yield _message_to_generate_content_response(
-                ChatCompletionAssistantMessage(
-                    role="assistant",
-                    content="",
-                    tool_calls=[
-                        ChatCompletionMessageToolCall(
-                            type="function",
-                            id=function_id,
-                            function=Function(
-                                name=function_name,
-                                arguments=function_args,
-                            ),
-                        )
-                    ],
+            aggregated_llm_response_with_tool_call = (
+                _message_to_generate_content_response(
+                    ChatCompletionAssistantMessage(
+                        role="assistant",
+                        content="",
+                        tool_calls=[
+                            ChatCompletionMessageToolCall(
+                                type="function",
+                                id=function_id,
+                                function=Function(
+                                    name=function_name,
+                                    arguments=function_args,
+                                ),
+                            )
+                        ],
+                    )
                 )
             )
             function_name = ""
             function_args = ""
             function_id = None
           elif finish_reason == "stop" and text:
-            yield _message_to_generate_content_response(
+            aggregated_llm_response = _message_to_generate_content_response(
                 ChatCompletionAssistantMessage(role="assistant", content=text)
             )
             text = ""
 
+      # waiting until streaming ends to yield the llm_response as litellm tends
+      # to send chunk that contains usage_metadata after the chunk with
+      # finish_reason set to tool_calls or stop.
+      if aggregated_llm_response:
+        if usage_metadata:
+          aggregated_llm_response.usage_metadata = usage_metadata
+          usage_metadata = None
+        yield aggregated_llm_response
+
+      if aggregated_llm_response_with_tool_call:
+        if usage_metadata:
+          aggregated_llm_response_with_tool_call.usage_metadata = usage_metadata
+        yield aggregated_llm_response_with_tool_call
+
     else:
       response = await self.llm_client.acompletion(**completion_args)
       yield _model_response_to_generate_content_response(response)

google/adk/models/llm_request.py

@@ -37,7 +37,7 @@ class LlmRequest(BaseModel):
   """
 
   model_config = ConfigDict(arbitrary_types_allowed=True)
-  """The model config."""
+  """The pydantic model config."""
 
   model: Optional[str] = None
   """The model name."""

google/adk/models/llm_response.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 from typing import Any, Optional
 
 from google.genai import types
+from pydantic import alias_generators
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
@@ -40,8 +41,12 @@ class LlmResponse(BaseModel):
     custom_metadata: The custom metadata of the LlmResponse.
   """
 
-  model_config = ConfigDict(extra='forbid')
-  """The model config."""
+  model_config = ConfigDict(
+      extra='forbid',
+      alias_generator=alias_generators.to_camel,
+      populate_by_name=True,
+  )
+  """The pydantic model config."""
 
   content: Optional[types.Content] = None
   """The content of the response."""
@@ -80,6 +85,9 @@ class LlmResponse(BaseModel):
   NOTE: the entire dict must be JSON serializable.
   """
 
+  usage_metadata: Optional[types.GenerateContentResponseUsageMetadata] = None
+  """The usage metadata of the LlmResponse"""
+
   @staticmethod
   def create(
       generate_content_response: types.GenerateContentResponse,
@@ -93,18 +101,20 @@ class LlmResponse(BaseModel):
     Returns:
       The LlmResponse.
     """
-
+    usage_metadata = generate_content_response.usage_metadata
     if generate_content_response.candidates:
       candidate = generate_content_response.candidates[0]
       if candidate.content and candidate.content.parts:
         return LlmResponse(
             content=candidate.content,
             grounding_metadata=candidate.grounding_metadata,
+            usage_metadata=usage_metadata,
         )
       else:
         return LlmResponse(
             error_code=candidate.finish_reason,
             error_message=candidate.finish_message,
+            usage_metadata=usage_metadata,
         )
     else:
       if generate_content_response.prompt_feedback:
@@ -112,9 +122,11 @@ class LlmResponse(BaseModel):
         return LlmResponse(
             error_code=prompt_feedback.block_reason,
             error_message=prompt_feedback.block_reason_message,
+            usage_metadata=usage_metadata,
         )
       else:
         return LlmResponse(
             error_code='UNKNOWN_ERROR',
             error_message='Unknown error.',
+            usage_metadata=usage_metadata,
         )

google/adk/models/registry.py

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
   from .base_llm import BaseLlm
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 _llm_registry_dict: dict[str, type[BaseLlm]] = {}

google/adk/runners.py

@@ -21,8 +21,8 @@ import threading
 from typing import AsyncGenerator
 from typing import Generator
 from typing import Optional
+import warnings
 
-from deprecated import deprecated
 from google.genai import types
 
 from .agents.active_streaming_tool import ActiveStreamingTool
@@ -32,7 +32,6 @@ from .agents.invocation_context import new_invocation_context_id
 from .agents.live_request_queue import LiveRequestQueue
 from .agents.llm_agent import LlmAgent
 from .agents.run_config import RunConfig
-from .agents.run_config import StreamingMode
 from .artifacts.base_artifact_service import BaseArtifactService
 from .artifacts.in_memory_artifact_service import InMemoryArtifactService
 from .events.event import Event
@@ -42,9 +41,9 @@ from .sessions.base_session_service import BaseSessionService
 from .sessions.in_memory_session_service import InMemorySessionService
 from .sessions.session import Session
 from .telemetry import tracer
-from .tools.built_in_code_execution_tool import built_in_code_execution
+from .tools._built_in_code_execution_tool import built_in_code_execution
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class Runner:
@@ -172,7 +171,7 @@ class Runner:
       The events generated by the agent.
     """
     with tracer.start_as_current_span('invocation'):
-      session = self.session_service.get_session(
+      session = await self.session_service.get_session(
           app_name=self.app_name, user_id=user_id, session_id=session_id
       )
       if not session:
@@ -186,7 +185,7 @@ class Runner:
       root_agent = self.agent
 
       if new_message:
-        self._append_new_message_to_session(
+        await self._append_new_message_to_session(
             session,
             new_message,
             invocation_context,
@@ -196,10 +195,10 @@ class Runner:
       invocation_context.agent = self._find_agent_to_run(session, root_agent)
       async for event in invocation_context.agent.run_async(invocation_context):
         if not event.partial:
-          self.session_service.append_event(session=session, event=event)
+          await self.session_service.append_event(session=session, event=event)
         yield event
 
-  def _append_new_message_to_session(
+  async def _append_new_message_to_session(
       self,
       session: Session,
       new_message: types.Content,
@@ -225,7 +224,7 @@ class Runner:
         if part.inline_data is None:
           continue
         file_name = f'artifact_{invocation_context.invocation_id}_{i}'
-        self.artifact_service.save_artifact(
+        await self.artifact_service.save_artifact(
             app_name=self.app_name,
             user_id=session.user_id,
             session_id=session.id,
@@ -241,30 +240,57 @@ class Runner:
         author='user',
         content=new_message,
     )
-    self.session_service.append_event(session=session, event=event)
+    await self.session_service.append_event(session=session, event=event)
 
   async def run_live(
       self,
       *,
-      session: Session,
+      user_id: Optional[str] = None,
+      session_id: Optional[str] = None,
       live_request_queue: LiveRequestQueue,
       run_config: RunConfig = RunConfig(),
+      session: Optional[Session] = None,
   ) -> AsyncGenerator[Event, None]:
     """Runs the agent in live mode (experimental feature).
 
     Args:
-        session: The session to use.
+        user_id: The user ID for the session. Required if `session` is None.
+        session_id: The session ID for the session. Required if `session` is
+          None.
         live_request_queue: The queue for live requests.
         run_config: The run config for the agent.
+        session: The session to use. This parameter is deprecated, please use
+          `user_id` and `session_id` instead.
 
     Yields:
-        The events generated by the agent.
+        AsyncGenerator[Event, None]: An asynchronous generator that yields
+        `Event`
+        objects as they are produced by the agent during its live execution.
 
     .. warning::
         This feature is **experimental** and its API or behavior may change
         in future releases.
+
+    .. note::
+        Either `session` or both `user_id` and `session_id` must be provided.
     """
-    # TODO: right now, only works for a single audio agent without FC.
+    if session is None and (user_id is None or session_id is None):
+      raise ValueError(
+          'Either session or user_id and session_id must be provided.'
+      )
+    if session is not None:
+      warnings.warn(
+          'The `session` parameter is deprecated. Please use `user_id` and'
+          ' `session_id` instead.',
+          DeprecationWarning,
+          stacklevel=2,
+      )
+    if not session:
+      session = self.session_service.get_session(
+          app_name=self.app_name, user_id=user_id, session_id=session_id
+      )
+      if not session:
+        raise ValueError(f'Session not found: {session_id}')
     invocation_context = self._new_invocation_context_for_live(
         session,
         live_request_queue=live_request_queue,
@@ -276,37 +302,29 @@ class Runner:
 
     invocation_context.active_streaming_tools = {}
     # TODO(hangfei): switch to use canonical_tools.
-    for tool in invocation_context.agent.tools:
-      # replicate a LiveRequestQueue for streaming tools that relis on
-      # LiveRequestQueue
-      from typing import get_type_hints
-
-      type_hints = get_type_hints(tool)
-      for arg_type in type_hints.values():
-        if arg_type is LiveRequestQueue:
-          if not invocation_context.active_streaming_tools:
-            invocation_context.active_streaming_tools = {}
-          active_streaming_tools = ActiveStreamingTool(
-              stream=LiveRequestQueue()
-          )
-          invocation_context.active_streaming_tools[tool.__name__] = (
-              active_streaming_tools
-          )
+    # for shell agents, there is no tools associated with it so we should skip.
+    if hasattr(invocation_context.agent, 'tools'):
+      for tool in invocation_context.agent.tools:
+        # replicate a LiveRequestQueue for streaming tools that relis on
+        # LiveRequestQueue
+        from typing import get_type_hints
+
+        type_hints = get_type_hints(tool)
+        for arg_type in type_hints.values():
+          if arg_type is LiveRequestQueue:
+            if not invocation_context.active_streaming_tools:
+              invocation_context.active_streaming_tools = {}
+            active_streaming_tools = ActiveStreamingTool(
+                stream=LiveRequestQueue()
+            )
+            invocation_context.active_streaming_tools[tool.__name__] = (
+                active_streaming_tools
+            )
 
     async for event in invocation_context.agent.run_live(invocation_context):
-      self.session_service.append_event(session=session, event=event)
+      await self.session_service.append_event(session=session, event=event)
       yield event
 
-  def close_session(self, session: Session):
-    """Closes a session and adds it to the memory service (experimental feature).
-
-    Args:
-        session: The session to close.
-    """
-    if self.memory_service:
-      self.memory_service.add_session_to_memory(session)
-    self.session_service.close_session(session=session)
-
   def _find_agent_to_run(
       self, session: Session, root_agent: BaseAgent
   ) -> BaseAgent:
@@ -391,7 +409,7 @@ class Runner:
             f'CFC is not supported for model: {model_name} in agent:'
             f' {self.agent.name}'
         )
-      if built_in_code_execution not in self.agent.canonical_tools:
+      if built_in_code_execution not in self.agent.canonical_tools():
         self.agent.tools.append(built_in_code_execution)
 
     return InvocationContext(
@@ -430,6 +448,9 @@ class Runner:
           run_config.output_audio_transcription = (
               types.AudioTranscriptionConfig()
           )
+      if not run_config.input_audio_transcription:
+        # need this input transcription for agent transferring in live mode.
+        run_config.input_audio_transcription = types.AudioTranscriptionConfig()
     return self._new_invocation_context(
         session,
         live_request_queue=live_request_queue,
@@ -448,9 +469,11 @@ class InMemoryRunner(Runner):
       agent: The root agent to run.
       app_name: The application name of the runner. Defaults to
         'InMemoryRunner'.
+      _in_memory_session_service: Deprecated. Please don't use. The in-memory
+        session service for the runner.
   """
 
-  def __init__(self, agent: LlmAgent, *, app_name: str = 'InMemoryRunner'):
+  def __init__(self, agent: BaseAgent, *, app_name: str = 'InMemoryRunner'):
     """Initializes the InMemoryRunner.
 
     Args:
@@ -458,10 +481,11 @@ class InMemoryRunner(Runner):
         app_name: The application name of the runner. Defaults to
           'InMemoryRunner'.
     """
+    self._in_memory_session_service = InMemorySessionService()
     super().__init__(
         app_name=app_name,
         agent=agent,
         artifact_service=InMemoryArtifactService(),
-        session_service=InMemorySessionService(),
+        session_service=self._in_memory_session_service,
         memory_service=InMemoryMemoryService(),
     )

google/adk/sessions/__init__.py

@@ -19,7 +19,7 @@ from .session import Session
 from .state import State
 from .vertex_ai_session_service import VertexAiSessionService
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 __all__ = [

google/adk/sessions/_session_util.py

@@ -1,3 +1,17 @@
+# 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.
+
 """Utility functions for session service."""
 
 import base64

google/adk/sessions/base_session_service.py

@@ -26,6 +26,7 @@ from .state import State
 
 class GetSessionConfig(BaseModel):
   """The configuration of getting a session."""
+
   num_recent_events: Optional[int] = None
   after_timestamp: Optional[float] = None
 
@@ -35,13 +36,8 @@ class ListSessionsResponse(BaseModel):
 
   The events and states are not set within each Session object.
   """
-  sessions: list[Session] = Field(default_factory=list)
-
 
-class ListEventsResponse(BaseModel):
-  """The response of listing events in a session."""
-  events: list[Event] = Field(default_factory=list)
-  next_page_token: Optional[str] = None
+  sessions: list[Session] = Field(default_factory=list)
 
 
 class BaseSessionService(abc.ABC):
@@ -51,7 +47,7 @@ class BaseSessionService(abc.ABC):
   """
 
   @abc.abstractmethod
-  def create_session(
+  async def create_session(
       self,
       *,
       app_name: str,
@@ -71,10 +67,9 @@ class BaseSessionService(abc.ABC):
     Returns:
       session: The newly created session instance.
     """
-    pass
 
   @abc.abstractmethod
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
@@ -83,39 +78,20 @@ class BaseSessionService(abc.ABC):
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
     """Gets a session."""
-    pass
 
   @abc.abstractmethod
-  def list_sessions(
+  async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
     """Lists all the sessions."""
-    pass
 
   @abc.abstractmethod
-  def delete_session(
+  async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     """Deletes a session."""
-    pass
-
-  @abc.abstractmethod
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    """Lists events in a session."""
-    pass
-
-  def close_session(self, *, session: Session):
-    """Closes a session."""
-    # TODO: determine whether we want to finalize the session here.
-    pass
 
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     """Appends an event to a session object."""
     if event.partial:
       return event
@@ -123,7 +99,7 @@ class BaseSessionService(abc.ABC):
     session.events.append(event)
     return event
 
-  def __update_session_state(self, session: Session, event: Event):
+  def __update_session_state(self, session: Session, event: Event) -> None:
     """Updates the session state based on the event."""
     if not event.actions or not event.actions.state_delta:
       return