try commit

app.py

@@ -1,191 +1,215 @@
-import streamlit as st
-import os
-import pickle
-import faiss
-import logging
-
-from multiprocessing import Lock
-from multiprocessing.managers import BaseManager
-from llama_index.callbacks import CallbackManager, LlamaDebugHandler
-from llama_index import VectorStoreIndex, Document,Prompt, SimpleDirectoryReader, ServiceContext, StorageContext, load_index_from_storage
-from llama_index.chat_engine import CondenseQuestionChatEngine;
-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.response_synthesizers import get_response_synthesizer
-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
-import tiktoken
-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"
-
-custom_prompt = Prompt("""\
-  以下はこれまでの会話履歴と、ドキュメントを検索して回答する必要がある、ユーザーからの会話文です。
-  会話と新しい会話文に基づいて、検索クエリを作成します。回答は日本語で行います。
-  新しい会話文が挨拶の場合、挨拶を返してください。
-  新しい会話文が質問の場合、検索した結果の回答を返してください。
-  答えがわからない場合は正直にわからないと回答してください。
-  会話履歴:
-  {chat_history}
-  新しい会話文:
-  {question}
-  Search query:
-""")
-
-chat_history = []
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("__name__")
-logger.debug("調査用ログ")
-
-def initialize_index():
-    logger.info("initialize_index start")
-    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)
-    # デバッグ用
-    llama_debug_handler = LlamaDebugHandler()
-    callback_manager = CallbackManager([llama_debug_handler])
-    service_context = ServiceContext.from_defaults(node_parser=node_parser,callback_manager=callback_manager)
-    lock = Lock()
-    with lock:
-        if os.path.exists(index_name):
-            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)
-            response_synthesizer = get_response_synthesizer(response_mode='refine')
-            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,
-                chat_history=chat_history,
-                verbose=True
-            )
-        else:
-            documents = SimpleDirectoryReader("./documents").load_data()
-            vector_store = FaissVectorStore(faiss_index=faiss_index)
-            storage_context = StorageContext.from_defaults(vector_store=vector_store)
-            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,service_context=service_context)
-            st.session_state.chat_engine = CondenseQuestionChatEngine.from_defaults(
-                query_engine=st.session_state.query_engine, 
-                condense_question_prompt=custom_prompt,
-                chat_history=chat_history,
-                verbose=True
-            )
-        if os.path.exists(pkl_name):
-            with open(pkl_name, "rb") as f:
-                st.session_state.stored_docs = pickle.load(f)
-        else:
-            st.session_state.stored_docs=list()
-
-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("サイドバーから利用するメニューをお選びください。")
+import streamlit as st
+import os
+import pickle
+import faiss
+import logging
+
+from multiprocessing import Lock
+from multiprocessing.managers import BaseManager
+from llama_index.callbacks import CallbackManager, LlamaDebugHandler
+from llama_index import VectorStoreIndex, Document,Prompt, SimpleDirectoryReader, ServiceContext, StorageContext, load_index_from_storage
+from llama_index.chat_engine import CondenseQuestionChatEngine;
+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.response_synthesizers import get_response_synthesizer
+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
+import tiktoken
+from streamlit import runtime
+from streamlit.runtime.scriptrunner import get_script_run_ctx
+import ipaddress
+
+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"
+
+custom_prompt = Prompt("""\
+  以下はこれまでの会話履歴と、ドキュメントを検索して回答する必要がある、ユーザーからの会話文です。
+  会話と新しい会話文に基づいて、検索クエリを作成します。回答は日本語で行います。
+  新しい会話文が挨拶の場合、挨拶を返してください。
+  新しい会話文が質問の場合、検索した結果の回答を返してください。
+  答えがわからない場合は正直にわからないと回答してください。
+  会話履歴:
+  {chat_history}
+  新しい会話文:
+  {question}
+  Search query:
+""")
+
+chat_history = []
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("__name__")
+logger.debug("調査用ログ")
+
+def initialize_index():
+    logger.info("initialize_index start")
+    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)
+    # デバッグ用
+    llama_debug_handler = LlamaDebugHandler()
+    callback_manager = CallbackManager([llama_debug_handler])
+    service_context = ServiceContext.from_defaults(node_parser=node_parser,callback_manager=callback_manager)
+    lock = Lock()
+    with lock:
+        if os.path.exists(index_name):
+            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)
+            response_synthesizer = get_response_synthesizer(response_mode='refine')
+            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,
+                chat_history=chat_history,
+                verbose=True
+            )
+        else:
+            documents = SimpleDirectoryReader("./documents").load_data()
+            vector_store = FaissVectorStore(faiss_index=faiss_index)
+            storage_context = StorageContext.from_defaults(vector_store=vector_store)
+            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,service_context=service_context)
+            st.session_state.chat_engine = CondenseQuestionChatEngine.from_defaults(
+                query_engine=st.session_state.query_engine, 
+                condense_question_prompt=custom_prompt,
+                chat_history=chat_history,
+                verbose=True
+            )
+        if os.path.exists(pkl_name):
+            with open(pkl_name, "rb") as f:
+                st.session_state.stored_docs = pickle.load(f)
+        else:
+            st.session_state.stored_docs=list()
+
+# 接続元IP取得
+def get_remote_ip():
+    ctx = get_script_run_ctx()
+    session_info = runtime.get_instance().get_client(ctx.session_id)
+    return session_info.request.remote_ip
+
+# 接続元IP許可判定
+def is_allow_ip_address():
+    remote_ip = get_remote_ip()
+
+    # localhost
+    if remote_ip == "::1":
+        return True
+
+    # プライベートIP
+    ipaddr = ipaddress.IPv4Address(remote_ip)
+    if ipaddr.is_private:
+        return True
+
+    # そ��他（許可リスト判定）
+    return remote_ip in ALLOW_IP_ADDRESS
+
+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 is_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()
+            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()
+            st.experimental_rerun()
+        st.text("ログイン済みです。")
+        st.text("サイドバーから利用するメニューをお選びください。")