Upload 3 files

app.py

@@ -1,4 +1,3 @@
-
 import streamlit as st
 import os
 import pickle
@@ -19,9 +18,47 @@ from llama_index.graph_stores import SimpleGraphStore
 from llama_index.storage.docstore import SimpleDocumentStore
 from llama_index.storage.index_store import SimpleIndexStore
 import tiktoken
-import streamlit_authenticator as stauth
-import yaml
-from logging import getLogger, StreamHandler, Formatter
+import requests
+
+from requests_oauthlib import OAuth2Session
+from time import time
+from dotenv import load_dotenv
+from streamlit import net_util
+
+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_PATH = os.environ["REDIRECT_PATH"]
+TOKEN_URL = f"{AUTHORITY}/oauth2/v2.0/token"
+AUTHORIZATION_URL = f"{AUTHORITY}/oauth2/v2.0/authorize"
+SCOPES = ["openid", "profile", "User.Read"]
+
+# 認証用URL取得
+def authorization_request():
+    oauth = OAuth2Session(CLIENT_ID, redirect_uri=REDIRECT_PATH, scope=SCOPES)
+    authorization_url, state = oauth.authorization_url(AUTHORIZATION_URL)
+    return authorization_url, state
+
+# 認証トークン取得
+def token_request(authorization_response, state):
+    oauth = OAuth2Session(CLIENT_ID, state=state)
+    token = oauth.fetch_token(
+        TOKEN_URL,
+        code=authorization_response[0],
+        authorization_response=authorization_response,
+        client_secret=CLIENT_SECRET,
+
+    )
+    return token
 
 index_name = "./data/storage"
 pkl_name = "./data/stored_documents.pkl"
@@ -51,14 +88,13 @@ def initialize_index():
       , chunk_overlap=DEFAULT_CHUNK_OVERLAP
       , tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode)
     node_parser = SimpleNodeParser(text_splitter=text_splitter)
-    service_context = ServiceContext.from_defaults(node_parser=node_parser)
     d = 1536
     k=2
     faiss_index = faiss.IndexFlatL2(d)
     # デバッグ用
     llama_debug_handler = LlamaDebugHandler()
     callback_manager = CallbackManager([llama_debug_handler])
-    service_context = ServiceContext.from_defaults(callback_manager=callback_manager)
+    service_context = ServiceContext.from_defaults(node_parser=node_parser,callback_manager=callback_manager)
     lock = Lock()
     with lock:
         if os.path.exists(index_name):
@@ -69,9 +105,8 @@ def initialize_index():
               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)
-            # index = load_index_from_storage(StorageContext.from_defaults(persist_dir=index_name), service_context=service_context)
             response_synthesizer = get_response_synthesizer(response_mode='refine')
