import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")

def chat_llama(chat_history):
    chat_completion = client.chat_completion(
        messages=chat_history,
        max_tokens=500,
        # stream=True
    )
    # Append the assistant's response to chat history
    chat_history.append({"role": "assistant", "content": chat_completion.choices[0].message.content})
    
    # Format chat_history as a list of tuples
    formatted_history = [(chat_history[i*2]["content"], chat_history[i*2+1]["content"]) 
                         for i in range(len(chat_history)//2)]
    
    return formatted_history

def chat_mem(message, chat_history):
    chat_history_role = [{"role": "system", "content": "You are a helpful assistant."}]
    
    if chat_history:
        for user_message, assistant_message in chat_history:
            chat_history_role.append({"role": "user", "content": user_message})
            chat_history_role.append({"role": "assistant", "content": assistant_message})
    
    chat_history_role.append({"role": "user", "content": message})
    
    chat_completion = client.chat_completion(
        messages=chat_history_role,
        max_tokens=500,
        # stream=True
    )
    assistant_message = chat_completion.choices[0].message.content
    chat_history_role.append({"role": "assistant", "content": assistant_message})
    
    # Format chat_history as a list of tuples
    formatted_history = [(chat_history_role[i*2]["content"], chat_history_role[i*2+1]["content"]) 
                         for i in range(len(chat_history_role)//2)]
    
    return "", formatted_history

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox(placeholder="Type a message...", interactive=True)
            with gr.Row():
                clear = gr.ClearButton([msg, chatbot], icon="https://img.icons8.com/?size=100&id=Xnx8cxDef16O&format=png&color=000000")
                send_btn = gr.Button("Send", variant='primary', icon="https://img.icons8.com/?size=100&id=g8ltXTwIfJ1n&format=png&color=000000")
            msg.submit(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])
            send_btn.click(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])
        with gr.Column():
            gr.Markdown("### API Testing")
            json_input = gr.Textbox(label="Input JSON", placeholder='Enter JSON here', lines=10)
            json_output = gr.Textbox(label="Output JSON", lines=10, interactive=False)
            test_btn = gr.Button("Test API")

            test_btn.click(fn=chat_llama, inputs=json_input, outputs=json_output)

if __name__ == "__main__":
    demo.launch()