Update app.py

app.py

@@ -13,6 +13,13 @@ from huggingface_hub import InferenceClient, hf_hub_download
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
+class LocalLLM:
+    def __init__(self, pipe):
+        self.pipe = pipe
+
+    def generate(self, prompt, **kwargs):
+        return self.pipe(prompt)[0]["generated_text"]
+
 @spaces.GPU
 def load_llm(hf_token):
     model_id = "meta-llama/Llama-3.1-8B-Instruct"
@@ -72,7 +79,7 @@ class BasicAgent:
         
         # Initialize the agent
         pipe = load_llm(hf_token = os.environ["HF_TOKEN"])
-        self.pipe = FunctionLLM(lambda prompt: pipe(prompt)[0]["generated_text"])
+        self.pipe = LocalLLM(pipe)
         self.agent = CodeAgent(llm=self.pipe, tools=[], add_base_tools=True)
         
         # self.agent = CodeAgent(