-            st.session_state.query_engine = st.session_state.index.as_query_engine(response_synthesizer=response_synthesizer)
+            st.session_state.query_engine = st.session_state.index.as_query_engine(response_synthesizer=response_synthesizer,service_context=service_context)
             st.session_state.chat_engine = CondenseQuestionChatEngine.from_defaults(
                 query_engine=st.session_state.query_engine, 
                 condense_question_prompt=custom_prompt,
@@ -85,7 +120,7 @@ def initialize_index():
             st.session_state.index = VectorStoreIndex.from_documents(documents, storage_context=storage_context,service_context=service_context)
             st.session_state.index.storage_context.persist(persist_dir=index_name)
             response_synthesizer = get_response_synthesizer(response_mode='refine')
-            st.session_state.query_engine = st.session_state.index.as_query_engine(response_synthesizer=response_synthesizer)
+            st.session_state.query_engine = st.session_state.index.as_query_engine(response_synthesizer=response_synthesizer,service_context=service_context)
             st.session_state.chat_engine = CondenseQuestionChatEngine.from_defaults(
                 query_engine=st.session_state.query_engine, 
                 condense_question_prompt=custom_prompt,
@@ -98,32 +133,59 @@ def initialize_index():
         else:
             st.session_state.stored_docs=list()
 
-with open('config.yaml') as file:
-    config = yaml.load(file, Loader=yaml.SafeLoader)
-
-authenticator = stauth.Authenticate(
-    config['credentials'],
-    config['cookie']['name'],
-    config['cookie']['key'],
-    config['cookie']['expiry_days'],
-    config['preauthorized'],
-)
-
-name, authentication_status, username = authenticator.login('Login', 'main')
-
-
-if 'authentication_status' not in st.session_state:
-    st.session_state['authentication_status'] = None
-
-if st.session_state["authentication_status"]:
-    authenticator.logout('Logout', 'main')
-    st.write(f'ログインに成功しました')
-    initialize_index()
-		# ここにログイン��の処理を書く。
-elif st.session_state["authentication_status"] is False:
-    st.error('ユーザ名またはパスワードが間違っています')
-elif st.session_state["authentication_status"] is None:
-    st.warning('ユーザ名やパスワードを入力してください')
-
-
-
+def logout():
+    st.session_state["token"] = None
+    st.session_state["token_expires"] = None
+    st.session_state["authorization_state"] = None
+
+# メイン
+def app():
+    # 初期化
+    st.session_state["token"] = None
+    st.session_state["token_expires"] = time()
+    st.session_state["authorization_state"] = None
+
+    # 接続元IP許可判定
+    if not net_util.get_external_ip() in ALLOW_IP_ADDRESS:
+        st.title("HTTP 403 Forbidden")
+        return
+
+    # 接続元OK
+    st.title("Azure AD Login with Streamlit")
+
+    # 認証後のリダイレクトのGETパラメータ値を取得
+    authorization_response = st.experimental_get_query_params().get("code")
+
+    # 認証OK、トークン無し
+    if authorization_response and st.session_state["token"] is None:
+        # トークン設定
+        token = token_request(authorization_response, st.session_state["authorization_state"])
+        st.session_state["token"] = token
+        st.session_state["token_expires"] = token["expires_at"]
+
+    # トークン無し or 期限切れ
+    if st.session_state["token"] is None or float(st.session_state["token_expires"]) <= time():
+        # 認証用リンク表示
+        authorization_url, st.session_state["authorization_state"] = authorization_request()
+        st.markdown(f'[Click here to log in]({authorization_url})', unsafe_allow_html=True)
+    else:
+        # 認証OK
+        st.markdown(f"Logged in successfully. Welcome, {st.session_state['token']['token_type']}!")
+        if st.button("logout",use_container_width=True):
+            logout()
+            st.experimental_set_query_params(test="test")
+            st.experimental_rerun()
+        st.text("サイドバーから利用するメニューをお選びください。")
+        initialize_index()
+
+if __name__ == "__main__":
+    if "token" not in st.session_state or st.session_state["token"] is None or float(st.session_state["token_expires"]) <= time():
+        app()
+    else:
+        st.title("Azure AD Login with Streamlit")
+        if st.button("logout",use_container_width=True):
+            logout()
+            st.experimental_set_query_params(test="test")
+            st.experimental_rerun()
+        st.text("ログイン済みです。")
+        st.text("サイドバーから利用するメニューをお選びください。")

common.py

@@ -2,15 +2,38 @@
 import extra_streamlit_components as stx
 import streamlit as st
 import logging
+import os
+
+from time import time
+from requests_oauthlib import OAuth2Session
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("__name__")
 logger.debug("調査用ログ")
 
+# 接続元制御
+ALLOW_IP_ADDRESS = os.environ["ALLOW_IP_ADDRESS"]
+
+# Azure AD app registration details
+CLIENT_ID = os.environ["CLIENT_ID"]
+TENANT_ID = os.environ["TENANT_ID"]
+
+# Azure API
+AUTHORITY = f"https://login.microsoftonline.com/{TENANT_ID}"
+REDIRECT_PATH = os.environ["REDIRECT_PATH"]
+AUTHORIZATION_URL = f"{AUTHORITY}/oauth2/v2.0/authorize"
+SCOPES = ["openid", "profile", "User.Read"]
+
+# 認証用URL取得
+def authorization_request():
+    oauth = OAuth2Session(CLIENT_ID, redirect_uri=REDIRECT_PATH, scope=SCOPES)
+    authorization_url, state = oauth.authorization_url(AUTHORIZATION_URL)
+    return authorization_url, state
+
 #ログインの確認
 def check_login():
-    if 'authentication_status' not in st.session_state:
-        st.session_state['authentication_status'] = None
-    if st.session_state["authentication_status"] is None or False:
-        st.warning("**ログインしてください**")
+    if "token" not in st.session_state or st.session_state["token"] is None or float(st.session_state["token_expires"]) <= time():
+        # 認証用リンク表示
+        authorization_url, st.session_state["authorization_state"] = authorization_request()
+        st.markdown(f'[Click here to log in]({authorization_url})', unsafe_allow_html=True)
         st.stop()

requirements.txt

@@ -8,4 +8,6 @@ pypdf==3.9.0
 faiss-cpu==1.7.4
 html2text
 streamlit-authenticator
-extra_streamlit_components
+extra_streamlit_components
+requests_oauthlib
+python-dotenv