add gemini

app.py

@@ -1,7 +1,12 @@
 import streamlit as st
-from langchain_pipeline import pipeline
+from langchain_pipeline import pipeline, model_names
 
-st.title("Composure AI")
+st.title("Canarie AI Prototype")
+st.subheader("Finding the canarie in the coal mine")
+
+option = st.selectbox(
+    "Model",
+    model_names())
 
 uploaded_file = st.file_uploader("Choose a file")
 if uploaded_file is not None:

langchain_pipeline.py

@@ -5,18 +5,29 @@ from langchain_astradb import AstraDBVectorStore
 from langchain_core.prompts import PromptTemplate
 from langchain_openai import OpenAIEmbeddings
 from langchain_anthropic import ChatAnthropic
+from langchain_google_genai import ChatGoogleGenerativeAI
 
 ASTRA_DB_API_ENDPOINT = os.environ["ASTRA_DB_API_ENDPOINT"]
 ASTRA_DB_APPLICATION_TOKEN = os.environ["ASTRA_DB_APPLICATION_TOKEN"]
 OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
 ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]
+GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
 
 collection_name = "ilj_test"
 
 embedding = OpenAIEmbeddings(model="text-embedding-ada-002")
-model = ChatAnthropic(model='claude-3-sonnet-20240229')
 
-def pipeline(bytes):
+models = {
+    "claude-3": ChatAnthropic(model='claude-3-sonnet-20240229'),
+    "gemini-pro": ChatGoogleGenerativeAI(model="gemini-pro")
+}
+
+
+def model_names():
+    return models.keys()
+
+
+def pipeline(bytes, model_name):
     disclosure_text = high_level.extract_text(bytes)
 #     disclosure_text = doc[0].page_content
 #
@@ -52,6 +63,6 @@ def pipeline(bytes):
     )
     val = prompt.format(context=related_docs, disclosure={disclosure_text})
 
-    chat_response = model.invoke(input=val)
+    chat_response = models[model_name].invoke(input=val)
 
     return chat_response.content