Update agent.py

agent.py

@@ -178,7 +178,29 @@ 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,LLM if other or have link  You're very serious,just give one word from given"
+    prompt = f"""
+      You are a strict classifier.  
+      According to the question below, choose **only one** of the following tools:
+      - WIKI → if the question is informative, factual, or science-related.
+      - WebInfo → if it requires up-to-date or current information (like news or dates).
+      - MATH → if it involves math, time, or numeric calculation.
+      - LLM → if it includes a link or anything else.
+
+      Question: "{state['question']}"
+
+      Rules:
+      - Output only one word: WIKI, WebInfo, MATH, or LLM.
+      - Do NOT write any explanation, punctuation, or extra text.
+
+      Examples:
+      Input: "What is the capital of France?"
+      Output: WIKI
+
+      Input: "What's the weather today in Cairo?"
+      Output: WebInfo
+
+      Now respond with one word only.
+      """
     model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
     state["answer_type"] = (model.invoke([HumanMessage(content=prompt)]).content)