Merge pull request #471 from MODSetter/dev

MODSetter · web-flow · commit a9c9e3fe8651 · 2025-11-06T13:33:47.000-08:00
feat: added top_k in chat Interface.
diff --git a/surfsense_backend/app/agents/researcher/configuration.py b/surfsense_backend/app/agents/researcher/configuration.py
@@ -27,6 +27,7 @@ class Configuration:
     search_mode: SearchMode
     document_ids_to_add_in_context: list[int]
     language: str | None = None
+    top_k: int = 10
 
     @classmethod
     def from_runnable_config(
diff --git a/surfsense_backend/app/agents/researcher/nodes.py b/surfsense_backend/app/agents/researcher/nodes.py
@@ -1366,8 +1366,8 @@ async def handle_qna_workflow(
         }
     )
 
-    # Use a reasonable top_k for QNA - not too many documents to avoid overwhelming the LLM
-    top_k = 5 if configuration.search_mode == SearchMode.DOCUMENTS else 20
+    # Use the top_k value from configuration
+    top_k = configuration.top_k
 
     relevant_documents = []
     user_selected_documents = []
diff --git a/surfsense_backend/app/routes/chats_routes.py b/surfsense_backend/app/routes/chats_routes.py
@@ -24,6 +24,7 @@
     validate_research_mode,
     validate_search_mode,
     validate_search_space_id,
+    validate_top_k,
 )
 
 router = APIRouter()
@@ -54,6 +55,7 @@ async def handle_chat_data(
         request_data.get("document_ids_to_add_in_context")
     )
     search_mode_str = validate_search_mode(request_data.get("search_mode"))
+    top_k = validate_top_k(request_data.get("top_k"))
     # print("RESQUEST DATA:", request_data)
     # print("SELECTED CONNECTORS:", selected_connectors)
 
@@ -123,6 +125,7 @@ async def handle_chat_data(
             search_mode_str,
             document_ids_to_add_in_context,
             language,
+            top_k,
         )
     )
 
diff --git a/surfsense_backend/app/tasks/stream_connector_search_results.py b/surfsense_backend/app/tasks/stream_connector_search_results.py
@@ -21,6 +21,7 @@ async def stream_connector_search_results(
     search_mode_str: str,
     document_ids_to_add_in_context: list[int],
     language: str | None = None,
+    top_k: int = 10,
 ) -> AsyncGenerator[str, None]:
     """
     Stream connector search results to the client
@@ -56,6 +57,7 @@ async def stream_connector_search_results(
             "search_mode": search_mode,
             "document_ids_to_add_in_context": document_ids_to_add_in_context,
             "language": language,  # Add language to the configuration
+            "top_k": top_k,  # Add top_k to the configuration
         }
     }
     # print(f"Researcher configuration: {config['configurable']}")  # Debug print
diff --git a/surfsense_backend/app/utils/validators.py b/surfsense_backend/app/utils/validators.py
@@ -241,6 +241,60 @@ def validate_search_mode(search_mode: Any) -> str:
     return normalized_mode
 
 
+def validate_top_k(top_k: Any) -> int:
+    """
+    Validate and convert top_k to integer.
+
+    Args:
+        top_k: The top_k value to validate
+
+    Returns:
+        int: Validated top_k value (defaults to 10 if None)
+
+    Raises:
+        HTTPException: If validation fails
+    """
+    if top_k is None:
+        return 10  # Default value
+
+    if isinstance(top_k, bool):
+        raise HTTPException(
+            status_code=400, detail="top_k must be an integer, not a boolean"
+        )
+
+    if isinstance(top_k, int):
+        if top_k <= 0:
+            raise HTTPException(
+                status_code=400, detail="top_k must be a positive integer"
+            )
+        if top_k > 100:
+            raise HTTPException(status_code=400, detail="top_k must not exceed 100")
+        return top_k
+
+    if isinstance(top_k, str):
+        if not top_k.strip():
+            raise HTTPException(status_code=400, detail="top_k cannot be empty")
+
+        if not re.match(r"^[1-9]\d*$", top_k.strip()):
+            raise HTTPException(
+                status_code=400, detail="top_k must be a valid positive integer"
+            )
+
+        value = int(top_k.strip())
+        if value <= 0:
+            raise HTTPException(
+                status_code=400, detail="top_k must be a positive integer"
+            )
+        if value > 100:
+            raise HTTPException(status_code=400, detail="top_k must not exceed 100")
+        return value
+
+    raise HTTPException(
+        status_code=400,
+        detail="top_k must be an integer or string representation of an integer",
+    )
+
+
 def validate_messages(messages: Any) -> list[dict]:
     """
     Validate messages structure.
