Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import base64 | |
| import logging | |
| from shillelagh.backends.apsw.db import connect | |
| class DatabaseManager: | |
| def __init__(self, spreadsheet_id, qa_guid, agents_config_guid, credentials_base64): | |
| """ | |
| Initialize the DatabaseManager with Google Sheets credentials and spreadsheet ID. | |
| """ | |
| self.spreadsheet_id = spreadsheet_id | |
| self.qa_table = f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit?gid={qa_guid}" | |
| self.agents_config_table = f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit?gid={agents_config_guid}" | |
| self.credentials_base64 = credentials_base64 | |
| self.connection = self._connect_to_sheets() | |
| self.cursor = self.connection.cursor() | |
| def _connect_to_sheets(self): | |
| """ | |
| Establish a connection to Google Sheets using the Shillelagh library. | |
| """ | |
| try: | |
| # Decode the Base64-encoded credentials | |
| credentials_json = base64.b64decode(self.credentials_base64).decode("utf-8") | |
| credentials_info = json.loads(credentials_json) | |
| return connect( | |
| ":memory:", | |
| adapter_kwargs={ | |
| "gsheetsapi": { | |
| "service_account_info": credentials_info, | |
| } | |
| } | |
| ) | |
| except Exception as e: | |
| logging.error(f"Error decoding credentials or connecting to Google Sheets: {e}") | |
| raise | |
| def read_configs(self): | |
| """ | |
| Read agent configurations from the AgentsConfig table. | |
| Returns: | |
| List of dictionaries containing agent configurations. | |
| """ | |
| try: | |
| query = f"SELECT * FROM '{self.agents_config_table}'" | |
| rows = self.cursor.execute(query).fetchall() | |
| column_names = [description[0] for description in self.cursor.description] | |
| configs = [dict(zip(column_names, row)) for row in rows] | |
| logging.info("Agent configurations retrieved successfully.") | |
| return configs | |
| except Exception as e: | |
| logging.error(f"Error reading agent configurations: {e}") | |
| return [] | |
| def write_qa(self, question, answer, timestamp, openai_link, result_object): | |
| """ | |
| Write a new Q&A entry to the QATable. | |
| """ | |
| try: | |
| query = f""" | |
| INSERT INTO '{self.qa_table}' | |
| (question, answer, timestamp, openai_link, result_object) | |
| VALUES (?, ?, ?, ?, ?) | |
| """ | |
| self.cursor.execute(query, ( | |
| question, answer, timestamp, openai_link, result_object | |
| )) | |
| logging.info("Q&A entry written successfully.") | |
| except Exception as e: | |
| logging.error(f"Error writing Q&A entry: {e}") | |
| def get_questions(self): | |
| """Retrieve a list of questions from the Questions table.""" | |
| try: | |
| questions_table = f"https://docs.google.com/spreadsheets/d/{self.spreadsheet_id}/edit?gid=253246580" | |
| query = f"SELECT * FROM '{questions_table}'" | |
| rows = self.cursor.execute(query).fetchall() | |
| column_names = [description[0] for description in self.cursor.description] | |
| questions = [dict(zip(column_names, row)) for row in rows] | |
| logging.info("Questions retrieved successfully.") | |
| return questions | |
| except Exception as e: | |
| logging.error(f"Error retrieving questions: {e}") | |
| return [] | |
| if __name__ == "__main__": | |
| logging.basicConfig(level=logging.INFO) | |
| # Replace with your actual Google Sheets ID and base64-encoded credentials | |
| SPREADSHEET_ID = "1gjhx54rsj6L9k_XssaVKF4fL1-Ov3me_HHBzs4mcimY" | |
| QA_GUID = "1674797002" | |
| AGENTS_CONFIG_GUID = "0" | |
| CREDENTIALS_BASE64 = os.getenv("GOOGLE_API_JSON_CONTENT") | |
| db_manager = DatabaseManager(SPREADSHEET_ID, QA_GUID, AGENTS_CONFIG_GUID, CREDENTIALS_BASE64) | |
| # Test reading configurations | |
| configs = db_manager.read_configs() | |
| print("Agent Configurations:", configs) | |
| # Test writing a Q&A entry | |
| db_manager.write_qa( | |
| question="What is AI?", | |
| answer="AI stands for Artificial Intelligence.", | |
| timestamp="2025-04-26T12:00:00", | |
| openai_link="https://platform.openai.com/logs/resp_1234567890abcdef", | |
| result_object=json.dumps({"example_key": "example_value"}) | |
| ) |