app.py

'''
LLM scanner streamlit app

streamlit run .\app.py

Functionality
- tokenize documents
- respond to queries
- generate new documents

Based on: 
1. https://huggingface.co/spaces/llamaindex/llama_index_vector_demo
2. https://github.com/logan-markewich/llama_index_starter_pack/blob/main/streamlit_term_definition/

TODO:
- customize to other [LLMs](https://gpt-index.readthedocs.io/en/latest/reference/llm_predictor.html#llama_index.llm_predictor.LLMPredictor) 
- guardrails on 
- prevent answers on facts outside the document (e.g. birthdate of Michael Jordan in the docs vs. the baseball player)
'''

import os
import streamlit as st
from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, ServiceContext, LLMPredictor, PromptHelper, readers
from llama_index import StorageContext, load_index_from_storage

from langchain import OpenAI, HuggingFaceHub

import app_constants

index_fpath = "./llamas_index"
documents_folder = "./documents" #initial documents - additional can be added via upload

if "dummy" not in st.session_state:
    st.session_state["dummy"] = "dummy"

#@st.cache_resource  #st makes this globally available for all users and sessions 
def initialize_index(index_name, documents_folder, persisted_to_storage=True):
    """
    creates an index of the documents in the folder
    if the index exists, skipped
    """
    # set maximum input size
    max_input_size = 4096
    # set number of output tokens
    num_outputs = 2000
    # set maximum chunk overlap
    max_chunk_overlap = 20
    # set chunk size limit
    chunk_size_limit = 600 

    llm_predictor = LLMPredictor(llm=OpenAI(openai_api_key=api_key, #from env
                                            temperature=0.5, 
                                            model_name="text-davinci-003", 
                                            max_tokens=num_outputs))    
    #wishlist: alternatives
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)
    if os.path.exists(index_name):
        storage_context = StorageContext.from_defaults(persist_dir=index_fpath)
        doc_index = load_index_from_storage(service_context=service_context, storage_context=storage_context)
    else:
        #st.info("Updating the document index")
        prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)
 
        documents = SimpleDirectoryReader(documents_folder).load_data()
        doc_index = GPTVectorStoreIndex.from_documents(
            documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper, 
            chunk_size_limit=512, service_context=service_context
        )
        if persisted_to_storage:
            doc_index.storage_context.persist(index_fpath)

    #avoid this side-effect: st.session_state["doc_index"] = "doc_index"
    return doc_index

#st returns data that's available for future caller
@st.cache_data(max_entries=200, persist=True)  
def query_index(_index, query_text):
    query_engine = _index.as_query_engine()
    response = query_engine.query(query_text)
    #response = _index.query(query_text)
    return str(response)


#page format is directly written her
st.title("LLM scanner")
st.markdown(
    (
        "This app allows you to query documents!\n\n"
        "Powered by [Llama Index](https://gpt-index.readthedocs.io/en/latest/index.html)"
    )
)

setup_tab, upload_tab, query_tab = st.tabs(
    ["Setup", "Index", "Query"]
)

with setup_tab:
    st.subheader("LLM Setup")
    api_key = st.text_input("Enter your OpenAI API key here", type="password")
    
    #wishlist llm_name = st.selectbox(
    #    "Which LLM?", ["text-davinci-003", "gpt-3.5-turbo", "gpt-4"]
    #)
    #repo_id = "google/flan-t5-xl" # See https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads for some other options
    #llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={"temperature":0, "max_length":64})
    
    #model_temperature = st.slider(
    #    "LLM Temperature", min_value=0.0, max_value=1.0, step=0.1
    #)

if api_key is not None and "doc_index" not in st.session_state: 
    st.session_state["doc_index"] = initialize_index(index_fpath, documents_folder, persisted_to_storage=False) 


with upload_tab:
    st.subheader("Upload documents")
    
    if st.button("Re-initialize index with pre-packaged documents"):
        st.session_state["doc_index"] = initialize_index(index_fpath, documents_folder, persisted_to_storage=False)   
        st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))

    if "doc_index" in st.session_state:
        doc_index = st.session_state["doc_index"]
        st.markdown(
            "Either upload a document, or enter the text manually."
        )
        uploaded_file = st.file_uploader(
            "Upload a document (pdf):", type=["pdf"]
        )
        document_text = st.text_area("Enter text")
        if st.button("Add document to index") and (uploaded_file or document_text):
            with st.spinner("Inserting (large files may be slow)..."):
                if document_text:
                    doc_index.refresh([readers.Document(text=document_text)]) #tokenizes new documents
                    st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))

                    st.session_state["doc_index"] = doc_index
                if uploaded_file:     
                    uploads_folder = "uploads/"
                    if not os.path.exists(uploads_folder):
                        os.mkdir(uploads_folder)
                    #file_details = {"FileName":uploaded_file.name,"FileType":uploaded_file.type}
                    with open(uploads_folder + "tmp.pdf", "wb") as f:
                        f.write(uploaded_file.getbuffer())
                    documents = SimpleDirectoryReader(uploads_folder).load_data()
                    doc_index.refresh(documents) #tokenizes new documents
                    st.session_state["doc_index"] = doc_index
                    st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))

                    st.session_state["doc_index"] = doc_index
                    os.remove(uploads_folder + "tmp.pdf")

with query_tab:
    st.subheader("Query Tab")
    st.write("Enter a query about the included documents. Find [documentation here](https://huggingface.co/spaces/agutfraind/llmscanner)")

    doc_index = None
    #api_key = st.text_input("Enter your OpenAI API key here:", type="password")
    if api_key:
        os.environ['OPENAI_API_KEY'] = api_key
        #doc_index = initialize_index(index_fpath, documents_folder)   

    if doc_index is None:
        if "doc_index" in st.session_state:
            doc_index = st.session_state["doc_index"]
            st.info('Documents in index: ' + str(doc_index.docstore.docs.__len__())) 
        else:
            st.warning("Doc index is not available - initialize or upload")
        #st.warning("Please enter your api key first.")

    if doc_index and api_key:
        select_type_your_own = 'type your own...'
        options_for_queries = app_constants.canned_questions + [select_type_your_own]
        query_selection = st.selectbox("Select option", options=options_for_queries)
        query_text = None
        
        if query_selection == select_type_your_own: 
            query_text = st.text_input("Query text")
        else:
            query_text = query_selection 

        if st.button("Run Query") and (doc_index is not None) and (query_text is not None):
            response = query_index(doc_index, query_text)
            st.markdown(response)
            
            llm_col, embed_col = st.columns(2)
            with llm_col:
                st.markdown(f"LLM Tokens Used: {doc_index.service_context.llm_predictor._last_token_usage}")
            
            with embed_col:
                st.markdown(f"Embedding Tokens Used: {doc_index.service_context.embed_model._last_token_usage}")