use session_state instead of st.cache

app.py

@@ -17,76 +17,6 @@ from openai.error import AuthenticationError
 import streamlit as st
 
 
-@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
 def setup_memory():
     msgs = StreamlitChatMessageHistory(key="basic_chat_app")
     memory = ConversationBufferWindowMemory(k=3, memory_key="chat_history", 
@@ -94,40 +24,49 @@ def setup_memory():
                                             return_messages=True)
     logging.info("setting up new chat memory")
     return memory
+    
 
-
-@st.cache_resource
-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(
-            "You are a nice chatbot having a conversation with a human."
-        ),
-        MessagesPlaceholder(variable_name="chat_history"),
-        HumanMessagePromptTemplate.from_template("{input}")
-    ]
-    )
-    conversation = LLMChain(
-                    llm=llm,
-                    prompt=prompt,
-                    verbose=True,
-                    memory=memory
-                )
-    return conversation
-
-
-
-
-            
+def use_existing_chain(model, provider, temp):
+    if "current_chain" in st.session_state:
+        current_chain = st.session_state.current_chain
+        if (current_chain.model == model) \
+            and (current_chain.provider == provider) \
+            and (current_chain.temp == temp):
+            return True
+    return False
+
+
+class CurrentChain():
+    def __init__(self, model, provider, memory, temp):
+        self.model = model
+        self.provider = provider
+        self.temp = temp
+        logging.info(f"setting up new chain with params {model_name}, {provider}, {temp}")
+        if provider == "OpenAI":                               
+            llm = ChatOpenAI(model_name=model, temperature=temp)
+        elif provider == "HuggingFace":
+            llm = HuggingFaceHub(repo_id=model,
+                                 model_kwargs={"temperature": temp, "max_length": 64})
+        prompt = ChatPromptTemplate(
+        messages=[
+            SystemMessagePromptTemplate.from_template(
+                "You are a nice chatbot having a conversation with a human."
+            ),
+            MessagesPlaceholder(variable_name="chat_history"),
+            HumanMessagePromptTemplate.from_template("{input}")
+        ]
+        )
+        self.conversation = LLMChain(
+                        llm=llm,
+                        prompt=prompt,
+                        verbose=True,
+                        memory=memory
+                    )
+        
+  
 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"):
@@ -144,22 +83,11 @@ if __name__ == "__main__":
                   ],
         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(
+    temp = st.sidebar.slider(
         label="Temperature",
         min_value=float(0),
         max_value=1.0,
@@ -168,29 +96,35 @@ if __name__ == "__main__":
         help="Set the decoding temperature"
         )
     ##########################
-    
     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, model_provider, memory, temperature)
-        if st.button("Clear history"):
-            chain.memory.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()
-        if text:
-            with st.chat_message("user"):
-                st.write(text)
-            try:
-                result = chain.predict(input=text)
-                with st.chat_message("assistant"):
-                    st.write(result)
-            except (AuthenticationError, ValueError):
-                st.warning("Enter a valid API key", icon="⚠️")
+    provider = model_name.split("(")[-1].split(")")[0]
+    if "session_memory" not in st.session_state:
+        st.session_state.session_memory = setup_memory()
+            
+    if use_existing_chain(model, provider, temp):
+        chain = st.session_state.current_chain
+    else:
+        chain =  CurrentChain(model, 
+                     provider, 
+                     st.session_state.session_memory, 
+                     temp)
+        st.session_state.current_chain = chain 
+
+    conversation = chain.conversation
+    if st.button("Clear history"):
+        conversation.memory.clear()
+    for message in conversation.memory.buffer:    # display chat history
+        st.chat_message(message.type).write(message.content)
+    text = st.chat_input()
+    if text:
+        with st.chat_message("user"):
+            st.write(text)
+        try:
+            result = conversation.predict(input=text)
+            with st.chat_message("assistant"):
+                st.write(result)
+        except (AuthenticationError, ValueError):
+            st.warning("Enter a valid API key", icon="⚠️")