Update app.py

app.py

@@ -160,7 +160,7 @@ def attack_sus(message, max_tokens, temperature, top_p):
         top_p=top_p,
     )
 
-    return output
+    return message.choices[0].delta.content, end=""
 
 def respond(
     message,
@@ -175,9 +175,9 @@ def respond(
     suspicious = getbool(clasif_client.predict(text=message, api_name="/predict"))
 
     if suspicious == "POSITIVE":
-        for message in attack_sus(message, max_tokens, temperature, top_p):
-            print("pos", message)
-            yield message
+        message = attack_sus(message, max_tokens, temperature, top_p):
+        print("pos", message)
+        return message
     elif suspicious == "NEGATIVE":
         messages = [{"role": "system", "content": system_message}]
 
@@ -189,15 +189,16 @@ def respond(
 
         messages.append({"role": "user", "content": message})
 
-        for message in llm_client.chat_completion(
-                messages,
-                model="meta-llama/Llama-3.2-11B-Vision-Instruct",
-                max_tokens=max_tokens,
-                stream=False,
-                temperature=temperature,
-                top_p=top_p):
-            print("neg: ", message)
-            yield message
+        message = llm_client.chat_completion(
+                    messages,
+                    model="meta-llama/Llama-3.2-11B-Vision-Instruct",
+                    max_tokens=max_tokens,
+                    stream=False,
+                    temperature=temperature,
+                    top_p=top_p
+                ):
+        print("neg: ", message.choices[0].delta.content, end="")
+        return message.choices[0].delta.content, end=""
     else:
         raise Exception(f"Invalid suspicious type:{suspicious}")