Update src.py

src.py

@@ -6,12 +6,14 @@ from PIL import Image
 from pandasai.llm import HuggingFaceTextGen
 from dotenv import load_dotenv
 from langchain_groq.chat_models import ChatGroq
+from langchain_google_genai import GoogleGenerativeAI
 
 load_dotenv()
 Groq_Token = os.environ["GROQ_API_KEY"]
-models = {"mixtral": "mixtral-8x7b-32768", "llama": "llama2-70b-4096", "gemma": "gemma-7b-it"}
+models = {"mixtral": "mixtral-8x7b-32768", "llama": "llama2-70b-4096", "gemma": "gemma-7b-it", "gemini-pro": "gemini-pro"}
 
 hf_token = os.getenv("HF_READ")
+gemini_token = os.getenv("GEMINI_TOKEN")
 
 def preprocess_and_load_df(path: str) -> pd.DataFrame:
     df = pd.read_csv(path)
@@ -27,7 +29,10 @@ def load_agent(df: pd.DataFrame, context: str, inference_server: str, name="mixt
     # top_k=5,
     # )
     # llm.client.headers = {"Authorization": f"Bearer {hf_token}"}
-    llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)
+    if name == "gemini-pro":
+        llm = GoogleGenerativeAI(model=model, google_api_key=gemini_token, temperature=0.1)
+    else:
+        llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)
     
     agent = Agent(df, config={"llm": llm, "enable_cache": False, "options": {"wait_for_model": True}})
     agent.add_message(context)
@@ -86,7 +91,10 @@ def show_response(st, response):
             return {"is_image": False}
 
 def ask_question(model_name, question):
-    llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0.1)
+    if model_name == "gemini-pro":
+        llm = GoogleGenerativeAI(model=model, google_api_key=os.environ.get("GOOGLE_API_KEY"), temperature=0)
+    else:
+        llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0.1)
 
     df_check = pd.read_csv("Data.csv")
     df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
@@ -121,11 +129,13 @@ df["Timestamp"] = pd.to_datetime(df["Timestamp"])
     {template}
 
     """
-
-    answer = llm.invoke(query)
+    if model_name == "gemini-pro":
+        answer = llm.invoke(query)
+    else:
+        answer = llm.invoke(query).content
     code = f"""
     {template.split("```python")[1].split("```")[0]}
-    {answer.content.split("```python")[1].split("```")[0]}
+    {answer.split("```python")[1].split("```")[0]}
     """
     # update variable `answer` when code is executed
     exec(code)