prompt fix

gemini_agent.py

@@ -294,6 +294,49 @@ def wikipedia_qa(question: str) -> str:
     out = qa_pipeline(question=question, context=page)
     return out["answer"]    
 
+import re
+
+def count_studio_albums(query: str) -> str:
+    """
+    解析维基 'Studio albums' 段落，统计 <artist> 在 start–end 年份之间出的专辑数。
+    query 格式示例："How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?"
+    """
+    # 1) 从问题里提取艺术家和年份区间
+    m = re.search(
+        r"studio albums were published by (.+?) between (\d{4}) and (\d{4})",
+        query, flags=re.IGNORECASE
+    )
+    if not m:
+        return ""  # 不匹配就空
+    artist, start, end = m.group(1), int(m.group(2)), int(m.group(3))
+
+    # 2) 抓维基全文，找 'Studio albums' 小节
+    try:
+        page = wikipedia.page(artist)
+    except Exception as e:
+        return f"Error fetching wiki page: {e}"
+    text = page.content
+
+    # 3) 拆段落
+    parts = re.split(r"\n==+\s*Studio albums\s*==+", text)
+    if len(parts) < 2:
+        return "0"
+    section = parts[1]
+
+    # 4) 从每行 `* Album Name (YYYY)` 抽年份、计数
+    years = []
+    for line in section.splitlines():
+        line = line.strip()
+        if not line.startswith("*"):
+            continue
+        y_m = re.search(r"\((\d{4})\)", line)
+        if y_m:
+            y = int(y_m.group(1))
+            years.append(y)
+
+    count = sum(1 for y in years if start <= y <= end)
+    return str(count)
+
 class GeminiAgent:
     def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
         # Suppress warnings
@@ -344,7 +387,13 @@ class GeminiAgent:
             name="wikipedia_qa",
             func=wikipedia_qa,
             description="给定问题 + 维基文章，直接抽取精确答案"
-)
+            ),
+            Tool(
+            name="count_studio_albums",
+            func=count_studio_albums,
+            description=("统计某位艺术家在指定年份区间内发布的 Studio albums 数量，"
+            "query 必须包含 “studio albums … between YYYY and YYYY”。")
+            )
         ]
         
         # Setup memory
@@ -589,7 +638,7 @@ Observation: the result of the action
 2. Action: wikipedia_qa
    Action Input: <original user question>
    Observation: <extracted fact or number>
-   
+
 When you have the final answer, you MUST use exactly one line with this format (no extra text, no lists):
 
 Thought: Do I need to use a tool? No