Update app.py

app.py

@@ -4,6 +4,7 @@ import requests
 import inspect
 import pandas as pd
 import wikipedia
+from transformers import pipeline
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -24,14 +25,27 @@ class EchoAgent:
 
 class WikipediaAgent:
     def __init__(self):
-        print("WikipediaAgent initialized.")
+        wikipedia.set_lang("en")
     def __call__(self, question: str) -> str:
         try:
-            summary = wikipedia.summary(question, sentences=1)
-            return summary
+            # Extract potential keyword
+            query = question.split(" ")[-1]  # naive keyword pick
+            page = wikipedia.page(query)
+            return wikipedia.summary(page.title, sentences=1)
         except Exception as e:
-            return f"Error: {str(e)}"
+            return "Could not answer"
+class TransformersAgent:
+    def __init__(self):
+        print("Loading transformers pipeline (flan-t5-base)...")
+        self.model = pipeline("text2text-generation", model="google/flan-t5-base", device=-1)  # CPU only
 
+    def __call__(self, question: str) -> str:
+        try:
+            response = self.model(question, max_new_tokens=100)
+            return response[0]['generated_text'].strip()
+        except Exception as e:
+            return f"Error: {str(e)}"
+            
 def run_and_submit_all(agent_name: str, profile: gr.OAuthProfile | None = None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
@@ -59,6 +73,8 @@ def run_and_submit_all(agent_name: str, profile: gr.OAuthProfile | None = None):
             agent = EchoAgent()
         elif agent_name == "WikipediaAgent":
             agent = WikipediaAgent()
+        elif agent_name == "TransformersAgent":
+            agent = TransformersAgent()
         else:
             return f"Unknown agent selected: {agent_name}", None
     except Exception as e:
@@ -167,7 +183,7 @@ with gr.Blocks() as demo:
     gr.LoginButton()
     
     agent_dropdown = gr.Dropdown(
-        choices=["BasicAgent", "EchoAgent", "WikipediaAgent"],
+        choices=["BasicAgent", "EchoAgent", "WikipediaAgent", "TransformersAgent"],
         label="Select Agent",
         value="BasicAgent"
     )