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='/Persian Chroma/')



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+context_text+'\n'+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=2)
    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
        #return response.strip()



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