import gradio as gr
from llama_cpp import Llama

llm = Llama(model_path="model.gguf", n_ctx=8000, n_threads=2, chat_format="chatml")
  
def generate(message, history, temperature=0.3, max_tokens=512):
    system_prompt = """You are a SQL virtual assistant, you will only create queries thinking step by step. Check that the syntax is perfect and don't miss any character. Pay attention to the names of the tables and fields. Do not make up fields that do not exist. I want you to give the query only. Do not speak and do not explain anything. Just provide the queries with no further words."""
    formatted_prompt = [{"role": "system", "content": system_prompt}]
    for user_prompt, bot_response in history:
        formatted_prompt.append({"role": "user", "content": user_prompt})
        formatted_prompt.append({"role": "assistant", "content": bot_response})
    formatted_prompt.append({"role": "user", "content": message})
    stream_response = llm.create_chat_completion(messages=formatted_prompt, temperature=temperature, max_tokens=max_tokens, stream=True)
    
    response = ""
    for chunk in stream_response:
        if len(chunk['choices'][0]["delta"]) != 0 and "content" in chunk['choices'][0]["delta"]:
            chunk_content = chunk['choices'][0]["delta"]["content"]
            stop_sequence_index = chunk_content.find(';---')  # Busca el índice del patrón de detención
            if stop_sequence_index != -1:  # Si se encuentra el patrón, recorta la respuesta
                response += chunk_content[:stop_sequence_index]
                break  # Termina el bucle ya que encontraste el punto de detención
            else:
                response += chunk_content
    yield response


mychatbot = gr.Chatbot(
avatar_images=["user.png", "botnb.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,)
        
iface = gr.ChatInterface(fn=generate, chatbot=mychatbot, retry_btn=None, undo_btn=None)

with gr.Blocks() as demo:
    gr.HTML("<center><h1>Natural SQL</h1></center>")
    iface.render()

demo.queue().launch(show_api=False, server_name="0.0.0.0")