Clean code & Add avatar

1. clean code
2. add the avatar base on sreamlit_chat: AI-Yash/st-chat#34

Author: Lin-jun-xiang

README.md

@@ -1,5 +1,7 @@
 # docGPT
 
+
+
 [English](./README.md) | [中文版](./README.zh-TW.md)
 
 - Table of Contents
@@ -16,6 +18,9 @@
 
 * Using this tool requires at least the `openai_api_key`. You can visit the [link](https://platform.openai.com/) to learn how to obtain the key.
 
+<p align="center">
+<img src='img/2023-08-24-15-02-11.png' width='100%'>
+</p>
 
 ---
 

app.py

@@ -12,7 +12,7 @@
 from streamlit_chat import message
 
 from agent import AgentHelper
-from docGPT import DocGPT
+from docGPT import DocGPT, OpenAiAPI, SerpAPI
 from model import PDFLoader
 
 langchain.llm_cache = InMemoryCache()
@@ -35,9 +35,6 @@ def theme():
         st.title('PDF Chatbot')
 
 
-theme()
-
-
 def load_api_key() -> None:
     with st.sidebar:
         if st.session_state.openai_api_key:
@@ -70,10 +67,7 @@ def load_api_key() -> None:
         os.environ['SERPAPI_API_KEY'] = SERPAPI_API_KEY
 
 
-load_api_key()
-
-
-with st.container():
+def upload_and_process_pdf():
     upload_file = st.file_uploader('#### Upload a PDF file:', type='pdf')
     if upload_file:
         temp_file = tempfile.NamedTemporaryFile(delete=False)
@@ -87,10 +81,33 @@ def load_api_key() -> None:
         if temp_file_path:
             os.remove(temp_file_path)
 
-        docGPT_tool, calculate_tool, search_tool, llm_tool = [None] * 4
+        return docs
+
+
+@lru_cache(maxsize=20)
+def get_response(query: str):
+    try:
+        if agent_.agent_ is not None:
+            response = agent_.query(query)
+            return response
+    except Exception as e:
+        app_logger.info(e)
+
+
+theme()
+load_api_key()
+
+doc_container = st.container()
+
+
+with doc_container:
+    docs = upload_and_process_pdf()
+
+    agent_, docGPT_tool, calculate_tool, search_tool, llm_tool = [None]*5
+    if OpenAiAPI.is_valid():
+        agent_ = AgentHelper()
 
-        try:
-            agent_ = AgentHelper()
+        if docs:
             docGPT = DocGPT(docs=docs)
             docGPT.create_qa_chain(
                 chain_type='refine',
@@ -100,59 +117,37 @@ def load_api_key() -> None:
             calculate_tool = agent_.get_calculate_chain
             llm_tool = agent_.create_llm_chain()
 
-        except Exception as e:
-            app_logger.info(e)
+    if SerpAPI.is_valid():
+        search_tool = agent_.get_searp_chain
 
-        try:
-            search_tool = agent_.get_searp_chain
-        except Exception as e:
-            app_logger.info(e)
-
-        try:
-            tools = [
-                docGPT_tool,
-                search_tool,
-                # llm_tool, # This will cause agent confuse
-                calculate_tool
-            ]
-            agent_.initialize(tools)
-        except Exception as e:
-            app_logger.info(e)
-
-
-if not st.session_state['openai_api_key']:
-    st.error('⚠️ :red[You have not pass OpenAPI key. (Or your api key cannot use.)] Necessary Pass')
-
-if not st.session_state['serpapi_api_key']:
-    st.warning('⚠️ You have not pass SEARPAPI key. (You cannot ask current events.)')
+    try:
+        tools = [
+            docGPT_tool,
+            search_tool,
+            # llm_tool, # This will cause agent confuse
+            calculate_tool
+        ]
+        agent_.initialize(tools)
+    except Exception as e:
+        app_logger.info(e)
 
-st.write('---')
+    st.write('---')
 
 if 'response' not in st.session_state:
     st.session_state['response'] = ['How can I help you?']
 
 if 'query' not in st.session_state:
     st.session_state['query'] = ['Hi']
 
-
-@lru_cache(maxsize=20)
-def get_response(query: str):
-    try:
-        if agent_.agent_ is not None:
-            response = agent_.query(query)
-            return response
-    except Exception as e:
-        app_logger.info(e)
-
-query = st.text_input(
-    "#### Question:",
-    placeholder='Enter your question'
-)
-
-response_container = st.container()
 user_container = st.container()
+response_container = st.container()
 
 with user_container:
+    query = st.text_input(
+        "#### Question:",
+        placeholder='Enter your question'
+    )
+
     if query and query != '':
         response = get_response(query)
         st.session_state.query.append(query)
@@ -161,5 +156,19 @@ def get_response(query: str):
 with response_container:
     if st.session_state['response']:
         for i in range(len(st.session_state['response'])-1, -1, -1):
-            message(st.session_state["response"][i], key=str(i))
-            message(st.session_state['query'][i], is_user=True, key=str(i) + '_user')
+            message(
+                st.session_state["response"][i], key=str(i),
+                logo=(
+                    'https://api.dicebear.com/6.x/bottts/svg?'
+                    'baseColor=fb8c00&eyes=bulging'
+                )    
+            )
+            message(
+                st.session_state['query'][i], is_user=True, key=str(i) + '_user',
+                logo=(
+                    'https://api.dicebear.com/6.x/adventurer/svg?'
+                    'hair=short16&hairColor=85c2c6&'
+                    'eyes=variant12&size=100&'
+                    'mouth=variant26&skinColor=f2d3b1'
+                )
+            )