Spaces:
Sleeping
Sleeping
| import io | |
| import json | |
| from typing import List, Dict, Literal, Optional | |
| from googleapiclient.discovery import build | |
| from googleapiclient.http import MediaIoBaseUpload, MediaIoBaseDownload | |
| from google.oauth2 import service_account | |
| class ChatUploader: | |
| def __init__( | |
| self, | |
| service_account_json: dict, | |
| root_folder_id: str = "1KtfVgL1Rg1iX-ZMHH4Im__ss-pgbaDM9", | |
| ): | |
| """ | |
| Initializes a new chat uploader instance using a service account JSON dict. | |
| By default writes into a fixed root folder. | |
| """ | |
| credentials = service_account.Credentials.from_service_account_info( | |
| service_account_json, | |
| scopes=["https://www.googleapis.com/auth/drive"], | |
| ) | |
| # cache_discovery=False avoids deprecation noise | |
| self.drive_service = build( | |
| "drive", "v3", credentials=credentials, cache_discovery=False | |
| ) | |
| self.root_folder_id = root_folder_id | |
| def _get_or_create_browser_folder(self, browser_id: str) -> str: | |
| """ | |
| Ensure a per-browser folder 'browser_{browser_id}' exists; return its file ID. | |
| """ | |
| folder_name = f"browser_{browser_id}" | |
| query = ( | |
| f"name = '{folder_name}' and '{self.root_folder_id}' in parents and " | |
| "mimeType = 'application/vnd.google-apps.folder' and trashed = false" | |
| ) | |
| results = self.drive_service.files().list(q=query, fields="files(id)").execute() | |
| folders = results.get("files", []) | |
| if folders: | |
| return folders[0]["id"] | |
| metadata = { | |
| "name": folder_name, | |
| "mimeType": "application/vnd.google-apps.folder", | |
| "parents": [self.root_folder_id], | |
| } | |
| folder = self.drive_service.files().create(body=metadata, fields="id").execute() | |
| return folder["id"] | |
| def _find_file(self, name: str, parent_id: str) -> Optional[str]: | |
| """ | |
| Return file ID for a JSON file with given name in parent, else None. | |
| """ | |
| query = ( | |
| f"name = '{name}' and '{parent_id}' in parents and " | |
| "mimeType = 'application/json' and trashed = false" | |
| ) | |
| results = self.drive_service.files().list(q=query, fields="files(id)").execute() | |
| files = results.get("files", []) | |
| return files[0]["id"] if files else None | |
| def upload_chat_history( | |
| self, | |
| chat_history: List[Dict[str, str]], | |
| browser_id: str, | |
| filename: str = "chat_log.json", | |
| mode: Literal["overwrite", "append"] = "overwrite", | |
| ) -> None: | |
| """ | |
| Write the chat log inside the browser's folder. | |
| - overwrite (default): REPLACE file contents with the provided chat_history | |
| (this is what you want to keep Drive in sync with the UI) | |
| - append: read existing JSON array and extend it with chat_history | |
| chat_history is expected to be the *complete* transcript you want stored | |
| (for overwrite), already normalized to [{role, content}, ...]. | |
| """ | |
| folder_id = self._get_or_create_browser_folder(browser_id) | |
| file_id = self._find_file(filename, folder_id) | |
| payload: List[Dict[str, str]] = chat_history | |
| if mode == "append" and file_id: | |
| # Load existing file and extend | |
| request = self.drive_service.files().get_media(fileId=file_id) | |
| existing_stream = io.BytesIO() | |
| downloader = MediaIoBaseDownload(existing_stream, request) | |
| done = False | |
| while not done: | |
| _, done = downloader.next_chunk() | |
| existing_stream.seek(0) | |
| try: | |
| existing_chat = json.loads(existing_stream.read()) | |
| if isinstance(existing_chat, list): | |
| payload = existing_chat + chat_history | |
| except json.JSONDecodeError: | |
| # Fall back to current chat_history only | |
| payload = chat_history | |
| content = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8") | |
| media = MediaIoBaseUpload(io.BytesIO(content), mimetype="application/json") | |
| if file_id: | |
| # REPLACE contents | |
| self.drive_service.files().update( | |
| fileId=file_id, media_body=media | |
| ).execute() | |
| else: | |
| metadata = { | |
| "name": filename, | |
| "parents": [folder_id], | |
| "mimeType": "application/json", | |
| } | |
| self.drive_service.files().create(body=metadata, media_body=media).execute() | |