diff --git a/surfsense_web/app/dashboard/[search_space_id]/researcher/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/researcher/[[...chat_id]]/page.tsx
@@ -27,6 +27,8 @@ export default function ResearcherPage() {
 		setSelectedConnectors,
 		selectedDocuments,
 		setSelectedDocuments,
+		topK,
+		setTopK,
 	} = useChatState({
 		search_space_id: search_space_id as string,
 		chat_id: chatIdParam,
@@ -66,6 +68,7 @@ export default function ResearcherPage() {
 		selectedConnectors: string[];
 		searchMode: "DOCUMENTS" | "CHUNKS";
 		researchMode: "QNA"; // Always QNA mode
+		topK: number;
 	}
 
 	const getChatStateStorageKey = (searchSpaceId: string, chatId: string) =>
@@ -105,6 +108,7 @@ export default function ResearcherPage() {
 				research_mode: researchMode,
 				search_mode: searchMode,
 				document_ids_to_add_in_context: documentIds,
+				top_k: topK,
 			},
 		},
 		onError: (error) => {
@@ -124,6 +128,7 @@ export default function ResearcherPage() {
 				selectedConnectors,
 				searchMode,
 				researchMode,
+				topK,
 			});
 			router.replace(`/dashboard/${search_space_id}/researcher/${newChatId}`);
 		}
@@ -145,10 +150,18 @@ export default function ResearcherPage() {
 				setSelectedDocuments(restoredState.selectedDocuments);
 				setSelectedConnectors(restoredState.selectedConnectors);
 				setSearchMode(restoredState.searchMode);
+				setTopK(restoredState.topK);
 				// researchMode is always "QNA", no need to restore
 			}
 		}
