fun 3 menos restrictivo

app.py

@@ -17,7 +17,7 @@ def format_prompt(message, history):
     return prompt
 
 def replace_variables_with_placeholders(query):
-    # Buscar patrones de variables en la cláusula WHERE y reemplazarlos por ?
+    # Reemplazar variables en la cláusula WHERE con ?
     return re.sub(r"= '.*?'", "= ?", query)
 
 def generate(prompt, history, temperature=0.2, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
@@ -47,13 +47,16 @@ def generate(prompt, history, temperature=0.2, max_new_tokens=256, top_p=0.95, r
 
     for response in stream:
         output += response.token.text
-        # Verificar si la respuesta contiene la consulta deseada
-        if "SELECT" in output.upper() and "FROM" in output.upper():
-            # Reemplazar variables en la cláusula WHERE con ?
-            output = replace_variables_with_placeholders(output)
-            break
 
-    yield output.strip()
+    # Reemplazar variables en la cláusula WHERE con ?
+    final_output = replace_variables_with_placeholders(output)
+    
+    # Verificar si la respuesta contiene la consulta deseada y devolver solo la consulta
+    if "SELECT" in final_output.upper() and "FROM" in final_output.upper():
+        final_output = final_output.strip().split("\n")[0]
+        final_output = "\n".join(line for line in output.split("\n") if line.strip().startswith(("SELECT", "FROM", "JOIN", "WHERE")))
+
+    yield final_output.strip()
 
 mychatbot = gr.Chatbot(
     avatar_images=["./user.png", "./botm.png"],