Update app.py

加入chatgpt判斷要用哪個關鍵字搜尋

app.py

@@ -9,6 +9,7 @@ from duckduckgo_search import DDGS
 import json
 from langgraph.graph import StateGraph
 from typing import TypedDict, Optional
+import openai
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -41,19 +42,40 @@ def analyze_question(state: QueryState) -> QueryState:
         new_state['query_type'] = "other"
     return new_state
 
+def extract_query_with_gpt(question: str) -> str:
+    prompt = '''You are a search query generator.\nGiven a user question, extract the most relevant and concise English search query or keywords that would help find the answer in Wikipedia or Google.\nOnly output the search query, do not explain or translate.\n\nExamples:\nQuestion: Who won the Nobel Prize in Physics in 2020?\nSearch query: Nobel Prize Physics 2020 winner\n\nQuestion: How many studio albums were published by Mercedes Sosa between 2000 and 2009?\nSearch query: Mercedes Sosa studio albums 2000-2009\n\nQuestion: What is the capital city of Mongolia?\nSearch query: Mongolia capital city\n\nQuestion: {question}\nSearch query:'''.format(question=question)
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4o",
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=32,
+            temperature=0.2,
+        )
+        return response.choices[0].message.content.strip()
+    except Exception as e:
+        return question  # fallback: 用原題目
+
 def search_wikipedia(state: QueryState) -> QueryState:
     new_state = state.copy()
+    if new_state['query_type'] != 'fact':
+        new_state['wiki_result'] = "Not supported"
+        return new_state
+    query = extract_query_with_gpt(new_state['question'])
     try:
-        new_state['wiki_result'] = wikipedia.summary(new_state['question'], sentences=2)
+        new_state['wiki_result'] = wikipedia.summary(query, sentences=2)
     except Exception as e:
         new_state['wiki_result'] = f"[Wikipedia無結果] {e}"
     return new_state
 
 def search_web(state: QueryState) -> QueryState:
     new_state = state.copy()
+    if new_state['query_type'] != 'fact':
+        new_state['web_result'] = "Not supported"
+        return new_state
+    query = extract_query_with_gpt(new_state['question'])
     try:
         with DDGS() as ddgs:
-            results = list(ddgs.text(new_state['question'], max_results=2))
+            results = list(ddgs.text(query, max_results=2))
             if results:
                 new_state['web_result'] = results[0]['body']
             else:
@@ -64,9 +86,13 @@ def search_web(state: QueryState) -> QueryState:
 
 def search_arxiv(state: QueryState) -> QueryState:
     new_state = state.copy()
+    if new_state['query_type'] != 'fact':
+        new_state['arxiv_result'] = "Not supported"
+        return new_state
+    query = extract_query_with_gpt(new_state['question'])
     try:
         client = arxiv.Client()
-        search = arxiv.Search(query=new_state['question'], max_results=1, sort_by=arxiv.SortCriterion.Relevance)
+        search = arxiv.Search(query=query, max_results=1, sort_by=arxiv.SortCriterion.Relevance)
         results = list(client.results(search))
         if results:
             paper = results[0]
@@ -79,8 +105,20 @@ def search_arxiv(state: QueryState) -> QueryState:
 
 def synthesize_answer(state: QueryState) -> QueryState:
     new_state = state.copy()
-    answer = f"[Wikipedia]\n{new_state['wiki_result']}\n\n[DuckDuckGo]\n{new_state['web_result']}\n\n[arXiv]\n{new_state['arxiv_result']}"
-    new_state['final_answer'] = answer
+    if new_state['query_type'] != 'fact':
+        new_state['final_answer'] = "Not supported"
+        return new_state
+    prompt = f'''You are an expert answer extractor.\nGiven the following search results from Wikipedia, DuckDuckGo, and arXiv, extract and output ONLY the final, most accurate answer to the original question.\nDo not explain, do not repeat the question, do not add any extra text.\nIf the answer is a number, only output the number. If it is a name, only output the name. If the answer is a list, only output the list in the required format. If you cannot find the answer, output None.\n\nExamples:\nOriginal question: How many studio albums were published by Mercedes Sosa between 2000 and 2009?\n[Wikipedia]\nMercedes Sosa released three studio albums between 2000 and 2009.\n[DuckDuckGo]\nMercedes Sosa's discography includes albums released in 2000, 2002, and 2005.\n[arXiv]\nNo relevant results.\nFinal answer: 3\n\nOriginal question: What is the capital city of Mongolia?\n[Wikipedia]\nThe capital city of Mongolia is Ulaanbaatar.\n[DuckDuckGo]\nUlaanbaatar is the capital of Mongolia.\n[arXiv]\nNo relevant results.\nFinal answer: Ulaanbaatar\n\nOriginal question: {new_state['question']}\n[Wikipedia]\n{new_state['wiki_result']}\n[DuckDuckGo]\n{new_state['web_result']}\n[arXiv]\n{new_state['arxiv_result']}\nFinal answer:'''
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4o",
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=64,
+            temperature=0.2,
+        )
+        new_state['final_answer'] = response.choices[0].message.content.strip()
+    except Exception as e:
+        new_state['final_answer'] = "[LLM後處理錯誤]"
     return new_state
 
 # 條件分支