Merge pull request #464 from MODSetter/dev

refactor: added batch commits and Increased task time limits in celery_app.py

surfsense_backend/app/celery_app.py

@@ -76,8 +76,8 @@ celery_app.conf.update(
     enable_utc=True,
     # Task execution settings
     task_track_started=True,
-    task_time_limit=3600,  # 1 hour hard limit
+    task_time_limit=28800,  # 8 hour hard limit
-    task_soft_time_limit=3000,  # 50 minutes soft limit
+    task_soft_time_limit=28200,  # 7 hours 50 minutes soft limit
     # Result backend settings
     result_expires=86400,  # Results expire after 24 hours
     result_extended=True,

surfsense_backend/app/connectors/notion_history.py

@@ -58,10 +58,23 @@ class NotionHistoryConnector:
                     "timestamp": "last_edited_time",
                 }
 
-        # First, get a list of all pages the integration has access to
+        # Paginate through all pages the integration has access to
-        search_results = await self.notion.search(**search_params)
+        pages = []
+        has_more = True
+        cursor = None
+
+        while has_more:
+            if cursor:
+                search_params["start_cursor"] = cursor
+
+            search_results = await self.notion.search(**search_params)
+
+            pages.extend(search_results["results"])
+            has_more = search_results.get("has_more", False)
+
+            if has_more:
+                cursor = search_results.get("next_cursor")
 
-        pages = search_results["results"]
         all_page_data = []
 
         for page in pages:

surfsense_backend/app/tasks/connector_indexers/airtable_indexer.py

@@ -390,6 +390,13 @@ async def index_airtable_records(
                                 f"Successfully indexed new Airtable record {summary_content}"
                             )
 
