import logging
import os
import pickle
import tempfile

import streamlit as st
from dotenv import load_dotenv
from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams
from ibm_watson_machine_learning.foundation_models.utils.enums import ModelTypes
from langchain.callbacks import StdOutCallbackHandler
from langchain.chains.question_answering import load_qa_chain
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import (HuggingFaceHubEmbeddings,
                                  HuggingFaceInstructEmbeddings)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS, Chroma
from PIL import Image

from langChainInterface import LangChainInterface

# Most GENAI logs are at Debug level.
logging.basicConfig(level=os.environ.get("LOGLEVEL", "DEBUG"))

st.set_page_config(
    page_title="Retrieval Augmented Generation",
    page_icon="🧊",
    layout="wide",
    initial_sidebar_state="expanded"
)
st.header("Retrieval Augmented Generation with watsonx.ai 💬")
# chunk_size=1500
# chunk_overlap = 200

load_dotenv()

handler = StdOutCallbackHandler()

api_key = os.getenv("API_KEY", None)
ibm_cloud_url = os.getenv("IBM_CLOUD_URL", None)
project_id = os.getenv("PROJECT_ID", None)

if api_key is None or ibm_cloud_url is None or project_id is None:
    print("Ensure you copied the .env file that you created earlier into the same directory as this notebook")
else:
    creds = {
        "url": ibm_cloud_url,
        "apikey": api_key 
    }

GEN_API_KEY = os.getenv("GENAI_KEY", None)

# Sidebar contents
with st.sidebar:
    st.title("RAG App")
    st.markdown('''
    ## About
    This app is an LLM-powered RAG built using:
    - [IBM Generative AI SDK](https://github.com/IBM/ibm-generative-ai/)
    - [HuggingFace](https://huggingface.co/)
    - [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai) LLM model
 
    ''')
    st.write('Powered by [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai)')
    image = Image.open('watsonxai.jpg')
    st.image(image, caption='Powered by watsonx.ai')
    max_new_tokens= st.number_input('max_new_tokens',1,1024,value=300)
    min_new_tokens= st.number_input('min_new_tokens',0,value=15)
    repetition_penalty = st.number_input('repetition_penalty',1,2,value=2)
    decoding = st.text_input(
            "Decoding",
            "greedy",
            key="placeholder",
        )
    
uploaded_files = st.file_uploader("Choose a PDF file", accept_multiple_files=True)

@st.cache_data
def read_pdf(uploaded_files,chunk_size =250,chunk_overlap=20):
    for uploaded_file in uploaded_files:
      bytes_data = uploaded_file.read()
      with tempfile.NamedTemporaryFile(mode='wb', delete=False) as temp_file:
      # Write content to the temporary file
          temp_file.write(bytes_data)
          filepath = temp_file.name
          with st.spinner('Waiting for the file to upload'):
             loader = PyPDFLoader(filepath)
             data = loader.load()
             text_splitter = RecursiveCharacterTextSplitter(chunk_size= chunk_size, chunk_overlap=chunk_overlap)
             docs = text_splitter.split_documents(data)
             return docs

@st.cache_data
def read_push_embeddings():
    embeddings = HuggingFaceHubEmbeddings(repo_id="sentence-transformers/all-MiniLM-L6-v2")
    if os.path.exists("db.pickle"):
        with open("db.pickle",'rb') as file_name:
            db = pickle.load(file_name)
    else:     
        db = FAISS.from_documents(docs, embeddings)
        with open('db.pickle','wb') as file_name  :
             pickle.dump(db,file_name)
        st.write("\n")
    return db

# show user input
if user_question := st.text_input(
    "Ask a question about your Policy Document:"
):
    docs = read_pdf(uploaded_files)
    db = read_push_embeddings()
    docs = db.similarity_search(user_question)
    params = {
        GenParams.DECODING_METHOD: "greedy",
        GenParams.MIN_NEW_TOKENS: 30,
        GenParams.MAX_NEW_TOKENS: 300,
        GenParams.TEMPERATURE: 0.0,
        # GenParams.TOP_K: 100,
        # GenParams.TOP_P: 1,
        GenParams.REPETITION_PENALTY: 1
    }
    model_llm = LangChainInterface(model=ModelTypes.LLAMA_2_70B_CHAT.value, credentials=creds, params=params, project_id=project_id)
    chain = load_qa_chain(model_llm, chain_type="stuff")

    response = chain.run(input_documents=docs, question=user_question)

    st.text_area(label="Model Response", value=response, height=100)
    st.write()

