google-adk 0.5.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

google/adk/cli/fast_api.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+from __future__ import annotations
+
 import asyncio
 from contextlib import asynccontextmanager
-import importlib
-import inspect
-import json
 import logging
 import os
 from pathlib import Path
-import re
-import sys
+import time
 import traceback
 import typing
 from typing import Any
@@ -30,7 +29,6 @@ from typing import Literal
 from typing import Optional
 
 import click
-from click import Tuple
 from fastapi import FastAPI
 from fastapi import HTTPException
 from fastapi import Query
@@ -48,16 +46,25 @@ from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
 from opentelemetry.sdk.trace import export
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace import TracerProvider
-from pydantic import BaseModel
+from pydantic import Field
 from pydantic import ValidationError
 from starlette.types import Lifespan
+from typing_extensions import override
 
 from ..agents import RunConfig
 from ..agents.live_request_queue import LiveRequest
 from ..agents.live_request_queue import LiveRequestQueue
 from ..agents.llm_agent import Agent
 from ..agents.run_config import StreamingMode
-from ..artifacts import InMemoryArtifactService
+from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
+from ..evaluation.eval_case import EvalCase
+from ..evaluation.eval_case import SessionInput
+from ..evaluation.eval_metrics import EvalMetric
+from ..evaluation.eval_metrics import EvalMetricResult
+from ..evaluation.eval_metrics import EvalMetricResultPerInvocation
+from ..evaluation.eval_result import EvalSetResult
+from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
+from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
 from ..events.event import Event
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
@@ -66,14 +73,15 @@ from ..sessions.in_memory_session_service import InMemorySessionService
 from ..sessions.session import Session
 from ..sessions.vertex_ai_session_service import VertexAiSessionService
 from .cli_eval import EVAL_SESSION_ID_PREFIX
-from .cli_eval import EvalMetric
-from .cli_eval import EvalMetricResult
 from .cli_eval import EvalStatus
+from .utils import cleanup
+from .utils import common
 from .utils import create_empty_state
 from .utils import envs
 from .utils import evals
+from .utils.agent_loader import AgentLoader
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 _EVAL_SET_FILE_EXTENSION = ".evalset.json"
 
@@ -103,7 +111,45 @@ class ApiServerSpanExporter(export.SpanExporter):
     return True
 
 
-class AgentRunRequest(BaseModel):
+class InMemoryExporter(export.SpanExporter):
+
+  def __init__(self, trace_dict):
+    super().__init__()
+    self._spans = []
+    self.trace_dict = trace_dict
+
+  @override
+  def export(
+      self, spans: typing.Sequence[ReadableSpan]
+  ) -> export.SpanExportResult:
+    for span in spans:
+      trace_id = span.context.trace_id
+      if span.name == "call_llm":
+        attributes = dict(span.attributes)
+        session_id = attributes.get("gcp.vertex.agent.session_id", None)
+        if session_id:
+          if session_id not in self.trace_dict:
+            self.trace_dict[session_id] = [trace_id]
+          else:
+            self.trace_dict[session_id] += [trace_id]
+    self._spans.extend(spans)
+    return export.SpanExportResult.SUCCESS
+
+  @override
+  def force_flush(self, timeout_millis: int = 30000) -> bool:
+    return True
+
+  def get_finished_spans(self, session_id: str):
+    trace_ids = self.trace_dict.get(session_id, None)
+    if trace_ids is None or not trace_ids:
+      return []
+    return [x for x in self._spans if x.context.trace_id in trace_ids]
+
+  def clear(self):
+    self._spans.clear()
+
+
+class AgentRunRequest(common.BaseModel):
   app_name: str
   user_id: str
   session_id: str
@@ -111,28 +157,41 @@ class AgentRunRequest(BaseModel):
   streaming: bool = False
 
 
-class AddSessionToEvalSetRequest(BaseModel):
+class AddSessionToEvalSetRequest(common.BaseModel):
   eval_id: str
   session_id: str
   user_id: str
 
 
-class RunEvalRequest(BaseModel):
+class RunEvalRequest(common.BaseModel):
   eval_ids: list[str]  # if empty, then all evals in the eval set are run.
   eval_metrics: list[EvalMetric]
 
 
-class RunEvalResult(BaseModel):
+class RunEvalResult(common.BaseModel):
+  eval_set_file: str
   eval_set_id: str
   eval_id: str
   final_eval_status: EvalStatus
