Create agent.py

agent.py

@@ -0,0 +1,85 @@
+# agent.py
+import os
+import re
+import google.generativeai as genai
+from dotenv import load_dotenv
+
+# 自動加載.env（如果有）
+load_dotenv()
+
+# 配置Gemini API Key（環境變數為 GOOGLE_API_KEY）
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+if GOOGLE_API_KEY is None:
+    raise ValueError("GOOGLE_API_KEY not set in environment variables!")
+
+genai.configure(api_key=GOOGLE_API_KEY)
+
+class Agent:
+    def __init__(self):
+        # 默認用flash，有就用pro
+        try:
+            self.model = genai.GenerativeModel('gemini-2.0-flash')
+        except Exception:
+            self.model = genai.GenerativeModel('gemini-1.5-pro')
+    
+    def __call__(self, question: str) -> str:
+        # 標準化prompt設計（已驗證能過30%分數線）
+        prompt = self._build_prompt(question)
+        response = self.model.generate_content(
+            prompt,
+            generation_config=genai.types.GenerationConfig(
+                temperature=0.1, max_output_tokens=512,
+            )
+        )
+        return self._clean_answer(response.text, question)
+    
+    def _build_prompt(self, question: str) -> str:
+        # 給LLM一個嚴格格式限制，最大限度提升自動答題正確率
+        return f"""You are a strict answer bot for GAIA evaluation.
+Only provide the FINAL concise answer to the following question.
+Rules:
+- No explanations, no extra info, no reasoning.
+- If the question asks for a number, answer only the number.
+- If the question is multiple choice, answer only the letter(s).
+- If the question is yes/no, answer 'yes' or 'no'.
+- If you are not sure, answer 'Unknown'.
+- For lists, answer with comma-separated values.
+End your answer with: FINAL ANSWER: [your answer]
+Question: {question}
+FINAL ANSWER:
+"""
+    def _clean_answer(self, answer: str, question: str) -> str:
+        # 從 LLM 返回文本中提取“FINAL ANSWER: ...”
+        if answer is None:
+            return "Unknown"
+        match = re.search(r"FINAL ANSWER:\s*([^\n]+)", answer, re.IGNORECASE)
+        if match:
+            ans = match.group(1).strip()
+        else:
+            # 如果沒這個格式，取最後一行或全部
+            ans = answer.strip().splitlines()[-1] if answer.strip() else "Unknown"
+        # 如果要求純數字、選項等，進行提純
+        q_lower = question.lower()
+        # 只要數字
+        if any(x in q_lower for x in ["how many", "number of", "count", "total", "amount"]):
+            num_match = re.search(r"\d+", ans)
+            if num_match:
+                return num_match.group(0)
+        # 只要選項字母
+        if re.search(r"[a-d]\)", question) or "option" in q_lower:
+            opt_match = re.findall(r"[a-d]", ans, re.IGNORECASE)
+            if opt_match:
+                return ",".join(x.lower() for x in opt_match)
+        # yes/no題
+        if q_lower.startswith("is ") or q_lower.startswith("are ") or q_lower.startswith("does "):
+            if "yes" in ans.lower():
+                return "yes"
+            if "no" in ans.lower():
+                return "no"
+        return ans
+
+# 測試
+if __name__ == "__main__":
+    agent = Agent()
+    q = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?"
+    print(agent(q))