langchain.chains.question_answering.load_qa_chain langchain.embeddings.HuggingFaceHubEmbeddings langchain.text_splitter.RecursiveCharacterTextSplitter langchain.callbacks.StdOutCallbackHandler langchain.vectorstores.FAISS.from_documents langchain.document_loaders.PyPDFLoader
[((861, 993), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""Retrieval Augmented Generation"""', 'page_icon': '"""🧊"""', 'layout': '"""wide"""', 'initial_sidebar_state': '"""expanded"""'}), "(page_title='Retrieval Augmented Generation', page_icon=\n    '🧊', layout='wide', initial_sidebar_state='expanded')\n", (879, 993), True, 'import streamlit as st\n'), ((1007, 1068), 'streamlit.header', 'st.header', (['"""Retrieval Augmented Generation with watsonx.ai 💬"""'], {}), "('Retrieval Augmented Generation with watsonx.ai 💬')\n", (1016, 1068), True, 'import streamlit as st\n'), ((1110, 1123), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (1121, 1123), False, 'from dotenv import load_dotenv\n'), ((1135, 1158), 'langchain.callbacks.StdOutCallbackHandler', 'StdOutCallbackHandler', ([], {}), '()\n', (1156, 1158), False, 'from langchain.callbacks import StdOutCallbackHandler\n'), ((1170, 1196), 'os.getenv', 'os.getenv', (['"""API_KEY"""', 'None'], {}), "('API_KEY', None)\n", (1179, 1196), False, 'import os\n'), ((1213, 1245), 'os.getenv', 'os.getenv', (['"""IBM_CLOUD_URL"""', 'None'], {}), "('IBM_CLOUD_URL', None)\n", (1222, 1245), False, 'import os\n'), ((1259, 1288), 'os.getenv', 'os.getenv', (['"""PROJECT_ID"""', 'None'], {}), "('PROJECT_ID', None)\n", (1268, 1288), False, 'import os\n'), ((1566, 1594), 'os.getenv', 'os.getenv', (['"""GENAI_KEY"""', 'None'], {}), "('GENAI_KEY', None)\n", (1575, 1594), False, 'import os\n'), ((2468, 2533), 'streamlit.file_uploader', 'st.file_uploader', (['"""Choose a PDF file"""'], {'accept_multiple_files': '(True)'}), "('Choose a PDF file', accept_multiple_files=True)\n", (2484, 2533), True, 'import streamlit as st\n'), ((1636, 1655), 'streamlit.title', 'st.title', (['"""RAG App"""'], {}), "('RAG App')\n", (1644, 1655), True, 'import streamlit as st\n'), ((1660, 1949), 'streamlit.markdown', 'st.markdown', (['"""\n    ## About\n    This app is an LLM-powered RAG built using:\n    - [IBM Generative AI SDK](https://github.com/IBM/ibm-generative-ai/)\n    - [HuggingFace](https://huggingface.co/)\n    - [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai) LLM model\n \n    """'], {}), '(\n    """\n    ## About\n    This app is an LLM-powered RAG built using:\n    - [IBM Generative AI SDK](https://github.com/IBM/ibm-generative-ai/)\n    - [HuggingFace](https://huggingface.co/)\n    - [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai) LLM model\n \n    """\n    )\n', (1671, 1949), True, 'import streamlit as st\n'), ((1944, 2029), 'streamlit.write', 'st.write', (['"""Powered by [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai)"""'], {}), "('Powered by [IBM watsonx.ai](https://www.ibm.com/products/watsonx-ai)'\n    )\n", (1952, 2029), True, 'import streamlit as st\n'), ((2037, 2064), 'PIL.Image.open', 'Image.open', (['"""watsonxai.jpg"""'], {}), "('watsonxai.jpg')\n", (2047, 2064), False, 'from PIL import Image\n'), ((2069, 2117), 'streamlit.image', 'st.image', (['image'], {'caption': '"""Powered by watsonx.ai"""'}), "(image, caption='Powered by watsonx.ai')\n", (2077, 2117), True, 'import streamlit as st\n'), ((2138, 2191), 'streamlit.number_input', 'st.number_input', (['"""max_new_tokens"""', '(1)', '(1024)'], {'value': '(300)'}), "('max_new_tokens', 1, 1024, value=300)\n", (2153, 2191), True, 'import streamlit as st\n'), ((2209, 2255), 'streamlit.number_input', 'st.number_input', (['"""min_new_tokens"""', '(0)'], {'value': '(15)'}), "('min_new_tokens', 0, value=15)\n", (2224, 2255), True, 'import streamlit as st\n'), ((2279, 2331), 'streamlit.number_input', 'st.number_input', (['"""repetition_penalty"""', '(1)', '(2)'], {'value': '(2)'}), "('repetition_penalty', 1, 2, value=2)\n", (2294, 2331), True, 'import streamlit as st\n'), ((2344, 2398), 'streamlit.text_input', 'st.text_input', (['"""Decoding"""', '"""greedy"""'], {'key': '"""placeholder"""'}), "('Decoding', 'greedy', key='placeholder')\n", (2357, 2398), True, 'import streamlit as st\n'), ((3284, 3358), 'langchain.embeddings.HuggingFaceHubEmbeddings', 'HuggingFaceHubEmbeddings', ([], {'repo_id': '"""sentence-transformers/all-MiniLM-L6-v2"""'}), "(repo_id='sentence-transformers/all-MiniLM-L6-v2')\n", (3308, 3358), False, 'from langchain.embeddings import HuggingFaceHubEmbeddings, HuggingFaceInstructEmbeddings\n'), ((3366, 3393), 'os.path.exists', 'os.path.exists', (['"""db.pickle"""'], {}), "('db.pickle')\n", (3380, 3393), False, 'import os\n'), ((3719, 3778), 'streamlit.text_input', 'st.text_input', (['"""Ask a question about your Policy Document:"""'], {}), "('Ask a question about your Policy Document:')\n", (3732, 3778), True, 'import streamlit as st\n'), ((4198, 4319), 'langChainInterface.LangChainInterface', 'LangChainInterface', ([], {'model': 'ModelTypes.LLAMA_2_70B_CHAT.value', 'credentials': 'creds', 'params': 'params', 'project_id': 'project_id'}), '(model=ModelTypes.LLAMA_2_70B_CHAT.value, credentials=\n    creds, params=params, project_id=project_id)\n', (4216, 4319), False, 'from langChainInterface import LangChainInterface\n'), ((4327, 4371), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['model_llm'], {'chain_type': '"""stuff"""'}), "(model_llm, chain_type='stuff')\n", (4340, 4371), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((4449, 4513), 'streamlit.text_area', 'st.text_area', ([], {'label': '"""Model Response"""', 'value': 'response', 'height': '(100)'}), "(label='Model Response', value=response, height=100)\n", (4461, 4513), True, 'import streamlit as st\n'), ((4518, 4528), 'streamlit.write', 'st.write', ([], {}), '()\n', (4526, 4528), True, 'import streamlit as st\n'), ((823, 858), 'os.environ.get', 'os.environ.get', (['"""LOGLEVEL"""', '"""DEBUG"""'], {}), "('LOGLEVEL', 'DEBUG')\n", (837, 858), False, 'import os\n'), ((3513, 3551), 'langchain.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['docs', 'embeddings'], {}), '(docs, embeddings)\n', (3533, 3551), False, 'from langchain.vectorstores import FAISS, Chroma\n'), ((3651, 3665), 'streamlit.write', 'st.write', (['"""\n"""'], {}), "('\\n')\n", (3659, 3665), True, 'import streamlit as st\n'), ((2705, 2757), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', ([], {'mode': '"""wb"""', 'delete': '(False)'}), "(mode='wb', delete=False)\n", (2732, 2757), False, 'import tempfile\n'), ((3462, 3484), 'pickle.load', 'pickle.load', (['file_name'], {}), '(file_name)\n', (3473, 3484), False, 'import pickle\n'), ((3617, 3643), 'pickle.dump', 'pickle.dump', (['db', 'file_name'], {}), '(db, file_name)\n', (3628, 3643), False, 'import pickle\n'), ((2905, 2949), 'streamlit.spinner', 'st.spinner', (['"""Waiting for the file to upload"""'], {}), "('Waiting for the file to upload')\n", (2915, 2949), True, 'import streamlit as st\n'), ((2973, 2994), 'langchain.document_loaders.PyPDFLoader', 'PyPDFLoader', (['filepath'], {}), '(filepath)\n', (2984, 2994), False, 'from langchain.document_loaders import PyPDFLoader\n'), ((3058, 3145), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=\n    chunk_overlap)\n', (3088, 3145), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n')]