change api key handling

app.py

@@ -1,7 +1,9 @@
+import logging
 import os
 
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
+from langchain.llms import HuggingFaceHub
 from langchain.prompts.chat import (
 ChatPromptTemplate,
 MessagesPlaceholder,
@@ -15,9 +17,73 @@ from openai.error import AuthenticationError
 import streamlit as st
 
 
-def set_api_key(USER_API_KEY):
-    os.environ["OPENAI_API_KEY"] = str(USER_API_KEY)
-    st.cache_resource.clear()  # clear existing chains with old api key
+@st.cache_resource
+class KeyManager():
+    """
+    Stores the original API keys from environment variables, which
+    can be overwritten if user supplies keys.
+    Also stores the currently active API key for each model provider and updates
+    these based on user input.
+    """
+    def __init__(self):
+        self.provider_names = {"OpenAI" : "OPENAI_API_KEY",
+             "HuggingFace" : "HUGGINGFACEHUB_API_TOKEN"} 
+        self.original_keys = {k : os.environ.get(v) for k, v
+                              in self.provider_names.items()}
+        self.current_keys = {k: os.environ.get(v) for k, v in self.provider_names.items()}
+        self.user_keys = {}  # most recent key supplied by user for each provider
+
+    def set_key(self, api_key, model_provider, user_entered=False):
+        self.current_keys[model_provider] = api_key
+        os.environ[self.provider_names[model_provider]] = api_key
+        if user_entered:
+            self.user_keys[model_provider] = api_key
+        get_chain.clear()
+
+    def list_keys(self):
+        """
+        For debugging purposes only. Do not use in deployed app. 
+        """
+        st.write("Active API keys:")
+        for k, v in self.provider_names.items():
+            st.write(k, " : ", os.environ.get(v))
+        st.write("Current API keys:")
+        for k, v in self.current_keys.items():
+            st.write(k, " : ", v)
+        st.write("User-supplied API keys:")
+        for k, v in self.user_keys.items():
+            st.write(k, " : ", v)
+        st.write("Original API keys:")
+        for k, v in self.original_keys.items():
+            st.write(k, " : ", v)
+
+    def configure_api_key(self, user_api_key, use_provided_key, model_provider):
+        """
+        Set the currently active API key(s) based on user input.
+        """
+        if user_api_key:
+            if use_provided_key:
+                st.warning("API key entered and 'use provided key' checked;"
+                           " using the key you entered", icon="⚠️")
+            self.set_key(str(user_api_key), model_provider, user_entered=True)
+            return True
+            
+        if use_provided_key:
+            self.set_key(self.original_keys[model_provider], model_provider)
+            return True
+            
+        if not user_api_key and not use_provided_key:
+            # check if user previously supplied a key for this provider
+            if model_provider in self.user_keys:
+                self.set_key(self.user_keys[model_provider], model_provider)
+                st.warning("No key entered and 'use provided key' not checked;"
+                           f" using previously entered {model_provider} key",  icon="⚠️")
+                return True
+
+            else:
+                st.warning("Enter an API key or check 'use provided key'"
+                           " to get started", icon="⚠️")
+                return False
 
 
 @st.cache_resource
@@ -26,12 +92,18 @@ def setup_memory():
     memory = ConversationBufferWindowMemory(k=3, memory_key="chat_history", 
                                             chat_memory=msgs,
                                             return_messages=True)
+    logging.info("setting up new chat memory")
     return memory
 
 
 @st.cache_resource
-def get_chain(model_name, _memory, temperature):
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
+def get_chain(model_name, model_provider, _memory, temperature):
+    logging.info(f"setting up new chain with params {model_name}, {model_provider}, {temperature}")
+    if model_provider == "OpenAI":                               
+        llm = ChatOpenAI(model_name=model_name, temperature=temperature)
+    elif model_provider == "HuggingFace":
+        llm = HuggingFaceHub(repo_id=model_name,
+                             model_kwargs={"temperature": temperature, "max_length": 64})
     prompt = ChatPromptTemplate(
     messages=[
         SystemMessagePromptTemplate.from_template(
@@ -50,31 +122,42 @@ def get_chain(model_name, _memory, temperature):
     return conversation
 
 
+
+
             
 if __name__ == "__main__":
-
+    logging.basicConfig(level=logging.INFO)
+    
     st.header("Basic chatbot")
-
     st.write("On small screens, click the `>` at top left to get started")
-
     with st.expander("How conversation history works"):
         st.write("To keep input lengths down and costs reasonable,"
                  " this bot only 'remembers' the past three turns of conversation.")
         st.write("To clear all memory and start fresh, click 'Clear history'" )
-
     st.sidebar.title("Choose options and enter API key")
 
-    USER_API_KEY = st.sidebar.text_input(
-        'API Key',
-        type='password',
-        help="Enter your OpenAI API key to use this app",
-        value=None)
-
+    #### USER INPUT ######
     model_name = st.sidebar.selectbox(
         label = "Choose a model",
-        options = ["gpt-3.5-turbo", "gpt-4"],
+        options = ["gpt-3.5-turbo (OpenAI)", 
+                   "bigscience/bloom (HuggingFace)"
+                  ],
         help="Which LLM to use",
-        )
+    )
+    
+    user_api_key = st.sidebar.text_input(
+        'Enter your API Key',
+        type='password',
+        help="Enter an API key for the appropriate model provider",
+        value="")
+
+    use_provided_key = st.sidebar.checkbox(
+        "Or use provided key",
+        help="If you don't have a key, you can use mine; usage limits apply.",
+    )
+
+    st.sidebar.write("Set the decoding temperature. Higher temperatures give "
+                     "more unpredictable outputs.")
     
     temperature = st.sidebar.slider(
         label="Temperature",
@@ -82,17 +165,20 @@ if __name__ == "__main__":
         max_value=1.0,
         step=0.1,
         value=0.9,
-        help="Set the decoding temperature. Lower temperatures give more predictable outputs."
+        help="Set the decoding temperature"
         )
-
+    ##########################
     
-    if USER_API_KEY is not None:
-        set_api_key(USER_API_KEY)
+    model = model_name.split("(")[0].rstrip()   # remove name of model provider
+    model_provider = model_name.split("(")[-1].split(")")[0]
+    key_manager = KeyManager()
+    if key_manager.configure_api_key(user_api_key, use_provided_key, model_provider):
+        # key_manager.list_keys()
         memory = setup_memory()
-        chain = get_chain(model_name, memory, temperature)
+        chain = get_chain(model, model_provider, memory, temperature)
         if st.button("Clear history"):
             chain.memory.clear()
-            st.cache_resource.clear()
+            # st.cache_resource.clear()
         for message in chain.memory.buffer:    # display chat history
             st.chat_message(message.type).write(message.content)
         text = st.chat_input()
@@ -103,11 +189,13 @@ if __name__ == "__main__":
                 result = chain.predict(input=text)
                 with st.chat_message("assistant"):
                     st.write(result)
-            except AuthenticationError:
+            except (AuthenticationError, ValueError):
                 st.warning("Enter a valid API key", icon="⚠️")
 
 
 
+
+