import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Access the value of OPENAI_API_KEY
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0,)
## Create embeddings and splitter

from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Create Embeddings
model_name = "BAAI/bge-large-en"

embedding = HuggingFaceBgeEmbeddings(
    model_name = model_name,
    model_kwargs = {'device':'cuda'},
    encode_kwargs = {'normalize_embeddings': True}
)

# Create Splitter
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,
)

from langchain_community.vectorstores import FAISS

persitsdirectory="/NLP/ATrad/ATrad/db/faiss_Test02_500_C_BGE_large"

vectorstore= FAISS.load_local(persitsdirectory, embedding) 

# Define a custom prompt for Unser manual
from langchain.prompts import PromptTemplate

qa_template = ("""
    You are the AI assistant of the IronOne Technologies which provide services for companies members and  novice users with learning with ATrad Aplication .
    You have provided context information below related to learning material.
    
    Context: {context}
    
    Given this information, please answer the question with the latest information.
    If you dont know the answer say you dont know, dont try to makeup answers.
    if context is not enough to answer the question, ask for more information.
    if context is not related to the question, say I dont know.
    
    each answer should start with code word ATrad Ai(QA):
    
    Question: {question}
    
    answer: let me think about it...""")

qa_template2 = ("""
    Welcome to IronOne Technologies' AI Assistant, designed to assist you in learning with the ATrad Application.
    
    Context: {context}
    
    As your AI assistant, I'm here to help you navigate through learning materials and provide guidance.
    Please provide me with any questions or concerns you have regarding the ATrad Application.
    If you're unsure about something or need more information, feel free to ask.
    
    Question: {question}
    
    ATrad Ai(QA): Let me think about it...""")


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


# Chain for Web 
from langchain.chains import RetrievalQA

Web_qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever = vectorstore.as_retriever(search_kwargs={"k": 4}),
    return_source_documents= True,
    input_key="question",
    chain_type_kwargs={"prompt": QA_PROMPT},
)