feat: enhance new chat functionality with document mentions support

- Updated the new chat routes to include handling for mentioned document IDs, allowing users to reference specific documents in their chat.
- Modified the NewChatRequest schema to accommodate optional document IDs.
- Implemented document mention formatting in the chat streaming service for improved context.
- Enhanced the frontend to manage document mentions, including a new atom for state management and UI updates for document selection.
- Refactored the DocumentsDataTable component for better integration with the new mention functionality.

surfsense_backend/app/routes/new_chat_routes.py

@@ -686,7 +686,7 @@ async def handle_new_chat(
         search_space = search_space_result.scalars().first()
 
         # TODO: Add new llm config arch then complete this
-        llm_config_id = -1  
+        llm_config_id = -4
 
         # Return streaming response
         return StreamingResponse(
@@ -698,6 +698,7 @@ async def handle_new_chat(
                 llm_config_id=llm_config_id,
                 messages=request.messages,
                 attachments=request.attachments,
+                mentioned_document_ids=request.mentioned_document_ids,
             ),
             media_type="text/event-stream",
             headers={

surfsense_backend/app/schemas/new_chat.py

@@ -160,3 +160,6 @@ class NewChatRequest(BaseModel):
     attachments: list[ChatAttachment] | None = (
         None  # Optional attachments with extracted content
     )
+    mentioned_document_ids: list[int] | None = (
+        None  # Optional document IDs mentioned with @ in the chat
+    )

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -9,6 +9,7 @@ import json
 from collections.abc import AsyncGenerator
 from langchain_core.messages import HumanMessage
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.future import select
 
 from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent
 from app.agents.new_chat.checkpointer import get_checkpointer
@@ -16,6 +17,7 @@ from app.agents.new_chat.llm_config import (
     create_chat_litellm_from_config,
     load_llm_config_from_yaml,
 )
+from app.db import Document
 from app.schemas.new_chat import ChatAttachment, ChatMessage
 from app.services.connector_service import ConnectorService
 from app.services.new_streaming_service import VercelStreamingService
@@ -38,6 +40,27 @@ def format_attachments_as_context(attachments: list[ChatAttachment]) -> str:
     return "\n".join(context_parts)
 
 
+def format_mentioned_documents_as_context(documents: list[Document]) -> str:
+    """Format mentioned documents as context for the agent."""
+    if not documents:
+        return ""
+
+    context_parts = ["<mentioned_documents>"]
+    context_parts.append(
+        "The user has explicitly mentioned the following documents from their knowledge base. "
+        "These documents are directly relevant to the query and should be prioritized as primary sources."
+    )
+    for i, doc in enumerate(documents, 1):
+        context_parts.append(
+            f"<document index='{i}' id='{doc.id}' title='{doc.title}' type='{doc.document_type.value}'>"
+        )
+        context_parts.append(f"<![CDATA[{doc.content}]]>")
+        context_parts.append("</document>")
+    context_parts.append("</mentioned_documents>")
+
+    return "\n".join(context_parts)
+
+
 async def stream_new_chat(
     user_query: str,
     search_space_id: int,
@@ -46,6 +69,7 @@ async def stream_new_chat(
     llm_config_id: int = -1,
     messages: list[ChatMessage] | None = None,
     attachments: list[ChatAttachment] | None = None,
+    mentioned_document_ids: list[int] | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Stream chat responses from the new SurfSense deep agent.
@@ -61,6 +85,8 @@ async def stream_new_chat(
         session: The database session
         llm_config_id: The LLM configuration ID (default: -1 for first global config)
         messages: Optional chat history from frontend (list of ChatMessage)
+        attachments: Optional attachments with extracted content
+        mentioned_document_ids: Optional list of document IDs mentioned with @ in the chat
 
     Yields:
         str: SSE formatted response strings
@@ -105,13 +131,30 @@ async def stream_new_chat(
         # Build input with message history from frontend
         langchain_messages = []
 
-        # Format the user query with attachment context if any
+        # Fetch mentioned documents if any
-        final_query = user_query
+        mentioned_documents: list[Document] = []
-        if attachments:
+        if mentioned_document_ids:
-            attachment_context = format_attachments_as_context(attachments)
+            result = await session.execute(
-            final_query = (
+                select(Document).filter(
-                f"{attachment_context}\n\n<user_query>{user_query}</user_query>"
+                    Document.id.in_(mentioned_document_ids),
+                    Document.search_space_id == search_space_id,
                 )
+            )
+            mentioned_documents = list(result.scalars().all())
+
+        # Format the user query with context (attachments + mentioned documents)
+        final_query = user_query
+        context_parts = []
+
+        if attachments:
+            context_parts.append(format_attachments_as_context(attachments))
+
+        if mentioned_documents:
+            context_parts.append(format_mentioned_documents_as_context(mentioned_documents))
+
+        if context_parts:
+            context = "\n\n".join(context_parts)
+            final_query = f"{context}\n\n<user_query>{user_query}</user_query>"
 
         # if messages:
         #     # Convert frontend messages to LangChain format

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -6,9 +6,11 @@ import {
 	type ThreadMessageLike,
 	useExternalStoreRuntime,
 } from "@assistant-ui/react";
+import { useAtomValue, useSetAtom } from "jotai";
 import { useParams, useRouter } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
+import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { Thread } from "@/components/assistant-ui/thread";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
 import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
@@ -107,6 +109,10 @@ export default function NewChatPage() {
 	>(new Map());
 	const abortControllerRef = useRef<AbortController | null>(null);
 
+	// Get mentioned document IDs from the composer
+	const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
+	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
+
 	// Create the attachment adapter for file processing
 	const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []);
 
@@ -372,6 +378,14 @@ export default function NewChatPage() {
 				// Extract attachment content to send with the request
 				const attachments = extractAttachmentContent(messageAttachments);
 
+				// Get mentioned document IDs for context
+				const documentIds = mentionedDocumentIds.length > 0 ? [...mentionedDocumentIds] : undefined;
+
+				// Clear mentioned documents after capturing them
+				if (mentionedDocumentIds.length > 0) {
+					setMentionedDocumentIds([]);
+				}
+
 				const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
 					method: "POST",
 					headers: {
@@ -384,6 +398,7 @@ export default function NewChatPage() {
 						search_space_id: searchSpaceId,
 						messages: messageHistory,
 						attachments: attachments.length > 0 ? attachments : undefined,
+						mentioned_document_ids: documentIds,
 					}),
 					signal: controller.signal,
 				});
@@ -546,7 +561,7 @@ export default function NewChatPage() {
 				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
-		[threadId, searchSpaceId, messages]
+		[threadId, searchSpaceId, messages, mentionedDocumentIds, setMentionedDocumentIds]
 	);
 
 	// Convert message (pass through since already in correct format)

surfsense_web/atoms/chat/mentioned-documents.atom.ts

@@ -0,0 +1,10 @@
+"use client";
+
+import { atom } from "jotai";
+
+/**
+ * Atom to store the IDs of documents mentioned in the current chat composer.
+ * This is used to pass document context to the backend when sending a message.
+ */
+export const mentionedDocumentIdsAtom = atom<number[]>([]);
+

surfsense_web/components/assistant-ui/thread.tsx

@@ -7,7 +7,6 @@ import {
 	MessagePrimitive,
 	ThreadPrimitive,
 	useAssistantState,
-	useMessage,
 } from "@assistant-ui/react";
 import {
 	ArrowDownIcon,
@@ -27,18 +26,20 @@ import {
 	Search,
 	Sparkles,
 	SquareIcon,
+	X,
 } from "lucide-react";
 import { useParams } from "next/navigation";
 import Link from "next/link";
-import { type FC, useState, useRef, useCallback, useEffect } from "react";
+import { type FC, useState, useRef, useCallback, useEffect, createContext, useContext, useMemo } from "react";
-import { useAtomValue } from "jotai";
+import { createPortal } from "react-dom";
+import { useAtomValue, useSetAtom } from "jotai";
+import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
 import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
 import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
-import { useRef, useState } from "react";
 import {
 	ComposerAddAttachment,
 	ComposerAttachments,
@@ -69,8 +70,6 @@ interface ThreadProps {
 }
 
 // Context to pass thinking steps to AssistantMessage
-import { createContext, useContext } from "react";
-
 const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
 
 /**
@@ -333,12 +332,15 @@ const getTimeBasedGreeting = (userEmail?: string): string => {
 const ThreadWelcome: FC = () => {
 	const { data: user } = useAtomValue(currentUserAtom);
 	
+	// Memoize greeting so it doesn't change on re-renders (only on user change)
+	const greeting = useMemo(() => getTimeBasedGreeting(user?.email), [user?.email]);
+	
 	return (
 		<div className="aui-thread-welcome-root mx-auto flex w-full max-w-(--thread-max-width) grow flex-col items-center px-4 relative">
 			{/* Greeting positioned above the composer - fixed position */}
 			<div className="aui-thread-welcome-message absolute bottom-[calc(50%+5rem)] left-0 right-0 flex flex-col items-center text-center z-10">
 			<h1 className="aui-thread-welcome-message-inner fade-in slide-in-from-bottom-2 animate-in text-5xl delay-100 duration-500 ease-out fill-mode-both">
-				{getTimeBasedGreeting(user?.email)}
+				{greeting}
 			</h1>
 			</div>
 			{/* Composer - top edge fixed, expands downward only */}
@@ -351,122 +353,155 @@ const ThreadWelcome: FC = () => {
 
 const Composer: FC = () => {
 	// ---- State for document mentions ----
-	const [allSelectedDocuments, setAllSelectedDocuments] = useState<Document[]>([]);
 	const [mentionedDocuments, setMentionedDocuments] = useState<Document[]>([]);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
-	const [inputValue, setInputValue] = useState("");
 	const inputRef = useRef<HTMLTextAreaElement | null>(null);
 	const { search_space_id } = useParams();
+	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
 
-	const handleInputOrKeyUp = (
+	// Sync mentioned document IDs to atom for use in chat request
-		e: React.FormEvent<HTMLTextAreaElement> | React.KeyboardEvent<HTMLTextAreaElement>
+	useEffect(() => {
-	) => {
+		setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
+	}, [mentionedDocuments, setMentionedDocumentIds]);
+
+	const handleKeyUp = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
 		const textarea = e.currentTarget;
 		const value = textarea.value;
-		setInputValue(value);
 
-		// Regex: finds all [title] occurrences
+		// Open document picker when user types '@'
-		const mentionRegex = /\[([^\]]+)\]/g;
+		if (e.key === "@" || (e.key === "2" && e.shiftKey)) {
-		const titlesMentioned: string[] = [];
+			setShowDocumentPopover(true);
-		let match;
-		while ((match = mentionRegex.exec(value)) !== null) {
-			titlesMentioned.push(match[1]);
 		}
 
-		// Use allSelectedDocuments to filter down for current chips
+		// Close popover if '@' is no longer in the input (user deleted it)
-		setMentionedDocuments(
+		if (showDocumentPopover && !value.includes("@")) {
-			allSelectedDocuments.filter((doc) => titlesMentioned.includes(doc.title))
-		);
-
-		const selectionStart = textarea.selectionStart;
-		// Only open if the last character before the caret is exactly '@'
-		if (
-			selectionStart !== null &&
-			value[selectionStart - 1] === "@" &&
-			value.length === selectionStart
-		) {
-			setShowDocumentPopover(true);
-		} else {
 			setShowDocumentPopover(false);
 		}
 	};
 
+	const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+		// Close popover on Escape
+		if (e.key === "Escape" && showDocumentPopover) {
+			e.preventDefault();
+			setShowDocumentPopover(false);
+			return;
+		}
+
+		// Remove last document chip when pressing backspace at the beginning of input
+		if (e.key === "Backspace" && mentionedDocuments.length > 0) {
+			const textarea = e.currentTarget;
+			const selectionStart = textarea.selectionStart;
+			const selectionEnd = textarea.selectionEnd;
+
+			// Only remove chip if cursor is at position 0 and nothing is selected
+			if (selectionStart === 0 && selectionEnd === 0) {
+				e.preventDefault();
+				// Remove the last document chip
+				setMentionedDocuments((prev) => prev.slice(0, -1));
+			}
+		}
+	};
+
 	const handleDocumentsMention = (documents: Document[]) => {
-		// Add newly selected docs to allSelectedDocuments
+		// Update mentioned documents (merge with existing, avoid duplicates)
-		setAllSelectedDocuments((prev) => {
+		setMentionedDocuments((prev) => {
-			const toAdd = documents.filter((doc) => !prev.find((p) => p.id === doc.id));
+			const existingIds = new Set(prev.map((d) => d.id));
-			return [...prev, ...toAdd];
+			const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
+			return [...prev, ...newDocs];
 		});
-		let newValue = inputValue;
+
-		documents.forEach((doc) => {
+		// Clean up the '@' trigger from input if present
-			const refString = `[${doc.title}]`;
+		if (inputRef.current) {
-			if (!newValue.includes(refString)) {
+			const input = inputRef.current;
-				if (newValue.trim() !== "" && !newValue.endsWith(" ")) {
+			const currentValue = input.value;
-					newValue += " ";
+			// Remove trailing @ if it exists
+			if (currentValue.endsWith("@")) {
+				// Use a native input event to properly update the controlled component
+				const nativeInputValueSetter = Object.getOwnPropertyDescriptor(
+					window.HTMLTextAreaElement.prototype,
+					"value"
+				)?.set;
+				if (nativeInputValueSetter) {
+					nativeInputValueSetter.call(input, currentValue.slice(0, -1));
+					input.dispatchEvent(new Event("input", { bubbles: true }));
 				}
-				newValue += refString;
 			}
-		});
+			// Focus the input so user can continue typing
-		setInputValue(newValue);
+			input.focus();
-		// Run the chip update as well right after change
-		const mentionRegex = /\[([^\]]+)\]/g;
-		const titlesMentioned: string[] = [];
-		let match;
-		while ((match = mentionRegex.exec(newValue)) !== null) {
-			titlesMentioned.push(match[1]);
 		}
-		setMentionedDocuments(
+	};
-			allSelectedDocuments.filter((doc) => titlesMentioned.includes(doc.title))
+
-		);
+	const handleRemoveDocument = (docId: number) => {
+		setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId));
 	};
 
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
 				<ComposerAttachments />
-				{/* -------- Input field w/ refs and handlers -------- */}
+				{/* -------- Input field with inline document chips -------- */}
+				<div className="aui-composer-input-wrapper flex flex-wrap items-center gap-1.5 px-3 pt-2 pb-6">
+					{/* Inline document chips */}
+					{mentionedDocuments.map((doc) => (
+						<span
+							key={doc.id}
+							className="inline-flex items-center gap-1 pl-2 pr-1 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20 shrink-0"
+							title={doc.title}
+						>
+							<span className="max-w-[120px] truncate">{doc.title}</span>
+							<button
+								type="button"
+								onClick={() => handleRemoveDocument(doc.id)}
+								className="size-4 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors"
+								aria-label={`Remove ${doc.title}`}
+							>
+								<X className="size-3" />
+							</button>
+						</span>
+					))}
+					{/* Text input */}
 					<ComposerPrimitive.Input
 						ref={inputRef}
-					value={inputValue}
+						onKeyUp={handleKeyUp}
-					onInput={handleInputOrKeyUp}
+						onKeyDown={handleKeyDown}
-					onKeyUp={handleInputOrKeyUp}
+						placeholder={mentionedDocuments.length > 0 ? "Ask about these documents..." : "Ask SurfSense (type @ to mention docs)"}
-					placeholder="Ask SurfSense"
+						className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1"
-					className="aui-composer-input mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0"
 						rows={1}
 						autoFocus
 						aria-label="Message input"
 					/>
+				</div>
 
-				{/* -------- Document mention popover (simple version) -------- */}
+				{/* -------- Document mention popover (rendered via portal) -------- */}
-				{showDocumentPopover && (
+				{showDocumentPopover && typeof document !== "undefined" && createPortal(
+					<>
+						{/* Backdrop */}
+						<button
+							type="button"
+							className="fixed inset-0 cursor-default"
+							style={{ zIndex: 9998 }}
+							onClick={() => setShowDocumentPopover(false)}
+							aria-label="Close document picker"
+						/>
+						{/* Popover positioned above input */}
 						<div
-						style={{ position: "absolute", bottom: "8rem", left: 0, zIndex: 50 }}
+							className="fixed shadow-2xl rounded-lg border border-border overflow-hidden"
-						className="shadow-lg rounded-md border bg-popover"
+							style={{ 
+								zIndex: 9999, 
+								backgroundColor: "#18181b",
+								bottom: inputRef.current ? `${window.innerHeight - inputRef.current.getBoundingClientRect().top + 8}px` : "200px",
+								left: inputRef.current ? `${inputRef.current.getBoundingClientRect().left}px` : "50%",
+							}}
 						>
-						<div className="p-2 max-h-96 w-full overflow-auto">
 							<DocumentsDataTable
 								searchSpaceId={Number(search_space_id)}
 								onSelectionChange={handleDocumentsMention}
 								onDone={() => setShowDocumentPopover(false)}
 								initialSelectedDocuments={mentionedDocuments}
-								viewOnly={true}
 							/>
 						</div>
-					</div>
+					</>,
-				)}
+					document.body
-				{/* ---- Mention chips for selected/mentioned documents ---- */}
-				{mentionedDocuments.length > 0 && (
-					<div className="aui-composer-mentioned-docs mx-2 flex flex-wrap gap-2">
-						{mentionedDocuments.map((doc) => (
-							<span
-								key={doc.id}
-								className="px-2 py-1 rounded bg-accent text-xs font-semibold max-w-xs truncate"
-								title={doc.title}
-							>
-								{doc.title}
-							</span>
-						))}
-					</div>
 				)}
 				<ComposerAction />
 			</ComposerPrimitive.AttachmentDropzone>
@@ -564,7 +599,7 @@ const ConnectorIndicator: FC = () => {
 						</div>
 						<div className="flex flex-wrap gap-2">
 							{/* Document types from the search space */}
-							{activeDocumentTypes.map(([docType, count]) => (
+							{activeDocumentTypes.map(([docType]) => (
 								<div
 									key={docType}
 									className="flex items-center gap-1.5 rounded-md bg-muted/80 px-2.5 py-1.5 text-xs border border-border/50"
@@ -707,7 +742,7 @@ const AssistantMessageInner: FC = () => {
 	const thinkingStepsMap = useContext(ThinkingStepsContext);
 	
 	// Get the current message ID to look up thinking steps
-	const messageId = useMessage((m) => m.id);
+	const messageId = useAssistantState(({ message }) => message?.id);
 	const thinkingSteps = thinkingStepsMap.get(messageId) || [];
 	
 	// Check if thread is still running (for stopping the spinner when cancelled)

surfsense_web/components/new-chat/DocumentsDataTable.tsx

@@ -1,49 +1,21 @@
 "use client";
 
 import { useQuery } from "@tanstack/react-query";
-import {
+import { FileText, Search } from "lucide-react";
-	type ColumnDef,
-	flexRender,
-	getCoreRowModel,
-	type SortingState,
-	useReactTable,
-} from "@tanstack/react-table";
-import { useAtomValue } from "jotai";
-import { ArrowUpDown, Calendar, FileText, Filter, Plus, Search } from "lucide-react";
-import { useRouter } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
-import { Button } from "@/components/ui/button";
-import { Checkbox } from "@/components/ui/checkbox";
 import { Input } from "@/components/ui/input";
-import { Label } from "@/components/ui/label";
+import { ScrollArea } from "@/components/ui/scroll-area";
-import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
-import {
-	Select,
-	SelectContent,
-	SelectItem,
-	SelectTrigger,
-	SelectValue,
-} from "@/components/ui/select";
-import {
-	Table,
-	TableBody,
-	TableCell,
-	TableHead,
-	TableHeader,
-	TableRow,
-} from "@/components/ui/table";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
-import type { Document, DocumentTypeEnum } from "@/contracts/types/document.types";
+import type { Document } from "@/contracts/types/document.types";
 import { documentsApiService } from "@/lib/apis/documents-api.service";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
+import { cn } from "@/lib/utils";
 
 interface DocumentsDataTableProps {
 	searchSpaceId: number;
 	onSelectionChange: (documents: Document[]) => void;
 	onDone: () => void;
 	initialSelectedDocuments?: Document[];
-	viewOnly?: boolean;
 }
 
 function useDebounced<T>(value: T, delay = 300) {
@@ -55,551 +27,190 @@ function useDebounced<T>(value: T, delay = 300) {
 	return debounced;
 }
 
-const columns: ColumnDef<Document>[] = [
-	{
-		id: "select",
-		header: ({ table }) => (
-			<Checkbox
-				checked={
-					table.getIsAllPageRowsSelected() || (table.getIsSomePageRowsSelected() && "indeterminate")
-				}
-				onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
-				aria-label="Select all"
-			/>
-		),
-		cell: ({ row }) => (
-			<Checkbox
-				checked={row.getIsSelected()}
-				onCheckedChange={(value) => row.toggleSelected(!!value)}
-				aria-label="Select row"
-			/>
-		),
-		enableSorting: false,
-		enableHiding: false,
-		size: 40,
-	},
-	{
-		accessorKey: "title",
-		header: ({ column }) => (
-			<Button
-				variant="ghost"
-				onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
-				className="h-8 px-1 sm:px-2 font-medium text-left justify-start"
-			>
-				<FileText className="mr-1 sm:mr-2 h-3 w-3 sm:h-4 sm:w-4 flex-shrink-0" />
-				<span className="hidden sm:inline">Title</span>
-				<span className="sm:hidden">Doc</span>
-				<ArrowUpDown className="ml-1 sm:ml-2 h-3 w-3 sm:h-4 sm:w-4 flex-shrink-0" />
-			</Button>
-		),
-		cell: ({ row }) => {
-			const title = row.getValue("title") as string;
-			return (
-				<div
-					className="font-medium max-w-[120px] sm:max-w-[250px] truncate text-xs sm:text-sm"
-					title={title}
-				>
-					{title}
-				</div>
-			);
-		},
-	},
-	{
-		accessorKey: "document_type",
-		header: "Type",
-		cell: ({ row }) => {
-			const type = row.getValue("document_type") as DocumentType;
-			return (
-				<div className="flex items-center gap-2" title={String(type)}>
-					<span className="text-primary">{getConnectorIcon(String(type))}</span>
-				</div>
-			);
-		},
-		size: 80,
-		meta: {
-			className: "hidden sm:table-cell",
-		},
-	},
-	{
-		accessorKey: "content",
-		header: "Preview",
-		cell: ({ row }) => {
-			const content = row.getValue("content") as string;
-			return (
-				<div
-					className="text-muted-foreground max-w-[150px] sm:max-w-[350px] truncate text-[10px] sm:text-sm"
-					title={content}
-				>
-					<span className="sm:hidden">{content.substring(0, 30)}...</span>
-					<span className="hidden sm:inline">{content.substring(0, 100)}...</span>
-				</div>
-			);
-		},
-		enableSorting: false,
-		meta: {
-			className: "hidden md:table-cell",
-		},
-	},
-	{
-		accessorKey: "created_at",
-		header: ({ column }) => (
-			<Button
-				variant="ghost"
-				onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
-				className="h-8 px-1 sm:px-2 font-medium"
-			>
-				<Calendar className="mr-1 sm:mr-2 h-3 w-3 sm:h-4 sm:w-4 flex-shrink-0" />
-				<span className="hidden sm:inline">Created</span>
-				<span className="sm:hidden">Date</span>
-				<ArrowUpDown className="ml-1 sm:ml-2 h-3 w-3 sm:h-4 sm:w-4 flex-shrink-0" />
-			</Button>
-		),
-		cell: ({ row }) => {
-			const date = new Date(row.getValue("created_at"));
-			return (
-				<div className="text-xs sm:text-sm whitespace-nowrap">
-					<span className="hidden sm:inline">
-						{date.toLocaleDateString("en-US", {
-							month: "short",
-							day: "numeric",
-							year: "numeric",
-						})}
-					</span>
-					<span className="sm:hidden">
-						{date.toLocaleDateString("en-US", {
-							month: "numeric",
-							day: "numeric",
-						})}
-					</span>
-				</div>
-			);
-		},
-		size: 80,
-	},
-];
-
 export function DocumentsDataTable({
 	searchSpaceId,
 	onSelectionChange,
 	onDone,
 	initialSelectedDocuments = [],
 }: DocumentsDataTableProps) {
-	const router = useRouter();
-	const [sorting, setSorting] = useState<SortingState>([]);
 	const [search, setSearch] = useState("");
 	const debouncedSearch = useDebounced(search, 300);
-	const [documentTypeFilter, setDocumentTypeFilter] = useState<DocumentTypeEnum[]>([]);
+	const [highlightedIndex, setHighlightedIndex] = useState(0);
-	const [pageIndex, setPageIndex] = useState(0);
+	const listRef = useRef<HTMLDivElement>(null);
-	const [pageSize, setPageSize] = useState(10);
+	const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
-	const { data: typeCounts } = useAtomValue(documentTypeCountsAtom);
 
 	const fetchQueryParams = useMemo(
 		() => ({
 			search_space_id: searchSpaceId,
-			page: pageIndex,
+			page: 0,
-			page_size: pageSize,
+			page_size: 20,
-			...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }),
 		}),
-		[searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch]
+		[searchSpaceId]
 	);
 
 	const searchQueryParams = useMemo(() => {
 		return {
 			search_space_id: searchSpaceId,
-			page: pageIndex,
+			page: 0,
-			page_size: pageSize,
+			page_size: 20,
-			...(documentTypeFilter.length > 0 && { document_types: documentTypeFilter }),
 			title: debouncedSearch,
 		};
-	}, [debouncedSearch, searchSpaceId, pageIndex, pageSize, documentTypeFilter, debouncedSearch]);
+	}, [debouncedSearch, searchSpaceId]);
 
 	// Use query for fetching documents
 	const { data: documents, isLoading: isDocumentsLoading } = useQuery({
 		queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
 		queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
-		staleTime: 3 * 60 * 1000, // 3 minutes
+		staleTime: 3 * 60 * 1000,
 		enabled: !!searchSpaceId && !debouncedSearch.trim(),
 	});
 
-	// Seaching
+	// Searching
 	const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
 		queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
 		queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
-		staleTime: 3 * 60 * 1000, // 3 minutes
+		staleTime: 3 * 60 * 1000,
 		enabled: !!searchSpaceId && !!debouncedSearch.trim(),
 	});
 
-	// Use query data when not searching, otherwise use hook data
 	const actualDocuments = debouncedSearch.trim()
 		? searchedDocuments?.items || []
 		: documents?.items || [];
-	const actualTotal = debouncedSearch.trim()
-		? searchedDocuments?.total || 0
-		: documents?.total || 0;
 	const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
 
-	// Memoize initial row selection to prevent infinite loops
+	// Track already selected document IDs
-	const initialRowSelection = useMemo(() => {
+	const selectedIds = useMemo(
-		if (!initialSelectedDocuments.length) return {};
+		() => new Set(initialSelectedDocuments.map((d) => d.id)),
-
+		[initialSelectedDocuments]
-		const selection: Record<string, boolean> = {};
-		initialSelectedDocuments.forEach((selectedDoc) => {
-			selection[selectedDoc.id] = true;
-		});
-		return selection;
-	}, [initialSelectedDocuments]);
-
-	const [rowSelection, setRowSelection] = useState<Record<string, boolean>>(
-		() => initialRowSelection
 	);
 
-	// Maintain a separate state for actually selected documents (across all pages)
+	// Filter out already selected documents for navigation
-	const [selectedDocumentsMap, setSelectedDocumentsMap] = useState<Map<number, Document>>(() => {
+	const selectableDocuments = useMemo(
-		const map = new Map<number, Document>();
+		() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
-		initialSelectedDocuments.forEach((doc) => map.set(doc.id, doc));
+		[actualDocuments, selectedIds]
-		return map;
+	);
-	});
 
-	// Track the last notified selection to avoid redundant parent calls
+	const handleSelectDocument = useCallback((doc: Document) => {
-	const lastNotifiedSelection = useRef<string>("");
+		onSelectionChange([...initialSelectedDocuments, doc]);
+		onDone();
+	}, [initialSelectedDocuments, onSelectionChange, onDone]);
 
-	// Update row selection only when initialSelectedDocuments changes (not rowSelection itself)
+	// Scroll highlighted item into view
 	useEffect(() => {
-		const initialKeys = Object.keys(initialRowSelection);
+		const item = itemRefs.current.get(highlightedIndex);
-		if (initialKeys.length === 0) return;
+		if (item) {
-
+			item.scrollIntoView({ block: "nearest", behavior: "smooth" });
-		const currentKeys = Object.keys(rowSelection);
-		// Quick length check before expensive comparison
-		if (currentKeys.length === initialKeys.length) {
-			// Check if all keys match (order doesn't matter for Sets)
-			const hasAllKeys = initialKeys.every((key) => rowSelection[key]);
-			if (hasAllKeys) return;
 		}
+	}, [highlightedIndex]);
 
-		setRowSelection(initialRowSelection);
+	// Handle keyboard navigation
-	}, [initialRowSelection]); // Remove rowSelection from dependencies to prevent loop
+	const handleKeyDown = useCallback((e: React.KeyboardEvent) => {
+		if (selectableDocuments.length === 0) return;
 
-	// Update the selected documents map when row selection changes
+		switch (e.key) {
-	useEffect(() => {
+			case "ArrowDown":
-		if (!actualDocuments || actualDocuments.length === 0) return;
+				e.preventDefault();
-
+				setHighlightedIndex((prev) => 
-		setSelectedDocumentsMap((prev) => {
+					prev < selectableDocuments.length - 1 ? prev + 1 : 0
-			const newMap = new Map(prev);
+				);
-			let hasChanges = false;
+				break;
-
+			case "ArrowUp":
-			// Process only current page documents
+				e.preventDefault();
-			for (const doc of actualDocuments) {
+				setHighlightedIndex((prev) => 
-				const docId = doc.id;
+					prev > 0 ? prev - 1 : selectableDocuments.length - 1
-				const isSelected = rowSelection[docId.toString()];
+				);
-				const wasInMap = newMap.has(docId);
+				break;
-
+			case "Enter":
-				if (isSelected && !wasInMap) {
+				e.preventDefault();
-					newMap.set(docId, doc);
+				if (selectableDocuments[highlightedIndex]) {
-					hasChanges = true;
+					handleSelectDocument(selectableDocuments[highlightedIndex]);
-				} else if (!isSelected && wasInMap) {
-					newMap.delete(docId);
-					hasChanges = true;
 				}
+				break;
+			case "Escape":
+				e.preventDefault();
+				onDone();
+				break;
 		}
-
+	}, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]);
-			// Return same reference if no changes to avoid unnecessary re-renders
-			return hasChanges ? newMap : prev;
-		});
-	}, [rowSelection, documents]);
-
-	// Memoize selected documents array
-	const selectedDocumentsArray = useMemo(() => {
-		return Array.from(selectedDocumentsMap.values());
-	}, [selectedDocumentsMap]);
-
-	// Notify parent of selection changes only when content actually changes
-	useEffect(() => {
-		// Create a stable string representation for comparison
-		const selectionKey = selectedDocumentsArray
-			.map((d) => d.id)
-			.sort()
-			.join(",");
-
-		// Skip if selection hasn't actually changed
-		if (selectionKey === lastNotifiedSelection.current) return;
-
-		lastNotifiedSelection.current = selectionKey;
-		onSelectionChange(selectedDocumentsArray);
-	}, [selectedDocumentsArray, onSelectionChange]);
-
-	const table = useReactTable({
-		data: actualDocuments || [],
-		columns,
-		getRowId: (row) => row.id.toString(),
-		onSortingChange: setSorting,
-		getCoreRowModel: getCoreRowModel(),
-		onRowSelectionChange: setRowSelection,
-		manualPagination: true,
-		pageCount: Math.ceil(actualTotal / pageSize),
-		state: { sorting, rowSelection, pagination: { pageIndex, pageSize } },
-	});
-
-	const handleClearAll = useCallback(() => {
-		setRowSelection({});
-		setSelectedDocumentsMap(new Map());
-	}, []);
-
-	const handleSelectPage = useCallback(() => {
-		const currentPageRows = table.getRowModel().rows;
-		const newSelection = { ...rowSelection };
-		currentPageRows.forEach((row) => {
-			newSelection[row.id] = true;
-		});
-		setRowSelection(newSelection);
-	}, [table, rowSelection]);
-
-	const handleToggleType = useCallback((type: DocumentTypeEnum, checked: boolean) => {
-		setDocumentTypeFilter((prev) => {
-			if (checked) {
-				return [...prev, type];
-			}
-			return prev.filter((t) => t !== type);
-		});
-		setPageIndex(0); // Reset to first page when filter changes
-	}, []);
-
-	const selectedCount = selectedDocumentsMap.size;
-
-	// Get available document types from type counts (memoized)
-	const availableTypes = useMemo(() => {
-		const types = typeCounts ? (Object.keys(typeCounts) as DocumentTypeEnum[]) : [];
-		return types.length > 0 ? types.sort() : [];
-	}, [typeCounts]);
 
 	return (
-		<div className="flex flex-col h-full space-y-3 md:space-y-4">
+		<div 
-			{/* Header Controls */}
+			className="flex flex-col w-[280px] sm:w-[320px] bg-zinc-900"
-			<div className="space-y-3 md:space-y-4 flex-shrink-0">
+			onKeyDown={handleKeyDown}
-				{/* Search and Filter Row */}
+			role="listbox"
-				<div className="flex flex-col sm:flex-row gap-3 sm:gap-4">
+			tabIndex={-1}
-					<div className="relative flex-1 max-w-full sm:max-w-sm">
+		>
-						<Search className="absolute left-3 top-1/2 transform -translate-y-1/2 h-4 w-4 text-muted-foreground" />
+			{/* Search */}
+			<div className="relative p-2 border-b">
+				<Search className="absolute left-4 top-1/2 transform -translate-y-1/2 h-4 w-4 text-muted-foreground" />
 				<Input
 					placeholder="Search documents..."
 					value={search}
-							onChange={(event) => {
+					onChange={(e) => {
-								setSearch(event.target.value);
+						setSearch(e.target.value);
-								setPageIndex(0); // Reset to first page on search
+						setHighlightedIndex(0);
 					}}
-							className="pl-10 text-sm"
+					className="pl-8 h-8 text-sm border-0 focus-visible:ring-0 focus-visible:ring-offset-0"
+					autoFocus
 				/>
 			</div>
-					<Popover>
-						<PopoverTrigger asChild>
-							<Button variant="outline" className="w-full sm:w-auto">
-								<Filter className="mr-2 h-4 w-4 opacity-60" />
-								Type
-								{documentTypeFilter.length > 0 && (
-									<span className="ml-2 inline-flex h-5 items-center rounded border border-border bg-background px-1.5 text-[0.625rem] font-medium text-muted-foreground/70">
-										{documentTypeFilter.length}
-									</span>
-								)}
-							</Button>
-						</PopoverTrigger>
-						<PopoverContent className="w-64 p-3" align="start">
-							<div className="space-y-3">
-								<div className="text-xs font-medium text-muted-foreground">Filter by Type</div>
-								<div className="space-y-2 max-h-[300px] overflow-y-auto">
-									{availableTypes.map((type) => (
-										<div key={type} className="flex items-center gap-2">
-											<Checkbox
-												id={`type-${type}`}
-												checked={documentTypeFilter.includes(type)}
-												onCheckedChange={(checked) => handleToggleType(type, !!checked)}
-											/>
-											<Label
-												htmlFor={`type-${type}`}
-												className="flex grow justify-between gap-2 font-normal text-sm cursor-pointer"
-											>
-												<span>{type.replace(/_/g, " ")}</span>
-												<span className="text-xs text-muted-foreground">{typeCounts?.[type]}</span>
-											</Label>
-										</div>
-									))}
-								</div>
-								{documentTypeFilter.length > 0 && (
-									<Button
-										variant="ghost"
-										size="sm"
-										className="w-full text-xs"
-										onClick={() => {
-											setDocumentTypeFilter([]);
-											setPageIndex(0);
-										}}
-									>
-										Clear Filters
-									</Button>
-								)}
-							</div>
-						</PopoverContent>
-					</Popover>
-				</div>
 
-				{/* Action Controls Row */}
+			{/* Document List */}
-				<div className="flex flex-col sm:flex-row sm:items-center sm:justify-between gap-3">
+			<ScrollArea className="h-[240px]">
-					<div className="flex flex-col sm:flex-row sm:items-center gap-2">
-						<span className="text-sm text-muted-foreground whitespace-nowrap">
-							{selectedCount} selected {actualLoading && "· Loading..."}
-						</span>
-						<div className="hidden sm:block h-4 w-px bg-border mx-2" />
-						<div className="flex items-center gap-2 flex-wrap">
-							<Button
-								variant="ghost"
-								size="sm"
-								onClick={handleClearAll}
-								disabled={selectedCount === 0}
-								className="text-xs sm:text-sm"
-							>
-								Clear All
-							</Button>
-							<Button
-								variant="ghost"
-								size="sm"
-								onClick={handleSelectPage}
-								className="text-xs sm:text-sm"
-								disabled={actualLoading}
-							>
-								Select Page
-							</Button>
-							<Select
-								value={pageSize.toString()}
-								onValueChange={(v) => {
-									setPageSize(Number(v));
-									setPageIndex(0);
-								}}
-							>
-								<SelectTrigger className="w-[100px] h-8 text-xs">
-									<SelectValue>{pageSize} per page</SelectValue>
-								</SelectTrigger>
-								<SelectContent>
-									{[10, 25, 50, 100].map((size) => (
-										<SelectItem key={size} value={size.toString()}>
-											{size} per page
-										</SelectItem>
-									))}
-								</SelectContent>
-							</Select>
-						</div>
-					</div>
-					<Button
-						onClick={onDone}
-						disabled={selectedCount === 0}
-						className="w-full sm:w-auto sm:min-w-[100px]"
-					>
-						Done ({selectedCount})
-					</Button>
-				</div>
-			</div>
-
-			{/* Table Container */}
-			<div className="border rounded-lg flex-1 min-h-0 overflow-hidden bg-background">
-				<div className="overflow-auto h-full">
 				{actualLoading ? (
-						<div className="flex items-center justify-center h-full">
+					<div className="flex items-center justify-center py-6">
-							<div className="text-center space-y-2">
+						<div className="animate-spin h-5 w-5 border-2 border-primary border-t-transparent rounded-full" />
-								<div className="animate-spin h-8 w-8 border-2 border-primary border-t-transparent rounded-full mx-auto" />
-								<p className="text-sm text-muted-foreground">Loading documents...</p>
 					</div>
+				) : actualDocuments.length === 0 ? (
+					<div className="flex flex-col items-center justify-center py-6 text-center px-4">
+						<FileText className="h-6 w-6 text-muted-foreground/50 mb-2" />
+						<p className="text-sm text-muted-foreground">No documents found</p>
 					</div>
 				) : (
-						<Table>
+					<div className="py-1" ref={listRef}>
-							<TableHeader className="sticky top-0 bg-background/95 backdrop-blur supports-[backdrop-filter]:bg-background/60 z-10">
+						{actualDocuments.map((doc) => {
-								{table.getHeaderGroups().map((headerGroup) => (
+							const isAlreadySelected = selectedIds.has(doc.id);
-									<TableRow key={headerGroup.id} className="border-b">
+							const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id);
-										{headerGroup.headers.map((header) => (
+							const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex;
-											<TableHead key={header.id} className="h-12 text-xs sm:text-sm">
-												{header.isPlaceholder
-													? null
-													: flexRender(header.column.columnDef.header, header.getContext())}
-											</TableHead>
-										))}
-									</TableRow>
-								))}
-							</TableHeader>
-							<TableBody>
-								{table.getRowModel().rows?.length ? (
-									table.getRowModel().rows.map((row) => (
-										<TableRow
-											key={row.id}
-											data-state={row.getIsSelected() && "selected"}
-											className="hover:bg-muted/30"
-										>
-											{row.getVisibleCells().map((cell) => (
-												<TableCell key={cell.id} className="py-3 text-xs sm:text-sm">
-													{flexRender(cell.column.columnDef.cell, cell.getContext())}
-												</TableCell>
-											))}
-										</TableRow>
-									))
-								) : (
-									<TableRow>
-										<TableCell colSpan={columns.length} className="h-64">
-											<div className="flex flex-col items-center justify-center gap-4 py-8">
-												<div className="rounded-full bg-muted p-3">
-													<FileText className="h-6 w-6 text-muted-foreground" />
-												</div>
-												<div className="space-y-2 text-center max-w-sm">
-													<h3 className="font-semibold">No documents found</h3>
-													<p className="text-sm text-muted-foreground">
-														Get started by adding your first data source to build your knowledge
-														base.
-													</p>
-												</div>
-												<Button
-													size="sm"
-													onClick={() => router.push(`/dashboard/${searchSpaceId}/sources/add`)}
-												>
-													<Plus className="mr-2 h-4 w-4" />
-													Add Sources
-												</Button>
-											</div>
-										</TableCell>
-									</TableRow>
-								)}
-							</TableBody>
-						</Table>
-					)}
-				</div>
-			</div>
 							
-			{/* Footer Pagination */}
+							return (
-			<div className="flex flex-col sm:flex-row sm:items-center sm:justify-between gap-3 text-xs sm:text-sm text-muted-foreground border-t pt-3 md:pt-4 flex-shrink-0">
+								<button
-				<div className="text-center sm:text-left">
+									key={doc.id}
-					Showing {pageIndex * pageSize + 1} to {Math.min((pageIndex + 1) * pageSize, actualTotal)}{" "}
+									ref={(el) => {
-					of {actualTotal} documents
+										if (el && selectableIndex >= 0) {
-				</div>
+											itemRefs.current.set(selectableIndex, el);
-				<div className="flex items-center justify-center sm:justify-end space-x-2">
+										}
-					<Button
+									}}
-						variant="outline"
+									type="button"
-						size="sm"
+									onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
-						onClick={() => setPageIndex((p) => Math.max(0, p - 1))}
+									onMouseEnter={() => {
-						disabled={pageIndex === 0 || actualLoading}
+										if (!isAlreadySelected && selectableIndex >= 0) {
-						className="text-xs sm:text-sm"
+											setHighlightedIndex(selectableIndex);
+										}
+									}}
+									disabled={isAlreadySelected}
+									className={cn(
+										"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
+										isAlreadySelected
+											? "opacity-50 cursor-not-allowed"
+											: "cursor-pointer",
+										isHighlighted && "bg-accent"
+									)}
 								>
-						Previous
+									{/* Type icon */}
-					</Button>
+									<span className="flex-shrink-0 text-muted-foreground text-sm">
-					<div className="flex items-center space-x-1 text-xs sm:text-sm">
+										{getConnectorIcon(doc.document_type)}
-						<span>Page</span>
+									</span>
-						<strong>{pageIndex + 1}</strong>
+									{/* Title */}
-						<span>of</span>
+									<span className="flex-1 text-sm truncate" title={doc.title}>
-						<strong>{Math.ceil(actualTotal / pageSize)}</strong>
+										{doc.title}
-					</div>
+									</span>
-					<Button
+								</button>
-						variant="outline"
+							);
-						size="sm"
+						})}
-						onClick={() => setPageIndex((p) => p + 1)}
-						disabled={pageIndex >= Math.ceil(actualTotal / pageSize) - 1 || actualLoading}
-						className="text-xs sm:text-sm"
-					>
-						Next
-					</Button>
-				</div>
 					</div>
+				)}
+			</ScrollArea>
 		</div>
 	);
 }