KnowledgeGPT works with Azure OpenAI

knowledge_gpt/components/sidebar.py

@@ -15,6 +15,7 @@ def sidebar():
             "2. Upload a pdf, docx, or txt file📄\n"
             "3. Ask a question about the document💬\n"
         )
+        """
         api_key_input = st.text_input(
             "OpenAI API Key",
             type="password",
@@ -25,6 +26,7 @@ def sidebar():
         )
 
         st.session_state["OPENAI_API_KEY"] = api_key_input
+        """
 
         st.markdown("---")
         st.markdown("# About")

knowledge_gpt/main.py

@@ -1,14 +1,24 @@
 import os
+os.environ["OPENAI_API_TYPE"] = "azure"  # configure API to Azure OpenAI
 
 import streamlit as st
 
+# add all secrets into environmental variables
+try:
+    for key, value in st.secrets.items():
+        # os.environ[key] = value
+        st.session_state[key] = value
+except FileNotFoundError as e:
+    print(e)
+    print("./streamlit/secrets.toml not found. Assuming secrets are already available" "as environmental variables...")
+
+
 from knowledge_gpt.components.sidebar import sidebar
 
 from knowledge_gpt.ui import (
     wrap_doc_in_html,
     is_query_valid,
     is_file_valid,
-    is_open_ai_key_valid,
     display_file_read_error,
 )
 
@@ -20,14 +30,7 @@ from knowledge_gpt.core.embedding import embed_files
 from knowledge_gpt.core.qa import query_folder
 from knowledge_gpt.core.utils import get_llm
 
-
-# add all secrets into environmental variables
-try:
-    for key, value in st.secrets.items():
-        os.environ[key] = value
-except FileNotFoundError as e:
-    print(e)
-    print("./streamlit/secrets.toml not found. Assuming secrets are already available" "as environmental variables...")
+from langchain.chat_models import AzureChatOpenAI
 
 
 def main():
@@ -38,8 +41,8 @@ def main():
     # Uncomment to enable debug mode
     # MODEL_LIST.insert(0, "debug")
 
-    st.set_page_config(page_title="KnowledgeGPT", page_icon="📖", layout="wide")
-    st.header("📖KnowledgeGPT")
+    st.set_page_config(page_title="ReferenceBot", page_icon="📖", layout="wide")
+    st.header("📖ReferenceBot")
 
     # Enable caching for expensive functions
     bootstrap_caching()
@@ -79,15 +82,17 @@ def main():
     if not is_file_valid(file):
         st.stop()
 
-    if not is_open_ai_key_valid(openai_api_key, model):
-        st.stop()
-
     with st.spinner("Indexing document... This may take a while⏳"):
         folder_index = embed_files(
             files=[chunked_file],
             embedding=EMBEDDING if model != "debug" else "debug",
             vector_store=VECTOR_STORE if model != "debug" else "debug",
-            openai_api_key=openai_api_key,
+            deployment=st.secrets["ENGINE_EMBEDDING"],
+            model=st.secrets["ENGINE"],
+            openai_api_key=st.secrets["OPENAI_API_KEY"],
+            openai_api_base=st.secrets["OPENAI_API_BASE"],
+            openai_api_type="azure",
+            chunk_size = 1,
         )
 
     with st.form(key="qa_form"):
@@ -106,13 +111,23 @@ def main():
         # Output Columns
         answer_col, sources_col = st.columns(2)
 
-        llm = get_llm(model=model, openai_api_key=openai_api_key, temperature=0)
-        result = query_folder(
-            folder_index=folder_index,
-            query=query,
-            return_all=return_all_chunks,
-            llm=llm,
-        )
+        with st.spinner("Setting up AzureChatOpenAI bot..."):
+            llm = AzureChatOpenAI(
+                openai_api_base=st.secrets["OPENAI_API_BASE"],
+                openai_api_version=st.secrets["OPENAI_API_VERSION"],
+                deployment_name=st.secrets["ENGINE"],
+                openai_api_key=st.secrets["OPENAI_API_KEY"],
+                openai_api_type="azure",
+                temperature=0,
+            )
+        
+        with st.spinner("Querying folder to get result..."):
+            result = query_folder(
+                folder_index=folder_index,
+                query=query,
+                return_all=return_all_chunks,
+                llm=llm,
+            )
 
         with answer_col:
             st.markdown("#### Answer")

knowledge_gpt/ui.py

@@ -2,7 +2,7 @@ from typing import List
 import streamlit as st
 from langchain.docstore.document import Document
 from knowledge_gpt.core.parsing import File
-import openai
+#import openai
 from streamlit.logger import get_logger
 from typing import NoReturn
 
@@ -51,11 +51,12 @@ def is_open_ai_key_valid(openai_api_key, model: str) -> bool:
         st.error("Please enter your OpenAI API key in the sidebar!")
         return False
     try:
-        openai.ChatCompletion.create(
-            model=model,
-            messages=[{"role": "user", "content": "test"}],
-            api_key=openai_api_key,
-        )
+        pass
+        #openai.ChatCompletion.create(
+        #    engine="chatbot-streamlit",#model,
+        #    messages=[{"role": "user", "content": "test"}],
+        #    #api_key=openai_api_key,
+        #)
     except Exception as e:
         st.error(f"{e.__class__.__name__}: {e}")
         logger.error(f"{e.__class__.__name__}: {e}")