lion-devs / otomati-app

Otomati APP help communicate with AI models in a more human-like way.
Apache License 2.0
1 stars 1 forks source link

[FEAT]: Add model selection #10

Closed JustinHung0407 closed 3 months ago

JustinHung0407 commented 3 months ago

Description

Description

We need to implement model selection in our Streamlit app. Users should be able to select a model (e.g., OpenAI, AFS, Ollama) and provide the corresponding API key. The app should dynamically handle the model and API key based on user selection.

1. config.py

Load API keys and model options from environment variables.

class Config:
    def __init__(self):
        self.models = ["OpenAI", "AFS", "Ollama"]
        self.load_from_env()

    def load_from_env(self):
        import os
        self.openai_api_key = os.getenv('OPENAI_API_KEY', "")
        self.afs_api_key = os.getenv('AFS_API_KEY', "")
        self.ollama_api_key = os.getenv('OLLAMA_API_KEY', "")
        self.model = os.getenv('MODEL', "OpenAI")

config = Config()

2. sidebar.py

Create a Sidebar class to handle model selection and API key input.

class Sidebar:
    def __init__(self):
        self.initialize_session_state()

    def initialize_session_state(self):
        if 'selected_model' not in st.session_state:
            st.session_state['selected_model'] = config.model
        if 'api_keys' not in st.session_state:
            st.session_state['api_keys'] = {
                "OpenAI": config.openai_api_key,
                "AFS": config.afs_api_key,
                "Ollama": config.ollama_api_key
            }

    def render_sidebar(self):
        with st.sidebar:
            st.title("Settings")
            selected_model = st.selectbox("Select Model", options=config.models, index=config.models.index(st.session_state['selected_model']))
            api_key_label = f"{selected_model} API Key"
            api_key = st.text_input(api_key_label, value=st.session_state['api_keys'][selected_model], key=f"{selected_model}_api_key_input", type="password")
            st.session_state['api_keys'][selected_model] = api_key
            if selected_model == "OpenAI":
                st.markdown("[Get an OpenAI API key](https://platform.openai.com/account/api-keys)")
            elif selected_model == "AFS":
                st.markdown("[Get an AFS API key](https://afs.example.com/account/api-keys)")
            elif selected_model == "Ollama":
                st.markdown("[Get an Ollama API key](https://ollama.example.com/account/api-keys)")
            with st.expander("Other Settings"):
                st.markdown("[View the source code](https://github.com/lion-devs/langchain-demo)")

otsidebar = Sidebar()

3. Home.py

Use the selected model and corresponding API key.

def generate_response(t, key, model):
    llm = ChatOpenAI(openai_api_key=key, model=model, temperature=0, max_retries=3)
    prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful assistant that tells jokes about {topic}."), ("human", "Tell me a joke about {topic}.")])
    response = (prompt | llm).invoke({"topic": t})
    st.info(response.content)

def main():
    load_dotenv()
    st.title("Otomati APP")
    otsidebar.render_sidebar()
    selected_model = st.session_state['selected_model']
    api_key = st.session_state['api_keys'][selected_model]
    if not api_key:
        st.warning(f"Please add your {selected_model} API key to continue.")
    with st.form("my_form"):
        topic = st.text_area("Enter a topic:", "programming")
        submitted = st.form_submit_button("Submit")
        if submitted and api_key:
            generate_response(topic, api_key, selected_model)

if __name__ == '__main__':
    main()

Additional Information

No response