Update app.py

app.py

@@ -18,36 +18,51 @@ class BasicAgent:
         print("Loading FLAN-T5 base model...")
         self.pipeline = pipeline(
             "text2text-generation",
-            model="google/flan-t5-base",
+            model = "google/flan-t5-base",
             max_new_tokens=128,
             temperature=0.3
         )
         print("Model loaded.")
 
-    def __call__(self, question: str) -> str:
-        print(f"Received question: {question[:60]}...")
-
-        few_shot_example = (
-            "Question: List just the vegetables from [milk, eggs, carrots, onions, cookies].\n"
-            "Answer: carrots, onions\n\n"
-        )
-
-        prompt = (
-            few_shot_example +
-            "Please solve the following step by step and give only the final answer:\n"
-            f"{question.strip()}"
-        )
-
-        response = self.pipeline(prompt)[0]["generated_text"]
-        answer = response.strip().split("Answer:")[-1].strip().split("\n")[0]
-        print(f"Generated answer: {answer}")
-        return answer
-
-
-
-
+    def __call__(self, question: str, task_id: str = None) -> str:
+        # Custom routing logic
+        if "wikipedia" in question.lower() or "how many" in question.lower():
+            return self.search_wikipedia(question)
+        elif "attached Excel" in question.lower():
+            return self.parse_excel(task_id)
+        elif "reverse" in question.lower() or "write the opposite" in question.lower():
+            return self.reverse_sentence(question)
+        else:
+            return self.model_response(question)
+        def model_response(self, question: str) -> str:
+        prompt = f"Please answer step by step:\n{question}"
+        output = self.pipeline(prompt)[0]["generated_text"]
+        return output.strip().split("Answer:")[-1].strip()
+
+    def search_wikipedia(self, question: str) -> str:
+        import wikipedia
+        try:
+            summary = wikipedia.summary(question, sentences=2)
+            return summary
+        except Exception as e:
+            return f"Could not find info: {e}"
 
+    def parse_excel(self, task_id: str) -> str:
+        import requests
+        import pandas as pd
+        try:
+            file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+            df = pd.read_excel(file_url)
+            # Your logic here, example:
+            food_sales = df[df['category'] == 'food']['sales'].sum()
+            return f"${food_sales:.2f}"
+        except Exception as e:
+            return f"Excel parsing error: {e}"
 
+    def reverse_sentence(self, question: str) -> str:
+        sentence = question.split("write")[0].strip()
+        reversed_text = sentence[::-1]
+        return reversed_text
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
@@ -111,7 +126,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            submitted_answer = agent(question_text, task_id=task_id)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e: