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, SystemMessagePromptTemplate, HumanMessagePromptTemplate, ) from langchain.memory import ConversationBufferWindowMemory from langchain.memory.chat_message_histories import StreamlitChatMessageHistory from langchain.schema import AIMessage, HumanMessage 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", chat_memory=msgs, 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 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 INPUT ###### model_name = st.sidebar.selectbox( label = "Choose a model", 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", min_value=float(0), max_value=1.0, step=0.1, value=0.9, 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="⚠️")