From bda3df24802d0456711a5cd05544aea54a13398d Mon Sep 17 00:00:00 2001 From: Alejandro Cruzado-Ruiz Date: Tue, 22 Jul 2025 14:57:56 -0700 Subject: [PATCH] feat: Refactor AgentLoader into base class and add InMemory impl alongside existing filesystem impl PiperOrigin-RevId: 786008518 --- src/google/adk/cli/cli_eval.py | 2 +- src/google/adk/cli/fast_api.py | 23 +------------ src/google/adk/cli/utils/agent_loader.py | 20 ++++++++++- src/google/adk/cli/utils/base_agent_loader.py | 34 +++++++++++++++++++ tests/unittests/cli/test_fast_api.py | 3 ++ 5 files changed, 58 insertions(+), 24 deletions(-) create mode 100644 src/google/adk/cli/utils/base_agent_loader.py diff --git a/src/google/adk/cli/cli_eval.py b/src/google/adk/cli/cli_eval.py index 42cc20b0..2f1d090c 100644 --- a/src/google/adk/cli/cli_eval.py +++ b/src/google/adk/cli/cli_eval.py @@ -27,7 +27,7 @@ import uuid from typing_extensions import deprecated -from ..agents import Agent +from ..agents.llm_agent import Agent from ..artifacts.base_artifact_service import BaseArtifactService from ..evaluation.base_eval_service import BaseEvalService from ..evaluation.base_eval_service import EvaluateConfig diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 05ed8fc4..09cd5d2e 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -27,7 +27,6 @@ import typing from typing import Any from typing import List from typing import Literal -from typing import Mapping from typing import Optional import click @@ -407,20 +406,7 @@ def get_fast_api_app( @app.get("/list-apps") def list_apps() -> list[str]: - base_path = Path.cwd() / agents_dir - if not base_path.exists(): - raise HTTPException(status_code=404, detail="Path not found") - if not base_path.is_dir(): - raise HTTPException(status_code=400, detail="Not a directory") - agent_names = [ - x - for x in os.listdir(base_path) - if os.path.isdir(os.path.join(base_path, x)) - and not x.startswith(".") - and x != "__pycache__" - ] - agent_names.sort() - return agent_names + return agent_loader.list_agents() @app.get("/debug/trace/{event_id}") def get_trace_dict(event_id: str) -> Any: @@ -525,13 +511,6 @@ def get_fast_api_app( return session - def _get_eval_set_file_path(app_name, agents_dir, eval_set_id) -> str: - return os.path.join( - agents_dir, - app_name, - eval_set_id + _EVAL_SET_FILE_EXTENSION, - ) - @app.post( "/apps/{app_name}/eval_sets/{eval_set_id}", response_model_exclude_none=True, diff --git a/src/google/adk/cli/utils/agent_loader.py b/src/google/adk/cli/utils/agent_loader.py index 1e206846..5b892487 100644 --- a/src/google/adk/cli/utils/agent_loader.py +++ b/src/google/adk/cli/utils/agent_loader.py @@ -17,20 +17,23 @@ from __future__ import annotations import importlib import logging import os +from pathlib import Path import sys from typing import Optional from pydantic import ValidationError +from typing_extensions import override from . import envs from ...agents import config_agent_utils from ...agents.base_agent import BaseAgent from ...utils.feature_decorator import working_in_progress +from .base_agent_loader import BaseAgentLoader logger = logging.getLogger("google_adk." + __name__) -class AgentLoader: +class AgentLoader(BaseAgentLoader): """Centralized agent loading with proper isolation, caching, and .env loading. Support loading agents from below folder/file structures: a) {agent_name}.agent as a module name: @@ -188,6 +191,7 @@ class AgentLoader: " exposed." ) + @override def load_agent(self, agent_name: str) -> BaseAgent: """Load an agent module (with caching & .env) and return its root_agent.""" if agent_name in self._agent_cache: @@ -199,6 +203,20 @@ class AgentLoader: self._agent_cache[agent_name] = agent return agent + @override + def list_agents(self) -> list[str]: + """Lists all agents available in the agent loader (sorted alphabetically).""" + base_path = Path.cwd() / self.agents_dir + agent_names = [ + x + for x in os.listdir(base_path) + if os.path.isdir(os.path.join(base_path, x)) + and not x.startswith(".") + and x != "__pycache__" + ] + agent_names.sort() + return agent_names + def remove_agent_from_cache(self, agent_name: str): # Clear module cache for the agent and its submodules keys_to_delete = [ diff --git a/src/google/adk/cli/utils/base_agent_loader.py b/src/google/adk/cli/utils/base_agent_loader.py new file mode 100644 index 00000000..015d450b --- /dev/null +++ b/src/google/adk/cli/utils/base_agent_loader.py @@ -0,0 +1,34 @@ +# 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. + +"""Base class for agent loaders.""" + +from __future__ import annotations + +from abc import ABC +from abc import abstractmethod + +from ...agents.base_agent import BaseAgent + + +class BaseAgentLoader(ABC): + """Abstract base class for agent loaders.""" + + @abstractmethod + def load_agent(self, agent_name: str) -> BaseAgent: + """Loads an instance of an agent with the given name.""" + + @abstractmethod + def list_agents(self) -> list[str]: + """Lists all agents available in the agent loader in alphabetical order.""" diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 8475b7e0..0c64cd0a 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -189,6 +189,9 @@ def mock_agent_loader(): def load_agent(self, app_name): return root_agent + def list_agents(self): + return ["test_app"] + return MockAgentLoader(".")