app.py

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="⚠️")