feat: Refactor AgentLoader into base class and add InMemory impl alongside existing filesystem impl

Poggecci · copybara-github · commit 5bccf8b68c97 · 2025-07-17T17:10:23.000-07:00
PiperOrigin-RevId: 779616089
diff --git a/src/google/adk/cli/cli.py b/src/google/adk/cli/cli.py
@@ -31,7 +31,7 @@
 from ..sessions.in_memory_session_service import InMemorySessionService
 from ..sessions.session import Session
 from .utils import envs
-from .utils.agent_loader import AgentLoader
+from .utils.file_system_agent_loader import FileSystemAgentLoader
 
 
 class InputFile(BaseModel):
@@ -137,7 +137,7 @@ async def run_cli(
   session = await session_service.create_session(
       app_name=agent_folder_name, user_id=user_id
   )
-  root_agent = AgentLoader(agents_dir=agent_parent_dir).load_agent(
+  root_agent = FileSystemAgentLoader(agents_dir=agent_parent_dir).load_agent(
       agent_folder_name
   )
   envs.load_dotenv_for_agent(agent_folder_name, agent_parent_dir)
diff --git a/src/google/adk/cli/cli_eval.py b/src/google/adk/cli/cli_eval.py
@@ -25,7 +25,7 @@
 from typing import Optional
 import uuid
 
-from ..agents import Agent
+from ..agents.llm_agent import Agent
 from ..artifacts.base_artifact_service import BaseArtifactService
 from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
 from ..evaluation.eval_case import EvalCase
diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py
@@ -87,7 +87,7 @@
 from .utils import create_empty_state
 from .utils import envs
 from .utils import evals
-from .utils.agent_loader import AgentLoader
+from .utils.file_system_agent_loader import FileSystemAgentLoader
 
 logger = logging.getLogger("google_adk." + __name__)
 
@@ -98,7 +98,7 @@
 
 class AgentChangeEventHandler(FileSystemEventHandler):
 
-  def __init__(self, agent_loader: AgentLoader):
+  def __init__(self, agent_loader: FileSystemAgentLoader):
     self.agent_loader = agent_loader
 
   def on_modified(self, event):
@@ -382,7 +382,7 @@ def _parse_agent_engine_resource_name(agent_engine_id_or_resource_name):
   credential_service = InMemoryCredentialService()
 
   # initialize Agent Loader
-  agent_loader = AgentLoader(agents_dir)
+  agent_loader = FileSystemAgentLoader(agents_dir)
 
   # Set up a file system watcher to detect changes in the agents directory.
   observer = Observer()
@@ -393,20 +393,7 @@ def _parse_agent_engine_resource_name(agent_engine_id_or_resource_name):
 
   @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:
@@ -511,13 +498,6 @@ async def create_session(
 
     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,
@@ -1042,7 +1022,6 @@ async def _get_runner_async(app_name: str) -> Runner:
       runner = runner_dict.pop(app_name, None)
       await cleanup.close_runners(list([runner]))
 
-    envs.load_dotenv_for_agent(os.path.basename(app_name), agents_dir)
     if app_name in runner_dict:
       return runner_dict[app_name]
     root_agent = agent_loader.load_agent(app_name)
diff --git a/src/google/adk/cli/utils/base_agent_loader.py b/src/google/adk/cli/utils/base_agent_loader.py
@@ -0,0 +1,36 @@
+# 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."""
+    raise NotImplementedError
+
+  @abstractmethod
+  def list_agents(self) -> list[str]:
+    """Lists all agents available in the agent loader in alphabetical order."""
+    raise NotImplementedError
diff --git a/src/google/adk/cli/utils/file_system_agent_loader.py b/src/google/adk/cli/utils/file_system_agent_loader.py
@@ -17,22 +17,26 @@
 import importlib
 import logging
 import os
+from pathlib import Path
 import sys
 from typing import Optional
 
+from google.adk.cli.utils import envs
 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 FileSystemAgentLoader(BaseAgentLoader):
   """Centralized agent loading with proper isolation, caching, and .env loading.
-  Support loading agents from below folder/file structures:
+
+  Supports loading agents from below folder/file structures:
   a)  {agent_name}.agent as a module name:
       agents_dir/{agent_name}/agent.py (with root_agent defined in the module)
   b)  {agent_name} as a module name
@@ -41,7 +45,6 @@ class AgentLoader:
       agents_dir/{agent_name}/__init__.py (with root_agent in the package)
   d)  {agent_name} as a YAML config folder:
       agents_dir/{agent_name}/root_agent.yaml defines the root agent
-
   """
 
   def __init__(self, agents_dir: str):
@@ -188,6 +191,7 @@ def _perform_load(self, agent_name: str) -> BaseAgent:
         " 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 @@ def load_agent(self, agent_name: str) -> BaseAgent:
     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/in_memory_agent_loader.py b/src/google/adk/cli/utils/in_memory_agent_loader.py
@@ -0,0 +1,45 @@
+# 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 logging
+from typing import Mapping
+
+from typing_extensions import override
+
+from ...agents.base_agent import BaseAgent
+from .base_agent_loader import BaseAgentLoader
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+class InMemoryAgentLoader(BaseAgentLoader):
+  """Simple agent loader that loads agent from in-memory dict."""
+
+  def __init__(self, agents: Mapping[str, BaseAgent]):
+    self._agent_cache: dict[str, BaseAgent] = {}
+    for app_name, agent in agents.items():
+      self._agent_cache[app_name] = agent
+    logger.debug("Added %d agents into memory.", len(self._agent_cache))
+
+  @override
+  def load_agent(self, agent_name: str) -> BaseAgent:
+    """Load an agent from in-memory dict."""
+    return self._agent_cache[agent_name]
+
+  @override
+  def list_agents(self) -> list[str]:
+    """List all agents available in the in-memory dict (sorted alphabetically)."""
+    return list(self._agent_cache.keys())
diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py
@@ -188,6 +188,9 @@ def __init__(self, agents_dir: str):
     def load_agent(self, app_name):
       return root_agent
 
+    def list_agents(self):
+      return ["test_app"]
+
   return MockAgentLoader(".")
 
 
@@ -445,7 +448,7 @@ def test_app(
           return_value=mock_memory_service,
       ),
       patch(
-          "google.adk.cli.fast_api.AgentLoader",
+          "google.adk.cli.fast_api.FileSystemAgentLoader",
           return_value=mock_agent_loader,
       ),
       patch(
@@ -596,7 +599,7 @@ def test_app_with_a2a(
           return_value=mock_memory_service,
       ),
       patch(
-          "google.adk.cli.fast_api.AgentLoader",
+          "google.adk.cli.fast_api.FileSystemAgentLoader",
           return_value=mock_agent_loader,
       ),
       patch(
diff --git a/tests/unittests/cli/utils/test_file_system_agent_loader.py b/tests/unittests/cli/utils/test_file_system_agent_loader.py
diff --git a/tests/unittests/cli/utils/test_in_memory_agent_loader.py b/tests/unittests/cli/utils/test_in_memory_agent_loader.py
