Update agent.py

agent.py

@@ -1,6 +1,8 @@
 from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel # HfApiModel, OpenAIServerModel
+from tools.final_answer import FinalAnswerTool
 import asyncio
 import os
+import yaml
 
 class MagAgent:
     def __init__(self):
@@ -16,11 +18,16 @@ class MagAgent:
 #            model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
 #            custom_role_conversions=None,
         )
+        # Load prompt templates
+        with open("prompts.yaml", 'r') as stream:
+            prompt_templates = yaml.safe_load(stream)
+        
         self.agent = CodeAgent(
             model= model,  # OpenAIServerModel(model_id="gpt-4o"), 
             tools=[
                 DuckDuckGoSearchTool(),
-                WikipediaSearchTool()
+                WikipediaSearchTool(),
+                FinalAnswerToolnswer()
             ]
         )
         print("MagAgent initialized.")
@@ -29,14 +36,22 @@ class MagAgent:
         """Process a question asynchronously using the MagAgent."""
         print(f"MagAgent received question (first 50 chars): {question[:50]}...")
         try:
-            # Run the agent asynchronously
+            # Define a task with fallback search logic
+            task = (
+                f"Answer the following question accurately and concisely: {question}\n"
+                "First, try searching Wikipedia with 'Mercedes Sosa'. If that fails, "
+                "use DuckDuckGo to search 'Mercedes Sosa discography 2000-2009'."
+            )
             response = await asyncio.to_thread(
                 self.agent.run,
-                task=f"Answer the following question accurately and concisely: {question}"
+                task=task
             )
+            if not response or "No Wikipedia page found" in response:
+                # Fallback response if search fails
+                response = "Unable to retrieve exact data. Please refine the question or check external sources."
             print(f"MagAgent response: {response[:50]}...")
             return response
         except Exception as e:
-            error_msg = f"Error processing question: {e}"
+            error_msg = f"Error processing question: {str(e)}. Check API key or network connectivity."
             print(error_msg)
             return error_msg