-  eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]]
+  eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]] = Field(
+      deprecated=True,
+      description=(
+          "This field is deprecated, use overall_eval_metric_results instead."
+      ),
+  )
+  overall_eval_metric_results: list[EvalMetricResult]
+  eval_metric_result_per_invocation: list[EvalMetricResultPerInvocation]
+  user_id: str
   session_id: str
 
 
+class GetEventGraphResult(common.BaseModel):
+  dot_src: str
+
+
 def get_fast_api_app(
     *,
-    agent_dir: str,
+    agents_dir: str,
     session_db_url: str = "",
     allow_origins: Optional[list[str]] = None,
     web: bool,
@@ -141,40 +200,42 @@ def get_fast_api_app(
 ) -> FastAPI:
   # InMemory tracing dict.
   trace_dict: dict[str, Any] = {}
+  session_trace_dict: dict[str, Any] = {}
 
   # Set up tracing in the FastAPI server.
   provider = TracerProvider()
   provider.add_span_processor(
       export.SimpleSpanProcessor(ApiServerSpanExporter(trace_dict))
   )
+  memory_exporter = InMemoryExporter(session_trace_dict)
+  provider.add_span_processor(export.SimpleSpanProcessor(memory_exporter))
   if trace_to_cloud:
-    envs.load_dotenv_for_agent("", agent_dir)
+    envs.load_dotenv_for_agent("", agents_dir)
     if project_id := os.environ.get("GOOGLE_CLOUD_PROJECT", None):
       processor = export.BatchSpanProcessor(
           CloudTraceSpanExporter(project_id=project_id)
       )
       provider.add_span_processor(processor)
     else:
-      logging.warning(
+      logger.warning(
           "GOOGLE_CLOUD_PROJECT environment variable is not set. Tracing will"
           " not be enabled."
       )
 
   trace.set_tracer_provider(provider)
 
-  exit_stacks = []
-
   @asynccontextmanager
   async def internal_lifespan(app: FastAPI):
-    if lifespan:
-      async with lifespan(app) as lifespan_context:
-        yield
 
-        if exit_stacks:
-          for stack in exit_stacks:
-            await stack.aclose()
-    else:
-      yield
+    try:
+      if lifespan:
+        async with lifespan(app) as lifespan_context:
+          yield lifespan_context
+      else:
+        yield
+    finally:
+      # Create tasks for all runner closures to run concurrently
+      await cleanup.close_runners(list(runner_dict.values()))
 
   # Run the FastAPI server.
   app = FastAPI(lifespan=internal_lifespan)
@@ -188,16 +249,15 @@ def get_fast_api_app(
         allow_headers=["*"],
     )
 
-  if agent_dir not in sys.path:
-    sys.path.append(agent_dir)
-
   runner_dict = {}
-  root_agent_dict = {}
 
   # Build the Artifact service
   artifact_service = InMemoryArtifactService()
   memory_service = InMemoryMemoryService()
 
+  eval_sets_manager = LocalEvalSetsManager(agents_dir=agents_dir)
+  eval_set_results_manager = LocalEvalSetResultsManager(agents_dir=agents_dir)
+
   # Build the Session service
   agent_engine_id = ""
   if session_db_url:
@@ -206,7 +266,7 @@ def get_fast_api_app(
       agent_engine_id = session_db_url.split("://")[1]
       if not agent_engine_id:
         raise click.ClickException("Agent engine id can not be empty.")
-      envs.load_dotenv_for_agent("", agent_dir)
+      envs.load_dotenv_for_agent("", agents_dir)
       session_service = VertexAiSessionService(
           os.environ["GOOGLE_CLOUD_PROJECT"],
           os.environ["GOOGLE_CLOUD_LOCATION"],
@@ -216,9 +276,12 @@ def get_fast_api_app(
   else:
     session_service = InMemorySessionService()
 
+  # initialize Agent Loader
+  agent_loader = AgentLoader(agents_dir)
+
   @app.get("/list-apps")
   def list_apps() -> list[str]:
-    base_path = Path.cwd() / agent_dir
+    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():
@@ -240,14 +303,34 @@ def get_fast_api_app(
       raise HTTPException(status_code=404, detail="Trace not found")
     return event_dict
 
+  @app.get("/debug/trace/session/{session_id}")
+  def get_session_trace(session_id: str) -> Any:
+    spans = memory_exporter.get_finished_spans(session_id)
+    if not spans:
+      return []
+    return [
+        {
+            "name": s.name,
+            "span_id": s.context.span_id,
+            "trace_id": s.context.trace_id,
+            "start_time": s.start_time,
+            "end_time": s.end_time,
+            "attributes": dict(s.attributes),
+            "parent_span_id": s.parent.span_id if s.parent else None,
+        }
+        for s in spans
+    ]
+
   @app.get(
       "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
       response_model_exclude_none=True,
   )
-  def get_session(app_name: str, user_id: str, session_id: str) -> Session:
+  async def get_session(
+      app_name: str, user_id: str, session_id: str
+  ) -> Session:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
     if not session:
@@ -258,14 +341,15 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions",
       response_model_exclude_none=True,
   )
-  def list_sessions(app_name: str, user_id: str) -> list[Session]:
+  async def list_sessions(app_name: str, user_id: str) -> list[Session]:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
+    list_sessions_response = await session_service.list_sessions(
+        app_name=app_name, user_id=user_id
+    )
     return [
         session
-        for session in session_service.list_sessions(
-            app_name=app_name, user_id=user_id
-        ).sessions
+        for session in list_sessions_response.sessions
         # Remove sessions that were generated as a part of Eval.
         if not session.id.startswith(EVAL_SESSION_ID_PREFIX)
     ]
@@ -274,7 +358,7 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
       response_model_exclude_none=True,
   )
-  def create_session_with_id(
+  async def create_session_with_id(
       app_name: str,
       user_id: str,
       session_id: str,
@@ -283,7 +367,7 @@ def get_fast_api_app(
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
     if (
-        session_service.get_session(
+        await session_service.get_session(
             app_name=app_name, user_id=user_id, session_id=session_id
         )
         is not None
@@ -292,9 +376,8 @@ def get_fast_api_app(
       raise HTTPException(
           status_code=400, detail=f"Session already exists: {session_id}"
       )
-
     logger.info("New session created: %s", session_id)
-    return session_service.create_session(
+    return await session_service.create_session(
         app_name=app_name, user_id=user_id, state=state, session_id=session_id
     )
 
@@ -302,22 +385,21 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions",
       response_model_exclude_none=True,
   )
-  def create_session(
+  async def create_session(
       app_name: str,
       user_id: str,
       state: Optional[dict[str, Any]] = None,
   ) -> Session:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-
     logger.info("New session created")
-    return session_service.create_session(
+    return await session_service.create_session(
         app_name=app_name, user_id=user_id, state=state
     )
 
-  def _get_eval_set_file_path(app_name, agent_dir, eval_set_id) -> str:
+  def _get_eval_set_file_path(app_name, agents_dir, eval_set_id) -> str:
     return os.path.join(
-        agent_dir,
+        agents_dir,
         app_name,
         eval_set_id + _EVAL_SET_FILE_EXTENSION,
     )
@@ -331,28 +413,13 @@ def get_fast_api_app(
       eval_set_id: str,
   ):
     """Creates an eval set, given the id."""
-    pattern = r"^[a-zA-Z0-9_]+$"
-    if not bool(re.fullmatch(pattern, eval_set_id)):
+    try:
+      eval_sets_manager.create_eval_set(app_name, eval_set_id)
+    except ValueError as ve:
       raise HTTPException(
           status_code=400,
-          detail=(
-              f"Invalid eval set id. Eval set id should have the `{pattern}`"
-              " format"
-          ),
-      )
-    # Define the file path
-    new_eval_set_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-
-    logger.info("Creating eval set file `%s`", new_eval_set_path)
-
-    if not os.path.exists(new_eval_set_path):
-      # Write the JSON string to the file
-      logger.info("Eval set file doesn't exist, we will create a new one.")
-      with open(new_eval_set_path, "w") as f:
-        empty_content = json.dumps([], indent=2)
-        f.write(empty_content)
+          detail=str(ve),
+      ) from ve
 
   @app.get(
       "/apps/{app_name}/eval_sets",
@@ -360,15 +427,7 @@ def get_fast_api_app(
   )
   def list_eval_sets(app_name: str) -> list[str]:
     """Lists all eval sets for the given app."""
-    eval_set_file_path = os.path.join(agent_dir, app_name)
-    eval_sets = []
-    for file in os.listdir(eval_set_file_path):
-      if file.endswith(_EVAL_SET_FILE_EXTENSION):
-        eval_sets.append(
-            os.path.basename(file).removesuffix(_EVAL_SET_FILE_EXTENSION)
-        )
-
-    return sorted(eval_sets)
+    return eval_sets_manager.list_eval_sets(app_name)
 
   @app.post(
       "/apps/{app_name}/eval_sets/{eval_set_id}/add_session",
@@ -377,54 +436,33 @@ def get_fast_api_app(
   async def add_session_to_eval_set(
       app_name: str, eval_set_id: str, req: AddSessionToEvalSetRequest
   ):
-    pattern = r"^[a-zA-Z0-9_]+$"
-    if not bool(re.fullmatch(pattern, req.eval_id)):
-      raise HTTPException(
-          status_code=400,
-          detail=f"Invalid eval id. Eval id should have the `{pattern}` format",
-      )
-
     # Get the session
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_name, user_id=req.user_id, session_id=req.session_id
     )
     assert session, "Session not found."
-    # Load the eval set file data
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    with open(eval_set_file_path, "r") as file:
-      eval_set_data = json.load(file)  # Load JSON into a list
-
-    if [x for x in eval_set_data if x["name"] == req.eval_id]:
-      raise HTTPException(
-          status_code=400,
-          detail=(
-              f"Eval id `{req.eval_id}` already exists in `{eval_set_id}`"
-              " eval set."
-          ),
-      )
 
-    # Convert the session data to evaluation format
-    test_data = evals.convert_session_to_eval_format(session)
+    # Convert the session data to eval invocations
+    invocations = evals.convert_session_to_eval_invocations(session)
 
     # Populate the session with initial session state.
     initial_session_state = create_empty_state(
-        await _get_root_agent_async(app_name)
+        agent_loader.load_agent(app_name)
+    )
+
+    new_eval_case = EvalCase(
+        eval_id=req.eval_id,
+        conversation=invocations,
+        session_input=SessionInput(
+            app_name=app_name, user_id=req.user_id, state=initial_session_state
+        ),
+        creation_timestamp=time.time(),
     )
 
-    eval_set_data.append({
-        "name": req.eval_id,
-        "data": test_data,
-        "initial_session": {
-            "state": initial_session_state,
-            "app_name": app_name,
-            "user_id": req.user_id,
-        },
-    })
-    # Serialize the test data to JSON and write to the eval set file.
-    with open(eval_set_file_path, "w") as f:
-      f.write(json.dumps(eval_set_data, indent=2))
+    try:
+      eval_sets_manager.add_eval_case(app_name, eval_set_id, new_eval_case)
+    except ValueError as ve:
+      raise HTTPException(status_code=400, detail=str(ve)) from ve
 
   @app.get(
       "/apps/{app_name}/eval_sets/{eval_set_id}/evals",
@@ -435,14 +473,9 @@ def get_fast_api_app(
       eval_set_id: str,
   ) -> list[str]:
     """Lists all evals in an eval set."""
-    # Load the eval set file data
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    with open(eval_set_file_path, "r") as file:
-      eval_set_data = json.load(file)  # Load JSON into a list
+    eval_set_data = eval_sets_manager.get_eval_set(app_name, eval_set_id)
 
-    return sorted([x["name"] for x in eval_set_data])
+    return sorted([x.eval_id for x in eval_set_data.eval_cases])
 
   @app.post(
       "/apps/{app_name}/eval_sets/{eval_set_id}/run_eval",
@@ -451,51 +484,89 @@ def get_fast_api_app(
   async def run_eval(
       app_name: str, eval_set_id: str, req: RunEvalRequest
   ) -> list[RunEvalResult]:
+    """Runs an eval given the details in the eval request."""
     from .cli_eval import run_evals
 
-    """Runs an eval given the details in the eval request."""
     # Create a mapping from eval set file to all the evals that needed to be
     # run.
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    eval_set_to_evals = {eval_set_file_path: req.eval_ids}
+    eval_set = eval_sets_manager.get_eval_set(app_name, eval_set_id)
 
-    if not req.eval_ids:
-      logger.info(
-          "Eval ids to run list is empty. We will all evals in the eval set."
-      )
-    root_agent = await _get_root_agent_async(app_name)
-    eval_results = list(
-        await run_evals(
-            eval_set_to_evals,
-            root_agent,
-            getattr(root_agent, "reset_data", None),
-            req.eval_metrics,
-            session_service=session_service,
-            artifact_service=artifact_service,
-        )
-    )
+    if req.eval_ids:
+      eval_cases = [e for e in eval_set.eval_cases if e.eval_id in req.eval_ids]
+      eval_set_to_evals = {eval_set_id: eval_cases}
+    else:
+      logger.info("Eval ids to run list is empty. We will run all eval cases.")
+      eval_set_to_evals = {eval_set_id: eval_set.eval_cases}
 
+    root_agent = agent_loader.load_agent(app_name)
     run_eval_results = []
-    for eval_result in eval_results:
+    eval_case_results = []
+    async for eval_case_result in run_evals(
+        eval_set_to_evals,
+        root_agent,
+        getattr(root_agent, "reset_data", None),
+        req.eval_metrics,
+        session_service=session_service,
+        artifact_service=artifact_service,
+    ):
       run_eval_results.append(
           RunEvalResult(
               app_name=app_name,
+              eval_set_file=eval_case_result.eval_set_file,
               eval_set_id=eval_set_id,
-              eval_id=eval_result.eval_id,
-              final_eval_status=eval_result.final_eval_status,
-              eval_metric_results=eval_result.eval_metric_results,
-              session_id=eval_result.session_id,
+              eval_id=eval_case_result.eval_id,
+              final_eval_status=eval_case_result.final_eval_status,
+              eval_metric_results=eval_case_result.eval_metric_results,
+              overall_eval_metric_results=eval_case_result.overall_eval_metric_results,
+              eval_metric_result_per_invocation=eval_case_result.eval_metric_result_per_invocation,
+              user_id=eval_case_result.user_id,
+              session_id=eval_case_result.session_id,
           )
       )
+      eval_case_result.session_details = await session_service.get_session(
+          app_name=app_name,
+          user_id=eval_case_result.user_id,
+          session_id=eval_case_result.session_id,
+      )
+      eval_case_results.append(eval_case_result)
+
+    eval_set_results_manager.save_eval_set_result(
+        app_name, eval_set_id, eval_case_results
+    )
+
     return run_eval_results
 
+  @app.get(
+      "/apps/{app_name}/eval_results/{eval_result_id}",
+      response_model_exclude_none=True,
+  )
+  def get_eval_result(
+      app_name: str,
+      eval_result_id: str,
+  ) -> EvalSetResult:
+    """Gets the eval result for the given eval id."""
+    try:
+      return eval_set_results_manager.get_eval_set_result(
+          app_name, eval_result_id
+      )
+    except ValueError as ve:
+      raise HTTPException(status_code=404, detail=str(ve)) from ve
+    except ValidationError as ve:
+      raise HTTPException(status_code=500, detail=str(ve)) from ve
+
+  @app.get(
+      "/apps/{app_name}/eval_results",
+      response_model_exclude_none=True,
+  )
+  def list_eval_results(app_name: str) -> list[str]:
+    """Lists all eval results for the given app."""
+    return eval_set_results_manager.list_eval_set_results(app_name)
+
   @app.delete("/apps/{app_name}/users/{user_id}/sessions/{session_id}")
-  def delete_session(app_name: str, user_id: str, session_id: str):
+  async def delete_session(app_name: str, user_id: str, session_id: str):
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-    session_service.delete_session(
+    await session_service.delete_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
 
@@ -589,9 +660,9 @@ def get_fast_api_app(
   @app.post("/run", response_model_exclude_none=True)
   async def agent_run(req: AgentRunRequest) -> list[Event]:
     # Connect to managed session if agent_engine_id is set.
-    app_id = agent_engine_id if agent_engine_id else req.app_name
-    session = session_service.get_session(
-        app_name=app_id, user_id=req.user_id, session_id=req.session_id
+    app_name = agent_engine_id if agent_engine_id else req.app_name
+    session = await session_service.get_session(
+        app_name=app_name, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
       raise HTTPException(status_code=404, detail="Session not found")
@@ -610,10 +681,10 @@ def get_fast_api_app(
   @app.post("/run_sse")
   async def agent_run_sse(req: AgentRunRequest) -> StreamingResponse:
     # Connect to managed session if agent_engine_id is set.
-    app_id = agent_engine_id if agent_engine_id else req.app_name
+    app_name = agent_engine_id if agent_engine_id else req.app_name
     # SSE endpoint
-    session = session_service.get_session(
-        app_name=app_id, user_id=req.user_id, session_id=req.session_id
+    session = await session_service.get_session(
+        app_name=app_name, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
       raise HTTPException(status_code=404, detail="Session not found")
@@ -652,9 +723,9 @@ def get_fast_api_app(
       app_name: str, user_id: str, session_id: str, event_id: str
   ):
     # Connect to managed session if agent_engine_id is set.
-    app_id = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
-        app_name=app_id, user_id=user_id, session_id=session_id
+    app_name = agent_engine_id if agent_engine_id else app_name
+    session = await session_service.get_session(
+        app_name=app_name, user_id=user_id, session_id=session_id
     )
     session_events = session.events if session else []
     event = next((x for x in session_events if x.id == event_id), None)
@@ -665,7 +736,7 @@ def get_fast_api_app(
 
     function_calls = event.get_function_calls()
     function_responses = event.get_function_responses()
-    root_agent = await _get_root_agent_async(app_name)
+    root_agent = agent_loader.load_agent(app_name)
     dot_graph = None
     if function_calls:
       function_call_highlights = []
@@ -673,7 +744,7 @@ def get_fast_api_app(
         from_name = event.author
         to_name = function_call.name
         function_call_highlights.append((from_name, to_name))
-        dot_graph = agent_graph.get_agent_graph(
+        dot_graph = await agent_graph.get_agent_graph(
             root_agent, function_call_highlights
         )
     elif function_responses:
@@ -682,17 +753,17 @@ def get_fast_api_app(
         from_name = function_response.name
         to_name = event.author
         function_responses_highlights.append((from_name, to_name))
-        dot_graph = agent_graph.get_agent_graph(
+        dot_graph = await agent_graph.get_agent_graph(
             root_agent, function_responses_highlights
         )
     else:
       from_name = event.author
       to_name = ""
-      dot_graph = agent_graph.get_agent_graph(
+      dot_graph = await agent_graph.get_agent_graph(
           root_agent, [(from_name, to_name)]
       )
     if dot_graph and isinstance(dot_graph, graphviz.Digraph):
-      return {"dot_src": dot_graph.source}
+      return GetEventGraphResult(dot_src=dot_graph.source)
     else:
       return {}
 
@@ -709,9 +780,9 @@ def get_fast_api_app(
     await websocket.accept()
 
     # Connect to managed session if agent_engine_id is set.
-    app_id = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
-        app_name=app_id, user_id=user_id, session_id=session_id
+    app_name = agent_engine_id if agent_engine_id else app_name
+    session = await session_service.get_session(
+        app_name=app_name, user_id=user_id, session_id=session_id
     )
     if not session:
       # Accept first so that the client is aware of connection establishment,
@@ -766,34 +837,12 @@ def get_fast_api_app(
       for task in pending:
         task.cancel()
 
-  async def _get_root_agent_async(app_name: str) -> Agent:
-    """Returns the root agent for the given app."""
-    if app_name in root_agent_dict:
-      return root_agent_dict[app_name]
-    agent_module = importlib.import_module(app_name)
-    if getattr(agent_module.agent, "root_agent"):
-      root_agent = agent_module.agent.root_agent
-    else:
-      raise ValueError(f'Unable to find "root_agent" from {app_name}.')
-
-    # Handle an awaitable root agent and await for the actual agent.
-    if inspect.isawaitable(root_agent):
-      try:
-        agent, exit_stack = await root_agent
-        exit_stacks.append(exit_stack)
-        root_agent = agent
-      except Exception as e:
-        raise RuntimeError(f"error getting root agent, {e}") from e
-
-    root_agent_dict[app_name] = root_agent
-    return root_agent
-
   async def _get_runner_async(app_name: str) -> Runner:
     """Returns the runner for the given app."""
-    envs.load_dotenv_for_agent(os.path.basename(app_name), agent_dir)
+    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 = await _get_root_agent_async(app_name)
+    root_agent = agent_loader.load_agent(app_name)
     runner = Runner(
         app_name=agent_engine_id if agent_engine_id else app_name,
         agent=root_agent,
@@ -809,14 +858,16 @@ def get_fast_api_app(
     ANGULAR_DIST_PATH = BASE_DIR / "browser"
 
     @app.get("/")
-    async def redirect_to_dev_ui():
-      return RedirectResponse("/dev-ui")
+    async def redirect_root_to_dev_ui():
+      return RedirectResponse("/dev-ui/")
 
     @app.get("/dev-ui")
-    async def dev_ui():
-      return FileResponse(BASE_DIR / "browser/index.html")
+    async def redirect_dev_ui_add_slash():
+      return RedirectResponse("/dev-ui/")
 
     app.mount(
-        "/", StaticFiles(directory=ANGULAR_DIST_PATH, html=True), name="static"
+        "/dev-ui/",
+        StaticFiles(directory=ANGULAR_DIST_PATH, html=True),
+        name="static",
     )
   return app