+                            # Batch commit every 10 documents
+                            if documents_indexed % 10 == 0:
+                                logger.info(
+                                    f"Committing batch: {documents_indexed} Airtable records processed so far"
+                                )
+                                await session.commit()
+
                         except Exception as e:
                             logger.error(
                                 f"Error processing the Airtable record {record.get('id', 'Unknown')}: {e!s}",
@@ -408,7 +415,10 @@ async def index_airtable_records(
                             session, connector, update_last_indexed
                         )
 
-                    # Commit all changes
+                    # Final commit for any remaining documents not yet committed in batches
+                    logger.info(
+                        f"Final commit: Total {documents_indexed} Airtable records processed"
+                    )
                     await session.commit()
                     logger.info(
                         "Successfully committed all Airtable document changes to database"

surfsense_backend/app/tasks/connector_indexers/clickup_indexer.py

@@ -354,6 +354,13 @@ async def index_clickup_tasks(
                     documents_indexed += 1
                     logger.info(f"Successfully indexed new task {task_name}")
 
+                    # Batch commit every 10 documents
+                    if documents_indexed % 10 == 0:
+                        logger.info(
+                            f"Committing batch: {documents_indexed} ClickUp tasks processed so far"
+                        )
+                        await session.commit()
+
                 except Exception as e:
                     logger.error(
                         f"Error processing task {task.get('name', 'Unknown')}: {e!s}",
@@ -366,6 +373,8 @@ async def index_clickup_tasks(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} ClickUp tasks processed")
         await session.commit()
 
         await task_logger.log_task_success(

surfsense_backend/app/tasks/connector_indexers/confluence_indexer.py

@@ -368,6 +368,13 @@ async def index_confluence_pages(
                 documents_indexed += 1
                 logger.info(f"Successfully indexed new page {page_title}")
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Confluence pages processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing page {page.get('title', 'Unknown')}: {e!s}",
@@ -384,7 +391,10 @@ async def index_confluence_pages(
         if update_last_indexed:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(
+            f"Final commit: Total {documents_indexed} Confluence pages processed"
+        )
         await session.commit()
         logger.info(
             "Successfully committed all Confluence document changes to database"

surfsense_backend/app/tasks/connector_indexers/discord_indexer.py

@@ -462,6 +462,13 @@ async def index_discord_messages(
                             f"Successfully indexed new channel {guild_name}#{channel_name} with {len(formatted_messages)} messages"
                         )
 
+                        # Batch commit every 10 documents
+                        if documents_indexed % 10 == 0:
+                            logger.info(
+                                f"Committing batch: {documents_indexed} Discord channels processed so far"
+                            )
+                            await session.commit()
+
                 except Exception as e:
                     logger.error(
                         f"Error processing guild {guild_name}: {e!s}", exc_info=True
@@ -476,6 +483,10 @@ async def index_discord_messages(
         if documents_indexed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(
+            f"Final commit: Total {documents_indexed} Discord channels processed"
+        )
         await session.commit()
 
         # Prepare result message

surfsense_backend/app/tasks/connector_indexers/github_indexer.py

@@ -381,13 +381,21 @@ async def index_github_repos(
                     session.add(document)
                     documents_processed += 1
 
+                    # Batch commit every 10 documents
+                    if documents_processed % 10 == 0:
+                        logger.info(
+                            f"Committing batch: {documents_processed} GitHub files processed so far"
+                        )
+                        await session.commit()
+
             except Exception as repo_err:
                 logger.error(
                     f"Failed to process repository {repo_full_name}: {repo_err}"
                 )
                 errors.append(f"Failed processing {repo_full_name}: {repo_err}")
 
-        # Commit all changes at the end
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_processed} GitHub files processed")
         await session.commit()
         logger.info(
             f"Finished GitHub indexing for connector {connector_id}. Processed {documents_processed} files."

surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py

@@ -407,6 +407,13 @@ async def index_google_calendar_events(
                 documents_indexed += 1
                 logger.info(f"Successfully indexed new event {event_summary}")
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Google Calendar events processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing event {event.get('summary', 'Unknown')}: {e!s}",
@@ -422,6 +429,10 @@ async def index_google_calendar_events(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(
+            f"Final commit: Total {documents_indexed} Google Calendar events processed"
+        )
         await session.commit()
 
         await task_logger.log_task_success(

surfsense_backend/app/tasks/connector_indexers/google_gmail_indexer.py

@@ -324,6 +324,13 @@ async def index_google_gmail_messages(
                 documents_indexed += 1
                 logger.info(f"Successfully indexed new email {summary_content}")
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Gmail messages processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing the email {message_id}: {e!s}",
@@ -338,7 +345,8 @@ async def index_google_gmail_messages(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} Gmail messages processed")
         await session.commit()
         logger.info(
             "Successfully committed all Google gmail document changes to database"

surfsense_backend/app/tasks/connector_indexers/jira_indexer.py

@@ -352,6 +352,13 @@ async def index_jira_issues(
                     f"Successfully indexed new issue {issue_identifier} - {issue_title}"
                 )
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Jira issues processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing issue {issue.get('identifier', 'Unknown')}: {e!s}",
@@ -368,7 +375,8 @@ async def index_jira_issues(
         if update_last_indexed:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} Jira issues processed")
         await session.commit()
         logger.info("Successfully committed all JIRA document changes to database")
 

surfsense_backend/app/tasks/connector_indexers/linear_indexer.py

@@ -371,6 +371,13 @@ async def index_linear_issues(
                     f"Successfully indexed new issue {issue_identifier} - {issue_title}"
                 )
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Linear issues processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing issue {issue.get('identifier', 'Unknown')}: {e!s}",
@@ -387,7 +394,8 @@ async def index_linear_issues(
         if update_last_indexed:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} Linear issues processed")
         await session.commit()
         logger.info("Successfully committed all Linear document changes to database")
 

surfsense_backend/app/tasks/connector_indexers/luma_indexer.py

@@ -438,6 +438,13 @@ async def index_luma_events(
                 documents_indexed += 1
                 logger.info(f"Successfully indexed new event {event_name}")
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} Luma events processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing event {event.get('name', 'Unknown')}: {e!s}",
@@ -453,6 +460,8 @@ async def index_luma_events(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} Luma events processed")
         await session.commit()
 
         await task_logger.log_task_success(

surfsense_backend/app/tasks/connector_indexers/notion_indexer.py

@@ -356,6 +356,14 @@ async def index_notion_pages(
 
                         documents_indexed += 1
                         logger.info(f"Successfully updated Notion page: {page_title}")
+
+                        # Batch commit every 10 documents
+                        if documents_indexed % 10 == 0:
+                            logger.info(
+                                f"Committing batch: {documents_indexed} documents processed so far"
+                            )
+                            await session.commit()
+
                         continue
 
                 # Document doesn't exist - create new one
@@ -406,6 +414,13 @@ async def index_notion_pages(
                 documents_indexed += 1
                 logger.info(f"Successfully indexed new Notion page: {page_title}")
 
+                # Batch commit every 10 documents
+                if documents_indexed % 10 == 0:
+                    logger.info(
+                        f"Committing batch: {documents_indexed} documents processed so far"
+                    )
+                    await session.commit()
+
             except Exception as e:
                 logger.error(
                     f"Error processing Notion page {page.get('title', 'Unknown')}: {e!s}",
@@ -423,7 +438,8 @@ async def index_notion_pages(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} documents processed")
         await session.commit()
 
         # Prepare result message

surfsense_backend/app/tasks/connector_indexers/slack_indexer.py

@@ -353,6 +353,14 @@ async def index_slack_messages(
 
                     session.add(document)
                     documents_indexed += 1
+
+                    # Batch commit every 10 documents
+                    if documents_indexed % 10 == 0:
+                        logger.info(
+                            f"Committing batch: {documents_indexed} Slack channels processed so far"
+                        )
+                        await session.commit()
+
                 logger.info(
                     f"Successfully indexed new channel {channel_name} with {len(formatted_messages)} messages"
                 )
@@ -376,7 +384,8 @@ async def index_slack_messages(
         if total_processed > 0:
             await update_connector_last_indexed(session, connector, update_last_indexed)
 
-        # Commit all changes
+        # Final commit for any remaining documents not yet committed in batches
+        logger.info(f"Final commit: Total {documents_indexed} Slack channels processed")
         await session.commit()
 
         # Prepare result message

surfsense_web/app/(home)/login/GoogleLoginButton.tsx

@@ -36,7 +36,7 @@ export function GoogleLoginButton() {
 				<h1 className="my-8 text-xl font-bold text-neutral-800 dark:text-neutral-100 md:text-4xl">
 					{t("welcome_back")}
 				</h1>
-{/* 
+				{/* 
 				<motion.div
 					initial={{ opacity: 0, y: -5 }}
 					animate={{ opacity: 1, y: 0 }}

surfsense_web/app/dashboard/page.tsx

@@ -337,7 +337,11 @@ const DashboardPage = () => {
 													</div>
 												</div>
 											</div>
-											<Link className="flex flex-1 flex-col p-4 cursor-pointer" href={`/dashboard/${space.id}/documents`} key={space.id}>
+											<Link
+												className="flex flex-1 flex-col p-4 cursor-pointer"
+												href={`/dashboard/${space.id}/documents`}
+												key={space.id}
+											>
 												<div className="flex flex-1 flex-col justify-between p-1">
 													<div>
 														<h3 className="font-medium text-lg">{space.name}</h3>

surfsense_web/components/chat/ChatInputGroup.tsx

@@ -26,6 +26,7 @@ import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import { useDocumentTypes } from "@/hooks/use-document-types";
 import type { Document } from "@/hooks/use-documents";
 import { useLLMConfigs, useLLMPreferences } from "@/hooks/use-llm-configs";
+import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
 
 const DocumentSelector = React.memo(
 	({
@@ -119,14 +120,31 @@ const ConnectorSelector = React.memo(
 			true
 		);
 
+		// Fetch live search connectors (non-indexable)
+		const {
+			connectors: searchConnectors,
+			isLoading: connectorsLoading,
+			isLoaded: connectorsLoaded,
+			fetchConnectors,
+		} = useSearchSourceConnectors(true, Number(search_space_id));
+
+		// Filter for non-indexable connectors (live search)
+		const liveSearchConnectors = React.useMemo(
+			() => searchConnectors.filter((connector) => !connector.is_indexable),
+			[searchConnectors]
+		);
+
 		const handleOpenChange = useCallback(
 			(open: boolean) => {
 				setIsOpen(open);
 				if (open && !isLoaded) {
 					fetchDocumentTypes(Number(search_space_id));
 				}
+				if (open && !connectorsLoaded) {
+					fetchConnectors(Number(search_space_id));
+				}
 			},
-			[fetchDocumentTypes, isLoaded, search_space_id]
+			[fetchDocumentTypes, isLoaded, fetchConnectors, connectorsLoaded, search_space_id]
 		);
 
 		const handleConnectorToggle = useCallback(
@@ -141,14 +159,18 @@ const ConnectorSelector = React.memo(
 		);
 
 		const handleSelectAll = useCallback(() => {
-			onSelectionChange?.(documentTypes.map((dt) => dt.type));
+			const allTypes = [
-		}, [documentTypes, onSelectionChange]);
+				...documentTypes.map((dt) => dt.type),
+				...liveSearchConnectors.map((c) => c.connector_type),
+			];
+			onSelectionChange?.(allTypes);
+		}, [documentTypes, liveSearchConnectors, onSelectionChange]);
 
 		const handleClearAll = useCallback(() => {
 			onSelectionChange?.([]);
 		}, [onSelectionChange]);
 
-		// Get display name for document type
+		// Get display name for connector type
 		const getDisplayName = (type: string) => {
 			return type
 				.split("_")
@@ -158,6 +180,13 @@ const ConnectorSelector = React.memo(
 
 		// Get selected document types with their counts
 		const selectedDocTypes = documentTypes.filter((dt) => selectedConnectors.includes(dt.type));
+		const selectedLiveConnectors = liveSearchConnectors.filter((c) =>
+			selectedConnectors.includes(c.connector_type)
+		);
+
+		// Total selected count
+		const totalSelectedCount = selectedDocTypes.length + selectedLiveConnectors.length;
+		const totalAvailableCount = documentTypes.length + liveSearchConnectors.length;
 
 		return (
 			<Dialog open={isOpen} onOpenChange={handleOpenChange}>
@@ -168,10 +197,10 @@ const ConnectorSelector = React.memo(
 						className="relative h-9 gap-2 px-3 border-dashed hover:border-solid hover:bg-accent/50 transition-all"
 					>
 						<div className="flex items-center gap-1.5">
-							{selectedDocTypes.length > 0 ? (
+							{totalSelectedCount > 0 ? (
 								<>
 									<div className="flex items-center -space-x-2">
-										{selectedDocTypes.slice(0, 3).map((docType) => (
+										{selectedDocTypes.slice(0, 2).map((docType) => (
 											<div
 												key={docType.type}
 												className="flex h-6 w-6 items-center justify-center rounded-full border-2 border-background bg-muted"
@@ -179,9 +208,19 @@ const ConnectorSelector = React.memo(
 												{getConnectorIcon(docType.type, "h-3 w-3")}
 											</div>
 										))}
+										{selectedLiveConnectors
+											.slice(0, 3 - selectedDocTypes.slice(0, 2).length)
+											.map((connector) => (
+												<div
+													key={connector.id}
+													className="flex h-6 w-6 items-center justify-center rounded-full border-2 border-background bg-muted"
+												>
+													{getConnectorIcon(connector.connector_type, "h-3 w-3")}
+												</div>
+											))}
 									</div>
 									<span className="text-xs font-medium">
-										{selectedDocTypes.length} {selectedDocTypes.length === 1 ? "source" : "sources"}
+										{totalSelectedCount} {totalSelectedCount === 1 ? "source" : "sources"}
 									</span>
 								</>
 							) : (
@@ -194,99 +233,167 @@ const ConnectorSelector = React.memo(
 					</Button>
 				</DialogTrigger>
 
-				<DialogContent className="sm:max-w-2xl">
+				<DialogContent className="sm:max-w-2xl max-h-[85vh] flex flex-col">
-					<div className="space-y-4">
+					<div className="space-y-4 flex-1 overflow-y-auto pr-2">
 						<div>
-							<DialogTitle className="text-xl">Select Document Types</DialogTitle>
+							<DialogTitle className="text-xl">Select Sources</DialogTitle>
 							<DialogDescription className="mt-1.5">
-								Choose which document types to include in your search
+								Choose indexed document types and live search connectors to include in your search
 							</DialogDescription>
 						</div>
 
-						{/* Document type selection grid */}
+						{isLoading || connectorsLoading ? (
-						<div className="grid grid-cols-2 gap-3">
+							<div className="flex justify-center py-8">
-							{isLoading ? (
+								<div className="animate-spin h-8 w-8 border-3 border-primary border-t-transparent rounded-full" />
-								<div className="col-span-2 flex justify-center py-8">
+							</div>
-									<div className="animate-spin h-8 w-8 border-3 border-primary border-t-transparent rounded-full" />
+						) : totalAvailableCount === 0 ? (
+							<div className="flex flex-col items-center justify-center py-12 text-center">
+								<div className="rounded-full bg-muted p-4 mb-4">
+									<Brain className="h-8 w-8 text-muted-foreground" />
 								</div>
-							) : documentTypes.length === 0 ? (
+								<h4 className="text-sm font-medium mb-1">No sources found</h4>
-								<div className="col-span-2 flex flex-col items-center justify-center py-12 text-center">
+								<p className="text-xs text-muted-foreground max-w-xs">
-									<div className="rounded-full bg-muted p-4 mb-4">
+									Add documents or configure search connectors for this search space
-										<Brain className="h-8 w-8 text-muted-foreground" />
+								</p>
-									</div>
+							</div>
-									<h4 className="text-sm font-medium mb-1">No documents found</h4>
+						) : (
-									<p className="text-xs text-muted-foreground max-w-xs">
+							<>
-										Add documents to this search space to enable filtering by type
+								{/* Live Search Connectors Section */}
-									</p>
+								{liveSearchConnectors.length > 0 && (
-								</div>
+									<div className="space-y-2">
-							) : (
+										<div className="flex items-center gap-2 pb-2">
-								documentTypes.map((docType) => {
+											<Zap className="h-4 w-4 text-primary" />
-									const isSelected = selectedConnectors.includes(docType.type);
+											<h3 className="text-sm font-semibold">Live Search Connectors</h3>
+											<Badge variant="outline" className="text-xs">
+												Real-time
+											</Badge>
+										</div>
+										<div className="grid grid-cols-2 gap-3">
+											{liveSearchConnectors.map((connector) => {
+												const isSelected = selectedConnectors.includes(connector.connector_type);
 
-									return (
+												return (
-										<button
+													<button
-											key={docType.type}
+														key={connector.id}
-											onClick={() => handleConnectorToggle(docType.type)}
+														onClick={() => handleConnectorToggle(connector.connector_type)}
-											type="button"
+														type="button"
-											className={`group relative flex items-center gap-3 p-3 rounded-lg border-2 transition-all ${
+														className={`group relative flex items-center gap-3 p-3 rounded-lg border-2 transition-all ${
-												isSelected
+															isSelected
-													? "border-primary bg-primary/5 shadow-sm"
+																? "border-primary bg-primary/5 shadow-sm"
-													: "border-border hover:border-primary/50 hover:bg-accent/50"
+																: "border-border hover:border-primary/50 hover:bg-accent/50"
-											}`}
+														}`}
-										>
+													>
-											<div
+														<div
-												className={`flex h-10 w-10 items-center justify-center rounded-md transition-colors ${
+															className={`flex h-10 w-10 items-center justify-center rounded-md transition-colors ${
-													isSelected ? "bg-primary/10" : "bg-muted group-hover:bg-primary/5"
+																isSelected ? "bg-primary/10" : "bg-muted group-hover:bg-primary/5"
-												}`}
+															}`}
-											>
+														>
-												{getConnectorIcon(
+															{getConnectorIcon(
-													docType.type,
+																connector.connector_type,
-													`h-5 w-5 ${isSelected ? "text-primary" : "text-muted-foreground group-hover:text-primary"}`
+																`h-5 w-5 ${isSelected ? "text-primary" : "text-muted-foreground group-hover:text-primary"}`
-												)}
+															)}
-											</div>
-											<div className="flex-1 text-left min-w-0">
-												<div className="flex items-center gap-2">
-													<p className="text-sm font-medium truncate">
-														{getDisplayName(docType.type)}
-													</p>
-													{isSelected && (
-														<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary">
-															<Check className="h-3 w-3 text-primary-foreground" />
 														</div>
-													)}
+														<div className="flex-1 text-left min-w-0">
-												</div>
+															<div className="flex items-center gap-2">
-												<p className="text-xs text-muted-foreground mt-0.5">
+																<p className="text-sm font-medium truncate">{connector.name}</p>
-													{docType.count} {docType.count === 1 ? "document" : "documents"}
+																{isSelected && (
-												</p>
+																	<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary">
-											</div>
+																		<Check className="h-3 w-3 text-primary-foreground" />
-										</button>
+																	</div>
-									);
+																)}
-								})
+															</div>
-							)}
+															<p className="text-xs text-muted-foreground mt-0.5 truncate">
-						</div>
+																{getDisplayName(connector.connector_type)}
+															</p>
+														</div>
+													</button>
+												);
+											})}
+										</div>
+									</div>
+								)}
 
-						{documentTypes.length > 0 && (
+								{/* Document Types Section */}
-							<DialogFooter className="flex flex-row justify-between items-center gap-2 pt-2">
+								{documentTypes.length > 0 && (
-								<Button
+									<div className="space-y-2">
-									variant="ghost"
+										<div className="flex items-center gap-2 pb-2">
-									size="sm"
+											<FolderOpen className="h-4 w-4 text-primary" />
-									onClick={handleClearAll}
+											<h3 className="text-sm font-semibold">Indexed Document Types</h3>
-									disabled={selectedConnectors.length === 0}
+											<Badge variant="outline" className="text-xs">
-									className="text-xs"
+												Stored
-								>
+											</Badge>
-									Clear All
+										</div>
-								</Button>
+										<div className="grid grid-cols-2 gap-3">
-								<Button
+											{documentTypes.map((docType) => {
-									size="sm"
+												const isSelected = selectedConnectors.includes(docType.type);
-									onClick={handleSelectAll}
+
-									disabled={selectedConnectors.length === documentTypes.length}
+												return (
-									className="text-xs"
+													<button
-								>
+														key={docType.type}
-									Select All ({documentTypes.length})
+														onClick={() => handleConnectorToggle(docType.type)}
-								</Button>
+														type="button"
-							</DialogFooter>
+														className={`group relative flex items-center gap-3 p-3 rounded-lg border-2 transition-all ${
+															isSelected
+																? "border-primary bg-primary/5 shadow-sm"
+																: "border-border hover:border-primary/50 hover:bg-accent/50"
+														}`}
+													>
+														<div
+															className={`flex h-10 w-10 items-center justify-center rounded-md transition-colors ${
+																isSelected ? "bg-primary/10" : "bg-muted group-hover:bg-primary/5"
+															}`}
+														>
+															{getConnectorIcon(
+																docType.type,
+																`h-5 w-5 ${isSelected ? "text-primary" : "text-muted-foreground group-hover:text-primary"}`
+															)}
+														</div>
+														<div className="flex-1 text-left min-w-0">
+															<div className="flex items-center gap-2">
+																<p className="text-sm font-medium truncate">
+																	{getDisplayName(docType.type)}
+																</p>
+																{isSelected && (
+																	<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary">
+																		<Check className="h-3 w-3 text-primary-foreground" />
+																	</div>
+																)}
+															</div>
+															<p className="text-xs text-muted-foreground mt-0.5">
+																{docType.count} {docType.count === 1 ? "document" : "documents"}
+															</p>
+														</div>
+													</button>
+												);
+											})}
+										</div>
+									</div>
+								)}
+							</>
 						)}
 					</div>
+
+					{totalAvailableCount > 0 && (
+						<DialogFooter className="flex flex-row justify-between items-center gap-2 pt-4 border-t">
+							<Button
+								variant="ghost"
+								size="sm"
+								onClick={handleClearAll}
+								disabled={selectedConnectors.length === 0}
+								className="text-xs"
+							>
+								Clear All
+							</Button>
+							<Button
+								size="sm"
+								onClick={handleSelectAll}
+								disabled={selectedConnectors.length === totalAvailableCount}
+								className="text-xs"
+							>
+								Select All ({totalAvailableCount})
+							</Button>
+						</DialogFooter>
+					)}
 				</DialogContent>
 			</Dialog>
 		);
@@ -483,7 +590,6 @@ const LLMSelector = React.memo(() => {
 												</p>
 											</div>
 										</div>
-										
 									</div>
 								</SelectItem>
 							))}

surfsense_web/components/sidebar/app-sidebar.tsx

@@ -1,8 +1,5 @@
 "use client";
 
-import Link from "next/link";
-import Image from "next/image";
-
 import {
 	AlertCircle,
 	BookOpen,
@@ -20,6 +17,8 @@ import {
 	Trash2,
 	Undo2,
 } from "lucide-react";
+import Image from "next/image";
+import Link from "next/link";
 import { memo, useMemo } from "react";
 
 import { Logo } from "@/components/Logo";
@@ -36,7 +35,6 @@ import {
 	SidebarMenuItem,
 } from "@/components/ui/sidebar";
 
-
 // Map of icon names to their components
 export const iconMap: Record<string, LucideIcon> = {
 	BookOpen,
@@ -218,21 +216,21 @@ export const AppSidebar = memo(function AppSidebar({
 				<SidebarMenu>
 					<SidebarMenuItem>
 						<SidebarMenuButton asChild size="lg">
-						<Link href="/" className="flex items-center gap-2 w-full">
+							<Link href="/" className="flex items-center gap-2 w-full">
-							<div className="bg-sidebar-primary text-sidebar-primary-foreground flex aspect-square size-8 items-center justify-center rounded-lg">
+								<div className="bg-sidebar-primary text-sidebar-primary-foreground flex aspect-square size-8 items-center justify-center rounded-lg">
-							<Image
+									<Image
-								src="/icon-128.png"
+										src="/icon-128.png"
-								alt="SurfSense logo"
+										alt="SurfSense logo"
-								width={32}
+										width={32}
-								height={32}
+										height={32}
-								className="rounded-lg"
+										className="rounded-lg"
-							/>
+									/>
-							</div>
+								</div>
-							<div className="grid flex-1 text-left text-sm leading-tight">
+								<div className="grid flex-1 text-left text-sm leading-tight">
-							<span className="truncate font-medium">SurfSense</span>
+									<span className="truncate font-medium">SurfSense</span>
-							<span className="truncate text-xs">beta v0.0.8</span>
+									<span className="truncate text-xs">beta v0.0.8</span>
-							</div>
+								</div>
-						</Link>
+							</Link>
 						</SidebarMenuButton>
 					</SidebarMenuItem>
 				</SidebarMenu>