try:
  from langchain_community.vectorstores import Chroma
except:
  from langchain_community.vectorstores import Chroma


from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
from langchain import PromptTemplate
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
import torch






import os
import requests  # Or your Groq library

groq_api_key = os.environ.get("my_groq_api_key")

# Initialize a ChatGroq object with a temperature of 0 and the "mixtral-8x7b-32768" model.
llm = ChatGroq(temperature=0, model_name="llama3-70b-8192",api_key=groq_api_key)

model_name = "BAAI/bge-m3"
device = "cuda" if torch.cuda.is_available() else "cpu"
embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': device})
# we run this cell every time
db = Chroma(embedding_function=embeddings, persist_directory='Master-Guidline')



memory = ConversationBufferWindowMemory(
    memory_key="history", k=3, return_only_outputs=True
)




template = """
فقط بر اساس محتوای داده شده، به سوال مطرح شده جواب بده. جواب باید به زبان فارسی باشه. اگه سوال مطرح شده خارج از محتوای داده شده هست، بگو که من فقط به سوالات مرتبط با تحصبلات تکمیلی جواب میدم
محتوا:
{history}

"""
s="""

\n سوال: {input}

\n جواب:""".strip()


prompt = PromptTemplate(input_variables=["history", "input"], template=template+s)






chain = ConversationChain(
    llm=llm,

    prompt=prompt,
    memory=memory,
    verbose=True,
)



# Generate a response from the Llama model
def get_llama_response(message: str, history: list) -> str:
    """
    Generates a conversational response from the Llama model.

    Parameters:
        message (str): User's input message.
        history (list): Past conversation history.

    Returns:
        str: Generated response from the Llama model.
    """
    query_text =message

    results = db.similarity_search_with_relevance_scores(query_text, k=3)
    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])




    template = """
    فقط بر اساس محتوای داده شده، به سوال مطرح شده جواب بده. جواب باید به زبان فارسی باشه. اگه سوال مطرح شده خارج از محتوای داده شده هست، بگو که من فقط به سوالات مرتبط با تحصبلات تکمیلی جواب میدم
    محتوا:
    {history}

    """
    s="""
    \n سوال: {input}

    \n جواب:""".strip()
    prompt = PromptTemplate(input_variables=["history", "input"], template=template+context_text+'\n'+s)
    

    #print(template)
    chain.prompt=prompt
    res = chain.predict(input=query_text)
    return res



import gradio as gr
iface = gr.Interface(fn=get_llama_response, inputs=gr.Textbox(),
             outputs="textbox")
iface.launch(share=True)