Spaces:
Build error
Build error
| # 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") |