import streamlit as st
import langchain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain import OpenAI, VectorDBQA
from langchain.chains import RetrievalQAWithSourcesChain
import PyPDF2

#This function will go through pdf and extract and return list of page texts.
def read_and_textify(files):
    text_list = []
    sources_list = []
    for file in files:
        pdfReader = PyPDF2.PdfReader(file)
        #print("Page Number:", len(pdfReader.pages))
        for i in range(len(pdfReader.pages)):
          pageObj = pdfReader.pages[i]
          text = pageObj.extract_text()
          pageObj.clear()
          text_list.append(text)
          sources_list.append(file.name + "_page_"+str(i))
    return [text_list,sources_list]
  
st.set_page_config(layout="centered", page_title="Multidoc_QnA")
st.header("Multidoc_QnA")
st.write("---")
  
#file uploader
uploaded_files = st.file_uploader("Upload documents",accept_multiple_files=True, type=["txt","pdf"])
st.write("---")

if uploaded_files is None:
  st.info(f"""Upload files to analyse""")
elif uploaded_files:
  st.write(str(len(uploaded_files)) + " document(s) loaded..")
  
  textify_output = read_and_textify(uploaded_files)
  
  documents = textify_output[0]
  sources = textify_output[1]
  
  #extract embeddings
  embeddings = OpenAIEmbeddings(openai_api_key = st.secrets["openai_api_key"])
  #vstore with metadata. Here we will store page numbers.
  vStore = Chroma.from_texts(documents, embeddings, metadatas=[{"source": s} for s in sources])
  #deciding model
  model_name = "gpt-3.5-turbo"
  # model_name = "gpt-4"

  retriever = vStore.as_retriever()
  retriever.search_kwargs = {'k':2}

  #initiate model
  llm = OpenAI(model_name=model_name, openai_api_key = st.secrets["openai_api_key"], streaming=True)
  model = RetrievalQAWithSourcesChain.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)
  
  st.header("Ask your data")
  user_q = st.text_area("Enter your questions here")
  
  if st.button("Get Response"):
    try:
      with st.spinner("Model is working on it..."):
        result = model({"question":user_q}, return_only_outputs=True)
        st.subheader('Your response:')
        st.write(result['answer'])
        st.subheader('Source pages:')
        st.write(result['sources'])
    except Exception as e:
      st.error(f"An error occurred: {e}")
      st.error('Oops, the GPT response resulted in an error :( Please try again with a different question.')
      
        
    
  
  
  
  
  
  
  
  
  

langchain.chains.RetrievalQAWithSourcesChain.from_chain_type langchain.vectorstores.Chroma.from_texts langchain.OpenAI langchain.embeddings.openai.OpenAIEmbeddings
[((868, 932), 'streamlit.set_page_config', 'st.set_page_config', ([], {'layout': '"""centered"""', 'page_title': '"""Multidoc_QnA"""'}), "(layout='centered', page_title='Multidoc_QnA')\n", (886, 932), True, 'import streamlit as st\n'), ((933, 958), 'streamlit.header', 'st.header', (['"""Multidoc_QnA"""'], {}), "('Multidoc_QnA')\n", (942, 958), True, 'import streamlit as st\n'), ((959, 974), 'streamlit.write', 'st.write', (['"""---"""'], {}), "('---')\n", (967, 974), True, 'import streamlit as st\n'), ((1010, 1100), 'streamlit.file_uploader', 'st.file_uploader', (['"""Upload documents"""'], {'accept_multiple_files': '(True)', 'type': "['txt', 'pdf']"}), "('Upload documents', accept_multiple_files=True, type=[\n    'txt', 'pdf'])\n", (1026, 1100), True, 'import streamlit as st\n'), ((1094, 1109), 'streamlit.write', 'st.write', (['"""---"""'], {}), "('---')\n", (1102, 1109), True, 'import streamlit as st\n'), ((1140, 1175), 'streamlit.info', 'st.info', (['f"""Upload files to analyse"""'], {}), "(f'Upload files to analyse')\n", (1147, 1175), True, 'import streamlit as st\n'), ((510, 532), 'PyPDF2.PdfReader', 'PyPDF2.PdfReader', (['file'], {}), '(file)\n', (526, 532), False, 'import PyPDF2\n'), ((1424, 1485), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'openai_api_key': "st.secrets['openai_api_key']"}), "(openai_api_key=st.secrets['openai_api_key'])\n", (1440, 1485), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((1557, 1645), 'langchain.vectorstores.Chroma.from_texts', 'Chroma.from_texts', (['documents', 'embeddings'], {'metadatas': "[{'source': s} for s in sources]"}), "(documents, embeddings, metadatas=[{'source': s} for s in\n    sources])\n", (1574, 1645), False, 'from langchain.vectorstores import Chroma\n'), ((1816, 1910), 'langchain.OpenAI', 'OpenAI', ([], {'model_name': 'model_name', 'openai_api_key': "st.secrets['openai_api_key']", 'streaming': '(True)'}), "(model_name=model_name, openai_api_key=st.secrets['openai_api_key'],\n    streaming=True)\n", (1822, 1910), False, 'from langchain import OpenAI, VectorDBQA\n'), ((1919, 2016), 'langchain.chains.RetrievalQAWithSourcesChain.from_chain_type', 'RetrievalQAWithSourcesChain.from_chain_type', ([], {'llm': 'llm', 'chain_type': '"""stuff"""', 'retriever': 'retriever'}), "(llm=llm, chain_type='stuff',\n    retriever=retriever)\n", (1962, 2016), False, 'from langchain.chains import RetrievalQAWithSourcesChain\n'), ((2018, 2044), 'streamlit.header', 'st.header', (['"""Ask your data"""'], {}), "('Ask your data')\n", (2027, 2044), True, 'import streamlit as st\n'), ((2056, 2097), 'streamlit.text_area', 'st.text_area', (['"""Enter your questions here"""'], {}), "('Enter your questions here')\n", (2068, 2097), True, 'import streamlit as st\n'), ((2106, 2131), 'streamlit.button', 'st.button', (['"""Get Response"""'], {}), "('Get Response')\n", (2115, 2131), True, 'import streamlit as st\n'), ((2153, 2192), 'streamlit.spinner', 'st.spinner', (['"""Model is working on it..."""'], {}), "('Model is working on it...')\n", (2163, 2192), True, 'import streamlit as st\n'), ((2272, 2302), 'streamlit.subheader', 'st.subheader', (['"""Your response:"""'], {}), "('Your response:')\n", (2284, 2302), True, 'import streamlit as st\n'), ((2311, 2337), 'streamlit.write', 'st.write', (["result['answer']"], {}), "(result['answer'])\n", (2319, 2337), True, 'import streamlit as st\n'), ((2346, 2375), 'streamlit.subheader', 'st.subheader', (['"""Source pages:"""'], {}), "('Source pages:')\n", (2358, 2375), True, 'import streamlit as st\n'), ((2384, 2411), 'streamlit.write', 'st.write', (["result['sources']"], {}), "(result['sources'])\n", (2392, 2411), True, 'import streamlit as st\n'), ((2445, 2480), 'streamlit.error', 'st.error', (['f"""An error occurred: {e}"""'], {}), "(f'An error occurred: {e}')\n", (2453, 2480), True, 'import streamlit as st\n'), ((2487, 2599), 'streamlit.error', 'st.error', (['"""Oops, the GPT response resulted in an error :( Please try again with a different question."""'], {}), "(\n    'Oops, the GPT response resulted in an error :( Please try again with a different question.'\n    )\n", (2495, 2599), True, 'import streamlit as st\n')]