Update utility.py

utility.py

@@ -7,18 +7,19 @@ from dotenv import load_dotenv
 import pandas as pd
 import ast
 import json
-from pandasai.responses.response_parser import  ResponseParser
+from pandasai.responses.response_parser import ResponseParser
 # Load environment variables and configure Firestore
 load_dotenv()
 db = firestore.Client.from_service_account_json("firestore-key.json")
+from langchain_google_genai import ChatGoogleGenerativeAI
+# Initialize the Gemini LLM client
+import google.generativeai as genai
 
-# Initialize the ChatSambaNovaCloud LLM client (using the latest model)
-from langchain_community.chat_models.sambanova import ChatSambaNovaCloud
-
-llm = ChatSambaNovaCloud(
-    model="Meta-Llama-3.1-70B-Instruct",
-    max_tokens=1024,
-    temperature=0.7,
+# Replace SambaNova with Google Gemini
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash-thinking-exp",
+    max_output_tokens=1024,
+    temperature=0.01,
     top_k=1,
     top_p=0.01,
 )
@@ -32,12 +33,9 @@ class FlaskResponse(ResponseParser):
       return result['value'].to_html()
 
   def format_plot(self, result):
-        # Save the plot using savefig
+      # Save the plot using savefig
       try:
-          
           img_path = result['value']
-      
-        
       except ValueError:
           img_path = str(result['value'])
           print("value error!", img_path)
@@ -48,7 +46,7 @@ class FlaskResponse(ResponseParser):
   def format_other(self, result):
       return str(result['value'])
       
-def generateResponse(prompt, model='Meta-Llama-3.1-70B-Instruct'):
+def generateResponse(prompt, model='gemini-2.0-flash-thinking-exp'):
     # Templates for extracting transaction information
     relevant_info_template = """
 Intent: The CRUD operation (create, read, update, delete)
@@ -69,24 +67,31 @@ Transaction 2:
 *Transaction Type*: Sale
 *Details*: - Item: Chair, - Quantity: 2, etc.
 """
-    response = openai.OpenAI(
-        api_key=os.environ.get("SAMBANOVA_API_KEY"),
-        base_url="https://api.sambanova.ai/v1",
-    ).chat.completions.create(
-        model=model,
-        messages=[
-            {"role": "system", "content": 
-             f"You are a helpful assistant that classifies transactions. Format your output with these guidelines: {relevant_info_template} "
-             f"Sample single transaction: {sample_single_transaction_template} "
-             f"Sample multi-transaction: {sample_multi_transaction_template}"},
-            {"role": "user", "content": prompt}
-        ]
+    
+    # Replace SambaNova with Google Gemini API
+    genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
+    
+    model_instance = genai.GenerativeModel(
+        model_name=model,
+        generation_config={
+            "temperature": 0.1,
+            "top_p": 0.01,
+            "top_k": 1,
+            "max_output_tokens": 1024,
+        }
     )
+    
+    system_prompt = f"You are a helpful assistant that classifies transactions. Format your output with these guidelines: {relevant_info_template} Sample single transaction: {sample_single_transaction_template} Sample multi-transaction: {sample_multi_transaction_template}"
+    
     try:
-        response_text = response.choices[0].message.content
+        response = model_instance.generate_content([
+            {"role": "user", "parts": [system_prompt, prompt]}
+        ])
+        response_text = response.text
     except Exception as e:
         print(f'An error occurred: {str(e)}')
         response_text = None
+    
     return response_text
 
 def parse_value(value):