Update agent.py

agent.py

@@ -197,30 +197,40 @@ def get_answer(state: InfoState) -> InfoState :
     state["final_answer"] = get_gpt_and_answer(prompt)
 
     return state
-
+    
 def get_type(state: InfoState) -> InfoState:
     """Choose which tool to use based on question type (WIKI, SEARCH, CODE)."""
-    print("Getting Type (Gemini)...")
-
-    prompt = f"""
-    You are a strict classifier.
-    Your job is to classify the following question into ONE of four categories:
-
-    - WIKI → informative, factual, or science question
-    - WebInfo → up-to-date, news, or current event question
-    - MATH → math, numeric, date, or time calculation
-    - LLM → all others, including links, reasoning, and file-related tasks
-
-    Question: "{state['question']}"
-
-    Rules:
-    - Reply with exactly one of these words: WIKI, WebInfo, MATH, or LLM
-    - Output nothing else. No punctuation, no quotes, no explanation.
-    - If unsure, default to LLM.
-    """
-    
-    state["answer_type"] = get_gpt_and_answer(prompt)
 
+    print("Getting Type (Gemini/OpenAI)...")
+    q = state["question"].lower()
+    if any(w in q for w in ["calculate", "sum", "add", "price", "how much", "date", "time"]):
+        state["answer_type"] = "MATH"
+    elif any(w in q for w in ["latest", "current", "today", "news"]):
+        state["answer_type"] = "WebInfo"
+    elif any(w in q for w in ["who", "where", "what", "when", "why", "how many", "which"]):
+        state["answer_type"] = "WIKI"
+    else:
+        try:
+            prompt = f"""
+            You are a strict classifier.
+            Your job is to classify the following question into ONE of four categories:
+        
+            - WIKI → informative, factual, or science question
+            - WebInfo → up-to-date, news, or current event question
+            - MATH → math, numeric, date, or time calculation
+            - LLM → all others, including links, reasoning, and file-related tasks
+        
+            Question: "{state['question']}"
+        
+            Rules:
+            - Reply with exactly one of these words: WIKI, WebInfo, MATH, or LLM
+            - Output nothing else. No punctuation, no quotes, no explanation.
+            - If unsure, default to LLM.
+            """
+            resp = get_gpt_and_answer(prompt)
+            state["answer_type"] = resp.split()[0].strip()
+        except:
+            state["answer_type"] = "LLM"
     return state