add gemini

app.py

@@ -13,6 +13,7 @@ from langchain_core.runnables import RunnablePassthrough
 from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
 from langchain_together import ChatTogether
+from langchain_google_genai import ChatGoogleGenerativeAI
 import streamlit as st
 
 import utils_mod
@@ -53,11 +54,18 @@ TOGETHER_CHAT_MODELS = {
         "cost": {"pmi": 5.00, "pmo": 5.00}
     },
 }
+GOOGLE_CHAT_MODELS = {
+    "gemini-1.5-flash": {"cost": {"pmi": 0.0, "pmo": 0.0}},
+    "gemini-1.5-pro": {"cost": {"pmi": 0.0, "pmo": 0.0}},
+    "gemini-1.5-pro-exp-0801": {"cost": {"pmi": 0.0, "pmo": 0.0}},
+}
+
 
 PROVIDER_MODELS = {
     "OpenAI": OPENAI_CHAT_MODELS,
     "Anthropic": ANTHROPIC_CHAT_MODELS,
     "Together": TOGETHER_CHAT_MODELS,
+    "Google": GOOGLE_CHAT_MODELS,
 }
 
 
@@ -215,6 +223,15 @@ def get_llm(gen_config: dict):
                 api_key=st.secrets["together_api_key"],
             )
 
+        case "Google":
+            llm = ChatGoogleGenerativeAI(
+                model=gen_config["model_name"],
+                temperature=gen_config["temperature"],
+                api_key=st.secrets["google_api_key"],
+                max_output_tokens=gen_config["max_output_tokens"],
+                top_p=gen_config["top_p"],
+            )
+            
         case _:
             raise ValueError()
 

usage_mod.py

@@ -43,6 +43,20 @@ def get_together_token_usage(response_metadata: dict, model_info: dict):
     }
 
 
+def get_google_token_usage(response_metadata: dict, model_info: dict):
+    input_tokens = 0
+    output_tokens = 0
+    cost = (
+        input_tokens * 1e-6 * model_info["cost"]["pmi"]
+        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
+    )
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cost": cost,
+    }
+
+
 def get_token_usage(response_metadata: dict, model_info: dict, provider: str):
     match provider:
         case "OpenAI":
@@ -51,6 +65,8 @@ def get_token_usage(response_metadata: dict, model_info: dict, provider: str):
             return get_anthropic_token_usage(response_metadata, model_info)
         case "Together":
             return get_together_token_usage(response_metadata, model_info)
+        case "Google":
+            return get_google_token_usage(response_metadata, model_info)
         case _:
             raise ValueError()