Prompts

app.py

@@ -9,7 +9,7 @@ from smolagents import (CodeAgent, InferenceClientModel,
                 PythonInterpreterTool, VisitWebpageTool)
 from smolagents.agents import PromptTemplates
 import yaml
-from inference import InferenceClient
+
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -19,27 +19,12 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 with open("prompts.yaml", "r", encoding="utf-8") as f:
     prompts = yaml.safe_load(f)
 
-class InferenceClientModelWithSystemPrompt(InferenceClientModel):
-    def __init__(self, model: str, system_prompt: str = ""):
-        self.client = InferenceClient(model)
-        self.system_prompt = system_prompt
-        self.name = model
-
-    def run(self, messages, **kwargs):
-        # Prepend the system prompt
-        if isinstance(messages, str):
-            messages = [{"role": "user", "content": messages}]
-        if messages[0]["role"] != "system":
-            messages = [{"role": "system", "content": self.system_prompt}] + messages
-        return self.client.chat(messages, **kwargs)
-
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
-        self.model = InferenceClientModelWithSystemPrompt(
-            model="meta-llama/Llama-3-8B-Instruct",
-            system_prompt=prompts["system"]
+        self.model = InferenceClientModel(
+            model="meta-llama/Llama-3-8B-Instruct"
         )
         search_tool = DuckDuckGoSearchTool()
         wiki_tool = WikipediaSearchTool()
@@ -51,11 +36,16 @@ class BasicAgent:
             model=self.model,
             tools=[search_tool, wiki_tool, python_tool, visit_tool],
         )
+        self.system_prompt = prompts["system"]
         print("BasicAgent initialized.")
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        answer = self.agent.run(question)
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": question},
+        ]
+        answer = self.agent.run(messages)
         print(f"Agent returning answer: {answer}")
         # Remove "FINAL ANSWER:" if present
         if isinstance(answer, str) and "FINAL ANSWER:" in answer: