feat: Add single row read option for Google Sheets

This change allows the user to specify a single row to be read from a spreadsheet, rather than reading the entire thing or a specified range.
This commit is contained in:
Toran Bruce Richards
2024-07-21 11:15:33 +01:00
parent c16b8763e7
commit 2bb37cc9c4

View File

@@ -2,6 +2,7 @@ from autogpt_server.data.block import Block, BlockSchema, BlockOutput
from autogpt_server.data.model import SchemaField
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from typing import Union, List
class GoogleSheetsWriter(Block):
class Input(BlockSchema):
@@ -79,38 +80,67 @@ class GoogleSheetsWriter(Block):
class GoogleSheetsReader(Block):
class Input(BlockSchema):
spreadsheet_id: str = SchemaField(
description="The ID of the Google Sheet to read",
description="Your Google Sheet ID (found in the sheet's URL)",
placeholder="1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms"
)
range: str = SchemaField(
description="The A1 notation of the range to read",
placeholder="Sheet1!A1:E10",
default= "Sheet1!A1:Z1000"
read_single_row: bool = SchemaField(
description="Do you want to read just one row? Select 'Yes' for a single row, 'No' for multiple rows.",
default=False
)
sheet_name: str = SchemaField(
description="The name of the sheet you want to read from (e.g., 'Sheet1')",
placeholder="Sheet1",
default="Sheet1"
)
row_number: int = SchemaField(
description="If reading a single row, which row number do you want? (e.g., 3 for the third row)",
placeholder="3"
)
cell_range: str = SchemaField(
description="If reading multiple rows, what range of cells do you want? (e.g., 'A1:E10' for the first 10 rows of columns A to E)",
placeholder="A1:E10",
default="A1:Z1000"
)
access_token: str = SchemaField(
description="Google OAuth2 access token",
description="Your Google OAuth2 access token (keep this secret!)",
secret=True
)
class Output(BlockSchema):
rows: list = SchemaField(description="The rows of from the specified Google Sheet range")
data: Union[List[str], List[List[str]]] = SchemaField(
description="The information read from your Google Sheet. For a single row, it's a simple list. For multiple rows, it's a list of lists."
)
def __init__(self):
super().__init__(
id="google-sheets-block",
id="google-sheets-reader",
input_schema=GoogleSheetsReader.Input,
output_schema=GoogleSheetsReader.Output,
)
def run(self, input_data: Input) -> BlockOutput:
credentials = Credentials(input_data.access_token)
service = build("sheets", "v4", credentials=credentials)
sheet = service.spreadsheets()
result = sheet.values().get(spreadsheetId=input_data.spreadsheet_id, range=input_data.range).execute()
values = result.get("values", [])
if not values:
yield "data", {"error": "No data found."}
else:
data = [row for row in values[0:]]
yield "rows", data
try:
credentials = Credentials(input_data.access_token)
service = build("sheets", "v4", credentials=credentials)
sheet = service.spreadsheets()
if input_data.read_single_row:
range_to_read = f"{input_data.sheet_name}!A{input_data.row_number}:ZZ{input_data.row_number}"
else:
range_to_read = f"{input_data.sheet_name}!{input_data.cell_range}"
result = sheet.values().get(spreadsheetId=input_data.spreadsheet_id, range=range_to_read).execute()
values = result.get("values", [])
if not values:
yield "data", {"message": "No data found in the specified range or row."}
else:
if input_data.read_single_row:
# Return a single list for a single row
yield "data", values[0] if values else []
else:
# Return a list of lists for multiple rows
yield "data", values
except Exception as e:
yield "data", {"error": f"Oops! Something went wrong: {str(e)}"}