import gradio as gr
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import FAISS
from langchain_community.chat_models.huggingface import ChatHuggingFace
from langchain.schema import SystemMessage, HumanMessage, AIMessage
from langchain_community.llms import HuggingFaceEndpoint

model_name = "sentence-transformers/all-mpnet-base-v2"
embedding_llm = SentenceTransformerEmbeddings(model_name=model_name)

db = FAISS.load_local("faiss_index", embedding_llm, allow_dangerous_deserialization=True)

# Set up Hugging Face model
llm = HuggingFaceEndpoint(
    repo_id="HuggingFaceH4/starchat2-15b-v0.1",
    task="text-generation",
    max_new_tokens=4096,
    temperature=0.6,
    top_p=0.9,
    top_k=40,
    repetition_penalty=1.2,
    do_sample=True,
)
chat_model = ChatHuggingFace(llm=llm)

messages = [
    SystemMessage(content="You are a helpful assistant."),
    HumanMessage(content="Hi AI, how are you today?"),
    AIMessage(content="I'm great thank you. How can I help you?")
]

def handle_query(mode: str, query: str):
    if mode == "chat":
        return chat_mode(query)
    elif mode == "web-search":
        return web_search(query)
    else:
        return "Invalid mode selected."

def chat_mode(query: str):
    prompt = HumanMessage(content=query)
    messages.append(prompt)
    response = chat_model.invoke(messages)
    messages.append(response.content)
    if len(messages) >= 6:
        messages = messages[-6:]
    return f"You: {query}\nIT-Assistant: {response.content}"

def web_search(query: str):
    similar_docs = db.similarity_search(query, k=3)
    if similar_docs:
        source_knowledge = "\n".join([x.page_content for x in similar_docs])
    else:
        source_knowledge = ""
    augmented_prompt = f"""
    Answer the next query using the provided Web Search.
    If the answer is not contained in the Web Search, ignore the web search and respond independently with your own knowledge.

    Query: {query}

    Web Search:
    {source_knowledge}
    """
    prompt = HumanMessage(content=augmented_prompt)
    messages.append(prompt)
    response = chat_model.invoke(messages)
    messages.append(response.content)
    if len(messages) >= 6:
        messages = messages[-6:]
    return f"You: {query}\nIT-Assistant: {response.content}"

demo = gr.Interface(fn=handle_query, inputs=["text", "text"], outputs="text", title="IT Assistant", description="Choose a mode and enter your message, then click submit to interact.", choices=["chat", "web-search"])
demo.launch()