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

google/adk/cli/cli.py

@@ -39,12 +39,12 @@ class InputFile(BaseModel):
 
 async def run_input_file(
     app_name: str,
+    user_id: str,
     root_agent: LlmAgent,
     artifact_service: BaseArtifactService,
-    session: Session,
     session_service: BaseSessionService,
     input_path: str,
-) -> None:
+) -> Session:
   runner = Runner(
       app_name=app_name,
       agent=root_agent,
@@ -55,9 +55,11 @@ async def run_input_file(
     input_file = InputFile.model_validate_json(f.read())
   input_file.state['_time'] = datetime.now()
 
-  session.state = input_file.state
+  session = session_service.create_session(
+      app_name=app_name, user_id=user_id, state=input_file.state
+  )
   for query in input_file.queries:
-    click.echo(f'user: {query}')
+    click.echo(f'[user]: {query}')
     content = types.Content(role='user', parts=[types.Part(text=query)])
     async for event in runner.run_async(
         user_id=session.user_id, session_id=session.id, new_message=content
@@ -65,23 +67,23 @@ async def run_input_file(
       if event.content and event.content.parts:
         if text := ''.join(part.text or '' for part in event.content.parts):
           click.echo(f'[{event.author}]: {text}')
+  return session
 
 
 async def run_interactively(
-    app_name: str,
     root_agent: LlmAgent,
     artifact_service: BaseArtifactService,
     session: Session,
     session_service: BaseSessionService,
 ) -> None:
   runner = Runner(
-      app_name=app_name,
+      app_name=session.app_name,
       agent=root_agent,
       artifact_service=artifact_service,
       session_service=session_service,
   )
   while True:
-    query = input('user: ')
+    query = input('[user]: ')
     if not query or not query.strip():
       continue
     if query == 'exit':
@@ -100,7 +102,8 @@ async def run_cli(
     *,
     agent_parent_dir: str,
     agent_folder_name: str,
-    json_file_path: Optional[str] = None,
+    input_file: Optional[str] = None,
+    saved_session_file: Optional[str] = None,
     save_session: bool,
 ) -> None:
   """Runs an interactive CLI for a certain agent.
@@ -109,8 +112,11 @@ async def run_cli(
     agent_parent_dir: str, the absolute path of the parent folder of the agent
       folder.
     agent_folder_name: str, the name of the agent folder.
-    json_file_path: Optional[str], the absolute path to the json file, either
-      *.input.json or *.session.json.
+    input_file: Optional[str], the absolute path to the json file that contains
+      the initial session state and user queries, exclusive with
+      saved_session_file.
+    saved_session_file: Optional[str], the absolute path to the json file that
+      contains a previously saved session, exclusive with input_file.
     save_session: bool, whether to save the session on exit.
   """
   if agent_parent_dir not in sys.path:
@@ -118,46 +124,50 @@ async def run_cli(
 
   artifact_service = InMemoryArtifactService()
   session_service = InMemorySessionService()
-  session = session_service.create_session(
-      app_name=agent_folder_name, user_id='test_user'
-  )
 
   agent_module_path = os.path.join(agent_parent_dir, agent_folder_name)
   agent_module = importlib.import_module(agent_folder_name)
+  user_id = 'test_user'
+  session = session_service.create_session(
+      app_name=agent_folder_name, user_id=user_id
+  )
   root_agent = agent_module.agent.root_agent
   envs.load_dotenv_for_agent(agent_folder_name, agent_parent_dir)
-  if json_file_path:
-    if json_file_path.endswith('.input.json'):
-      await run_input_file(
-          app_name=agent_folder_name,
-          root_agent=root_agent,
-          artifact_service=artifact_service,
-          session=session,
-          session_service=session_service,
-          input_path=json_file_path,
-      )
-    elif json_file_path.endswith('.session.json'):
-      with open(json_file_path, 'r') as f:
-        session = Session.model_validate_json(f.read())
-      for content in session.get_contents():
-        if content.role == 'user':
-          print('user: ', content.parts[0].text)
+  if input_file:
+    session = await run_input_file(
+        app_name=agent_folder_name,
+        user_id=user_id,
+        root_agent=root_agent,
+        artifact_service=artifact_service,
+        session_service=session_service,
+        input_path=input_file,
+    )
+  elif saved_session_file:
+
+    loaded_session = None
+    with open(saved_session_file, 'r') as f:
+      loaded_session = Session.model_validate_json(f.read())
+
+    if loaded_session:
+      for event in loaded_session.events:
+        session_service.append_event(session, event)
+        content = event.content
+        if not content or not content.parts or not content.parts[0].text:
+          continue
+        if event.author == 'user':
+          click.echo(f'[user]: {content.parts[0].text}')
         else:
-          print(content.parts[0].text)
-      await run_interactively(
-          agent_folder_name,
-          root_agent,
-          artifact_service,
-          session,
-          session_service,
-      )
-    else:
-      print(f'Unsupported file type: {json_file_path}')
-      exit(1)
+          click.echo(f'[{event.author}]: {content.parts[0].text}')
+
+    await run_interactively(
+        root_agent,
+        artifact_service,
+        session,
+        session_service,
+    )
   else:
-    print(f'Running agent {root_agent.name}, type exit to exit.')
+    click.echo(f'Running agent {root_agent.name}, type exit to exit.')
     await run_interactively(
-        agent_folder_name,
         root_agent,
         artifact_service,
         session,
@@ -165,11 +175,8 @@ async def run_cli(
     )
 
   if save_session:
-    if json_file_path:
-      session_path = json_file_path.replace('.input.json', '.session.json')
-    else:
-      session_id = input('Session ID to save: ')
-      session_path = f'{agent_module_path}/{session_id}.session.json'
+    session_id = input('Session ID to save: ')
+    session_path = f'{agent_module_path}/{session_id}.session.json'
 
     # Fetch the session again to get all the details.
     session = session_service.get_session(

google/adk/cli/cli_deploy.py

@@ -54,7 +54,7 @@ COPY "agents/{app_name}/" "/app/agents/{app_name}/"
 
 EXPOSE {port}
 
-CMD adk {command} --port={port} {trace_to_cloud_option} "/app/agents"
+CMD adk {command} --port={port} {session_db_option} {trace_to_cloud_option} "/app/agents"
 """
 
 
@@ -85,6 +85,7 @@ def to_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
+    session_db_url: str,
 ):
   """Deploys an agent to Google Cloud Run.
 
@@ -112,6 +113,7 @@ def to_cloud_run(
     trace_to_cloud: Whether to enable Cloud Trace.
     with_ui: Whether to deploy with UI.
     verbosity: The verbosity level of the CLI.
+    session_db_url: The database URL to connect the session.
   """
   app_name = app_name or os.path.basename(agent_folder)
 
@@ -144,6 +146,9 @@ def to_cloud_run(
         port=port,
         command='web' if with_ui else 'api_server',
         install_agent_deps=install_agent_deps,
+        session_db_option=f'--session_db_url={session_db_url}'
+        if session_db_url
+        else '',
         trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
     )
     dockerfile_path = os.path.join(temp_folder, 'Dockerfile')

google/adk/cli/cli_eval.py

@@ -256,7 +256,7 @@ def run_evals(
         )
 
         if final_eval_status == EvalStatus.PASSED:
-          result = "✅ Passsed"
+          result = "✅ Passed"
         else:
           result = "❌ Failed"
 

google/adk/cli/cli_tools_click.py

@@ -96,6 +96,23 @@ def cli_create_cmd(
   )
 
 
+def validate_exclusive(ctx, param, value):
+  # Store the validated parameters in the context
+  if not hasattr(ctx, "exclusive_opts"):
+    ctx.exclusive_opts = {}
+
+  # If this option has a value and we've already seen another exclusive option
+  if value is not None and any(ctx.exclusive_opts.values()):
+    exclusive_opt = next(key for key, val in ctx.exclusive_opts.items() if val)
+    raise click.UsageError(
+        f"Options '{param.name}' and '{exclusive_opt}' cannot be set together."
+    )
+
+  # Record this option's value
+  ctx.exclusive_opts[param.name] = value is not None
+  return value
+
+
 @main.command("run")
 @click.option(
     "--save_session",
@@ -105,13 +122,43 @@ def cli_create_cmd(
     default=False,
     help="Optional. Whether to save the session to a json file on exit.",
 )
+@click.option(
+    "--replay",
+    type=click.Path(
+        exists=True, dir_okay=False, file_okay=True, resolve_path=True
+    ),
+    help=(
+        "The json file that contains the initial state of the session and user"
+        " queries. A new session will be created using this state. And user"
+        " queries are run againt the newly created session. Users cannot"
+        " continue to interact with the agent."
+    ),
+    callback=validate_exclusive,
+)
+@click.option(
+    "--resume",
+    type=click.Path(
+        exists=True, dir_okay=False, file_okay=True, resolve_path=True
+    ),
+    help=(
+        "The json file that contains a previously saved session (by"
+        "--save_session option). The previous session will be re-displayed. And"
+        " user can continue to interact with the agent."
+    ),
+    callback=validate_exclusive,
+)
 @click.argument(
     "agent",
     type=click.Path(
         exists=True, dir_okay=True, file_okay=False, resolve_path=True
     ),
 )
-def cli_run(agent: str, save_session: bool):
+def cli_run(
+    agent: str,
+    save_session: bool,
+    replay: Optional[str],
+    resume: Optional[str],
+):
   """Runs an interactive CLI for a certain agent.
 
   AGENT: The path to the agent source code folder.
@@ -129,6 +176,8 @@ def cli_run(agent: str, save_session: bool):
       run_cli(
           agent_parent_dir=agent_parent_folder,
           agent_folder_name=agent_folder_name,
+          input_file=replay,
+          saved_session_file=resume,
           save_session=save_session,
       )
   )
@@ -245,12 +294,13 @@ def cli_eval(
 @click.option(
     "--session_db_url",
     help=(
-        "Optional. The database URL to store the session.\n\n  - Use"
-        " 'agentengine://<agent_engine_resource_id>' to connect to Vertex"
-        " managed session service.\n\n  - Use 'sqlite://<path_to_sqlite_file>'"
-        " to connect to a SQLite DB.\n\n  - See"
-        " https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls"
-        " for more details on supported DB URLs."
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
     ),
 )
 @click.option(
@@ -366,12 +416,13 @@ def cli_web(
 @click.option(
     "--session_db_url",
     help=(
-        "Optional. The database URL to store the session.\n\n  - Use"
-        " 'agentengine://<agent_engine_resource_id>' to connect to Vertex"
-        " managed session service.\n\n  - Use 'sqlite://<path_to_sqlite_file>'"
-        " to connect to a SQLite DB.\n\n  - See"
-        " https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls"
-        " for more details on supported DB URLs."
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
     ),
 )
 @click.option(
@@ -541,6 +592,18 @@ def cli_api_server(
     default="WARNING",
     help="Optional. Override the default verbosity level.",
 )
+@click.option(
+    "--session_db_url",
+    help=(
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
+    ),
+)
 @click.argument(
     "agent",
     type=click.Path(
@@ -558,6 +621,7 @@ def cli_deploy_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
+    session_db_url: str,
 ):
   """Deploys an agent to Cloud Run.
 
@@ -579,6 +643,7 @@ def cli_deploy_cloud_run(
         trace_to_cloud=trace_to_cloud,
         with_ui=with_ui,
         verbosity=verbosity,
+        session_db_url=session_db_url,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

google/adk/cli/fast_api.py

@@ -756,6 +756,12 @@ def get_fast_api_app(
     except Exception as e:
       logger.exception("Error during live websocket communication: %s", e)
       traceback.print_exc()
+      WEBSOCKET_INTERNAL_ERROR_CODE = 1011
+      WEBSOCKET_MAX_BYTES_FOR_REASON = 123
+      await websocket.close(
+          code=WEBSOCKET_INTERNAL_ERROR_CODE,
+          reason=str(e)[:WEBSOCKET_MAX_BYTES_FOR_REASON],
+      )
     finally:
       for task in pending:
         task.cancel()

google/adk/evaluation/agent_evaluator.py

@@ -55,7 +55,7 @@ def load_json(file_path: str) -> Union[Dict, List]:
 
 
 class AgentEvaluator:
-  """An evaluator for Agents, mainly intented for helping with test cases."""
+  """An evaluator for Agents, mainly intended for helping with test cases."""
 
   @staticmethod
   def find_config_for_test_file(test_file: str):
@@ -91,7 +91,7 @@ class AgentEvaluator:
         look for 'root_agent' in the loaded module.
       eval_dataset: The eval data set. This can be either a string representing
         full path to the file containing eval dataset, or a directory that is
-        recusively explored for all files that have a `.test.json` suffix.
+        recursively explored for all files that have a `.test.json` suffix.
       num_runs: Number of times all entries in the eval dataset should be
         assessed.
       agent_name: The name of the agent.

google/adk/evaluation/response_evaluator.py

@@ -35,7 +35,7 @@ class ResponseEvaluator:
     Args:
       raw_eval_dataset: The dataset that will be evaluated.
       evaluation_criteria: The evaluation criteria to be used. This method
-        support two criterias, `response_evaluation_score` and
+        support two criteria, `response_evaluation_score` and
         `response_match_score`.
       print_detailed_results: Prints detailed results on the console. This is
         usually helpful during debugging.
@@ -56,7 +56,7 @@ class ResponseEvaluator:
         Value range: [0, 5], where 0 means that the agent's response is not
         coherent, while 5 means it is . High values are good.
     A note on raw_eval_dataset:
-      The dataset should be a list session, where each sesssion is represented
+      The dataset should be a list session, where each session is represented
       as a list of interaction that need evaluation. Each evaluation is
       represented as a dictionary that is expected to have values for the
       following keys:

google/adk/evaluation/trajectory_evaluator.py

@@ -31,10 +31,9 @@ class TrajectoryEvaluator:
   ):
     r"""Returns the mean tool use accuracy of the eval dataset.
 
-    Tool use accuracy is calculated by comparing the expected and actuall tool
-    use trajectories. An exact match scores a 1, 0 otherwise. The final number
-    is an
-    average of these individual scores.
+    Tool use accuracy is calculated by comparing the expected and the actual
+    tool use trajectories. An exact match scores a 1, 0 otherwise. The final
+    number is an average of these individual scores.
 
     Value range: [0, 1], where 0 is means none of the too use entries aligned,
     and 1 would mean all of them aligned. Higher value is good.
@@ -45,7 +44,7 @@ class TrajectoryEvaluator:
         usually helpful during debugging.
 
     A note on eval_dataset:
-      The dataset should be a list session, where each sesssion is represented
+      The dataset should be a list session, where each session is represented
       as a list of interaction that need evaluation. Each evaluation is
       represented as a dictionary that is expected to have values for the
       following keys:

google/adk/events/event_actions.py

@@ -48,8 +48,13 @@ class EventActions(BaseModel):
   """The agent is escalating to a higher level agent."""
 
   requested_auth_configs: dict[str, AuthConfig] = Field(default_factory=dict)
-  """Will only be set by a tool response indicating tool request euc.
-  dict key is the function call id since one function call response (from model)
-  could correspond to multiple function calls.
-  dict value is the required auth config.
+  """Authentication configurations requested by tool responses.
+
+  This field will only be set by a tool response event indicating tool request
+  auth credential.
+  - Keys: The function call id. Since one function response event could contain
+  multiple function responses that correspond to multiple function calls. Each
+  function call could request different auth configs. This id is used to
+  identify the function call.
+  - Values: The requested auth config.
   """

google/adk/flows/llm_flows/agent_transfer.py

@@ -94,7 +94,7 @@ can answer it.
 
 If another agent is better for answering the question according to its
 description, call `{_TRANSFER_TO_AGENT_FUNCTION_NAME}` function to transfer the
-question to that agent. When transfering, do not generate any text other than
+question to that agent. When transferring, do not generate any text other than
 the function call.
 """
 

google/adk/flows/llm_flows/base_llm_flow.py

@@ -115,7 +115,7 @@ class BaseLlmFlow(ABC):
           yield event
           # send back the function response
           if event.get_function_responses():
-            logger.debug('Sending back last function resonse event: %s', event)
+            logger.debug('Sending back last function response event: %s', event)
             invocation_context.live_request_queue.send_content(event.content)
           if (
               event.content

google/adk/flows/llm_flows/contents.py

@@ -15,9 +15,7 @@
 from __future__ import annotations
 
 import copy
-from typing import AsyncGenerator
-from typing import Generator
-from typing import Optional
+from typing import AsyncGenerator, Generator, Optional
 
 from google.genai import types
 from typing_extensions import override
@@ -111,7 +109,7 @@ def _rearrange_events_for_latest_function_response(
   """Rearrange the events for the latest function_response.
 
   If the latest function_response is for an async function_call, all events
-  bewteen the initial function_call and the latest function_response will be
+  between the initial function_call and the latest function_response will be
   removed.
 
   Args:
@@ -202,8 +200,14 @@ def _get_contents(
   # Parse the events, leaving the contents and the function calls and
   # responses from the current agent.
   for event in events:
-    if not event.content or not event.content.role:
-      # Skip events without content, or generated neither by user nor by model.
+    if (
+        not event.content
+        or not event.content.role
+        or not event.content.parts
+        or event.content.parts[0].text == ''
+    ):
+      # Skip events without content, or generated neither by user nor by model
+      # or has empty text.
       # E.g. events purely for mutating session states.
       continue
     if not _is_event_belongs_to_branch(current_branch, event):

google/adk/flows/llm_flows/functions.py

@@ -151,28 +151,33 @@ async def handle_function_calls_async(
     # do not use "args" as the variable name, because it is a reserved keyword
     # in python debugger.
     function_args = function_call.args or {}
-    function_response = None
-    # Calls the tool if before_tool_callback does not exist or returns None.
+    function_response: Optional[dict] = None
+
+    # before_tool_callback (sync or async)
     if agent.before_tool_callback:
       function_response = agent.before_tool_callback(
           tool=tool, args=function_args, tool_context=tool_context
       )
+      if inspect.isawaitable(function_response):
+        function_response = await function_response
 
     if not function_response:
       function_response = await __call_tool_async(
           tool, args=function_args, tool_context=tool_context
       )
 
-    # Calls after_tool_callback if it exists.
+    # after_tool_callback (sync or async)
     if agent.after_tool_callback:
-      new_response = agent.after_tool_callback(
+      altered_function_response = agent.after_tool_callback(
           tool=tool,
           args=function_args,
           tool_context=tool_context,
           tool_response=function_response,
       )
-      if new_response:
-        function_response = new_response
+      if inspect.isawaitable(altered_function_response):
+        altered_function_response = await altered_function_response
+      if altered_function_response is not None:
+        function_response = altered_function_response
 
     if tool.is_long_running:
       # Allow long running function to return None to not provide function response.
@@ -223,11 +228,17 @@ async def handle_function_calls_live(
     # in python debugger.
     function_args = function_call.args or {}
     function_response = None
-    # Calls the tool if before_tool_callback does not exist or returns None.
+    # # Calls the tool if before_tool_callback does not exist or returns None.
+    # if agent.before_tool_callback:
+    #   function_response = agent.before_tool_callback(
+    #       tool, function_args, tool_context
+    #   )
     if agent.before_tool_callback:
       function_response = agent.before_tool_callback(
-          tool, function_args, tool_context
+          tool=tool, args=function_args, tool_context=tool_context
       )
+      if inspect.isawaitable(function_response):
+        function_response = await function_response
 
     if not function_response:
       function_response = await _process_function_live_helper(
@@ -235,15 +246,26 @@ async def handle_function_calls_live(
       )
 
     # Calls after_tool_callback if it exists.
+    # if agent.after_tool_callback:
+    #   new_response = agent.after_tool_callback(
+    #       tool,
+    #       function_args,
+    #       tool_context,
+    #       function_response,
+    #   )
+    #   if new_response:
+    #     function_response = new_response
     if agent.after_tool_callback:
-      new_response = agent.after_tool_callback(
-          tool,
-          function_args,
-          tool_context,
-          function_response,
+      altered_function_response = agent.after_tool_callback(
+          tool=tool,
+          args=function_args,
+          tool_context=tool_context,
+          tool_response=function_response,
       )
-      if new_response:
-        function_response = new_response
+      if inspect.isawaitable(altered_function_response):
+        altered_function_response = await altered_function_response
+      if altered_function_response is not None:
+        function_response = altered_function_response
 
     if tool.is_long_running:
       # Allow async function to return None to not provide function response.
@@ -310,9 +332,7 @@ async def _process_function_live_helper(
       function_response = {
           'status': f'No active streaming function named {function_name} found'
       }
-  elif inspect.isasyncgenfunction(tool.func):
-    print('is async')
-
+  elif hasattr(tool, "func") and inspect.isasyncgenfunction(tool.func):
     # for streaming tool use case
     # we require the function to be a async generator function
     async def run_tool_and_update_queue(tool, function_args, tool_context):

google/adk/flows/llm_flows/instructions.py

@@ -52,7 +52,7 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
     # Appends global instructions if set.
     if (
         isinstance(root_agent, LlmAgent) and root_agent.global_instruction
-    ):  # not emtpy str
+    ):  # not empty str
       raw_si = root_agent.canonical_global_instruction(
           ReadonlyContext(invocation_context)
       )
@@ -60,7 +60,7 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
       llm_request.append_instructions([si])
 
     # Appends agent instructions if set.
-    if agent.instruction:  # not emtpy str
+    if agent.instruction:  # not empty str
       raw_si = agent.canonical_instruction(ReadonlyContext(invocation_context))
       si = _populate_values(raw_si, invocation_context)
       llm_request.append_instructions([si])

google/adk/models/gemini_llm_connection.py

@@ -152,7 +152,7 @@ class GeminiLlmConnection(BaseLlmConnection):
         ):
           # TODO: Right now, we just support output_transcription without
           # changing interface and data protocol. Later, we can consider to
-          # support output_transcription as a separete field in LlmResponse.
+          # support output_transcription as a separate field in LlmResponse.
 
           # Transcription is always considered as partial event
           # We rely on other control signals to determine when to yield the
@@ -179,7 +179,7 @@ class GeminiLlmConnection(BaseLlmConnection):
         # in case of empty content or parts, we sill surface it
         # in case it's an interrupted message, we merge the previous partial
         # text. Other we don't merge. because content can be none when model
-        # safty threshold is triggered
+        # safety threshold is triggered
         if message.server_content.interrupted and text:
           yield self.__build_full_text_response(text)
           text = ''