Update megumin_agent/chat.py

megumin_agent/chat.py

@@ -12,9 +12,9 @@ from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.genai import types
 
-from .agent import DATASET_DIR
+from .agent import FACT_RETRIEVER
+from .agent import PERSONA_RETRIEVER
 from .agent import root_agent
-from .retrieval import JsonQaRetriever
 
 
 APP_NAME = "megumin_rag_app"
@@ -84,7 +84,8 @@ def _trim_session_history(
 
 
 def create_chat_services() -> ChatServices:
-    JsonQaRetriever(DATASET_DIR).warmup()
+    PERSONA_RETRIEVER.warmup()
+    FACT_RETRIEVER.warmup()
     session_service = InMemorySessionService()
     runner = Runner(
         agent=root_agent,