Update agent.py

agent.py

@@ -31,7 +31,7 @@ class InfoState(TypedDict):
 def get_wiki_relate(state: InfoState) -> InfoState:
     """
     Tool to Get the wikipedia info from keywords extracted from preprocessing at main_parts.
-    
+
     Uses: Wikipedia API
     Returns: tool_answer (summary)
     """
@@ -85,17 +85,17 @@ def get_code(state:InfoState) -> InfoState:
     response = model.invoke([HumanMessage(content=prompt)]).content.strip()
 
     state["answer_code"] = response
-    
+
 
     return state
 
 
 
-def preprocess_text(state: dict) -> InfoState: 
-    
+def preprocess_text(state: dict) -> InfoState:
+
     """
     Preprocess text to get the keywords to help get results directly from wikipedia.
-    
+
     Input: raw question
     Output: main_parts (list of keywords)
     """
@@ -150,7 +150,7 @@ def preprocess_text(state: dict) -> InfoState:
 def get_answer(state: InfoState) -> InfoState :
     """
     Final Node that returns the final answer organized.
-    
+
     Combines: tool_answer → final_answer
     """
     print("Getting Answer (Gemini)...")
@@ -165,7 +165,9 @@ def get_answer(state: InfoState) -> InfoState :
         "- 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"
+        "- If the question asks about an attached file or audio, reply briefly that you cannot access attachments or audio files."
+        "- If absolutely nothing is relevant, reply: I'm not sure because the question depends on a file or missing data.\n\n"
+        "- Final answer should be complete text not part of answer"
         "Final Answer:"
     )
 
@@ -179,28 +181,22 @@ def get_type(state: InfoState) -> InfoState:
     print("Getting Type (Gemini)...")
 
     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']}"
+    You are a strict classifier.
+    Your job is to classify the following question into ONE of four categories:
 
-      Rules:
-      - Output only one word: WIKI, WebInfo, MATH, or LLM.
-      - Do NOT write any explanation, punctuation, or extra text.
+    - 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
 
-      Examples:
-      Input: "What is the capital of France?"
-      Output: WIKI
+    Question: "{state['question']}"
 
-      Input: "What's the weather today in Cairo?"
-      Output: WebInfo
+    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.
+    """
 
-      Now respond with one word only.
-      """
     model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
     state["answer_type"] = (model.invoke([HumanMessage(content=prompt)]).content)
 
@@ -260,7 +256,7 @@ def get_graph():
         )
 
         # Add final edges
-        graph.add_edge("get_search_results", "get_answer")  
+        graph.add_edge("get_search_results", "get_answer")
         #graph.add_edge("get_wiki_relate", "get_answer")
 
         graph.add_edge("get_code", "execute_code")
@@ -279,4 +275,4 @@ def ask(compiled_graph,question):
 
         })
 
-        return legitimate_result['final_answer']
+        return legitimate_result['final_answer'],legitimate_result