commit

app.py

@@ -26,6 +26,13 @@ You MUST:
 - No extra text.
 """
 
+RULES = {
+    "number": "Answer with ONLY the number. No words.",
+    "date": "Answer with ONLY the year or date. No words.",
+    "yesno": "Answer with ONLY Yes or No.",
+    "fact": "Answer with ONLY the final answer."
+}
+
 def clean_answer(text: str) -> str:
     if not text:
         return ""
@@ -35,8 +42,22 @@ def clean_answer(text: str) -> str:
     s = s.strip('"\'`')
     if len(s) > 1 and s.endswith("."):
         s = s[:-1].strip()
+
+    if s.lower() in ["yes", "no"]:
+        s = s.capitalize()
+        
     return s
 
+def classify_question(q: str) -> str:
+    ql = q.lower().strip()
+    if any(k in ql for k in ["how many", "how much", "number of", "population"]):
+        return "number"
+    if any(k in ql for k in ["what year", "when", "date"]):
+        return "date"
+    if ql.startswith(("is ", "are ", "was ", "were ", "did ", "does ")):
+        return "yesno"
+    return "fact"
+
 # -------------------------------
 # State
 # -------------------------------
@@ -67,25 +88,38 @@ class BasicAgent:
     def __call__(self, question: str) -> str:
         print(f"Question: {question[:80]}...")
 
-        # 1) Search
+        qtype = classify_question(question)
+        rule = RULES[qtype]
+
+        # 검색 쿼리 보정 (숫자/날짜는 Wikipedia 우선)
+        search_query = question
+        if qtype in ["number", "date"]:
+            search_query = question + " site:wikipedia.org"
+
+        # 검색 (실패 시 1회 재시도)
         try:
-            search_result = search_tool.run(question)
-        except Exception as e:
-            print("Search error:", e)
-            search_result = ""
+            search_result = search_tool.run(search_query)
+            if not search_result:
+                raise RuntimeError("empty search")
+        except Exception:
+            try:
+                search_result = search_tool.run(question)
+            except Exception:
+                search_result = ""
 
-        # 2) Prompt with evidence
         prompt = f"""
-{SYSTEM_PROMPT}
+    {SYSTEM_PROMPT}
+
+    Additional Rules:
+    - {rule}
 
-Question:
-{question}
+    Question:
+    {question}
 
-Search Results:
-{search_result}
-""".strip()
+    Search Results:
+    {search_result}
+    """.strip()
 
-        # 3) LLM Answer
         response = llm.invoke([HumanMessage(content=prompt)])
         answer = clean_answer(response.content)