import arxiv
import openai
import langchain
import pinecone
from langchain_community.document_loaders import ArxivLoader
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Pinecone
from langchain.llms import OpenAI
from langchain import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.chains.summarize import load_summarize_chain
from langchain.chains.question_answering import load_qa_chain
from langchain import OpenAI
from utils import *
import streamlit as st

import os
from dotenv import load_dotenv
load_dotenv()

openai_api_key = os.getenv('OPENAI_API_KEY')
pinecone_api_key = os.getenv('PINECONE_API_KEY')
environment = os.getenv('PINECONE_ENV')

llm_summary = ChatOpenAI(temperature=0.3, model_name="gpt-3.5-turbo-0125")    
llm = OpenAI(model_name="gpt-3.5-turbo-0125", temperature=0.6, api_key=openai_api_key)

if 'summary' not in st.session_state:
    st.session_state.summary = None

if 'documents' not in st.session_state:
    st.session_state.documents = None


st.title('Arxiv Paper Summarizer and Interactive Q&A')

paper_id_input = st.text_input('Enter Arxiv Paper ID', '')

if st.button('Summarize Paper') and paper_id_input:
    with st.spinner('Fetching and summarizing the paper...'):
        try:
            doc = arxiv_loader(paper_id=paper_id_input)
            st.session_state.documents = chunk_data(docs=doc)
#            st.write(st.session_state.documents)
            chain = load_summarize_chain(
                llm=llm_summary,
                    chain_type='map_reduce',
                    verbose=False
                    )
            summary = chain.run(st.session_state.documents)
            st.subheader('Summary')
            st.write(summary)
        except Exception as e:
            st.error(f"An error occurred: {e}")
            

def initialize_index(index_name='arxiv-summarizer'):
#    documents = chunk_data(docs=doc)
    embeddings = OpenAIEmbeddings(api_key=openai_api_key)
    index_name = index_name
    # Make sure environment is correctly spelled (there was a typo in your provided code)
    pinecone.Pinecone(
        api_key=pinecone_api_key,
        environment=environment
    )
    if st.session_state.documents:
        index = Pinecone.from_documents(st.session_state.documents, embeddings, index_name=index_name)
    else:
        index = None
    return index

index = initialize_index()

def retrieve_query(query, k=2):
    matching_results = index.similarity_search(query, k=k)
    return matching_results

def retrieve_answers(query):
    chain = load_qa_chain(llm, chain_type='stuff')
    doc_search = retrieve_query(query)
    print(doc_search)
    response = chain.run(input_documents=doc_search, question=query)
    return response

if paper_id_input:
    user_query = st.text_input("Ask a question about the paper:", '')
    
    if user_query:
        if st.button('Get Answer'):
            with st.spinner('Retrieving your answer...'):
                try:
                    answer = retrieve_answers(user_query)
                    st.subheader('Answer')
                    st.write(answer)
                except Exception as e:
                    st.error(f"An error occurred while retrieving the answer: {e}")



