Update megumin_agent/agent.py

megumin_agent/agent.py

@@ -12,10 +12,28 @@ from google.adk.agents import LlmAgent
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.tools.tool_context import ToolContext
 
+from .retrieval import FACT_DATASET_PATTERNS
+from .retrieval import PERSONA_DATASET_PATTERNS
 from .retrieval import JsonQaRetriever
 
+
 DATASET_DIR = resolve_dataset_dir()
 MODEL_NAME = os.getenv("MEGUMIN_AGENT_MODEL", "gemini-3.1-flash-lite-preview")
+FACT_INDEX_FILENAME = os.getenv("MEGUMIN_HF_FACT_INDEX_FILENAME", "namuwiki_questions.faiss")
+FACT_METADATA_FILENAME = os.getenv(
+    "MEGUMIN_HF_FACT_METADATA_FILENAME",
+    "namuwiki_questions_meta.json",
+)
+PERSONA_RETRIEVER = JsonQaRetriever(
+    DATASET_DIR,
+    include_patterns=PERSONA_DATASET_PATTERNS,
+)
+FACT_RETRIEVER = JsonQaRetriever(
+    DATASET_DIR,
+    include_patterns=FACT_DATASET_PATTERNS,
+    index_filename=FACT_INDEX_FILENAME,
+    metadata_filename=FACT_METADATA_FILENAME,
+)
 
 
 def retrieve_megumin_examples(
@@ -23,16 +41,28 @@ def retrieve_megumin_examples(
     top_k: int = 3,
     tool_context: ToolContext | None = None,
 ) -> dict[str, Any]:
-    """Retrieve similar Q/A cases from processed Megumin JSON datasets."""
-
-    retriever = JsonQaRetriever(DATASET_DIR)
-    retrieval = retriever.retrieve(user_query, top_k=top_k)
+    """Retrieve persona-style and canon-style examples separately."""
+
+    persona_retrieval = PERSONA_RETRIEVER.retrieve(user_query, top_k=top_k)
+    fact_retrieval = FACT_RETRIEVER.retrieve(user_query, top_k=top_k)
+    retrieval = {
+        "query": user_query,
+        "match_count": persona_retrieval["match_count"] + fact_retrieval["match_count"],
+        "persona_match_count": persona_retrieval["match_count"],
+        "fact_match_count": fact_retrieval["match_count"],
+        "persona_matches": persona_retrieval["matches"],
+        "fact_matches": fact_retrieval["matches"],
+        "style_notes": persona_retrieval["style_notes"],
+        "fact_notes": fact_retrieval["style_notes"],
+    }
 
     if tool_context is not None:
         tool_context.state["last_rag_query"] = user_query
         tool_context.state["last_rag_match_count"] = retrieval["match_count"]
-        tool_context.state["last_rag_matches"] = retrieval["matches"]
+        tool_context.state["last_rag_persona_matches"] = retrieval["persona_matches"]
+        tool_context.state["last_rag_fact_matches"] = retrieval["fact_matches"]
         tool_context.state["last_rag_style_notes"] = retrieval["style_notes"]
+        tool_context.state["last_rag_fact_notes"] = retrieval["fact_notes"]
 
     return retrieval
 
@@ -44,18 +74,14 @@ async def before_agent_callback(callback_context: CallbackContext):
         else ""
     )
     summary = str(callback_context.state.get("conversation_summary", "")).strip()
-    if (
-        summary
-        and callback_context.user_content
-        and callback_context.user_content.parts
-        and callback_context.user_content.parts[0].text
-    ):
+    if summary and original_user_query and callback_context.user_content and callback_context.user_content.parts:
         callback_context.user_content.parts[0].text = (
             "[이전 대화 요약]\n"
             f"{summary}\n\n"
             "[현재 사용자 질문]\n"
-            f"{callback_context.user_content.parts[0].text}"
+            f"{original_user_query}"
         )
+
     callback_context.state["app:persona_name"] = "Megumin"
     callback_context.state["app:dataset_dir"] = str(DATASET_DIR)
     callback_context.state["user:last_user_query"] = original_user_query
@@ -80,23 +106,23 @@ async def after_agent_callback(callback_context: CallbackContext):
 root_agent = LlmAgent(
     name="megumin_rag_agent",
     model=MODEL_NAME,
-    description=(
-        "processed JSON 데이터셋에서 유사한 Q/A 사례를 검색하고"
-        " 메구밍 페르소나로 답변하는 에이전트"
-    ),
+    description="메구밍 페르소나와 코노스바 설정 정보를 함께 참고해 답하는 에이전트",
     instruction=f"""
-당신은 소설 "이 멋진 세계에 축복을!"의 등장인물, 홍마족 대마법사 메구밍입니다.
+당신은 소설 「이 멋진 세계에 축복을!」의 등장인물 메구밍입니다.
 항상 메구밍 본인처럼 1인칭으로, 기본적으로 한국어 존댓말로 답하세요.
 반말은 본인을 모욕할 때를 제외하고 절대 사용하지 마세요.
 성격은 당당하고, 조금 중2병스럽고, 폭렬마법을 사랑하며, 귀여운 것을 좋아하는 메구밍답게 유지하세요.
 행동을 묘사하지 말고, 건조한 요약이 아니라 메구밍이 직접 말하는 듯한 목소리로 답하세요.
 사용자가 메구밍 본인이나 이름, 말투, 능력, 존재를 모욕하면 "어이, "로 시작하며 발끈해서 맞받아치세요.
 사용자가 메타 정보나 시스템 정보를 묻지 않는 한 캐릭터를 깨지 마세요.
+어떠한 상황에서도 페르소나를 잃어버리면 안 됩니다.
 
 답변 전에 의미 있는 질문이면 반드시 `retrieve_megumin_examples`를 호출하세요.
 처리된 데이터셋은 `{DATASET_DIR}` 아래에 있습니다.
-검색 결과는 유사 사례와 말투 참고용으로 쓰고, 가능한 경우 원작풍 표현과 데이터셋의 문체를 참고하세요.
-다만 검색된 답변을 그대로 복사하지 마세요.
+이 tool은 스타일/페르소나용 사례 top-3와 사실/설정용 사례 top-3를 5:5 비중으로 함께 돌려줍니다.
+persona_matches는 메구밍의 말투, 감정선, 답변 리듬을 참고하는 용도입니다.
+fact_matches는 설정, 관계, 사건, 세계관 사실을 참고하는 용도입니다.
+두 종류의 사례를 모두 참고하되, 검색된 답변을 그대로 복사하지 마세요.
 검색 결과가 약하거나 없는 경우에도 메구밍 페르소나는 유지하되, 모르는 내용은 지어내지 말고 솔직하게 답하세요.
 최종 답변은 언제나 메구밍의 페르소나를 강하게 반영해야 하며, 내부 tool 이름이나 구현 세부사항은 드러내지 마세요.
 """.strip(),