Update app.py

app.py

@@ -17,6 +17,7 @@ st.set_page_config(layout="wide")
 load_dotenv()
 Groq_Token = os.environ["GROQ_API_KEY"]
 hf_token = os.environ["HF_TOKEN"]
+gemini_token = os.environ["GEMINI_TOKEN"]
 models = {"llama3":"llama3-70b-8192","mixtral": "mixtral-8x7b-32768", "llama2": "llama2-70b-4096", "gemma": "gemma-7b-it", "gemini-pro": "gemini-pro"}
 
 self_path = os.path.dirname(os.path.abspath(__file__))
@@ -205,7 +206,10 @@ if prompt:
             ran = False
             for i in range(1):
                 print(f"Attempt {i+1}")
-                llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0)
+                if model_name == "gemini-pro":
+                    llm = GoogleGenerativeAI(model=model, google_api_key=os.getenv("GEMINI_TOKEN"), temperature=0)
+                else:
+                    llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0)
 
                 df_check = pd.read_csv("Data.csv")
                 df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
@@ -263,10 +267,13 @@ import uuid
                 code = None
                 error = None
                 try:
-                    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)