Update app.py

app.py

@@ -116,13 +116,7 @@ class BasicAgent:
         allowed = {"max_new_tokens", "temperature", "top_k", "top_p"}
         gen_kwargs = {k: v for k, v in kwargs.items() if k in allowed}
         
-        # 2. Optionally map SmolAgents’ stop_sequences → HF pipeline’s 'stop'
-        if stop_sequences is not None:
-            # if your pipeline accepts `stop`, you can do:
-            gen_kwargs["stop"] = stop_sequences
-            # otherwise just ignore it 
-        
-        # 3. Serialize the message
+        # 2. Serialize the message and get the response
         prompt_str = (
             self._serialize_messages(question)
             if isinstance(question, list)
@@ -131,6 +125,14 @@ class BasicAgent:
         outputs = self.pipe(prompt_str, **gen_kwargs)
         response = outputs[0]["generated_text"]
         # response = self.agent.run(question)
+
+        # 3. Optionally map SmolAgents’ stop_sequences → HF pipeline’s 'stop'
+        if stop_sequences:
+            # find the earliest occurrence of any stop token
+            cuts = [response.find(s) for s in stop_sequences if response.find(s) != -1]
+            if cuts:
+                response = response[: min(cuts)]
+        
         print(f"Agent returning its generated answer: {response}")
         
         # wrap back into a chat message dict