app.py

import streamlit as st
import os
import pickle
import faiss
import common
import glob
from multiprocessing import Lock
from multiprocessing.managers import BaseManager
from pathlib import Path
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index import Document,VectorStoreIndex, SimpleDirectoryReader, ServiceContext, StorageContext, load_index_from_storage
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
from llama_index.constants import DEFAULT_CHUNK_OVERLAP
from llama_index.vector_stores.faiss import FaissVectorStore
from llama_index.graph_stores import SimpleGraphStore
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from msal_streamlit_authentication import msal_authentication
from llama_hub.file.cjk_pdf.base import CJKPDFReader 
from llama_hub.file.pptx.base import PptxReader
from llama_hub.file.pandas_excel.base import PandasExcelReader
from llama_hub.file.docx.base import DocxReader 
from llama_index.llms import OpenAI
import tiktoken
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from dotenv import load_dotenv

load_dotenv()

# 接続元制御
ALLOW_IP_ADDRESS = os.environ["ALLOW_IP_ADDRESS"]

# Azure AD app registration details
CLIENT_ID = os.environ["CLIENT_ID"]
CLIENT_SECRET = os.environ["CLIENT_SECRET"]
TENANT_ID = os.environ["TENANT_ID"]

# Azure API
AUTHORITY = f"https://login.microsoftonline.com/{TENANT_ID}"
REDIRECT_URI = os.environ["REDIRECT_URI"]
SCOPES = ["openid", "profile", "User.Read"]

INDEX_NAME = os.environ["INDEX_NAME"]
PKL_NAME = os.environ["PKL_NAME"]
st.session_state.llama_debug_handler = LlamaDebugHandler()
from log import logger

def initialize_index():
    logger.info("initialize_index start")
    llm = OpenAI(model='gpt-3.5-turbo', temperature=0.8, max_tokens=256)
    text_splitter = TokenTextSplitter(separator="。",chunk_size=1500
      , chunk_overlap=DEFAULT_CHUNK_OVERLAP
      , tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode)
    node_parser = SimpleNodeParser(text_splitter=text_splitter)
    d = 1536
    k=2
    faiss_index = faiss.IndexFlatL2(d)
    # デバッグ用
    callback_manager = CallbackManager([st.session_state.llama_debug_handler])
    service_context = ServiceContext.from_defaults(llm=llm,node_parser=node_parser,callback_manager=callback_manager)
    lock = Lock()
    with lock:
        if os.path.exists(INDEX_NAME):
            logger.info("start import index")
            storage_context = StorageContext.from_defaults(
              docstore=SimpleDocumentStore.from_persist_dir(persist_dir=INDEX_NAME),
              graph_store=SimpleGraphStore.from_persist_dir(persist_dir=INDEX_NAME),
              vector_store=FaissVectorStore.from_persist_dir(persist_dir=INDEX_NAME),
              index_store=SimpleIndexStore.from_persist_dir(persist_dir=INDEX_NAME),
            )
            st.session_state.index = load_index_from_storage(storage_context=storage_context,service_context=service_context)
            with open(PKL_NAME, "rb") as f:
                st.session_state.stored_docs = pickle.load(f)
            common.setChatEngine()
        else:
            logger.info("start create index")
            documents = list()
            files = glob.glob("./documents/*")
            vector_store = FaissVectorStore(faiss_index=faiss_index)
            storage_context = StorageContext.from_defaults(vector_store=vector_store)
            st.session_state.stored_docs=list()
            for file in files:
                loader=None
                noextpath,extension = os.path.splitext(file)
                logger.info(file)
                document = Document()
                if extension == ".txt" or extension ==".md":
                    document = SimpleDirectoryReader(input_files=[file],  filename_as_id=True).load_data()[0]
                else:
                    if extension == ".pdf":
                        loader = CJKPDFReader()
                    elif extension == ".pptx":
                        loader = PptxReader()
                    elif extension == ".xlsx":
                        loader = PandasExcelReader(pandas_config={"header": 0})
                    elif extension == ".docx":
                        loader = DocxReader()
                    else:
                        logger.error("Can`t read file:" + file)
                        continue
                    document = loader.load_data(file=Path(file))[0]
                document.metadata={'filename': os.path.basename(file)}
                documents.append(document)
                st.session_state.stored_docs.append(os.path.basename(file)) 
            st.session_state.index = VectorStoreIndex.from_documents( documents=documents,storage_context=storage_context,service_context=service_context)
            st.session_state.index.storage_context.persist(persist_dir=INDEX_NAME)
            with open(PKL_NAME, "wb") as f:
                print("pickle")
                pickle.dump(st.session_state.stored_docs, f)
            common.setChatEngine()

def logout():
    st.session_state["login_token"] = None

# メイン
st.session_state["login_token"] = msal_authentication(
    auth={
        "clientId": CLIENT_ID,
        "authority": AUTHORITY,
        "redirectUri": REDIRECT_URI,
        "postLogoutRedirectUri": ""
    }, # Corresponds to the 'auth' configuration for an MSAL Instance
    cache={
        "cacheLocation": "sessionStorage",
        "storeAuthStateInCookie": False
    }, # Corresponds to the 'cache' configuration for an MSAL Instance
    login_request={
        "scopes": SCOPES
    }, # Optional
    logout_request={}, # Optional
    login_button_text="Login", # Optional, defaults to "Login"
    logout_button_text="Logout", # Optional, defaults to "Logout"
    class_name="css_button_class_selector", # Optional, defaults to None. Corresponds to HTML class.
    html_id="html_id_for_button", # Optional, defaults to None. Corresponds to HTML id.
    #key=1 # Optional if only a single instance is needed
)
# st.write("Recevied login token:", st.session_state.login_token)

if st.session_state.login_token:
    initialize_index()
    st.write("ようこそ", st.session_state.login_token["account"]["name"])
    st.write("サイドメニューからファイルインポート又はChatbotへの質問を開始してください。")
    st.markdown("""
        ## 使い方
        - **Chatbot**  
            初期からインポートされているファイルとImportXXFileでインポートしたファイルの内容に関する質問に対して、GenerativeAIが回答します。  
            ※返答が正常に帰ってこない場合があります。参照ファイルを記載しているので、判断の目安にしてください。
        
        - **ChatbotWebRead**  
            入力したURLのサイトの情報に関して、GenerativeAIが回答します。  
            スクレイピングが禁止されているサイトは入力しないでください。  
            ImportAllFileの内容は登録されていません。
        
        - **ImportAllFile**  
            テキストファイル,mdファイル,Excel,PDF,PowerPoint,Wordをインポートできます。
    """)