add delete_from_kb support to trash google drive file tool

surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py

@@ -20,6 +20,7 @@ def create_trash_google_drive_file_tool(
     @tool
     async def trash_google_drive_file(
         file_name: str,
+        delete_from_kb: bool = False,
     ) -> dict[str, Any]:
         """Move a Google Drive file to trash.
 
@@ -28,11 +29,15 @@ def create_trash_google_drive_file_tool(
 
         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:
@@ -41,13 +46,13 @@ def create_trash_google_drive_file_tool(
             - 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 the action.
+              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"trash_google_drive_file called: file_name='{file_name}'")
+        logger.info(f"trash_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 {
@@ -71,6 +76,7 @@ def create_trash_google_drive_file_tool(
 
             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:
@@ -80,7 +86,7 @@ def create_trash_google_drive_file_tool(
                 }
 
             logger.info(
-                f"Requesting approval for trashing Google Drive file: '{file_name}' (file_id={file_id})"
+                f"Requesting approval for trashing Google Drive file: '{file_name}' (file_id={file_id}, delete_from_kb={delete_from_kb})"
             )
             approval = interrupt(
                 {
@@ -90,6 +96,7 @@ def create_trash_google_drive_file_tool(
                         "params": {
                             "file_id": file_id,
                             "connector_id": connector_id_from_context,
+                            "delete_from_kb": delete_from_kb,
                         },
                     },
                     "context": context,
@@ -124,6 +131,7 @@ def create_trash_google_drive_file_tool(
 
             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."}
@@ -167,12 +175,44 @@ def create_trash_google_drive_file_tool(
                 raise
 
             logger.info(f"Google Drive file trashed: file_id={final_file_id}")
-            return {
+
+            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
 

surfsense_web/components/tool-ui/google-drive/trash-file.tsx

@@ -49,6 +49,14 @@ interface SuccessResult {
 	status: "success";
 	file_id: string;
 	message?: string;
+	deleted_from_kb?: boolean;
+}
+
+interface WarningResult {
+	status: "success";
+	warning: string;
+	file_id?: string;
+	message?: string;
 }
 
 interface ErrorResult {
@@ -70,6 +78,7 @@ interface InsufficientPermissionsResult {
 type TrashGoogleDriveFileResult =
 	| InterruptResult
 	| SuccessResult
+	| WarningResult
 	| ErrorResult
 	| NotFoundResult
 	| InsufficientPermissionsResult;
@@ -101,6 +110,17 @@ function isNotFoundResult(result: unknown): result is NotFoundResult {
 	);
 }
 
+function isWarningResult(result: unknown): result is WarningResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as WarningResult).status === "success" &&
+		"warning" in result &&
+		typeof (result as WarningResult).warning === "string"
+	);
+}
+
 function isInsufficientPermissionsResult(result: unknown): result is InsufficientPermissionsResult {
 	return (
 		typeof result === "object" &&
@@ -130,6 +150,7 @@ function ApprovalCard({
 	const [decided, setDecided] = useState<"approve" | "reject" | null>(
 		interruptData.__decided__ ?? null
 	);
+	const [deleteFromKb, setDeleteFromKb] = useState(false);
 
 	const account = interruptData.context?.account;
 	const file = interruptData.context?.file;
@@ -218,6 +239,26 @@ function ApprovalCard({
 				</div>
 			)}
 
+			{/* Checkbox for deleting from knowledge base */}
+			{!decided && (
+				<div className="px-4 py-3 border-b border-border bg-muted/20">
+					<label className="flex items-start gap-2 cursor-pointer">
+						<input
+							type="checkbox"
+							checked={deleteFromKb}
+							onChange={(e) => setDeleteFromKb(e.target.checked)}
+							className="mt-0.5"
+						/>
+						<div className="flex-1">
+							<span className="text-sm text-foreground">Also remove from knowledge base</span>
+							<p className="text-xs text-muted-foreground mt-1">
+								⚠️ This will permanently delete the file from your knowledge base (cannot be undone)
+							</p>
+						</div>
+					</label>
+				</div>
+			)}
+
 			{/* Action buttons */}
 			<div
 				className={`flex items-center gap-2 border-t ${
@@ -252,6 +293,7 @@ function ApprovalCard({
 										args: {
 											file_id: file?.file_id,
 											connector_id: account?.id,
+											delete_from_kb: deleteFromKb,
 										},
 									},
 								});
@@ -328,6 +370,24 @@ function InsufficientPermissionsCard({ result }: { result: InsufficientPermissio
 	);
 }
 
+function WarningCard({ result }: { result: WarningResult }) {
+	return (
+		<div className="my-4 max-w-md overflow-hidden rounded-xl border border-amber-500/50 bg-card">
+			<div className="flex items-center gap-3 border-b border-amber-500/50 px-4 py-3">
+				<div className="flex size-9 shrink-0 items-center justify-center rounded-lg bg-amber-500/10">
+					<AlertTriangleIcon className="size-4 text-amber-500" />
+				</div>
+				<div className="min-w-0 flex-1">
+					<p className="text-sm font-medium text-amber-600 dark:text-amber-500">Partial success</p>
+				</div>
+			</div>
+			<div className="space-y-2 px-4 py-3 text-xs">
+				<p className="text-sm text-muted-foreground">{result.warning}</p>
+			</div>
+		</div>
+	);
+}
+
 function ErrorCard({ result }: { result: ErrorResult }) {
 	return (
 		<div className="my-4 max-w-md overflow-hidden rounded-xl border border-destructive/50 bg-card">
@@ -364,7 +424,7 @@ function NotFoundCard({ result }: { result: NotFoundResult }) {
 function SuccessCard({ result }: { result: SuccessResult }) {
 	return (
 		<div className="my-4 max-w-md overflow-hidden rounded-xl border border-border bg-card">
-			<div className="flex items-center gap-3 px-4 py-3">
+			<div className="flex items-center gap-3 border-b border-border px-4 py-3">
 				<div className="flex size-9 shrink-0 items-center justify-center rounded-lg bg-green-500/10">
 					<CheckIcon className="size-4 text-green-500" />
 				</div>
@@ -374,12 +434,19 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 					</p>
 				</div>
 			</div>
+			{result.deleted_from_kb && (
+				<div className="px-4 py-3 text-xs">
+					<span className="text-green-600 dark:text-green-500">
+						✓ Also removed from knowledge base
+					</span>
+				</div>
+			)}
 		</div>
 	);
 }
 
 export const TrashGoogleDriveFileToolUI = makeAssistantToolUI<
-	{ file_name: string },
+	{ file_name: string; delete_from_kb?: boolean },
 	TrashGoogleDriveFileResult
 >({
 	toolName: "trash_google_drive_file",
@@ -421,6 +488,7 @@ export const TrashGoogleDriveFileToolUI = makeAssistantToolUI<
 			return <InsufficientPermissionsCard result={result} />;
 
 		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+		if (isWarningResult(result)) return <WarningCard result={result} />;
 		if (isErrorResult(result)) return <ErrorCard result={result} />;
 
 		return <SuccessCard result={result as SuccessResult} />;