Update app.py

app.py

@@ -1,7 +1,7 @@
 import yaml
 import re
-from smolagents import HfApiModel  # Removed CodeAgent import.
-from tools.final_answer import FinalAnswerTool  # Retained in case you need it.
+from smolagents import HfApiModel
+from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
 # Updated system prompt: Only output the final, direct advice in plain text.
@@ -28,11 +28,14 @@ def remove_code_snippets(text):
     text = re.sub(r"\[Step \d+: Duration .*", "", text)
     # Remove lines that mention code snippet instructions.
     text = re.sub(r"Make sure to include code with the correct pattern.*", "", text)
-    # Remove any remaining lines that seem to be debug logs.
+    # Finally, remove any remaining lines that seem to be debug logs.
     lines = text.splitlines()
     cleaned_lines = [line for line in lines if not re.search(r"Step \d+|Duration", line)]
     return "\n".join(cleaned_lines).strip()
 
+# Use only the final_answer tool.
+final_answer = FinalAnswerTool()
+
 # Set up the model with a reduced token limit.
 model = HfApiModel(
     max_tokens=1024,
@@ -41,41 +44,34 @@ model = HfApiModel(
     custom_role_conversions=None,
 )
 
-# Load prompt templates if needed (unused here but kept for compatibility).
+# Load prompt templates from YAML.
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-def run_agent(user_input, stream=False, reset=False, additional_args=None):
+# Ensure the final_answer key exists in prompt_templates to prevent KeyError.
+if "final_answer" not in prompt_templates:
+    prompt_templates["final_answer"] = {"pre_messages": "", "post_messages": ""}
+
+# Initialize CodeAgent with a low verbosity level to reduce extra debug output.
+
+def run_agent(user_input):
     """
-    Constructs a prompt from the system prompt and the user input,
-    then directly calls the model to generate an answer.
-    After generation, it removes any chain-of-thought, debug logs, or code snippets.
-    Accepts extra keyword arguments to be compatible with Gradio's expectations.
+    Runs the agent, then removes any internal chain-of-thought, step logs, and code snippets
+    before returning the final plain-text answer.
     """
-    prompt = f"{system_prompt}\nUser Input: {user_input}\nFinal Answer:"
-    try:
-        raw_response = model.run(prompt)
-    except AttributeError:
-        raw_response = model(prompt)
-    
-    print("Raw Model Response:", raw_response)  # Debug output (optional)
+    raw_response = agent.run(user_input)
+    print("Raw Agent Response:", raw_response)  # Debugging output (optional)
     
     if not raw_response.strip():
         return "I'm sorry, but I couldn't generate a response. Please try again."
+    if "final_answer" not in raw_response.lower():
+        return "Error: The response did not use the `final_answer` tool. Please try again."
     
     clean_response = remove_code_snippets(raw_response)
     words = clean_response.split()
-    if len(set(words)) < 5:  # Check for repetition or overly brief output.
+    if len(set(words)) < 5:
         return "I'm unable to generate a meaningful response. Please refine your query."
     return clean_response
 
-class SimpleAgent:
-    def __init__(self, run_function):
-        self.run_function = run_function
-
-    def run(self, *args, **kwargs):
-        return self.run_function(*args, **kwargs)
-
-agent = SimpleAgent(run_agent)
-
+# Launch the Gradio UI.
 GradioUI(agent).launch()