-	}, [chatIdParam, search_space_id, setSelectedDocuments, setSelectedConnectors, setSearchMode]);
+	}, [
+		chatIdParam,
+		search_space_id,
+		setSelectedDocuments,
+		setSelectedConnectors,
+		setSearchMode,
+		setTopK,
+	]);
 
 	// Set all sources as default for new chats
 	useEffect(() => {
@@ -234,6 +247,8 @@ export default function ResearcherPage() {
 			selectedConnectors={selectedConnectors}
 			searchMode={searchMode}
 			onSearchModeChange={setSearchMode}
+			topK={topK}
+			onTopKChange={setTopK}
 		/>
 	);
 }
diff --git a/surfsense_web/components/chat/ChatInputGroup.tsx b/surfsense_web/components/chat/ChatInputGroup.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { ChatInput } from "@llamaindex/chat-ui";
-import { Brain, Check, FolderOpen, Zap } from "lucide-react";
+import { Brain, Check, FolderOpen, Minus, Plus, Zap } from "lucide-react";
 import { useParams } from "next/navigation";
 import React, { Suspense, useCallback, useState } from "react";
 import { DocumentsDataTable } from "@/components/chat/DocumentsDataTable";
@@ -15,13 +15,15 @@ import {
 	DialogTitle,
 	DialogTrigger,
 } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
 import {
 	Select,
 	SelectContent,
 	SelectItem,
 	SelectTrigger,
 	SelectValue,
 } from "@/components/ui/select";
+import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import { useDocumentTypes } from "@/hooks/use-document-types";
 import type { Document } from "@/hooks/use-documents";
@@ -447,6 +449,119 @@ const SearchModeSelector = React.memo(
 
 SearchModeSelector.displayName = "SearchModeSelector";
 
+const TopKSelector = React.memo(
+	({ topK = 10, onTopKChange }: { topK?: number; onTopKChange?: (topK: number) => void }) => {
+		const MIN_VALUE = 1;
+		const MAX_VALUE = 100;
+
+		const handleIncrement = React.useCallback(() => {
+			if (topK < MAX_VALUE) {
+				onTopKChange?.(topK + 1);
+			}
+		}, [topK, onTopKChange]);
+
+		const handleDecrement = React.useCallback(() => {
+			if (topK > MIN_VALUE) {
+				onTopKChange?.(topK - 1);
+			}
+		}, [topK, onTopKChange]);
+
+		const handleInputChange = React.useCallback(
+			(e: React.ChangeEvent<HTMLInputElement>) => {
+				const value = e.target.value;
+				// Allow empty input for editing
+				if (value === "") {
+					return;
+				}
+				const numValue = parseInt(value, 10);
+				if (!isNaN(numValue) && numValue >= MIN_VALUE && numValue <= MAX_VALUE) {
+					onTopKChange?.(numValue);
+				}
+			},
+			[onTopKChange]
+		);
+
+		const handleInputBlur = React.useCallback(
+			(e: React.FocusEvent<HTMLInputElement>) => {
+				const value = e.target.value;
+				if (value === "") {
+					// Reset to default if empty
+					onTopKChange?.(10);
+					return;
+				}
+				const numValue = parseInt(value, 10);
+				if (isNaN(numValue) || numValue < MIN_VALUE) {
+					onTopKChange?.(MIN_VALUE);
+				} else if (numValue > MAX_VALUE) {
+					onTopKChange?.(MAX_VALUE);
+				}
+			},
+			[onTopKChange]
+		);
+
+		return (
+			<TooltipProvider>
+				<Tooltip delayDuration={200}>
+					<TooltipTrigger asChild>
+						<div className="flex items-center h-8 border rounded-md bg-background hover:bg-accent/50 transition-colors">
+							<Button
+								type="button"
+								variant="ghost"
+								size="icon"
+								className="h-full w-7 rounded-l-md rounded-r-none hover:bg-accent border-r"
+								onClick={handleDecrement}
+								disabled={topK <= MIN_VALUE}
+							>
+								<Minus className="h-3.5 w-3.5" />
+							</Button>
+							<div className="flex flex-col items-center justify-center px-2 min-w-[60px]">
+								<Input
+									type="number"
+									value={topK}
+									onChange={handleInputChange}
+									onBlur={handleInputBlur}
+									min={MIN_VALUE}
+									max={MAX_VALUE}
+									className="h-5 w-full px-1 text-center text-sm font-semibold border-0 bg-transparent focus-visible:ring-0 focus-visible:ring-offset-0 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none"
+								/>
+								<span className="text-[10px] text-muted-foreground leading-none">Results</span>
+							</div>
+							<Button
+								type="button"
+								variant="ghost"
+								size="icon"
+								className="h-full w-7 rounded-r-md rounded-l-none hover:bg-accent border-l"
+								onClick={handleIncrement}
+								disabled={topK >= MAX_VALUE}
+							>
+								<Plus className="h-3.5 w-3.5" />
+							</Button>
+						</div>
+					</TooltipTrigger>
+					<TooltipContent side="top" className="max-w-xs">
+						<div className="space-y-2">
+							<p className="text-sm font-semibold">Results per Source</p>
+							<p className="text-xs text-muted-foreground leading-relaxed">
+								Control how many results to fetch from each data source. Set a higher number to get
+								more information, or a lower number for faster, more focused results.
+							</p>
+							<div className="flex items-center gap-2 text-xs text-muted-foreground pt-1 border-t">
+								<span>Recommended: 5-20</span>
+								<span>•</span>
+								<span>
+									Range: {MIN_VALUE}-{MAX_VALUE}
+								</span>
+							</div>
+						</div>
+					</TooltipContent>
+				</Tooltip>
+			</TooltipProvider>
+		);
+	}
+);
+
+TopKSelector.displayName = "TopKSelector";
+
 const LLMSelector = React.memo(() => {
 	const { search_space_id } = useParams();
 	const searchSpaceId = Number(search_space_id);
@@ -604,13 +719,17 @@ const CustomChatInputOptions = React.memo(
 		selectedConnectors,
 		searchMode,
 		onSearchModeChange,
+		topK,
+		onTopKChange,
 	}: {
 		onDocumentSelectionChange?: (documents: Document[]) => void;
 		selectedDocuments?: Document[];
 		onConnectorSelectionChange?: (connectorTypes: string[]) => void;
 		selectedConnectors?: string[];
 		searchMode?: "DOCUMENTS" | "CHUNKS";
 		onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
+		topK?: number;
+		onTopKChange?: (topK: number) => void;
 	}) => {
 		// Memoize the loading fallback to prevent recreation
 		const loadingFallback = React.useMemo(
@@ -637,6 +756,8 @@ const CustomChatInputOptions = React.memo(
 				<div className="h-4 w-px bg-border hidden sm:block" />
 				<SearchModeSelector searchMode={searchMode} onSearchModeChange={onSearchModeChange} />
 				<div className="h-4 w-px bg-border hidden sm:block" />
+				<TopKSelector topK={topK} onTopKChange={onTopKChange} />
+				<div className="h-4 w-px bg-border hidden sm:block" />
 				<LLMSelector />
 			</div>
 		);
@@ -653,13 +774,17 @@ export const ChatInputUI = React.memo(
 		selectedConnectors,
 		searchMode,
 		onSearchModeChange,
+		topK,
+		onTopKChange,
 	}: {
 		onDocumentSelectionChange?: (documents: Document[]) => void;
 		selectedDocuments?: Document[];
 		onConnectorSelectionChange?: (connectorTypes: string[]) => void;
 		selectedConnectors?: string[];
 		searchMode?: "DOCUMENTS" | "CHUNKS";
 		onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
+		topK?: number;
+		onTopKChange?: (topK: number) => void;
 	}) => {
 		return (
 			<ChatInput>
@@ -674,6 +799,8 @@ export const ChatInputUI = React.memo(
 					selectedConnectors={selectedConnectors}
 					searchMode={searchMode}
 					onSearchModeChange={onSearchModeChange}
+					topK={topK}
+					onTopKChange={onTopKChange}
 				/>
 			</ChatInput>
 		);
diff --git a/surfsense_web/components/chat/ChatInterface.tsx b/surfsense_web/components/chat/ChatInterface.tsx
diff --git a/surfsense_web/hooks/use-chat.ts b/surfsense_web/hooks/use-chat.ts
diff --git a/surfsense_web/package.json b/surfsense_web/package.json

Original file line number	Diff line number	Diff line change
`@@ -1366,8 +1366,8 @@ async def handle_qna_workflow(`
`1366`	`1366`	`}`
`1367`	`1367`	`)`
`1368`	`1368`
`1369`		`- # Use a reasonable top_k for QNA - not too many documents to avoid overwhelming the LLM`
`1370`		`- top_k = 5 if configuration.search_mode == SearchMode.DOCUMENTS else 20`
	`1369`	`+ # Use the top_k value from configuration`
	`1370`	`+ top_k = configuration.top_k`
`1371`	`1371`
`1372`	`1372`	`relevant_documents = []`
`1373`	`1373`	`user_selected_documents = []`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@`
`24`	`24`	`validate_research_mode,`
`25`	`25`	`validate_search_mode,`
`26`	`26`	`validate_search_space_id,`
	`27`	`+ validate_top_k,`
`27`	`28`	`)`
`28`	`29`
`29`	`30`	`router = APIRouter()`
`@@ -54,6 +55,7 @@ async def handle_chat_data(`
`54`	`55`	`request_data.get("document_ids_to_add_in_context")`
`55`	`56`	`)`
`56`	`57`	`search_mode_str = validate_search_mode(request_data.get("search_mode"))`
	`58`	`+ top_k = validate_top_k(request_data.get("top_k"))`
`57`	`59`	`# print("RESQUEST DATA:", request_data)`
`58`	`60`	`# print("SELECTED CONNECTORS:", selected_connectors)`
`59`	`61`
`@@ -123,6 +125,7 @@ async def handle_chat_data(`
`123`	`125`	`search_mode_str,`
`124`	`126`	`document_ids_to_add_in_context,`
`125`	`127`	`language,`
	`128`	`+ top_k,`
`126`	`129`	`)`
`127`	`130`	`)`
`128`	`131`