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

google/adk/flows/llm_flows/base_llm_flow.py

@@ -0,0 +1,559 @@
+# 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 abc import ABC
+import asyncio
+import logging
+from typing import AsyncGenerator
+from typing import cast
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from websockets.exceptions import ConnectionClosedOK
+
+from ...agents.base_agent import BaseAgent
+from ...agents.callback_context import CallbackContext
+from ...agents.invocation_context import InvocationContext
+from ...agents.live_request_queue import LiveRequestQueue
+from ...agents.run_config import StreamingMode
+from ...agents.transcription_entry import TranscriptionEntry
+from ...events.event import Event
+from ...models.base_llm_connection import BaseLlmConnection
+from ...models.llm_request import LlmRequest
+from ...models.llm_response import LlmResponse
+from ...telemetry import trace_call_llm
+from ...telemetry import trace_send_data
+from ...telemetry import tracer
+from ...tools.tool_context import ToolContext
+from . import functions
+
+if TYPE_CHECKING:
+  from ...agents.llm_agent import LlmAgent
+  from ...models.base_llm import BaseLlm
+  from ._base_llm_processor import BaseLlmRequestProcessor
+  from ._base_llm_processor import BaseLlmResponseProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class BaseLlmFlow(ABC):
+  """A basic flow that calls the LLM in a loop until a final response is generated.
+
+  This flow ends when it transfer to another agent.
+  """
+
+  def __init__(self):
+    self.request_processors: list[BaseLlmRequestProcessor] = []
+    self.response_processors: list[BaseLlmResponseProcessor] = []
+
+  async def run_live(
+      self,
+      invocation_context: InvocationContext,
+  ) -> AsyncGenerator[Event, None]:
+    """Runs the flow using live api."""
+    llm_request = LlmRequest()
+    event_id = Event.new_id()
+
+    # Preprocess before calling the LLM.
+    async for event in self._preprocess_async(invocation_context, llm_request):
+      yield event
+    if invocation_context.end_invocation:
+      return
+
+    llm = self.__get_llm(invocation_context)
+    logger.info(
+        'Establishing live connection for agent: %s with llm request: %s',
+        invocation_context.agent.name,
+        llm_request,
+    )
+    async with llm.connect(llm_request) as llm_connection:
+      if llm_request.contents:
+        # Sends the conversation history to the model.
+        with tracer.start_as_current_span('send_data'):
+
+          if invocation_context.transcription_cache:
+            from . import audio_transcriber
+
+            audio_transcriber = audio_transcriber.AudioTranscriber()
+            contents = audio_transcriber.transcribe_file(invocation_context)
+            logger.debug('Sending history to model: %s', contents)
+            await llm_connection.send_history(contents)
+            invocation_context.transcription_cache = None
+            trace_send_data(invocation_context, event_id, contents)
+          else:
+            await llm_connection.send_history(llm_request.contents)
+            trace_send_data(invocation_context, event_id, llm_request.contents)
+
+      send_task = asyncio.create_task(
+          self._send_to_model(llm_connection, invocation_context)
+      )
+
+      try:
+        async for event in self._receive_from_model(
+            llm_connection,
+            event_id,
+            invocation_context,
+            llm_request,
+        ):
+          # Empty event means the queue is closed.
+          if not event:
+            break
+          logger.debug('Receive new event: %s', event)
+          yield event
+          # send back the function response
+          if event.get_function_responses():
+            logger.debug('Sending back last function resonse event: %s', event)
+            invocation_context.live_request_queue.send_content(event.content)
+          if (
+              event.content
+              and event.content.parts
+              and event.content.parts[0].function_response
+              and event.content.parts[0].function_response.name
+              == 'transfer_to_agent'
+          ):
+            await asyncio.sleep(1)
+            # cancel the tasks that belongs to the closed connection.
+            send_task.cancel()
+            await llm_connection.close()
+      finally:
+        # Clean up
+        if not send_task.done():
+          send_task.cancel()
+        try:
+          await send_task
+        except asyncio.CancelledError:
+          pass
+
+  async def _send_to_model(
+      self,
+      llm_connection: BaseLlmConnection,
+      invocation_context: InvocationContext,
+  ):
+    """Sends data to model."""
+    while True:
+      live_request_queue = invocation_context.live_request_queue
+      try:
+        # Streamlit's execution model doesn't preemptively yield to the event
+        # loop. Therefore, we must explicitly introduce timeouts to allow the
+        # event loop to process events.
+        # TODO: revert back(remove timeout) once we move off streamlit.
+        live_request = await asyncio.wait_for(
+            live_request_queue.get(), timeout=0.25
+        )
+        # duplicate the live_request to all the active streams
+        logger.debug(
+            'Sending live request %s to active streams: %s',
+            live_request,
+            invocation_context.active_streaming_tools,
+        )
+        if invocation_context.active_streaming_tools:
+          for active_streaming_tool in (
+              invocation_context.active_streaming_tools
+          ).values():
+            if active_streaming_tool.stream:
+              active_streaming_tool.stream.send(live_request)
+        await asyncio.sleep(0)
+      except asyncio.TimeoutError:
+        continue
+      if live_request.close:
+        await llm_connection.close()
+        return
+      if live_request.blob:
+        # Cache audio data here for transcription
+        if not invocation_context.transcription_cache:
+          invocation_context.transcription_cache = []
+        invocation_context.transcription_cache.append(
+            TranscriptionEntry(role='user', data=live_request.blob)
+        )
+        await llm_connection.send_realtime(live_request.blob)
+      if live_request.content:
+        await llm_connection.send_content(live_request.content)
+
+  async def _receive_from_model(
+      self,
+      llm_connection: BaseLlmConnection,
+      event_id: str,
+      invocation_context: InvocationContext,
+      llm_request: LlmRequest,
+  ) -> AsyncGenerator[Event, None]:
+    """Receive data from model and process events using BaseLlmConnection."""
+    assert invocation_context.live_request_queue
+    try:
+      while True:
+        async for llm_response in llm_connection.receive():
+          model_response_event = Event(
+              id=Event.new_id(),
+              invocation_id=invocation_context.invocation_id,
+              author=invocation_context.agent.name,
+          )
+          async for event in self._postprocess_live(
+              invocation_context,
+              llm_request,
+              llm_response,
+              model_response_event,
+          ):
+            if (
+                event.content
+                and event.content.parts
+                and event.content.parts[0].text
+                and not event.partial
+            ):
+              if not invocation_context.transcription_cache:
+                invocation_context.transcription_cache = []
+              invocation_context.transcription_cache.append(
+                  TranscriptionEntry(role='model', data=event.content)
+              )
+            yield event
+        # Give opportunity for other tasks to run.
+        await asyncio.sleep(0)
+    except ConnectionClosedOK:
+      pass
+
+  async def run_async(
+      self, invocation_context: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    """Runs the flow."""
+    while True:
+      last_event = None
+      async for event in self._run_one_step_async(invocation_context):
+        last_event = event
+        yield event
+      if not last_event or last_event.is_final_response():
+        break
+
+  async def _run_one_step_async(
+      self,
+      invocation_context: InvocationContext,
+  ) -> AsyncGenerator[Event, None]:
+    """One step means one LLM call."""
+    llm_request = LlmRequest()
+
+    # Preprocess before calling the LLM.
+    async for event in self._preprocess_async(invocation_context, llm_request):
+      yield event
+    if invocation_context.end_invocation:
+      return
+
+    # Calls the LLM.
+    model_response_event = Event(
+        id=Event.new_id(),
+        invocation_id=invocation_context.invocation_id,
+        author=invocation_context.agent.name,
+        branch=invocation_context.branch,
+    )
+    async for llm_response in self._call_llm_async(
+        invocation_context, llm_request, model_response_event
+    ):
+      # Postprocess after calling the LLM.
+      async for event in self._postprocess_async(
+          invocation_context, llm_request, llm_response, model_response_event
+      ):
+        yield event
+
+  async def _preprocess_async(
+      self, invocation_context: InvocationContext, llm_request: LlmRequest
+  ) -> AsyncGenerator[Event, None]:
+    from ...agents.llm_agent import LlmAgent
+
+    agent = invocation_context.agent
+    if not isinstance(agent, LlmAgent):
+      return
+
+    # Runs processors.
+    for processor in self.request_processors:
+      async for event in processor.run_async(invocation_context, llm_request):
+        yield event
+
+    # Run processors for tools.
+    for tool in agent.canonical_tools:
+      tool_context = ToolContext(invocation_context)
+      await tool.process_llm_request(
+          tool_context=tool_context, llm_request=llm_request
+      )
+
+  async def _postprocess_async(
+      self,
+      invocation_context: InvocationContext,
+      llm_request: LlmRequest,
+      llm_response: LlmResponse,
+      model_response_event: Event,
+  ) -> AsyncGenerator[Event, None]:
+    """Postprocess after calling the LLM.
+
+    Args:
+      invocation_context: The invocation context.
+      llm_request: The original LLM request.
+      llm_response: The LLM response from the LLM call.
+      model_response_event: A mutable event for the LLM response.
+
+    Yields:
+      A generator of events.
+    """
+
+    # Runs processors.
+    async for event in self._postprocess_run_processors_async(
+        invocation_context, llm_response
+    ):
+      yield event
+
+    # Skip the model response event if there is no content and no error code.
+    # This is needed for the code executor to trigger another loop.
+    if (
+        not llm_response.content
+        and not llm_response.error_code
+        and not llm_response.interrupted
+    ):
+      return
+
+    # Builds the event.
+    model_response_event = self._finalize_model_response_event(
+        llm_request, llm_response, model_response_event
+    )
+    yield model_response_event
+
+    # Handles function calls.
+    if model_response_event.get_function_calls():
+      async for event in self._postprocess_handle_function_calls_async(
+          invocation_context, model_response_event, llm_request
+      ):
+        yield event
+
+  async def _postprocess_live(
+      self,
+      invocation_context: InvocationContext,
+      llm_request: LlmRequest,
+      llm_response: LlmResponse,
+      model_response_event: Event,
+  ) -> AsyncGenerator[Event, None]:
+    """Postprocess after calling the LLM asynchronously.
+
+    Args:
+      invocation_context: The invocation context.
+      llm_request: The original LLM request.
+      llm_response: The LLM response from the LLM call.
+      model_response_event: A mutable event for the LLM response.
+
+    Yields:
+      A generator of events.
+    """
+
+    # Runs processors.
+    async for event in self._postprocess_run_processors_async(
+        invocation_context, llm_response
+    ):
+      yield event
+
+    # Skip the model response event if there is no content and no error code.
+    # This is needed for the code executor to trigger another loop.
+    # But don't skip control events like turn_complete.
+    if (
+        not llm_response.content
+        and not llm_response.error_code
+        and not llm_response.interrupted
+        and not llm_response.turn_complete
+    ):
+      return
+
+    # Builds the event.
+    model_response_event = self._finalize_model_response_event(
+        llm_request, llm_response, model_response_event
+    )
+    yield model_response_event
+
+    # Handles function calls.
+    if model_response_event.get_function_calls():
+      function_response_event = await functions.handle_function_calls_live(
+          invocation_context, model_response_event, llm_request.tools_dict
+      )
+      yield function_response_event
+
+      transfer_to_agent = function_response_event.actions.transfer_to_agent
+      if transfer_to_agent:
+        agent_to_run = self._get_agent_to_run(
+            invocation_context, transfer_to_agent
+        )
+        async for item in agent_to_run.run_live(invocation_context):
+          yield item
+
+  async def _postprocess_run_processors_async(
+      self, invocation_context: InvocationContext, llm_response: LlmResponse
+  ) -> AsyncGenerator[Event, None]:
+    for processor in self.response_processors:
+      async for event in processor.run_async(invocation_context, llm_response):
+        yield event
+
+  async def _postprocess_handle_function_calls_async(
+      self,
+      invocation_context: InvocationContext,
+      function_call_event: Event,
+      llm_request: LlmRequest,
+  ) -> AsyncGenerator[Event, None]:
+    if function_response_event := await functions.handle_function_calls_async(
+        invocation_context, function_call_event, llm_request.tools_dict
+    ):
+      auth_event = functions.generate_auth_event(
+          invocation_context, function_response_event
+      )
+      if auth_event:
+        yield auth_event
+
+      yield function_response_event
+      transfer_to_agent = function_response_event.actions.transfer_to_agent
+      if transfer_to_agent:
+        agent_to_run = self._get_agent_to_run(
+            invocation_context, transfer_to_agent
+        )
+        async for event in agent_to_run.run_async(invocation_context):
+          yield event
+
+  def _get_agent_to_run(
+      self, invocation_context: InvocationContext, transfer_to_agent
+  ) -> BaseAgent:
+    root_agent = invocation_context.agent.root_agent
+    agent_to_run = root_agent.find_agent(transfer_to_agent)
+    if not agent_to_run:
+      raise ValueError(
+          f'Agent {transfer_to_agent} not found in the agent tree.'
+      )
+    return agent_to_run
+
+  async def _call_llm_async(
+      self,
+      invocation_context: InvocationContext,
+      llm_request: LlmRequest,
+      model_response_event: Event,
+  ) -> AsyncGenerator[LlmResponse, None]:
+    # Runs before_model_callback if it exists.
+    if response := self._handle_before_model_callback(
+        invocation_context, llm_request, model_response_event
+    ):
+      yield response
+      return
+
+    # Calls the LLM.
+    llm = self.__get_llm(invocation_context)
+    with tracer.start_as_current_span('call_llm'):
+      if invocation_context.run_config.support_cfc:
+        invocation_context.live_request_queue = LiveRequestQueue()
+        async for llm_response in self.run_live(invocation_context):
+          # Runs after_model_callback if it exists.
+          if altered_llm_response := self._handle_after_model_callback(
+              invocation_context, llm_response, model_response_event
+          ):
+            llm_response = altered_llm_response
+          # only yield partial response in SSE streaming mode
+          if (
+              invocation_context.run_config.streaming_mode == StreamingMode.SSE
+              or not llm_response.partial
+          ):
+            yield llm_response
+          if llm_response.turn_complete:
+            invocation_context.live_request_queue.close()
+      else:
+        # Check if we can make this llm call or not. If the current call pushes
+        # the counter beyond the max set value, then the execution is stopped
+        # right here, and exception is thrown.
+        invocation_context.increment_llm_call_count()
+        async for llm_response in llm.generate_content_async(
+            llm_request,
+            stream=invocation_context.run_config.streaming_mode
+            == StreamingMode.SSE,
+        ):
+          trace_call_llm(
+              invocation_context,
+              model_response_event.id,
+              llm_request,
+              llm_response,
+          )
+          # Runs after_model_callback if it exists.
+          if altered_llm_response := self._handle_after_model_callback(
+              invocation_context, llm_response, model_response_event
+          ):
+            llm_response = altered_llm_response
+
+          yield llm_response
+
+  def _handle_before_model_callback(
+      self,
+      invocation_context: InvocationContext,
+      llm_request: LlmRequest,
+      model_response_event: Event,
+  ) -> Optional[LlmResponse]:
+    from ...agents.llm_agent import LlmAgent
+
+    agent = invocation_context.agent
+    if not isinstance(agent, LlmAgent):
+      return
+
+    if not agent.before_model_callback:
+      return
+
+    callback_context = CallbackContext(
+        invocation_context, event_actions=model_response_event.actions
+    )
+    return agent.before_model_callback(
+        callback_context=callback_context, llm_request=llm_request
+    )
+
+  def _handle_after_model_callback(
+      self,
+      invocation_context: InvocationContext,
+      llm_response: LlmResponse,
+      model_response_event: Event,
+  ) -> Optional[LlmResponse]:
+    from ...agents.llm_agent import LlmAgent
+
+    agent = invocation_context.agent
+    if not isinstance(agent, LlmAgent):
+      return
+
+    if not agent.after_model_callback:
+      return
+
+    callback_context = CallbackContext(
+        invocation_context, event_actions=model_response_event.actions
+    )
+    return agent.after_model_callback(
+        callback_context=callback_context, llm_response=llm_response
+    )
+
+  def _finalize_model_response_event(
+      self,
+      llm_request: LlmRequest,
+      llm_response: LlmResponse,
+      model_response_event: Event,
+  ) -> Event:
+    model_response_event = Event.model_validate({
+        **model_response_event.model_dump(exclude_none=True),
+        **llm_response.model_dump(exclude_none=True),
+    })
+
+    if model_response_event.content:
+      function_calls = model_response_event.get_function_calls()
+      if function_calls:
+        functions.populate_client_function_call_id(model_response_event)
+        model_response_event.long_running_tool_ids = (
+            functions.get_long_running_function_calls(
+                function_calls, llm_request.tools_dict
+            )
+        )
+
+    return model_response_event
+
+  def __get_llm(self, invocation_context: InvocationContext) -> BaseLlm:
+    from ...agents.llm_agent import LlmAgent
+
+    return cast(LlmAgent, invocation_context.agent).canonical_model

google/adk/flows/llm_flows/basic.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.
+
+"""Handles basic information to build the LLM request."""
+
+from __future__ import annotations
+
+from typing import AsyncGenerator
+from typing import Generator
+
+from google.genai import types
+from typing_extensions import override
+
+from ...agents.invocation_context import InvocationContext
+from ...events.event import Event
+from ...models.llm_request import LlmRequest
+from ._base_llm_processor import BaseLlmRequestProcessor
+
+
+class _BasicLlmRequestProcessor(BaseLlmRequestProcessor):
+
+  @override
+  async def run_async(
+      self, invocation_context: InvocationContext, llm_request: LlmRequest
+  ) -> AsyncGenerator[Event, None]:
+    from ...agents.llm_agent import LlmAgent
+
+    agent = invocation_context.agent
+    if not isinstance(agent, LlmAgent):
+      return
+
+    llm_request.model = (
+        agent.canonical_model
+        if isinstance(agent.canonical_model, str)
+        else agent.canonical_model.model
+    )
+    llm_request.config = (
+        agent.generate_content_config.model_copy(deep=True)
+        if agent.generate_content_config
+        else types.GenerateContentConfig()
+    )
+    if agent.output_schema:
+      llm_request.set_output_schema(agent.output_schema)
+
+    llm_request.live_connect_config.response_modalities = (
+        invocation_context.run_config.response_modalities
+    )
+    llm_request.live_connect_config.speech_config = (
+        invocation_context.run_config.speech_config
+    )
+    llm_request.live_connect_config.output_audio_transcription = (
+        invocation_context.run_config.output_audio_transcription
+    )
+
+    # TODO: handle tool append here, instead of in BaseTool.process_llm_request.
+
+    return
+    yield  # Generator requires yield statement in function body.
+
+
+request_processor = _BasicLlmRequestProcessor()