Update app.py

app.py

@@ -8,21 +8,26 @@ from transformers import pipeline
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
-class BasicAgent:
+class GeneralAgent:
     def __init__(self):
-        print("Loading QA pipeline...")
+        print("Initializing general QA agent...")
         self.qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
     
     def __call__(self, question: str, context: str = None) -> str:
-        try:
-            if context is None:
-                context = question
-            result = self.qa_pipeline(question=question, context=context)
-            answer = result["answer"]
-        except Exception as e:
-            print(f"Error during QA: {e}")
-            answer = "Error processing question."
-        return answer
+        """
+        Process the question and return an answer based on the given context.
+        Uses a prompt template to provide a clear structure for the answer.
+        """
+        prompt = """
+You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+        """
+        question_with_prompt = f"{prompt}\nQuestion: {question}\nContext: {context if context else 'No context provided.'}\n"
+        
+        result = self.qa_pipeline(question=question_with_prompt, context=context)
+        answer = result["answer"]
+        
+        # This is where you can modify the final format of the answer if needed.
+        return f"FINAL ANSWER: {answer}"
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")