sheet_extractor.py

import math
import json
import regex
import inspect
from ast import literal_eval
from transformers import Tool

from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build


class DataExtractorTool(Tool):
    def __init__(self):
        super().__init__()

    def fetch_data(self, source_link: str) -> str:
        """
        Abstract method to fetch raw data from the source.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    def __call__(self, source_link: str) -> dict:
        raw_data = self.fetch_data(source_link)
        return raw_data


class GoogleSheetExtractorTool(DataExtractorTool):
    name = "google_sheet_extractor_tool"
    description = """
    Tool to extract data from Google Sheets.
    Input is source_link which is a str of a url or google sheets id
    e.x. source_link='https://docs.google.com/spreadsheets/d/SHEETS_ID/'.
    Output is a string.
    """
    inputs = ["text"]
    outputs = ["text"]

    def fetch_data(self, source_link: str) -> str:
        # Set up the credentials
        scope = ["https://www.googleapis.com/auth/spreadsheets.readonly"]
        creds = Credentials.from_authorized_user_file('token.json', scope)
        service = build('sheets', 'v4', credentials=creds)

        # Open the spreadsheet and get all values
        source_link = source_link.rstrip("/")
        spreadsheet_id = source_link.split('/')[-1]
        range_name = 'Sheet1!A1:Z1000'
        sheet = service.spreadsheets()
        result = sheet.values().get(spreadsheetId=spreadsheet_id,
                                    range=range_name).execute()
        data = result.get('values', [])
        
        # Convert the data to a string representation
        raw_data = '\n'.join([','.join(row) for row in data])
        print(raw_data)
        return raw_data