diff --git a/surfsense_backend/app/agents/new_chat/tools/google_drive/__init__.py b/surfsense_backend/app/agents/new_chat/tools/google_drive/__init__.py new file mode 100644 index 000000000..9c63bceb1 --- /dev/null +++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/__init__.py @@ -0,0 +1,11 @@ +from app.agents.new_chat.tools.google_drive.create_file import ( + create_create_google_drive_file_tool, +) +from app.agents.new_chat.tools.google_drive.trash_file import ( + create_delete_google_drive_file_tool, +) + +__all__ = [ + "create_create_google_drive_file_tool", + "create_delete_google_drive_file_tool", +] diff --git a/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py b/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py new file mode 100644 index 000000000..0dd683f7e --- /dev/null +++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py @@ -0,0 +1,227 @@ +import logging +from typing import Any, Literal + +from googleapiclient.errors import HttpError +from langchain_core.tools import tool +from langgraph.types import interrupt +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.google_drive.client import GoogleDriveClient +from app.connectors.google_drive.file_types import GOOGLE_DOC, GOOGLE_SHEET +from app.services.google_drive import GoogleDriveToolMetadataService + +logger = logging.getLogger(__name__) + +_MIME_MAP: dict[str, str] = { + "google_doc": GOOGLE_DOC, + "google_sheet": GOOGLE_SHEET, +} + + +def create_create_google_drive_file_tool( + db_session: AsyncSession | None = None, + search_space_id: int | None = None, + user_id: str | None = None, +): + @tool + async def create_google_drive_file( + name: str, + file_type: Literal["google_doc", "google_sheet"], + content: str | None = None, + ) -> dict[str, Any]: + """Create a new Google Doc or Google Sheet in Google Drive. + + Use this tool when the user explicitly asks to create a new document + or spreadsheet in Google Drive. + + Args: + name: The file name (without extension). + file_type: Either "google_doc" or "google_sheet". + content: Optional initial content. For google_doc, provide markdown text. + For google_sheet, provide CSV-formatted text. + + Returns: + Dictionary with: + - status: "success", "rejected", or "error" + - file_id: Google Drive file ID (if success) + - name: File name (if success) + - web_view_link: URL to open the file (if success) + - message: Result message + + IMPORTANT: + - If status is "rejected", the user explicitly declined the action. + Respond with a brief acknowledgment and do NOT retry or suggest alternatives. + - If status is "insufficient_permissions", the connector lacks the required OAuth scope. + Inform the user they need to re-authenticate and do NOT retry the action. + + Examples: + - "Create a Google Doc called 'Meeting Notes'" + - "Create a spreadsheet named 'Budget 2026' with some sample data" + """ + logger.info(f"create_google_drive_file called: name='{name}', type='{file_type}'") + + if db_session is None or search_space_id is None or user_id is None: + return { + "status": "error", + "message": "Google Drive tool not properly configured. Please contact support.", + } + + if file_type not in _MIME_MAP: + return { + "status": "error", + "message": f"Unsupported file type '{file_type}'. Use 'google_doc' or 'google_sheet'.", + } + + try: + metadata_service = GoogleDriveToolMetadataService(db_session) + context = await metadata_service.get_creation_context(search_space_id, user_id) + + if "error" in context: + logger.error(f"Failed to fetch creation context: {context['error']}") + return {"status": "error", "message": context["error"]} + + logger.info( + f"Requesting approval for creating Google Drive file: name='{name}', type='{file_type}'" + ) + approval = interrupt( + { + "type": "google_drive_file_creation", + "action": { + "tool": "create_google_drive_file", + "params": { + "name": name, + "file_type": file_type, + "content": content, + "connector_id": None, + "parent_folder_id": None, + }, + }, + "context": context, + } + ) + + decisions_raw = approval.get("decisions", []) if isinstance(approval, dict) else [] + decisions = decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": + return { + "status": "rejected", + "message": "User declined. The file was not created. Do not ask again or suggest alternatives.", + } + + final_params: dict[str, Any] = {} + edited_action = decision.get("edited_action") + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_name = final_params.get("name", name) + final_file_type = final_params.get("file_type", file_type) + final_content = final_params.get("content", content) + final_connector_id = final_params.get("connector_id") + final_parent_folder_id = final_params.get("parent_folder_id") + + if not final_name or not final_name.strip(): + return {"status": "error", "message": "File name cannot be empty."} + + mime_type = _MIME_MAP.get(final_file_type) + if not mime_type: + return { + "status": "error", + "message": f"Unsupported file type '{final_file_type}'.", + } + + from sqlalchemy.future import select + + from app.db import SearchSourceConnector, SearchSourceConnectorType + + if final_connector_id is not None: + result = await db_session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == final_connector_id, + SearchSourceConnector.search_space_id == search_space_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + return { + "status": "error", + "message": "Selected Google Drive connector is invalid or has been disconnected.", + } + actual_connector_id = connector.id + else: + result = await db_session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.search_space_id == search_space_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + return { + "status": "error", + "message": "No Google Drive connector found. Please connect Google Drive in your workspace settings.", + } + actual_connector_id = connector.id + + logger.info( + f"Creating Google Drive file: name='{final_name}', type='{final_file_type}', connector={actual_connector_id}" + ) + client = GoogleDriveClient(session=db_session, connector_id=actual_connector_id) + try: + created = await client.create_file( + name=final_name, + mime_type=mime_type, + parent_folder_id=final_parent_folder_id, + content=final_content, + ) + except HttpError as http_err: + if http_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {actual_connector_id}: {http_err}" + ) + return { + "status": "insufficient_permissions", + "connector_id": actual_connector_id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate.", + } + raise + + logger.info(f"Google Drive file created: id={created.get('id')}, name={created.get('name')}") + return { + "status": "success", + "file_id": created.get("id"), + "name": created.get("name"), + "web_view_link": created.get("webViewLink"), + "message": f"Successfully created '{created.get('name')}' in Google Drive.", + } + + except Exception as e: + from langgraph.errors import GraphInterrupt + + if isinstance(e, GraphInterrupt): + raise + + logger.error(f"Error creating Google Drive file: {e}", exc_info=True) + return { + "status": "error", + "message": "Something went wrong while creating the file. Please try again.", + } + + return create_google_drive_file diff --git a/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py b/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py new file mode 100644 index 000000000..600aae983 --- /dev/null +++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py @@ -0,0 +1,228 @@ +import logging +from typing import Any + +from googleapiclient.errors import HttpError +from langchain_core.tools import tool +from langgraph.types import interrupt +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.google_drive.client import GoogleDriveClient +from app.services.google_drive import GoogleDriveToolMetadataService + +logger = logging.getLogger(__name__) + + +def create_delete_google_drive_file_tool( + db_session: AsyncSession | None = None, + search_space_id: int | None = None, + user_id: str | None = None, +): + @tool + async def delete_google_drive_file( + file_name: str, + delete_from_kb: bool = False, + ) -> dict[str, Any]: + """Move a Google Drive file to trash. + + Use this tool when the user explicitly asks to delete, remove, or trash + a file in Google Drive. + + Args: + file_name: The exact name of the file to trash (as it appears in Drive). + delete_from_kb: Whether to also remove the file from the knowledge base. + Default is False. + Set to True to remove from both Google Drive and knowledge base. + + Returns: + Dictionary with: + - status: "success", "rejected", "not_found", or "error" + - file_id: Google Drive file ID (if success) + - deleted_from_kb: whether the document was removed from the knowledge base + - message: Result message + + IMPORTANT: + - If status is "rejected", the user explicitly declined. Respond with a brief + acknowledgment and do NOT retry or suggest alternatives. + - If status is "not_found", relay the exact message to the user and ask them + to verify the file name or check if it has been indexed. + - If status is "insufficient_permissions", the connector lacks the required OAuth scope. + Inform the user they need to re-authenticate and do NOT retry this tool. + + Examples: + - "Delete the 'Meeting Notes' file from Google Drive" + - "Trash the 'Old Budget' spreadsheet" + """ + logger.info(f"delete_google_drive_file called: file_name='{file_name}', delete_from_kb={delete_from_kb}") + + if db_session is None or search_space_id is None or user_id is None: + return { + "status": "error", + "message": "Google Drive tool not properly configured. Please contact support.", + } + + try: + metadata_service = GoogleDriveToolMetadataService(db_session) + context = await metadata_service.get_trash_context( + search_space_id, user_id, file_name + ) + + if "error" in context: + error_msg = context["error"] + if "not found" in error_msg.lower(): + logger.warning(f"File not found: {error_msg}") + return {"status": "not_found", "message": error_msg} + logger.error(f"Failed to fetch trash context: {error_msg}") + return {"status": "error", "message": error_msg} + + file = context["file"] + file_id = file["file_id"] + document_id = file.get("document_id") + connector_id_from_context = context["account"]["id"] + + if not file_id: + return { + "status": "error", + "message": "File ID is missing from the indexed document. Please re-index the file and try again.", + } + + logger.info( + f"Requesting approval for deleting Google Drive file: '{file_name}' (file_id={file_id}, delete_from_kb={delete_from_kb})" + ) + approval = interrupt( + { + "type": "google_drive_file_trash", + "action": { + "tool": "delete_google_drive_file", + "params": { + "file_id": file_id, + "connector_id": connector_id_from_context, + "delete_from_kb": delete_from_kb, + }, + }, + "context": context, + } + ) + + decisions_raw = approval.get("decisions", []) if isinstance(approval, dict) else [] + decisions = decisions_raw if isinstance(decisions_raw, list) else [decisions_raw] + decisions = [d for d in decisions if isinstance(d, dict)] + if not decisions: + logger.warning("No approval decision received") + return {"status": "error", "message": "No approval decision received"} + + decision = decisions[0] + decision_type = decision.get("type") or decision.get("decision_type") + logger.info(f"User decision: {decision_type}") + + if decision_type == "reject": + return { + "status": "rejected", + "message": "User declined. The file was not trashed. Do not ask again or suggest alternatives.", + } + + edited_action = decision.get("edited_action") + final_params: dict[str, Any] = {} + if isinstance(edited_action, dict): + edited_args = edited_action.get("args") + if isinstance(edited_args, dict): + final_params = edited_args + elif isinstance(decision.get("args"), dict): + final_params = decision["args"] + + final_file_id = final_params.get("file_id", file_id) + final_connector_id = final_params.get("connector_id", connector_id_from_context) + final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) + + if not final_connector_id: + return {"status": "error", "message": "No connector found for this file."} + + from sqlalchemy.future import select + + from app.db import SearchSourceConnector, SearchSourceConnectorType + + result = await db_session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == final_connector_id, + SearchSourceConnector.search_space_id == search_space_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + return { + "status": "error", + "message": "Selected Google Drive connector is invalid or has been disconnected.", + } + + logger.info( + f"Deleting Google Drive file: file_id='{final_file_id}', connector={final_connector_id}" + ) + client = GoogleDriveClient(session=db_session, connector_id=connector.id) + try: + await client.trash_file(file_id=final_file_id) + except HttpError as http_err: + if http_err.resp.status == 403: + logger.warning( + f"Insufficient permissions for connector {connector.id}: {http_err}" + ) + return { + "status": "insufficient_permissions", + "connector_id": connector.id, + "message": "This Google Drive account needs additional permissions. Please re-authenticate.", + } + raise + + logger.info(f"Google Drive file deleted (moved to trash): file_id={final_file_id}") + + trash_result: dict[str, Any] = { + "status": "success", + "file_id": final_file_id, + "message": f"Successfully moved '{file['name']}' to trash.", + } + + deleted_from_kb = False + if final_delete_from_kb and document_id: + try: + from app.db import Document + + doc_result = await db_session.execute( + select(Document).filter(Document.id == document_id) + ) + document = doc_result.scalars().first() + if document: + await db_session.delete(document) + await db_session.commit() + deleted_from_kb = True + logger.info(f"Deleted document {document_id} from knowledge base") + else: + logger.warning(f"Document {document_id} not found in KB") + except Exception as e: + logger.error(f"Failed to delete document from KB: {e}") + await db_session.rollback() + trash_result["warning"] = ( + f"File moved to trash, but failed to remove from knowledge base: {e!s}" + ) + + trash_result["deleted_from_kb"] = deleted_from_kb + if deleted_from_kb: + trash_result["message"] = ( + f"{trash_result.get('message', '')} (also removed from knowledge base)" + ) + + return trash_result + + except Exception as e: + from langgraph.errors import GraphInterrupt + + if isinstance(e, GraphInterrupt): + raise + + logger.error(f"Error deleting Google Drive file: {e}", exc_info=True) + return { + "status": "error", + "message": "Something went wrong while trashing the file. Please try again.", + } + + return delete_google_drive_file diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 4c6345bc3..01342e920 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/surfsense_backend/app/agents/new_chat/tools/registry.py @@ -55,6 +55,10 @@ from .linear import ( ) from .link_preview import create_link_preview_tool from .mcp_tool import load_mcp_tools +from .google_drive import ( + create_create_google_drive_file_tool, + create_delete_google_drive_file_tool, +) from .notion import ( create_create_notion_page_tool, create_delete_notion_page_tool, @@ -292,6 +296,29 @@ BUILTIN_TOOLS: list[ToolDefinition] = [ ), requires=["db_session", "search_space_id", "user_id"], ), + # ========================================================================= + # GOOGLE DRIVE TOOLS - create files, delete files + # ========================================================================= + ToolDefinition( + name="create_google_drive_file", + description="Create a new Google Doc or Google Sheet in Google Drive", + factory=lambda deps: create_create_google_drive_file_tool( + db_session=deps["db_session"], + search_space_id=deps["search_space_id"], + user_id=deps["user_id"], + ), + requires=["db_session", "search_space_id", "user_id"], + ), + ToolDefinition( + name="delete_google_drive_file", + description="Move an indexed Google Drive file to trash", + factory=lambda deps: create_delete_google_drive_file_tool( + db_session=deps["db_session"], + search_space_id=deps["search_space_id"], + user_id=deps["user_id"], + ), + requires=["db_session", "search_space_id", "user_id"], + ), ] diff --git a/surfsense_backend/app/connectors/google_drive/client.py b/surfsense_backend/app/connectors/google_drive/client.py index a001e42be..2910320b2 100644 --- a/surfsense_backend/app/connectors/google_drive/client.py +++ b/surfsense_backend/app/connectors/google_drive/client.py @@ -1,12 +1,15 @@ """Google Drive API client.""" +import io from typing import Any from googleapiclient.discovery import build from googleapiclient.errors import HttpError +from googleapiclient.http import MediaIoBaseUpload from sqlalchemy.ext.asyncio import AsyncSession from .credentials import get_valid_credentials +from .file_types import GOOGLE_DOC, GOOGLE_SHEET class GoogleDriveClient: @@ -179,3 +182,65 @@ class GoogleDriveClient: return None, f"HTTP error exporting file: {e.resp.status}" except Exception as e: return None, f"Error exporting file: {e!s}" + + async def create_file( + self, + name: str, + mime_type: str, + parent_folder_id: str | None = None, + content: str | None = None, + ) -> dict[str, Any]: + service = await self.get_service() + + body: dict[str, Any] = {"name": name, "mimeType": mime_type} + if parent_folder_id: + body["parents"] = [parent_folder_id] + + media: MediaIoBaseUpload | None = None + if content: + if mime_type == GOOGLE_DOC: + import markdown as md_lib + + html = md_lib.markdown(content) + media = MediaIoBaseUpload( + io.BytesIO(html.encode("utf-8")), + mimetype="text/html", + resumable=False, + ) + elif mime_type == GOOGLE_SHEET: + media = MediaIoBaseUpload( + io.BytesIO(content.encode("utf-8")), + mimetype="text/csv", + resumable=False, + ) + + if media: + return ( + service.files() + .create( + body=body, + media_body=media, + fields="id,name,mimeType,webViewLink", + supportsAllDrives=True, + ) + .execute() + ) + + return ( + service.files() + .create( + body=body, + fields="id,name,mimeType,webViewLink", + supportsAllDrives=True, + ) + .execute() + ) + + async def trash_file(self, file_id: str) -> bool: + service = await self.get_service() + service.files().update( + fileId=file_id, + body={"trashed": True}, + supportsAllDrives=True, + ).execute() + return True diff --git a/surfsense_backend/app/routes/google_drive_add_connector_route.py b/surfsense_backend/app/routes/google_drive_add_connector_route.py index 6b4159d29..9fa83f3a2 100644 --- a/surfsense_backend/app/routes/google_drive_add_connector_route.py +++ b/surfsense_backend/app/routes/google_drive_add_connector_route.py @@ -76,9 +76,9 @@ def get_token_encryption() -> TokenEncryption: # Google Drive OAuth scopes SCOPES = [ - "https://www.googleapis.com/auth/drive.readonly", # Read-only access to Drive - "https://www.googleapis.com/auth/userinfo.email", # User email - "https://www.googleapis.com/auth/userinfo.profile", # User profile + "https://www.googleapis.com/auth/drive", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", "openid", ] @@ -151,6 +151,75 @@ async def connect_drive(space_id: int, user: User = Depends(current_active_user) ) from e +@router.get("/auth/google/drive/connector/reauth") +async def reauth_drive( + space_id: int, + connector_id: int, + return_url: str | None = None, + user: User = Depends(current_active_user), + session: AsyncSession = Depends(get_async_session), +): + """ + Initiate Google Drive re-authentication to upgrade OAuth scopes. + + Query params: + space_id: Search space ID the connector belongs to + connector_id: ID of the existing connector to re-authenticate + + Returns: + JSON with auth_url to redirect user to Google authorization + """ + try: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == connector_id, + SearchSourceConnector.user_id == user.id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + raise HTTPException( + status_code=404, + detail="Google Drive connector not found or access denied", + ) + + if not config.SECRET_KEY: + raise HTTPException( + status_code=500, detail="SECRET_KEY not configured for OAuth security." + ) + + flow = get_google_flow() + + state_manager = get_state_manager() + extra: dict = {"connector_id": connector_id} + if return_url and return_url.startswith("/"): + extra["return_url"] = return_url + state_encoded = state_manager.generate_secure_state(space_id, user.id, **extra) + + auth_url, _ = flow.authorization_url( + access_type="offline", + prompt="consent", + include_granted_scopes="true", + state=state_encoded, + ) + + logger.info( + f"Initiating Google Drive re-auth for user {user.id}, connector {connector_id}" + ) + return {"auth_url": auth_url} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Failed to initiate Google Drive re-auth: {e!s}", exc_info=True) + raise HTTPException( + status_code=500, detail=f"Failed to initiate Google re-auth: {e!s}" + ) from e + + @router.get("/auth/google/drive/connector/callback") async def drive_callback( request: Request, @@ -214,6 +283,8 @@ async def drive_callback( user_id = UUID(data["user_id"]) space_id = data["space_id"] + reauth_connector_id = data.get("connector_id") + reauth_return_url = data.get("return_url") logger.info( f"Processing Google Drive callback for user {user_id}, space {space_id}" @@ -253,7 +324,45 @@ async def drive_callback( # Mark that credentials are encrypted for backward compatibility creds_dict["_token_encrypted"] = True - # Check for duplicate connector (same account already connected) + if reauth_connector_id: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == reauth_connector_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + db_connector = result.scalars().first() + if not db_connector: + raise HTTPException( + status_code=404, + detail="Connector not found or access denied during re-auth", + ) + + existing_start_page_token = db_connector.config.get("start_page_token") + db_connector.config = { + **creds_dict, + "start_page_token": existing_start_page_token, + } + from sqlalchemy.orm.attributes import flag_modified + + flag_modified(db_connector, "config") + await session.commit() + await session.refresh(db_connector) + + logger.info( + f"Re-authenticated Google Drive connector {db_connector.id} for user {user_id}" + ) + if reauth_return_url and reauth_return_url.startswith("/"): + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}{reauth_return_url}" + ) + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/new-chat?modal=connectors&tab=all&success=true&connector=google-drive-connector&connectorId={db_connector.id}" + ) + is_duplicate = await check_duplicate_connector( session, SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, diff --git a/surfsense_backend/app/services/google_drive/__init__.py b/surfsense_backend/app/services/google_drive/__init__.py new file mode 100644 index 000000000..5958a1761 --- /dev/null +++ b/surfsense_backend/app/services/google_drive/__init__.py @@ -0,0 +1,11 @@ +from app.services.google_drive.tool_metadata_service import ( + GoogleDriveAccount, + GoogleDriveFile, + GoogleDriveToolMetadataService, +) + +__all__ = [ + "GoogleDriveAccount", + "GoogleDriveFile", + "GoogleDriveToolMetadataService", +] diff --git a/surfsense_backend/app/services/google_drive/tool_metadata_service.py b/surfsense_backend/app/services/google_drive/tool_metadata_service.py new file mode 100644 index 000000000..8bc60ecbc --- /dev/null +++ b/surfsense_backend/app/services/google_drive/tool_metadata_service.py @@ -0,0 +1,149 @@ +from dataclasses import dataclass + +from sqlalchemy import and_, func +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + +from app.db import ( + Document, + DocumentType, + SearchSourceConnector, + SearchSourceConnectorType, +) + + +@dataclass +class GoogleDriveAccount: + id: int + name: str + + @classmethod + def from_connector(cls, connector: SearchSourceConnector) -> "GoogleDriveAccount": + return cls(id=connector.id, name=connector.name) + + def to_dict(self) -> dict: + return {"id": self.id, "name": self.name} + + +@dataclass +class GoogleDriveFile: + file_id: str + name: str + mime_type: str + web_view_link: str + connector_id: int + document_id: int + + @classmethod + def from_document(cls, document: Document) -> "GoogleDriveFile": + meta = document.document_metadata or {} + return cls( + file_id=meta.get("google_drive_file_id", ""), + name=meta.get("google_drive_file_name", document.title), + mime_type=meta.get("google_drive_mime_type", ""), + web_view_link=meta.get("web_view_link", ""), + connector_id=document.connector_id, + document_id=document.id, + ) + + def to_dict(self) -> dict: + return { + "file_id": self.file_id, + "name": self.name, + "mime_type": self.mime_type, + "web_view_link": self.web_view_link, + "connector_id": self.connector_id, + "document_id": self.document_id, + } + + +class GoogleDriveToolMetadataService: + def __init__(self, db_session: AsyncSession): + self._db_session = db_session + + async def get_creation_context(self, search_space_id: int, user_id: str) -> dict: + accounts = await self._get_google_drive_accounts(search_space_id, user_id) + + if not accounts: + return { + "accounts": [], + "supported_types": [], + "error": "No Google Drive account connected", + } + + return { + "accounts": [acc.to_dict() for acc in accounts], + "supported_types": ["google_doc", "google_sheet"], + } + + async def get_trash_context( + self, search_space_id: int, user_id: str, file_name: str + ) -> dict: + result = await self._db_session.execute( + select(Document) + .join( + SearchSourceConnector, Document.connector_id == SearchSourceConnector.id + ) + .filter( + and_( + Document.search_space_id == search_space_id, + Document.document_type == DocumentType.GOOGLE_DRIVE_FILE, + func.lower(Document.title) == func.lower(file_name), + SearchSourceConnector.user_id == user_id, + ) + ) + ) + document = result.scalars().first() + + if not document: + return { + "error": ( + f"File '{file_name}' not found in your indexed Google Drive files. " + "This could mean: (1) the file doesn't exist, (2) it hasn't been indexed yet, " + "or (3) the file name is different." + ) + } + + if not document.connector_id: + return {"error": "Document has no associated connector"} + + result = await self._db_session.execute( + select(SearchSourceConnector).filter( + and_( + SearchSourceConnector.id == document.connector_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + ) + connector = result.scalars().first() + + if not connector: + return {"error": "Connector not found or access denied"} + + account = GoogleDriveAccount.from_connector(connector) + file = GoogleDriveFile.from_document(document) + + return { + "account": account.to_dict(), + "file": file.to_dict(), + } + + async def _get_google_drive_accounts( + self, search_space_id: int, user_id: str + ) -> list[GoogleDriveAccount]: + result = await self._db_session.execute( + select(SearchSourceConnector) + .filter( + and_( + SearchSourceConnector.search_space_id == search_space_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + .order_by(SearchSourceConnector.last_indexed_at.desc()) + ) + connectors = result.scalars().all() + return [GoogleDriveAccount.from_connector(c) for c in connectors] diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index ecf04ce08..f5df5862b 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -840,6 +840,8 @@ async def _stream_agent_events( "create_linear_issue", "update_linear_issue", "delete_linear_issue", + "create_google_drive_file", + "delete_google_drive_file", ): yield streaming_service.format_tool_output_available( tool_call_id, diff --git a/surfsense_backend/pyproject.toml b/surfsense_backend/pyproject.toml index 3df84141d..5f79c3154 100644 --- a/surfsense_backend/pyproject.toml +++ b/surfsense_backend/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "kokoro>=0.9.4", "linkup-sdk>=0.2.4", "llama-cloud-services>=0.6.25", + "Markdown>=3.7", "markdownify>=0.14.1", "notion-client>=2.3.0", "numpy>=1.24.0", diff --git a/surfsense_backend/uv.lock b/surfsense_backend/uv.lock index 50ed66617..68e7d1f1b 100644 --- a/surfsense_backend/uv.lock +++ b/surfsense_backend/uv.lock @@ -6892,6 +6892,7 @@ dependencies = [ { name = "linkup-sdk" }, { name = "litellm" }, { name = "llama-cloud-services" }, + { name = "markdown" }, { name = "markdownify" }, { name = "mcp" }, { name = "notion-client" }, @@ -6963,6 +6964,7 @@ requires-dist = [ { name = "linkup-sdk", specifier = ">=0.2.4" }, { name = "litellm", specifier = ">=1.80.10" }, { name = "llama-cloud-services", specifier = ">=0.6.25" }, + { name = "markdown", specifier = ">=3.7" }, { name = "markdownify", specifier = ">=0.14.1" }, { name = "mcp", specifier = ">=1.25.0" }, { name = "notion-client", specifier = ">=2.3.0" }, diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 8720078cc..d1d98bbf6 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -38,6 +38,10 @@ import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; import { DisplayImageToolUI } from "@/components/tool-ui/display-image"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; import { GenerateReportToolUI } from "@/components/tool-ui/generate-report"; +import { + CreateGoogleDriveFileToolUI, + DeleteGoogleDriveFileToolUI, +} from "@/components/tool-ui/google-drive"; import { CreateLinearIssueToolUI, DeleteLinearIssueToolUI, @@ -152,6 +156,8 @@ const TOOLS_WITH_UI = new Set([ "create_linear_issue", "update_linear_issue", "delete_linear_issue", + "create_google_drive_file", + "delete_google_drive_file", "execute", // "write_todos", // Disabled for now ]); @@ -1666,6 +1672,8 @@ export default function NewChatPage() { + + {/* Disabled for now */}
diff --git a/surfsense_web/components/tool-ui/google-drive/create-file.tsx b/surfsense_web/components/tool-ui/google-drive/create-file.tsx new file mode 100644 index 000000000..d6f08653d --- /dev/null +++ b/surfsense_web/components/tool-ui/google-drive/create-file.tsx @@ -0,0 +1,582 @@ +"use client"; + +import { makeAssistantToolUI } from "@assistant-ui/react"; +import { + AlertTriangleIcon, + CheckIcon, + FileIcon, + Loader2Icon, + PencilIcon, + RefreshCwIcon, + XIcon, +} from "lucide-react"; +import { useParams } from "next/navigation"; +import { useMemo, useState } from "react"; +import { toast } from "sonner"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; +import { authenticatedFetch } from "@/lib/auth-utils"; + +interface GoogleDriveAccount { + id: number; + name: string; +} + +interface InterruptResult { + __interrupt__: true; + __decided__?: "approve" | "reject" | "edit"; + action_requests: Array<{ + name: string; + args: Record; + }>; + review_configs: Array<{ + action_name: string; + allowed_decisions: Array<"approve" | "edit" | "reject">; + }>; + context?: { + accounts?: GoogleDriveAccount[]; + supported_types?: string[]; + error?: string; + }; +} + +interface SuccessResult { + status: "success"; + file_id: string; + name: string; + web_view_link?: string; + message?: string; +} + +interface ErrorResult { + status: "error"; + message: string; +} + +interface InsufficientPermissionsResult { + status: "insufficient_permissions"; + connector_id: number; + message: string; +} + +type CreateGoogleDriveFileResult = + | InterruptResult + | SuccessResult + | ErrorResult + | InsufficientPermissionsResult; + +function isInterruptResult(result: unknown): result is InterruptResult { + return ( + typeof result === "object" && + result !== null && + "__interrupt__" in result && + (result as InterruptResult).__interrupt__ === true + ); +} + +function isErrorResult(result: unknown): result is ErrorResult { + return ( + typeof result === "object" && + result !== null && + "status" in result && + (result as ErrorResult).status === "error" + ); +} + +function isInsufficientPermissionsResult(result: unknown): result is InsufficientPermissionsResult { + return ( + typeof result === "object" && + result !== null && + "status" in result && + (result as InsufficientPermissionsResult).status === "insufficient_permissions" + ); +} + +const FILE_TYPE_LABELS: Record = { + google_doc: "Google Doc", + google_sheet: "Google Sheet", +}; + +function ApprovalCard({ + args, + interruptData, + onDecision, +}: { + args: { name: string; file_type: string; content?: string }; + interruptData: InterruptResult; + onDecision: (decision: { + type: "approve" | "reject" | "edit"; + message?: string; + edited_action?: { name: string; args: Record }; + }) => void; +}) { + const [decided, setDecided] = useState<"approve" | "reject" | "edit" | null>( + interruptData.__decided__ ?? null + ); + const [isEditing, setIsEditing] = useState(false); + const [editedName, setEditedName] = useState(args.name ?? ""); + const [editedContent, setEditedContent] = useState(args.content ?? ""); + const [committedArgs, setCommittedArgs] = useState<{ + name: string; + file_type: string; + content?: string | null; + } | null>(null); + + const accounts = interruptData.context?.accounts ?? []; + + const defaultAccountId = useMemo(() => { + if (accounts.length === 1) return String(accounts[0].id); + return ""; + }, [accounts]); + + const [selectedAccountId, setSelectedAccountId] = useState(defaultAccountId); + const [selectedFileType, setSelectedFileType] = useState(args.file_type ?? "google_doc"); + const [parentFolderId, setParentFolderId] = useState(""); + + const isNameValid = useMemo( + () => (isEditing ? editedName.trim().length > 0 : args.name?.trim().length > 0), + [isEditing, editedName, args.name] + ); + + const canApprove = !!selectedAccountId && isNameValid; + + const reviewConfig = interruptData.review_configs[0]; + const allowedDecisions = reviewConfig?.allowed_decisions ?? ["approve", "reject"]; + const canEdit = allowedDecisions.includes("edit"); + + function buildFinalArgs() { + return { + name: isEditing ? editedName : args.name, + file_type: selectedFileType, + content: isEditing ? editedContent || null : (args.content ?? null), + connector_id: selectedAccountId ? Number(selectedAccountId) : null, + parent_folder_id: parentFolderId.trim() || null, + }; + } + + return ( +
+ {/* Header */} +
+
+ +
+
+

Create Google Drive File

+

+ {isEditing ? "You can edit the arguments below" : "Requires your approval to proceed"} +

+
+
+ + {/* Context section */} + {!decided && interruptData.context && ( +
+ {interruptData.context.error ? ( +

{interruptData.context.error}

+ ) : ( + <> + {accounts.length > 0 && ( +
+
+ Google Drive Account * +
+ +
+ )} + +
+
+ File Type * +
+ +
+ +
+
+ Parent Folder ID (optional) +
+ setParentFolderId(e.target.value)} + placeholder="Leave blank to create at Drive root" + /> +

+ Paste a Google Drive folder ID to place the file in a specific folder. +

+
+ + )} +
+ )} + + {/* Display mode */} + {!isEditing && ( +
+
+

Name

+

{committedArgs?.name ?? args.name}

+
+
+

Type

+

+ {FILE_TYPE_LABELS[committedArgs?.file_type ?? args.file_type] ?? committedArgs?.file_type ?? args.file_type} +

+
+ {(committedArgs?.content ?? args.content) && ( +
+

Content

+

+ {committedArgs?.content ?? args.content} +

+
+ )} +
+ )} + + {/* Edit mode */} + {isEditing && !decided && ( +
+
+ + setEditedName(e.target.value)} + placeholder="Enter file name" + className={!isNameValid ? "border-destructive" : ""} + /> + {!isNameValid &&

Name is required

} +
+
+ +