option to switch models

app.py

@@ -15,10 +15,27 @@ set_initial_state()
 sidebar()
 
 st.write("# Get the summaries of latest top Hacker News posts 🧡")
+if st.session_state.get("model") == None:
+    mistral, openai = st.columns(2)
 
-if st.session_state.get("HF_TGI_TOKEN"):
-    pipeline = start_haystack(st.session_state.get("HF_TGI_TOKEN"))
-    st.session_state["api_key_configured"] = True
+    with mistral:
+        mistral_pressed = st.button("Mistral")
+        if mistral_pressed:
+            st.session_state["model"] = "Mistral"
+    with openai:
+        openai_pressed = st.button("OpenAI")
+        if openai_pressed:
+            st.session_state["model"] = "GPT-4"
+
+if st.session_state.get("model") and (st.session_state.get("HF_TGI_TOKEN") or st.session_state.get("OPENAI_API_KEY")):
+    if st.session_state.get("HF_TGI_TOKEN"):
+        pipeline = start_haystack(st.session_state.get("HF_TGI_TOKEN"), st.session_state.get("model"))
+        st.session_state["api_key_configured"] = True
+    
+    elif st.session_state.get("OPENAI_API_KEY"):
+        pipeline = start_haystack(st.session_state.get("OPENAI_API_KEY"), st.session_state.get("model"))
+        st.session_state["api_key_configured"] = True
+    
     search_bar, button = st.columns(2)
     # Search bar
     with search_bar: 
@@ -29,7 +46,7 @@ if st.session_state.get("HF_TGI_TOKEN"):
         st.write("")
         run_pressed = st.button("Get summaries")
 else:
-    st.write("Please provide your Hugging Face Token to start using the application")
+    st.write("Please provide your Hugging Face or OpenAI key to start using the application")
     st.write("If you are using a smaller screen, open the sidebar from the top left to provide your token 🙌")
     
 if st.session_state.get("api_key_configured"):

utils/haystack.py

@@ -1,10 +1,10 @@
 import streamlit as st
 from haystack import Pipeline
 from haystack.components.builders.prompt_builder import PromptBuilder
-from haystack.components.generators import HuggingFaceTGIGenerator
+from haystack.components.generators import HuggingFaceTGIGenerator, OpenAIGenerator
 from .hackernews_fetcher import HackernewsFetcher
 
-def start_haystack(hf_token):
+def start_haystack(key, model):
     prompt_template = """
 You will be provided one or more top HakcerNews posts, followed by their URL.
 For the posts you have, provide a short summary followed by the URL that the post can be found at.
@@ -18,7 +18,10 @@ Summaries:
 """
 
     prompt_builder = PromptBuilder(template=prompt_template)
-    llm = HuggingFaceTGIGenerator("mistralai/Mixtral-8x7B-Instruct-v0.1", token=hf_token)
+    if model == "Mistral":
+        llm = HuggingFaceTGIGenerator("mistralai/Mixtral-8x7B-Instruct-v0.1", token=key)
+    elif model == "GPT-4":
+        llm = OpenAIGenerator(api_key=key, model="gpt-4")
     fetcher = HackernewsFetcher()
 
     pipe = Pipeline()
@@ -34,9 +37,14 @@ Summaries:
 @st.cache_data(show_spinner=True)
 def query(top_k, _pipeline):
     try:
-        replies = _pipeline.run(data={"hackernews_fetcher": {"top_k": top_k}, 
-                                      "llm": {"generation_kwargs": {"max_new_tokens": 600}}
-                                      })
+        run_args = {"hackernews_fetcher": {"top_k": top_k}}
+        
+        if st.session_state.get("model") == "Mistral":
+            run_args = {"hackernews_fetcher": {"top_k": top_k}, 
+                        "llm": {"generation_kwargs": {"max_new_tokens": 600}}
+                        }
+
+        replies = _pipeline.run(data=run_args)
         
         result = replies['llm']['replies']
     except Exception as e:

utils/ui.py

@@ -9,14 +9,18 @@ def set_initial_state():
     set_state_if_absent("top_k", "How many of the top posts would you like a summary for?")
     set_state_if_absent("result", None)
     set_state_if_absent("haystack_started", False)
+    set_state_if_absent("model", None)
 
 def reset_results(*args):
     st.session_state.result = None
     st.session_state.top_k = None
 
-def set_openai_api_key(api_key: str):
+def set_hf_token(api_key: str):
     st.session_state["HF_TGI_TOKEN"] = api_key
 
+def set_openai_key(api_key: str):
+    st.session_state["OPENAI_API_KEY"] = api_key
+
 def sidebar():
     with st.sidebar:
         # image = Image.open('logo/haystack-logo-colored.png')
@@ -33,16 +37,28 @@ def sidebar():
             "3. Enjoy 🤗\n"
         )
 
-        api_key_input = st.text_input(
-            "Hugging Face Token",
-            type="password",
-            placeholder="Paste your Hugging Face TGI Token",
-            help="You can get your API key from https://platform.openai.com/account/api-keys.",
-            value=st.session_state.get("HF_TGI_TOKEN", ""),
-        )
-
-        if api_key_input:
-            set_openai_api_key(api_key_input)
+        if st.session_state.model == "Mistral":
+            api_key_input = st.text_input(
+                "Hugging Face Token",
+                type="password",
+                placeholder="Paste your Hugging Face TGI Token",
+                help="You can get your API key from https://platform.openai.com/account/api-keys.",
+                value=st.session_state.get("HF_TGI_TOKEN", ""),
+            )
+            if api_key_input:
+                set_hf_token(api_key_input)
+        
+        elif st.session_state.model == "GPT-4":
+            api_key_input = st.text_input(
+                "OpenAI API Key",
+                type="password",
+                placeholder="Paste your OpenAI API Key",
+                help="You can get your API key from https://platform.openai.com/account/api-keys.",
+                value=st.session_state.get("OPENAI_API_KEY", ""),
+            )
+            if api_key_input:
+                set_openai_key(api_key_input)
+        
 
         st.markdown("---")
         st.markdown(