from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFaceHub
import os
from dotenv import load_dotenv
from langchain import PromptTemplate

from langchain_openai import OpenAI


from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain.smith import RunEvalConfig

from langchain.globals import set_debug
import streamlit as st

set_debug(True)



load_dotenv()

# Load, chunk and index the contents of the blog.
raw_documents = TextLoader('data.txt').load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
documents = text_splitter.split_documents(raw_documents)


vectorstore = Chroma.from_documents(documents=documents, embedding=HuggingFaceEmbeddings())

# Retrieve and generate using the relevant snippets of the blog.
retriever = vectorstore.as_retriever()


# query it
#query = "comment gagner des points ?"
#docs = vectorstore.similarity_search(query)




repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
llm = HuggingFaceHub(
    repo_id=repo_id,
    model_kwargs={"temperature": 0.6, "top_k": 2, "max_tokens":2},
    huggingfacehub_api_token=os.getenv("HUGGINGFACE_API_KEY"),
)
llm = OpenAI(openai_api_key=os.getenv("OPENAI_API_KEY"), model_name="gpt-3.5-turbo-instruct")



template = """

Vous êtes un chatbot de service client. Utilisez le morceau de contexte suivant pour répondre à la question. Si vous ne connaissez pas la réponse, dites simplement que vous ne savez pas. Gardez la réponse concise.

Contexte: {context}
Question: {question}



"""

prompt = PromptTemplate(template=template, input_variables=["context", "question"])


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)





def main():
    st.set_page_config(page_title="Kenzup Chatbot")
    with st.sidebar:
        st.title('Kenzup Chatbot')

    # Function for generating LLM response
    def generate_response(input):
        result = rag_chain.invoke(input)
        return result

    # Store LLM generated responses
    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": "Bonjour ! En quoi puis-je vous aider aujourd'hui ?"}]

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    # User-provided prompt
    if input := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": input})
        with st.chat_message("user"):
            st.write(input)

    # Generate a new response if last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant"):
            with st.spinner("Getting your answer from mystery stuff.."):
                response = generate_response(input) 
                st.write(response) 
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)



main()