Untitled

import streamlit as st
from streamlit_chat import message
from langchain.chat_models import ChatOpenAI
from langchain.indexes import VectorstoreIndexCreator
from langchain.document_loaders import TextLoader
from llama_index import LLMPredictor, GPTSimpleVectorIndex, \
    GithubRepositoryReader, SimpleDirectoryReader, Document
    
    
import os
import zipfile
from io import StringIO
from IPython.display import Markdown, display

os.environ['OPENAI_API_KEY'] = "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"

# Load the LangChain
# documents = SimpleDirectoryReader('./sample').load_data()
# index = GPTSimpleVectorIndex(documents, chunk_size_limit=512)

index = GPTSimpleVectorIndex.load_from_disk('./index.json')
llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"))

github_template = "https://github.com/{}/{}"

def load_data(file_or_text):
    """
    Input raw code file; compress file (zip); github link or pdf documents
    and return model's ready indices
    """
    # try:
    if type(file_or_text) == str:
        if 'https://github.com/' in file_or_text:
            link = str(file_or_text).replace('https://github.com/', '')
            link = link.split('/', 1)
            author_name = link[0]
            repo_name = link[1]
            index = GithubRepositoryReader(
                author_name,
                repo_name,
                github_token='',
            ).load_data(branch='main')
            
        elif os.path.exists(file_or_text):
            index = GPTSimpleVectorIndex.load_from_disk(file_or_text)

        else:
            index = GPTSimpleVectorIndex([Document(file_or_text)])
    else:
        # Input file
        file_name = file_or_text.name
            
        if str(file_name).endswith('.zip'):
            tmp_save_path = f'/tmp/{file_name}'
            with open(tmp_save_path, "wb") as f:
                f.write(file_or_text.getbuffer())

            # extract
            with zipfile.ZipFile(tmp_save_path, 'r') as zip_ref:
                tmp_save_path = f'/tmp/{str(file_name)[:-4]}'
                print(tmp_save_path)
                zip_ref.extractall(tmp_save_path)
            index = SimpleDirectoryReader(tmp_save_path).load_data()
        else:
            index = GPTSimpleVectorIndex([Document(file_or_text)])

    return index


def on_input_change():
    user_input = st.session_state.user_input
    st.session_state.past.append(user_input)
    st.session_state.generated.append("The `messages` from Bot")
    st.session_state.user_input = ''


def on_btn_click():
    del st.session_state.past[:]
    del st.session_state.generated[:]
    

def main_app():
    st.session_state.setdefault('past', [])
    st.session_state.setdefault('generated', [])
    source_data = None
    select_data = st.empty()
    
    with select_data.container():
        st.title("Create New Chatbot")
        st.warning('We support raw filetext, pdf, zip code or github link for create chatbot', icon="⚠️")

        uploaded_files = st.file_uploader("Select a file")
        with st.expander("Advanced Options"):
            txt = st.text_area("Github link or Raw txt")

        if uploaded_files or txt:
            source_data = uploaded_files or txt
            source_data = load_data(source_data)
            print(source_data)
            select_data.empty()
        
    if source_data:
        st.title("ChatBot")
        message("Hi, I am an AI assistant. I have read your project, ask me anything.")

        chat_placeholder = st.empty()

        with chat_placeholder.container():
            for i in range(len(st.session_state['generated'])):                
                message(st.session_state['past'][i], is_user=True, key=f"{i}_user")
                message(st.session_state['generated'][i], key=f"{i}")

        with st.container():
            st.text_input("User message:", on_change=on_input_change, key="user_input")

        with st.sidebar:
            st.header("Chatbot Options:")
            st.write('Reset Conversation')
            st.button('Reset', key="reset-butt", on_click=on_btn_click)
            with st.expander("Advanced Option"):
                st.warning('Feature in development', icon="⚠️")

if __name__ == '__main__':
    # st.set_page_config(page_title="QA Bot", page_icon=":robot:")
    st.set_page_config(
        page_title="Question Answering Bot",
        page_icon="🧊",
        initial_sidebar_state="expanded"
    )
    main_app()