__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.messages import HumanMessage
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import gradio as gr

load_dotenv()

# OpenAI's GPT model to use
GPT_VERSION = "gpt-3.5-turbo-0125"

# Prompt used to contextualize a question inside the context of a conversation
CONTEXTUALIZATION_PROMPT = """Given a chat history and the latest user question \
which might reference context in the chat history, formulate a standalone question \
which can be understood without the chat history. Do NOT answer the question, \
just reformulate it if needed and otherwise return it as is."""

# Prompt used for building an answer to a question using the retrieved context
SYSTEM_PROMPT = """You are a marketing assistant for question-answering tasks for a company called Tryolabs. \
You must answer the question using pieces of retrieved context consisting of blogposts written by Tryolabs. \
Answer the question in a friendly manner, trying to refer to Tryolabs' content to engage the user and encourage \
to make contact with the company. \
When you can, include properly formatted hyperlinks to the content. TRY providing references to different content \
within the same response. Extract the hyperlink from the context. \
If you don't know the answer to the question, encourage the user to contact us for more information. \
Tryolabs' contact email is hello@tryolabs.com. \
Keep your answers concise. \
ALWAYS respond with Markdown. \
You MUST only use information found on the pieces of context given. You MUST NOT attribute any inventions to \
Tryolabs unless explicitly mentioned on the context. \

{context}"""

llm = ChatOpenAI(model=GPT_VERSION)

vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=OpenAIEmbeddings())

retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.3})

contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", CONTEXTUALIZATION_PROMPT),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)

custom_rag_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", SYSTEM_PROMPT),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}")
    ]
)

document_prompt = PromptTemplate(
    input_variables=['page_content', 'title', 'link'], 
    template="Title: {title}, Link to blogpost: {link}, Content: {page_content}"
)
qa_chain = create_stuff_documents_chain(llm, custom_rag_prompt, document_prompt=document_prompt)
rag_chain = create_retrieval_chain(history_aware_retriever, qa_chain)


def predict(message, history):
    chat_history = []

    for m, a in history:
        chat_history.append(HumanMessage(content=m))
        chat_history.append(a)

    ai_response = rag_chain.invoke({"input": message, "chat_history": chat_history})

    return ai_response["answer"]

theme = gr.themes.Default().set(
    button_primary_background_fill_dark="#26C8A1",
    button_primary_background_fill_hover_dark="#3cd6b2",
    button_primary_border_color_dark="#26C8A1",
    button_primary_text_color_dark="#0B0723",
    background_fill_primary_dark="#0B0723",
)

gr.ChatInterface(
    fn=predict,
    examples=[
        "What experience do you have working with pricing optimization?",
        "How can I build my own MLOps pipeline?",
        "How can I fine-tune my own LLMs?"
    ],
    theme=theme
).launch()