SurfSense/surfsense_web/components/chat/ChatInputGroup.tsx
2025-08-03 12:16:40 +02:00

576 lines
18 KiB
TypeScript

"use client";
import { ChatInput } from "@llamaindex/chat-ui";
import { Brain, Check, FolderOpen, Zap } from "lucide-react";
import { useParams } from "next/navigation";
import React, { Suspense, useCallback, useState } from "react";
import type { ResearchMode } from "@/components/chat";
import {
ConnectorButton as ConnectorButtonComponent,
getConnectorIcon,
} from "@/components/chat/ConnectorComponents";
import { DocumentsDataTable } from "@/components/chat/DocumentsDataTable";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogTitle,
DialogTrigger,
} from "@/components/ui/dialog";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { type Document, useDocuments } from "@/hooks/use-documents";
import { useLLMConfigs, useLLMPreferences } from "@/hooks/use-llm-configs";
import { useSearchSourceConnectors } from "@/hooks/useSearchSourceConnectors";
const DocumentSelector = React.memo(
({
onSelectionChange,
selectedDocuments = [],
}: {
onSelectionChange?: (documents: Document[]) => void;
selectedDocuments?: Document[];
}) => {
const { search_space_id } = useParams();
const [isOpen, setIsOpen] = useState(false);
const { documents, loading, isLoaded, fetchDocuments } = useDocuments(
Number(search_space_id),
true
);
const handleOpenChange = useCallback(
(open: boolean) => {
setIsOpen(open);
if (open && !isLoaded) {
fetchDocuments();
}
},
[fetchDocuments, isLoaded]
);
const handleSelectionChange = useCallback(
(documents: Document[]) => {
onSelectionChange?.(documents);
},
[onSelectionChange]
);
const handleDone = useCallback(() => {
setIsOpen(false);
}, []);
const selectedCount = React.useMemo(() => selectedDocuments.length, [selectedDocuments.length]);
return (
<Dialog open={isOpen} onOpenChange={handleOpenChange}>
<DialogTrigger asChild>
<Button variant="outline" className="relative">
<FolderOpen className="w-4 h-4" />
{selectedCount > 0 && (
<span className="absolute -top-2 -right-2 h-5 w-5 rounded-full bg-primary text-primary-foreground text-xs flex items-center justify-center">
{selectedCount}
</span>
)}
</Button>
</DialogTrigger>
<DialogContent className="max-w-[95vw] md:max-w-5xl h-[90vh] md:h-[85vh] p-0 flex flex-col">
<div className="flex flex-col h-full">
<div className="px-4 md:px-6 py-4 border-b flex-shrink-0">
<DialogTitle className="text-lg md:text-xl">Select Documents</DialogTitle>
<DialogDescription className="mt-1 text-sm">
Choose documents to include in your research context
</DialogDescription>
</div>
<div className="flex-1 min-h-0 p-4 md:p-6">
{loading ? (
<div className="flex items-center justify-center h-full">
<div className="text-center space-y-2">
<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>
</div>
) : isLoaded ? (
<DocumentsDataTable
documents={documents}
onSelectionChange={handleSelectionChange}
onDone={handleDone}
initialSelectedDocuments={selectedDocuments}
/>
) : null}
</div>
</div>
</DialogContent>
</Dialog>
);
}
);
DocumentSelector.displayName = "DocumentSelector";
const ConnectorSelector = React.memo(
({
onSelectionChange,
selectedConnectors = [],
}: {
onSelectionChange?: (connectorTypes: string[]) => void;
selectedConnectors?: string[];
}) => {
const [isOpen, setIsOpen] = useState(false);
const { connectorSourceItems, isLoading, isLoaded, fetchConnectors } =
useSearchSourceConnectors(true);
const handleOpenChange = useCallback(
(open: boolean) => {
setIsOpen(open);
if (open && !isLoaded) {
fetchConnectors();
}
},
[fetchConnectors, isLoaded]
);
const handleConnectorToggle = useCallback(
(connectorType: string) => {
const isSelected = selectedConnectors.includes(connectorType);
const newSelection = isSelected
? selectedConnectors.filter((type) => type !== connectorType)
: [...selectedConnectors, connectorType];
onSelectionChange?.(newSelection);
},
[selectedConnectors, onSelectionChange]
);
const handleSelectAll = useCallback(() => {
onSelectionChange?.(connectorSourceItems.map((c) => c.type));
}, [connectorSourceItems, onSelectionChange]);
const handleClearAll = useCallback(() => {
onSelectionChange?.([]);
}, [onSelectionChange]);
return (
<Dialog open={isOpen} onOpenChange={handleOpenChange}>
<DialogTrigger asChild>
<ConnectorButtonComponent
selectedConnectors={selectedConnectors}
onClick={() => setIsOpen(true)}
connectorSources={connectorSourceItems}
/>
</DialogTrigger>
<DialogContent className="sm:max-w-md">
<DialogTitle>Select Connectors</DialogTitle>
<DialogDescription>
Choose which data sources to include in your research
</DialogDescription>
{/* Connector selection grid */}
<div className="grid grid-cols-2 gap-4 py-4">
{isLoading ? (
<div className="col-span-2 flex justify-center py-4">
<div className="animate-spin h-6 w-6 border-2 border-primary border-t-transparent rounded-full" />
</div>
) : (
connectorSourceItems.map((connector) => {
const isSelected = selectedConnectors.includes(connector.type);
return (
<Button
key={connector.id}
className={`flex items-center gap-2 p-2 rounded-md border cursor-pointer transition-colors`}
onClick={() => handleConnectorToggle(connector.type)}
variant={isSelected ? "default" : "outline"}
size="sm"
type="button"
>
{getConnectorIcon(connector.type)}
<span className="flex-1 text-sm truncate font-medium">{connector.name}</span>
</Button>
);
})
)}
</div>
<DialogFooter className="flex justify-between items-center">
<div className="flex gap-2">
<Button variant="outline" onClick={handleClearAll}>
Clear All
</Button>
<Button onClick={handleSelectAll}>Select All</Button>
</div>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
);
ConnectorSelector.displayName = "ConnectorSelector";
const SearchModeSelector = React.memo(
({
searchMode,
onSearchModeChange,
}: {
searchMode?: "DOCUMENTS" | "CHUNKS";
onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
}) => {
const handleDocumentsClick = React.useCallback(() => {
onSearchModeChange?.("DOCUMENTS");
}, [onSearchModeChange]);
const handleChunksClick = React.useCallback(() => {
onSearchModeChange?.("CHUNKS");
}, [onSearchModeChange]);
return (
<div className="flex items-center gap-1 sm:gap-2">
<span className="text-xs text-muted-foreground hidden sm:block">Scope:</span>
<div className="flex rounded-md border border-border overflow-hidden">
<Button
variant={searchMode === "DOCUMENTS" ? "default" : "ghost"}
size="sm"
className="rounded-none border-r h-8 px-2 sm:px-3 text-xs transition-all duration-200 hover:bg-muted/80"
onClick={handleDocumentsClick}
>
<span className="hidden sm:inline">Documents</span>
<span className="sm:hidden">Docs</span>
</Button>
<Button
variant={searchMode === "CHUNKS" ? "default" : "ghost"}
size="sm"
className="rounded-none h-8 px-2 sm:px-3 text-xs transition-all duration-200 hover:bg-muted/80"
onClick={handleChunksClick}
>
Chunks
</Button>
</div>
</div>
);
}
);
SearchModeSelector.displayName = "SearchModeSelector";
const ResearchModeSelector = React.memo(
({
researchMode,
onResearchModeChange,
}: {
researchMode?: ResearchMode;
onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
const handleValueChange = React.useCallback(
(value: string) => {
onResearchModeChange?.(value as ResearchMode);
},
[onResearchModeChange]
);
// Memoize mode options to prevent recreation
const modeOptions = React.useMemo(
() => [
{ value: "QNA", label: "Q&A", shortLabel: "Q&A" },
{
value: "REPORT_GENERAL",
label: "General Report",
shortLabel: "General",
},
{
value: "REPORT_DEEP",
label: "Deep Report",
shortLabel: "Deep",
},
{
value: "REPORT_DEEPER",
label: "Deeper Report",
shortLabel: "Deeper",
},
],
[]
);
return (
<div className="flex items-center gap-1 sm:gap-2">
<span className="text-xs text-muted-foreground hidden sm:block">Mode:</span>
<Select value={researchMode} onValueChange={handleValueChange}>
<SelectTrigger className="w-auto min-w-[80px] sm:min-w-[120px] h-8 text-xs border-border bg-background hover:bg-muted/50 transition-colors duration-200 focus:ring-2 focus:ring-primary/20">
<SelectValue placeholder="Mode" className="text-xs" />
</SelectTrigger>
<SelectContent align="end" className="min-w-[140px]">
<div className="px-2 py-1.5 text-xs font-medium text-muted-foreground border-b bg-muted/30">
Research Mode
</div>
{modeOptions.map((option) => (
<SelectItem
key={option.value}
value={option.value}
className="px-3 py-2 cursor-pointer hover:bg-accent/50 focus:bg-accent"
>
<span className="hidden sm:inline">{option.label}</span>
<span className="sm:hidden">{option.shortLabel}</span>
</SelectItem>
))}
</SelectContent>
</Select>
</div>
);
}
);
ResearchModeSelector.displayName = "ResearchModeSelector";
const LLMSelector = React.memo(() => {
const { llmConfigs, loading: llmLoading, error } = useLLMConfigs();
const { preferences, updatePreferences, loading: preferencesLoading } = useLLMPreferences();
const isLoading = llmLoading || preferencesLoading;
// Memoize the selected config to avoid repeated lookups
const selectedConfig = React.useMemo(() => {
if (!preferences.fast_llm_id || !llmConfigs.length) return null;
return llmConfigs.find((config) => config.id === preferences.fast_llm_id) || null;
}, [preferences.fast_llm_id, llmConfigs]);
// Memoize the display value for the trigger
const displayValue = React.useMemo(() => {
if (!selectedConfig) return null;
return (
<div className="flex items-center gap-1">
<span className="font-medium text-xs">{selectedConfig.provider}</span>
<span className="text-muted-foreground"></span>
<span className="hidden sm:inline text-muted-foreground text-xs truncate max-w-[60px]">
{selectedConfig.name}
</span>
</div>
);
}, [selectedConfig]);
const handleValueChange = React.useCallback(
(value: string) => {
const llmId = value ? parseInt(value, 10) : undefined;
updatePreferences({ fast_llm_id: llmId });
},
[updatePreferences]
);
// Loading skeleton
if (isLoading) {
return (
<div className="h-8 min-w-[100px] sm:min-w-[120px]">
<div className="h-8 rounded-md bg-muted animate-pulse flex items-center px-3">
<div className="w-3 h-3 rounded bg-muted-foreground/20 mr-2" />
<div className="h-3 w-16 rounded bg-muted-foreground/20" />
</div>
</div>
);
}
// Error state
if (error) {
return (
<div className="h-8 min-w-[100px] sm:min-w-[120px]">
<Button
variant="outline"
size="sm"
className="h-8 px-3 text-xs text-destructive border-destructive/50 hover:bg-destructive/10"
disabled
>
<span className="text-xs">Error</span>
</Button>
</div>
);
}
return (
<div className="h-8 min-w-0">
<Select
value={preferences.fast_llm_id?.toString() || ""}
onValueChange={handleValueChange}
disabled={isLoading}
>
<SelectTrigger className="h-8 w-auto min-w-[100px] sm:min-w-[120px] px-3 text-xs border-border bg-background hover:bg-muted/50 transition-colors duration-200 focus:ring-2 focus:ring-primary/20">
<div className="flex items-center gap-2 min-w-0">
<Zap className="h-3 w-3 text-primary flex-shrink-0" />
<SelectValue placeholder="Fast LLM" className="text-xs">
{displayValue || <span className="text-muted-foreground">Select LLM</span>}
</SelectValue>
</div>
</SelectTrigger>
<SelectContent align="end" className="w-[300px] max-h-[400px]">
<div className="px-3 py-2 text-xs font-medium text-muted-foreground border-b bg-muted/30">
<div className="flex items-center gap-2">
<Zap className="h-3 w-3" />
Fast LLM Selection
</div>
</div>
{llmConfigs.length === 0 ? (
<div className="px-4 py-6 text-center">
<div className="mx-auto w-12 h-12 rounded-full bg-muted flex items-center justify-center mb-3">
<Brain className="h-5 w-5 text-muted-foreground" />
</div>
<h4 className="text-sm font-medium mb-1">No LLM configurations</h4>
<p className="text-xs text-muted-foreground mb-3">
Configure AI models to get started
</p>
<Button
variant="outline"
size="sm"
className="text-xs"
onClick={() => window.open("/settings", "_blank")}
>
Open Settings
</Button>
</div>
) : (
<div className="py-1">
{llmConfigs.map((config) => (
<SelectItem
key={config.id}
value={config.id.toString()}
className="px-3 py-2 cursor-pointer hover:bg-accent/50 focus:bg-accent"
>
<div className="flex items-center justify-between w-full min-w-0">
<div className="flex items-center gap-3 min-w-0 flex-1">
<div className="flex h-8 w-8 items-center justify-center rounded-md bg-primary/10 flex-shrink-0">
<Brain className="h-4 w-4 text-primary" />
</div>
<div className="min-w-0 flex-1">
<div className="flex items-center gap-2 mb-1">
<span className="font-medium text-sm truncate">{config.name}</span>
<Badge variant="outline" className="text-xs px-1.5 py-0.5 flex-shrink-0">
{config.provider}
</Badge>
</div>
<p className="text-xs text-muted-foreground font-mono truncate">
{config.model_name}
</p>
</div>
</div>
{preferences.fast_llm_id === config.id && (
<div className="flex h-5 w-5 items-center justify-center rounded-full bg-primary ml-2 flex-shrink-0">
<Check className="h-3 w-3 text-primary-foreground" />
</div>
)}
</div>
</SelectItem>
))}
</div>
)}
</SelectContent>
</Select>
</div>
);
});
LLMSelector.displayName = "LLMSelector";
const CustomChatInputOptions = React.memo(
({
onDocumentSelectionChange,
selectedDocuments,
onConnectorSelectionChange,
selectedConnectors,
searchMode,
onSearchModeChange,
researchMode,
onResearchModeChange,
}: {
onDocumentSelectionChange?: (documents: Document[]) => void;
selectedDocuments?: Document[];
onConnectorSelectionChange?: (connectorTypes: string[]) => void;
selectedConnectors?: string[];
searchMode?: "DOCUMENTS" | "CHUNKS";
onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
researchMode?: ResearchMode;
onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
// Memoize the loading fallback to prevent recreation
const loadingFallback = React.useMemo(
() => <div className="h-8 min-w-[100px] animate-pulse bg-muted rounded-md" />,
[]
);
return (
<div className="flex flex-wrap gap-2 sm:gap-3 items-center justify-start">
<Suspense fallback={loadingFallback}>
<DocumentSelector
onSelectionChange={onDocumentSelectionChange}
selectedDocuments={selectedDocuments}
/>
</Suspense>
<Suspense fallback={loadingFallback}>
<ConnectorSelector
onSelectionChange={onConnectorSelectionChange}
selectedConnectors={selectedConnectors}
/>
</Suspense>
<SearchModeSelector searchMode={searchMode} onSearchModeChange={onSearchModeChange} />
<ResearchModeSelector
researchMode={researchMode}
onResearchModeChange={onResearchModeChange}
/>
<LLMSelector />
</div>
);
}
);
CustomChatInputOptions.displayName = "CustomChatInputOptions";
export const ChatInputUI = React.memo(
({
onDocumentSelectionChange,
selectedDocuments,
onConnectorSelectionChange,
selectedConnectors,
searchMode,
onSearchModeChange,
researchMode,
onResearchModeChange,
}: {
onDocumentSelectionChange?: (documents: Document[]) => void;
selectedDocuments?: Document[];
onConnectorSelectionChange?: (connectorTypes: string[]) => void;
selectedConnectors?: string[];
searchMode?: "DOCUMENTS" | "CHUNKS";
onSearchModeChange?: (mode: "DOCUMENTS" | "CHUNKS") => void;
researchMode?: ResearchMode;
onResearchModeChange?: (mode: ResearchMode) => void;
}) => {
return (
<ChatInput>
<ChatInput.Form className="flex gap-2">
<ChatInput.Field className="flex-1" />
<ChatInput.Submit />
</ChatInput.Form>
<CustomChatInputOptions
onDocumentSelectionChange={onDocumentSelectionChange}
selectedDocuments={selectedDocuments}
onConnectorSelectionChange={onConnectorSelectionChange}
selectedConnectors={selectedConnectors}
searchMode={searchMode}
onSearchModeChange={onSearchModeChange}
researchMode={researchMode}
onResearchModeChange={onResearchModeChange}
/>
</ChatInput>
);
}
);
ChatInputUI.displayName = "ChatInputUI";