add configuration options (#13)

app.py

@@ -24,6 +24,11 @@ def set_api_key():
     st.session_state.api_message = gm.api_message(openai.api_key)
 
 
+def click_wk_button():
+    """Set the OpenAI API key."""
+    st.session_state.wk_button = not st.session_state.wk_button
+
+
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 if "api_message" not in st.session_state:
@@ -44,16 +49,28 @@ with st.sidebar:
     st.write(
         "You can find your API key at https://platform.openai.com/account/api-keys"
     )
+    if "wk_button" not in st.session_state:
+        st.session_state.wk_button = False
+
+    st.checkbox(
+        "Use Wikipedia", on_change=click_wk_button, value=st.session_state.wk_button
+    )
+    st.subheader("Creativity")
+    st.write("The higher the value, the crazier the text.")
+    st.slider(
+        "Temperature",
+        min_value=0.0,
+        max_value=2.0,
+        value=0.9,
+        step=0.01,
+        key="temperature",
+    )
+
 
 # Build settings
 chroma_db = ChromaDB(openai.api_key)
 openai_client, collection = settings.build(chroma_db)
 
-# Create Agent
-openai_api_key = openai.api_key
-llm = ChatOpenAI(temperature=0.9, model="gpt-3.5-turbo-16k", api_key=openai_api_key)
-agent = PDFExplainer(llm, chroma_db).agent
-
 # Main
 st.title("PDF Explainer")
 st.subheader("Create your knowledge base")
@@ -81,18 +98,26 @@ else:
     st.write("Please upload a file of type: pdf")
 
 st.subheader("Search on your knowledge base")
-# if st.button("Chroma data collection"):
-#     st.write(collection)
-
-# if st.button("Delete Chroma Collection"):
-#     try:
-#         chroma_db.client.delete_collection(collection.name)
-#     except AttributeError:
-#         st.error("Collection erased.")
 
 prompt = st.chat_input()
 
 if prompt:
+    # Create Agent
+    try:
+        openai_api_key = openai.api_key
+        llm = ChatOpenAI(
+            temperature=st.session_state.temperature,
+            model="gpt-3.5-turbo-16k",
+            api_key=openai.api_key,
+        )
+        agent = PDFExplainer(
+            llm,
+            chroma_db,
+            extra_tools=st.session_state.wk_button,
+        ).agent
+    except Exception:  # pylint: disable=broad-exception-caught
+        st.warning("Missing OpenAI API Key.")
+
     st.chat_message("user").write(prompt)
     with st.chat_message("assistant"):
         st_callback = StreamlitCallbackHandler(st.container())

requirements.txt

@@ -3,6 +3,7 @@ tiktoken
 langchain
 pymupdf
 pypdf
+duckduckgo-search
 chromadb>='0.4.18'
 sentence_transformers
 streamlit

src/agent.py

@@ -1,5 +1,5 @@
 """An Langchain Agent that uses ChromaDB as a query tool"""
-from langchain.agents import AgentType, initialize_agent
+from langchain.agents import AgentType, initialize_agent, load_tools
 from langchain.tools import Tool
 from src.search import Search
 
@@ -7,7 +7,7 @@ from src.search import Search
 class PDFExplainer:
     """An Agent that uses ChromaDB as a query tool"""
 
-    def __init__(self, llm, chroma_db):
+    def __init__(self, llm, chroma_db, extra_tools=False):
         """Initialize the Agent"""
         search = Search(chroma_db)
 
@@ -15,11 +15,14 @@ class PDFExplainer:
             Tool.from_function(
                 func=search.run,
                 name="Search DB",
-                description="Useful when you need more context about a specific topic.",
+                description="Useful when you need more context for answering a question.",
                 handle_parsing_errors=True,
             )
         ]
 
+        if extra_tools:
+            self.tools.extend(load_tools(["wikipedia"]))
+
         self.agent = initialize_agent(
             self.tools,
             llm,