Update agent.py

agent.py

@@ -33,6 +33,17 @@ class BraveSearchTool:
         except Exception as e:
             return f"BraveSearchTool ERROR: {str(e)}"
 
+def classify_question_type(question: str) -> str:
+    q = question.lower()
+    if any(k in q for k in ["spreadsheet", "excel", "csv", "table", "data", "json", "file attached"]):
+        return "file"
+    elif any(k in q for k in ["calculate", "total", "sum", "difference", "convert", "how many", "what is the number"]):
+        return "math"
+    elif any(k in q for k in ["wikipedia", "who", "what", "where", "when", "name", "define", "explain"]):
+        return "knowledge"
+    else:
+        return "search"
+
 class Agent:
     def __init__(self):
         gemini_key = os.getenv("GEMINI_API_KEY")
@@ -79,6 +90,9 @@ class Agent:
         file_names = input_data.get("file_names", [])
         task_id = input_data.get("task_id", "")
 
+        task_type = classify_question_type(question)
+
+        type_prefix = f"[Task Type: {task_type.upper()}]\n\n"
         system_prompt = (
             "You are a member of a multidisciplinary research institute, tackling complex and ambiguous problems across knowledge, reasoning, and vision.\n\n"
             "You have access to tools like search engines, calculators, and data analysis environments. Your task is to solve the following question carefully and completely.\n\n"
@@ -126,7 +140,7 @@ class Agent:
                 else:
                     file_summary = "(Unsupported file type — skipping file content.)"
 
-            full_prompt = system_prompt + file_summary + f"\n\nTASK:\n{question}"
+            full_prompt = type_prefix + system_prompt + file_summary + f"\n\nTASK:\n{question}"
             result = self.agent.run(full_prompt)
             return result.strip()
         except Exception as e: