Upload agent.py

agent.py

@@ -46,9 +46,14 @@ def get_wiki_relate(state: InfoState) -> InfoState:
     # Make a tool from it
     wiki_tool = WikipediaQueryRun(api_wrapper=wiki)
 
-    wiki_answer = wiki_tool.run(" ".join(state["main_parts"]))
-    state['tool_answer'] = wiki_answer
-    return state
+    try:
+      wiki_answer = wiki_tool.run(" ".join(state["main_parts"]) + " full wikipedia article about this topic")
+      state['tool_answer'] = wiki_answer
+      return state
+    except Exception as e:
+      print("Rate limit Exception")
+      state['tool_answer'] = ""
+      return state
 
 
 def execute_code(state: InfoState) -> InfoState :
@@ -96,17 +101,28 @@ def preprocess_text(state: dict) -> InfoState:
     """
     print("Preprocess text (Gemini)...")
     # 1️⃣ Prepare the prompt
-    prompt = (
-    "Extract the most important content words (nouns, proper names, key concepts) from this question that would help find the best-matching Wikipedia article. "
-    "If the question is not in English; translate key terms to English for Wikipedia's English edition. "
-    "Ignore stopwords (like 'who', 'is', 'the', 'of', 'in', 'current', 'what'), filler words, and typos. "
-    "Focus on entities and topics that exist as Wikipedia page titles. "
-    "Correct obvious spelling mistakes and expand common abbreviations if needed for better Wikipedia matching.\n\n"
-    "Question: '" + state["question"] + "'\n\n"
-    "Output ONLY a valid JSON list of 1–4 corrected, title-cased strings (e.g. [\"President of the United States\", \"Joe Biden\"]). "
-    "No explanations, no markdown, no extra text, no quotes around words, no trailing commas."
+    prompt =  (
+    "We want to find the best-matching English Wikipedia pages for a factual question, "
+    "so we must extract only the essential topic names or entities that Wikipedia likely has pages for. "
+    "These should include the main subject (e.g., a person, event, place, or concept) and any directly relevant subtopic "
+    "(like 'Discography', 'Career', or 'History') if they help narrow the search.\n\n"
+
+    "Rules:\n"
+    "- Output 1 to 3 items maximum.\n"
+    "- Use English Wikipedia title format (capitalize each main word).\n"
+    "- Translate non-English names or terms to English.\n"
+    "- Exclude question words, pronouns, and filler terms.\n"
+    "- Fix spelling errors if necessary.\n"
+    "- Prefer specific Wikipedia topics over vague ones.\n\n"
+
+    "Example:\n"
+    "Q: 'Who built the Eiffel Tower?'\n"
+    "A: [\"Eiffel Tower\", \"Gustave Eiffel\"]\n\n"
+    f"Question: '{state['question']}'\n\n"
+    "Output ONLY a valid JSON list as described — no explanations, markdown, or extra formatting."
     )
 
+
     # 2️⃣ Call Gemini
     model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
     response = model.invoke([HumanMessage(content=prompt)]).content.strip()
@@ -138,12 +154,19 @@ def get_answer(state: InfoState) -> InfoState :
     print("Getting Answer (Gemini)...")
 
     prompt = (
-        "Answer the question based on the context below. "
-        #"If the question cannot be answered using the information provided, answer with 'I don't know'. "
-        "Question: " + state["question"] +
-        "\nContext: " + str(state.get("tool_answer")) +
-        "\nRewrite answer in clearer, simple way."
+        "You are a knowledgeable assistant that answers questions based on context and common factual knowledge.\n"
+        "Use the context first, but if it clearly lacks the needed details, you may rely on well-known public facts "
+        "(such as from Wikipedia) that logically complete the answer.\n\n"
+        f"Question: {state['question']}\n"
+        f"Context:\n{state.get('tool_answer')}\n\n"
+        "Instructions:\n"
+        "- Focus on producing one short factual answer.\n"
+        "- Do not include tool names, prefixes, or metadata.\n"
+        "- If the context contains partial hints, you can infer the answer from general knowledge of the same topic.\n"
+        "- If absolutely nothing is relevant, reply: I don't know.\n\n"
+        "Final Answer:"
     )
+
     model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
     state["final_answer"] = (model.invoke([HumanMessage(content=prompt)]).content)
 
@@ -153,7 +176,7 @@ def get_type(state: InfoState) -> InfoState:
     """Choose which tool to use based on question type (WIKI, SEARCH, CODE)."""
     print("Getting Type (Gemini)...")
 
-    prompt = "According to the Question " +state["question"] + " Select the best tool to answer WIKI if it's informatative or science question, WebInfo if it need up to data news,MATH if math or date or time  You're very serious,just give one word from given"
+    prompt = "According to the Question " +state["question"] + " Select the best tool to answer WIKI if it's informatative or science question, WebInfo if it need up to data news,MATH if math or date or time,LLM if other or have link  You're very serious,just give one word from given"
     model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
     state["answer_type"] = (model.invoke([HumanMessage(content=prompt)]).content)
 
@@ -168,10 +191,12 @@ def get_search_results(state: InfoState) -> InfoState:
 
     search = DuckDuckGoSearchRun()
 
-    # Simple text result
-    state['tool_answer'] = search.run(state["question"])   #" " .join(state["main_parts"]))
-
-    return state
+    try:
+      state['tool_answer'] = search.run(state["question"])   #" " .join(state["main_parts"]))
+      return state
+    except Exception:
+      state['tool_answer'] = ""
+      return state
 
 
 def route(state: InfoState):
@@ -185,7 +210,7 @@ def get_graph():
         graph = StateGraph(InfoState)
 
         # Add nodes
-        graph.add_node("get_wiki_relate", get_wiki_relate)
+        #graph.add_node("get_wiki_relate", get_wiki_relate)
         graph.add_node("preprocess_text", preprocess_text)
         graph.add_node("get_answer", get_answer)
         graph.add_node("get_type", get_type)
@@ -204,14 +229,15 @@ def get_graph():
             route,
             {
                 "WebInfo": "get_search_results",
-                "WIKI": "get_wiki_relate",
-                "MATH": "get_code"
+                "WIKI": "get_search_results",#"get_wiki_relate",
+                "MATH": "get_code",
+                "LLM": "get_answer"
             }
         )
 
         # Add final edges
         graph.add_edge("get_search_results", "get_answer")  
-        graph.add_edge("get_wiki_relate", "get_answer")
+        #graph.add_edge("get_wiki_relate", "get_answer")
 
         graph.add_edge("get_code", "execute_code")
         graph.add_edge("execute_code", "get_answer")
@@ -229,4 +255,4 @@ def ask(compiled_graph,question):
 
         })
 
-        return legitimate_result['final_answer']
+        return legitimate_result['final_answer'],legitimate_result