langchain.chains.question_answering.load_qa_chain langchain.chains.summarize.load_summarize_chain langchain.vectorstores.Pinecone.from_documents langchain.chat_models.ChatOpenAI langchain.OpenAI langchain.embeddings.openai.OpenAIEmbeddings
[((690, 703), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (701, 703), False, 'from dotenv import load_dotenv\n'), ((722, 749), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (731, 749), False, 'import os\n'), ((769, 798), 'os.getenv', 'os.getenv', (['"""PINECONE_API_KEY"""'], {}), "('PINECONE_API_KEY')\n", (778, 798), False, 'import os\n'), ((813, 838), 'os.getenv', 'os.getenv', (['"""PINECONE_ENV"""'], {}), "('PINECONE_ENV')\n", (822, 838), False, 'import os\n'), ((854, 914), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.3)', 'model_name': '"""gpt-3.5-turbo-0125"""'}), "(temperature=0.3, model_name='gpt-3.5-turbo-0125')\n", (864, 914), False, 'from langchain.chat_models import ChatOpenAI\n'), ((925, 1010), 'langchain.OpenAI', 'OpenAI', ([], {'model_name': '"""gpt-3.5-turbo-0125"""', 'temperature': '(0.6)', 'api_key': 'openai_api_key'}), "(model_name='gpt-3.5-turbo-0125', temperature=0.6, api_key=openai_api_key\n    )\n", (931, 1010), False, 'from langchain import OpenAI\n'), ((1162, 1216), 'streamlit.title', 'st.title', (['"""Arxiv Paper Summarizer and Interactive Q&A"""'], {}), "('Arxiv Paper Summarizer and Interactive Q&A')\n", (1170, 1216), True, 'import streamlit as st\n'), ((1235, 1276), 'streamlit.text_input', 'st.text_input', (['"""Enter Arxiv Paper ID"""', '""""""'], {}), "('Enter Arxiv Paper ID', '')\n", (1248, 1276), True, 'import streamlit as st\n'), ((1281, 1309), 'streamlit.button', 'st.button', (['"""Summarize Paper"""'], {}), "('Summarize Paper')\n", (1290, 1309), True, 'import streamlit as st\n'), ((2076, 2116), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'api_key': 'openai_api_key'}), '(api_key=openai_api_key)\n', (2092, 2116), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((2239, 2307), 'pinecone.Pinecone', 'pinecone.Pinecone', ([], {'api_key': 'pinecone_api_key', 'environment': 'environment'}), '(api_key=pinecone_api_key, environment=environment)\n', (2256, 2307), False, 'import pinecone\n'), ((2706, 2744), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['llm'], {'chain_type': '"""stuff"""'}), "(llm, chain_type='stuff')\n", (2719, 2744), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((2932, 2984), 'streamlit.text_input', 'st.text_input', (['"""Ask a question about the paper:"""', '""""""'], {}), "('Ask a question about the paper:', '')\n", (2945, 2984), True, 'import streamlit as st\n'), ((1339, 1390), 'streamlit.spinner', 'st.spinner', (['"""Fetching and summarizing the paper..."""'], {}), "('Fetching and summarizing the paper...')\n", (1349, 1390), True, 'import streamlit as st\n'), ((2381, 2472), 'langchain.vectorstores.Pinecone.from_documents', 'Pinecone.from_documents', (['st.session_state.documents', 'embeddings'], {'index_name': 'index_name'}), '(st.session_state.documents, embeddings, index_name=\n    index_name)\n', (2404, 2472), False, 'from langchain.vectorstores import Pinecone\n'), ((3020, 3043), 'streamlit.button', 'st.button', (['"""Get Answer"""'], {}), "('Get Answer')\n", (3029, 3043), True, 'import streamlit as st\n'), ((1593, 1670), 'langchain.chains.summarize.load_summarize_chain', 'load_summarize_chain', ([], {'llm': 'llm_summary', 'chain_type': '"""map_reduce"""', 'verbose': '(False)'}), "(llm=llm_summary, chain_type='map_reduce', verbose=False)\n", (1613, 1670), False, 'from langchain.chains.summarize import load_summarize_chain\n'), ((1821, 1844), 'streamlit.subheader', 'st.subheader', (['"""Summary"""'], {}), "('Summary')\n", (1833, 1844), True, 'import streamlit as st\n'), ((1857, 1874), 'streamlit.write', 'st.write', (['summary'], {}), '(summary)\n', (1865, 1874), True, 'import streamlit as st\n'), ((1918, 1953), 'streamlit.error', 'st.error', (['f"""An error occurred: {e}"""'], {}), "(f'An error occurred: {e}')\n", (1926, 1953), True, 'import streamlit as st\n'), ((3062, 3101), 'streamlit.spinner', 'st.spinner', (['"""Retrieving your answer..."""'], {}), "('Retrieving your answer...')\n", (3072, 3101), True, 'import streamlit as st\n'), ((3202, 3224), 'streamlit.subheader', 'st.subheader', (['"""Answer"""'], {}), "('Answer')\n", (3214, 3224), True, 'import streamlit as st\n'), ((3245, 3261), 'streamlit.write', 'st.write', (['answer'], {}), '(answer)\n', (3253, 3261), True, 'import streamlit as st\n'), ((3321, 3384), 'streamlit.error', 'st.error', (['f"""An error occurred while retrieving the answer: {e}"""'], {}), "(f'An error occurred while retrieving the answer: {e}')\n", (3329, 3384), True, 'import streamlit as st\n')]