Ai-Email-Assistant / google_sheets_utils.py
0Learn's picture
Update google_sheets_utils.py
7b73282 verified
# google_sheets_utils.py
# Filepath: ai-email-assistant/google_sheets_utils.py
# Manages interactions with Google Sheets for profile and template storage
import gspread
import json
import os
import logging
from google.oauth2.service_account import Credentials
import uuid
import time
from google.auth.exceptions import TransportError
from dotenv import load_dotenv
logging.basicConfig(level=logging.INFO)
# Constants
SCOPES = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive']
SPREADSHEET_NAME = "AI Email Assistant Data"
WORKSHEET_NAMES = ["Sender Profiles", "Receiver Profiles", "Email Templates"]
HEADERS = {
"Sender Profiles": ["Type", "ID", "Name", "Email", "Position", "Company", "Context"],
"Receiver Profiles": ["Type", "ID", "Name", "Email", "Position", "Company", "Context"],
"Email Templates": ["Type", "ID", "Name", "Subject", "Body"]
}
SAMPLE_DATA = {
"Sender Profiles": ["Sender Profile", "PROF_1", "John Doe", "john.doe@example.com", "Manager", "Example Corp", "Sample context for John Doe"],
"Receiver Profiles": ["Receiver Profile", "PROF_2", "Jane Smith", "jane.smith@example.com", "Developer", "Example Inc", "Sample context for Jane Smith"],
"Email Templates": ["Template", "TEMP_1", "Welcome Email", "Welcome to our service", "Hello, welcome to our service. We are glad to have you."]
}
def load_config():
load_dotenv()
return {
'google_sheets_credentials_file': os.getenv('GOOGLE_SHEETS_CREDENTIALS_FILE')
}
def get_credentials():
config = load_config()
if not config:
return None
creds_json = config.get('google_sheets_credentials_file')
if not creds_json:
logging.error("Google Sheets credentials file path not found in .env file")
return None
try:
creds_dict = json.loads(creds_json)
return Credentials.from_service_account_info(creds_dict, scopes=SCOPES)
except Exception as e:
logging.error("Error creating credentials: %s", str(e))
return None
def get_client():
creds = get_credentials()
if not creds:
logging.error("Failed to obtain credentials")
return None
return gspread.authorize(creds)
def get_or_create_spreadsheet():
gc = get_client()
if not gc:
logging.error("Failed to get Google Sheets client")
return None
max_retries = 3
retry_delay = 5 # seconds
for attempt in range(max_retries):
try:
sh = gc.open(SPREADSHEET_NAME)
logging.info(f"Opened existing spreadsheet: {SPREADSHEET_NAME}")
return sh
except gspread.SpreadsheetNotFound:
try:
sh = gc.create(SPREADSHEET_NAME)
logging.info(f"Created new spreadsheet: {SPREADSHEET_NAME}")
return sh
except Exception as e:
logging.error(f"Failed to create spreadsheet: {str(e)}")
return None
except (gspread.exceptions.APIError, TransportError) as e:
logging.error(f"API Error (attempt {attempt + 1}/{max_retries}): {str(e)}")
if attempt < max_retries - 1:
logging.info(f"Retrying in {retry_delay} seconds...")
time.sleep(retry_delay)
else:
logging.error("Max retries reached. Unable to connect to Google Sheets.")
return None
# Ensure all required worksheets exist
existing_worksheets = [worksheet.title for worksheet in sh.worksheets()]
for worksheet_name in WORKSHEET_NAMES:
if worksheet_name not in existing_worksheets:
try:
worksheet = sh.add_worksheet(title=worksheet_name, rows=1, cols=len(HEADERS[worksheet_name]))
worksheet.append_row(HEADERS[worksheet_name])
worksheet.append_row(SAMPLE_DATA[worksheet_name])
logging.info(f"Created worksheet: {worksheet_name}")
except Exception as e:
logging.error(f"Failed to create worksheet {worksheet_name}: {str(e)}")
return sh
def get_profiles_from_sheet(worksheet):
try:
records = worksheet.get_all_records()
return records
except Exception as e:
logging.error("Error getting profiles from worksheet: %s", str(e))
return []
def save_profile(worksheet, profile_data):
try:
worksheet.append_row(profile_data)
logging.info("Profile saved successfully")
return True
except Exception as e:
logging.error("Error saving profile: %s", str(e))
return False
def delete_profile_from_sheet(worksheet, profile_id):
try:
cell = worksheet.find(profile_id)
worksheet.delete_rows(cell.row) # Use delete_rows instead of delete_row
logging.info(f"Profile {profile_id} deleted successfully")
return True
except Exception as e:
logging.error(f"Error deleting profile {profile_id}: {str(e)}")
return False
def get_templates(worksheet):
try:
records = worksheet.get_all_records()
return records
except Exception as e:
logging.error(f"Error getting templates: {str(e)}")
return []
def save_template(worksheet, template_id, name, template_type, subject, body):
try:
template_data = [template_type, template_id, name, subject, body]
worksheet.append_row(template_data)
logging.info(f"Template saved successfully with ID: {template_id}")
return f"Template saved successfully with ID: {template_id}"
except Exception as e:
logging.error(f"Error saving template: {str(e)}")
return f"Error saving template: {str(e)}"
def delete_template_from_sheet(worksheet, template_id):
try:
cell = worksheet.find(template_id)
worksheet.delete_rows(cell.row) # Use delete_rows instead of delete_row
logging.info(f"Template {template_id} deleted successfully")
return f"Template {template_id} deleted successfully"
except Exception as e:
logging.error(f"Error deleting template {template_id}: {str(e)}")
return f"Error deleting template {template_id}: {str(e)}"
def update_profile_in_sheet(worksheet, profile_id, updated_data):
try:
cell = worksheet.find(profile_id)
for idx, value in enumerate(updated_data):
worksheet.update_cell(cell.row, idx + 1, value)
logging.info(f"Profile {profile_id} updated successfully")
return True
except Exception as e:
logging.error(f"Error updating profile {profile_id}: {str(e)}")
return False
def get_templates_from_sheet(worksheet):
try:
# Get all records from the worksheet
records = worksheet.get_all_records()
# Convert records to a list of dictionaries
templates = []
for record in records:
template = {
"Name": record.get("Name", ""),
"Type": record.get("Type", ""),
"ID": record.get("ID", ""),
"Subject": record.get("Subject", ""),
"Body": record.get("Body", "")
}
templates.append(template)
return templates
except Exception as e:
print(f"An error occurred while fetching templates: {str(e)}")
return []
def get_profile_summaries(worksheet, profile_type):
try:
profiles = worksheet.get_all_records()
return [f"{profile['Name']} ({profile_type})" for profile in profiles]
except Exception as e:
logging.error(f"Error getting profile summaries: {str(e)}")
return []
def get_template_summaries(worksheet):
try:
records = worksheet.get_all_records()
summaries = [{"Name": record["Name"], "ID": record["ID"]} for record in records]
return summaries
except Exception as e:
logging.error(f"Error getting template summaries: {str(e)}")
return []
# Additional utility functions can be added here as needed
def generate_template_id():
return f"TEMP_{uuid.uuid4().hex[:8].upper()}"
def generate_profile_id():
return f"PROF_{uuid.uuid4().hex[:8].upper()}"
def update_template(worksheet, template_id, name, template_type, subject, body):
try:
cell = worksheet.find(template_id)
row = cell.row
worksheet.update(f'A{row}:E{row}', [[template_type, template_id, name, subject, body]])
logging.info(f"Template {template_id} updated successfully")
return True
except Exception as e:
logging.error(f"Error updating template {template_id}: {str(e)}")
return False
def get_profile_by_id(worksheet, profile_id):
try:
cell = worksheet.find(profile_id)
row = worksheet.row_values(cell.row)
profile_data = dict(zip(HEADERS[worksheet.title], row))
return profile_data
except Exception as e:
logging.error(f"Error fetching profile {profile_id}: {str(e)}")
return None
def get_template_by_id(worksheet, template_id):
try:
cell = worksheet.find(template_id)
row = worksheet.row_values(cell.row)
template_data = dict(zip(HEADERS[worksheet.title], row))
return template_data
except Exception as e:
logging.error(f"Error fetching template {template_id}: {str(e)}")
return None
def update_template_in_sheet(worksheet, template_id, updated_data):
try:
cell = worksheet.find(template_id)
for idx, value in enumerate(updated_data):
worksheet.update_cell(cell.row, idx + 1, value)
logging.info(f"Template {template_id} updated successfully")
return True
except Exception as e:
logging.error(f"Error updating template {template_id}: {str(e)}")
return False
if __name__ == "__main__":
# Test the functions
print("Testing Google Sheets integration...")
sh = get_or_create_spreadsheet()
if sh:
print(f"Successfully connected to spreadsheet: {sh.title}")
print("Worksheets:", [worksheet.title for worksheet in sh.worksheets()])
else:
print("Failed to connect to Google Sheets")