# import os
# import gradio as gr
# from langchain.chat_models import ChatOpenAI
# from langchain.prompts import PromptTemplate
# from langchain.chains import LLMChain
# from langchain.memory import ConversationBufferMemory

# # Set OpenAI API Key
# OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

# # Define the template for the chatbot's response
# template = """You are a helpful assistant to answer all user queries.
# {chat_history}
# User: {user_message}
# Chatbot:"""

# # Define the prompt template
# prompt = PromptTemplate(
#     input_variables=["chat_history", "user_message"], 
#     template=template
# )

# # Initialize conversation memory
# memory = ConversationBufferMemory(memory_key="chat_history")

# # Define the LLM chain with the ChatOpenAI model and conversation memory
# llm_chain = LLMChain(
#     llm=ChatOpenAI(temperature=0.5, model="gpt-3.5-turbo"),  # Use 'model' instead of 'model_name'
#     prompt=prompt,
#     verbose=True,
#     memory=memory,
# )

# # Function to get chatbot response
# def get_text_response(user_message, history):
#     response = llm_chain.predict(user_message=user_message)
#     return response

# # Create a Gradio chat interface
# demo = gr.Interface(fn=get_text_response, inputs="text", outputs="text")

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

# import os
# import gradio as gr
# from langchain.chat_models import ChatOpenAI
# from langchain.schema import AIMessage, HumanMessage

# # Set OpenAI API Key
# os.environ["OPENAI_API_KEY"] = "sk-3_mJiR5z9Q3XN-D33cgrAIYGffmMvHfu5Je1U0CW1ZT3BlbkFJA2vfSvDqZAVUyHo2JIcU91XPiAq424OSS8ci29tWMA"  # Replace with your key

# # Initialize the ChatOpenAI model
# llm = ChatOpenAI(temperature=1.0, model="gpt-3.5-turbo-0613")

# # Function to predict response
# def get_text_response(message, history=None):
#     # Ensure history is a list
#     if history is None:
#         history = []
    
#     # Convert the Gradio history format to LangChain message format
#     history_langchain_format = []
#     for human, ai in history:
#         history_langchain_format.append(HumanMessage(content=human))
#         history_langchain_format.append(AIMessage(content=ai))
    
#     # Add the new user message to the history
#     history_langchain_format.append(HumanMessage(content=message))

#     # Get the model's response
#     gpt_response = llm(history_langchain_format)

#     # Append AI response to history
#     history.append((message, gpt_response.content))

#     # Return the response and updated history
#     return gpt_response.content, history

# # Create a Gradio chat interface
# demo = gr.ChatInterface(
#     fn=get_text_response, 
#     inputs=["text", "state"], 
#     outputs=["text", "state"]
# )

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

# import os  # Import the os module
# import time
# import gradio as gr
# from langchain_community.chat_models import ChatOpenAI  # Updated import based on deprecation warning
# from langchain.schema import AIMessage, HumanMessage
# import openai

# # Set your OpenAI API key
# os.environ["OPENAI_API_KEY"] = "sk-3_mJiR5z9Q3XN-D33cgrAIYGffmMvHfu5Je1U0CW1ZT3BlbkFJA2vfSvDqZAVUyHo2JIcU91XPiAq424OSS8ci29tWMA"  # Replace with your OpenAI key

# # Initialize ChatOpenAI
# llm = ChatOpenAI(temperature=1.0, model='gpt-3.5-turbo-0613')

# def predict(message, history):
#     # Reformat history for LangChain
#     history_langchain_format = []
#     for human, ai in history:
#         history_langchain_format.append(HumanMessage(content=human))
#         history_langchain_format.append(AIMessage(content=ai))
    
#     # Add latest human message
#     history_langchain_format.append(HumanMessage(content=message))
    
#     # Get response from the model
#     gpt_response = llm(history_langchain_format)
    
#     # Return response
#     return gpt_response.content

# # Using ChatInterface to create a chat-style UI
# demo = gr.ChatInterface(fn=predict, type="messages")
# if __name__ == "__main__":
#     demo.launch()
import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


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