mirror of
https://github.com/NVIDIA/dgx-spark-playbooks.git
synced 2026-04-23 02:23:53 +00:00
359 lines
15 KiB
TypeScript
359 lines
15 KiB
TypeScript
//
|
|
// SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
import { useState, useEffect } from "react";
|
|
import { Triple } from "@/types/graph";
|
|
import { Search as SearchIcon, Zap, Database, Cpu } from "lucide-react";
|
|
import { LLMSelectorCompact } from "./llm-selector-compact";
|
|
|
|
interface RagQueryProps {
|
|
onQuerySubmit: (query: string, params: RagParams) => Promise<void>;
|
|
clearResults: () => void;
|
|
isLoading: boolean;
|
|
error?: string | null;
|
|
className?: string;
|
|
vectorEnabled?: boolean;
|
|
}
|
|
|
|
type QueryMode = 'traditional' | 'vector-search' | 'pure-rag';
|
|
|
|
export interface RagParams {
|
|
kNeighbors: number;
|
|
fanout: number;
|
|
numHops: number;
|
|
topK: number;
|
|
useVectorSearch?: boolean;
|
|
usePureRag?: boolean;
|
|
queryMode?: QueryMode;
|
|
}
|
|
|
|
export function RagQuery({
|
|
onQuerySubmit,
|
|
clearResults,
|
|
isLoading,
|
|
error,
|
|
className,
|
|
vectorEnabled
|
|
}: RagQueryProps) {
|
|
const [query, setQuery] = useState("");
|
|
const [showAdvanced, setShowAdvanced] = useState(false);
|
|
const [queryMode, setQueryMode] = useState<QueryMode>('traditional');
|
|
const [params, setParams] = useState<RagParams>({
|
|
kNeighbors: 4096,
|
|
fanout: 400,
|
|
numHops: 2,
|
|
topK: 40, // Increased to 40 for multi-hop paths (was 20, originally 5)
|
|
useVectorSearch: false,
|
|
usePureRag: false,
|
|
queryMode: 'traditional'
|
|
});
|
|
|
|
// Update query mode when vectorEnabled changes - but don't override user selection
|
|
useEffect(() => {
|
|
// Only set default mode on initial load, not when user selects Graph Search
|
|
if (vectorEnabled && queryMode === 'traditional') {
|
|
console.log('Not forcing mode change - respecting user selection');
|
|
// No longer forcing mode change - let user select what they want
|
|
}
|
|
}, [vectorEnabled, queryMode]);
|
|
|
|
const handleSubmit = async (e: React.FormEvent) => {
|
|
e.preventDefault();
|
|
if (!query.trim()) return;
|
|
|
|
// Set the appropriate params based on the selected query mode
|
|
const updatedParams = { ...params };
|
|
|
|
if (queryMode === 'pure-rag') {
|
|
updatedParams.usePureRag = true;
|
|
updatedParams.useVectorSearch = false;
|
|
} else if (queryMode === 'vector-search') {
|
|
updatedParams.usePureRag = false;
|
|
updatedParams.useVectorSearch = true;
|
|
} else {
|
|
// traditional
|
|
updatedParams.usePureRag = false;
|
|
updatedParams.useVectorSearch = false;
|
|
}
|
|
|
|
// Pass the current queryMode to the parent component
|
|
await onQuerySubmit(query, {
|
|
...updatedParams,
|
|
queryMode: queryMode // Add the query mode to the params
|
|
});
|
|
};
|
|
|
|
const handleReset = () => {
|
|
setQuery("");
|
|
clearResults();
|
|
};
|
|
|
|
const updateParam = (key: keyof RagParams, value: number | boolean) => {
|
|
setParams(prev => ({ ...prev, [key]: value }));
|
|
};
|
|
|
|
// Handle changing the query mode
|
|
const handleQueryModeChange = (mode: QueryMode) => {
|
|
console.log(`Changing query mode to: ${mode}`);
|
|
setQueryMode(mode);
|
|
|
|
// Force update params based on the selected mode
|
|
const updatedParams = { ...params };
|
|
if (mode === 'pure-rag') {
|
|
updatedParams.usePureRag = true;
|
|
updatedParams.useVectorSearch = false;
|
|
} else if (mode === 'vector-search') {
|
|
updatedParams.usePureRag = false;
|
|
updatedParams.useVectorSearch = true;
|
|
} else {
|
|
// traditional mode
|
|
updatedParams.usePureRag = false;
|
|
updatedParams.useVectorSearch = false;
|
|
}
|
|
setParams(updatedParams);
|
|
};
|
|
|
|
return (
|
|
<div className={`nvidia-build-card ${className}`}>
|
|
<div className="mb-6">
|
|
<div className="flex items-center gap-3 mb-3">
|
|
<div className="w-8 h-8 rounded-lg bg-nvidia-green/15 flex items-center justify-center">
|
|
<SearchIcon className="h-4 w-4 text-nvidia-green" />
|
|
</div>
|
|
<h2 className="text-xl font-semibold text-foreground">RAG Query Engine</h2>
|
|
</div>
|
|
<p className="text-sm text-muted-foreground leading-relaxed">
|
|
Query your knowledge graph using natural language
|
|
</p>
|
|
</div>
|
|
|
|
{/* Query Type Selection */}
|
|
<div className="mb-5">
|
|
<h3 className="text-xs font-semibold text-foreground mb-3">Select Query Type</h3>
|
|
<div className="grid grid-cols-1 md:grid-cols-3 gap-3">
|
|
<button
|
|
type="button"
|
|
onClick={() => handleQueryModeChange('pure-rag')}
|
|
disabled={!vectorEnabled}
|
|
className={`relative flex flex-col items-center p-3 border rounded-lg transition-all duration-200 hover:shadow-md ${
|
|
queryMode === 'pure-rag'
|
|
? 'border-nvidia-green bg-nvidia-green/10 text-nvidia-green shadow-sm'
|
|
: vectorEnabled
|
|
? 'border-border/40 hover:border-border/60 hover:bg-muted/20'
|
|
: 'border-border/30 opacity-50 cursor-not-allowed'
|
|
}`}
|
|
>
|
|
<div className="w-5 h-5 rounded-md bg-nvidia-green/15 flex items-center justify-center mb-1.5">
|
|
<Zap className="h-2.5 w-2.5 text-nvidia-green" />
|
|
</div>
|
|
<span className="text-sm font-semibold">Pure RAG</span>
|
|
<span className="text-[10px] mt-0.5 text-center text-muted-foreground leading-tight">
|
|
Vector DB + LLM
|
|
</span>
|
|
{queryMode === 'pure-rag' && (
|
|
<div className="absolute top-2 right-2 w-1.5 h-1.5 bg-nvidia-green rounded-full"></div>
|
|
)}
|
|
</button>
|
|
|
|
<button
|
|
type="button"
|
|
onClick={() => handleQueryModeChange('traditional')}
|
|
className={`relative flex flex-col items-center p-3 border rounded-lg transition-all duration-200 hover:shadow-md ${
|
|
queryMode === 'traditional'
|
|
? 'border-nvidia-green bg-nvidia-green/10 text-nvidia-green shadow-sm'
|
|
: 'border-border/40 hover:border-border/60 hover:bg-muted/20'
|
|
}`}
|
|
>
|
|
<div className="w-5 h-5 rounded-md bg-nvidia-green/15 flex items-center justify-center mb-1.5">
|
|
<Database className="h-2.5 w-2.5 text-nvidia-green" />
|
|
</div>
|
|
<span className="text-sm font-semibold">Graph Search</span>
|
|
<span className="text-[10px] mt-0.5 text-center text-muted-foreground leading-tight">
|
|
Graph DB + LLM
|
|
</span>
|
|
{queryMode === 'traditional' && (
|
|
<div className="absolute top-2 right-2 w-1.5 h-1.5 bg-nvidia-green rounded-full"></div>
|
|
)}
|
|
</button>
|
|
|
|
<button
|
|
type="button"
|
|
onClick={() => handleQueryModeChange('vector-search')}
|
|
disabled={true}
|
|
title="GraphRAG requires a GNN model (not yet available)"
|
|
className="relative flex flex-col items-center p-3 border rounded-lg border-border/30 opacity-50 cursor-not-allowed"
|
|
>
|
|
<div className="w-5 h-5 rounded-md bg-muted/15 flex items-center justify-center mb-1.5">
|
|
<Cpu className="h-2.5 w-2.5 text-muted-foreground" />
|
|
</div>
|
|
<span className="text-sm font-semibold text-muted-foreground">GraphRAG</span>
|
|
<span className="text-[10px] mt-0.5 text-center text-muted-foreground leading-tight">
|
|
RAG + GNN
|
|
</span>
|
|
<div className="text-[9px] px-1.5 py-0.5 bg-amber-500/20 text-amber-700 dark:text-amber-500 rounded mt-1 font-medium">
|
|
COMING SOON
|
|
</div>
|
|
</button>
|
|
</div>
|
|
</div>
|
|
|
|
{/* LLM Selection */}
|
|
<div className="mb-6 flex items-center justify-between">
|
|
<div>
|
|
<h3 className="text-sm font-semibold text-foreground mb-1">LLM Model</h3>
|
|
<p className="text-xs text-muted-foreground">Select the LLM for answer generation</p>
|
|
</div>
|
|
<LLMSelectorCompact />
|
|
</div>
|
|
|
|
<form onSubmit={handleSubmit} className="space-y-6">
|
|
<div className="flex gap-3">
|
|
<div className="relative flex-1">
|
|
<input
|
|
type="text"
|
|
placeholder={
|
|
queryMode === 'pure-rag'
|
|
? "Ask a question using RAG..."
|
|
: queryMode === 'vector-search'
|
|
? "Ask a question about your vector-enhanced knowledge graph..."
|
|
: "Ask a question about your knowledge graph..."
|
|
}
|
|
value={query}
|
|
onChange={(e) => setQuery(e.target.value)}
|
|
className="w-full p-3 pl-11 rounded-xl border border-border/60 bg-background text-foreground focus:border-nvidia-green/50 focus:ring-2 focus:ring-nvidia-green/20 transition-colors"
|
|
disabled={isLoading}
|
|
/>
|
|
<SearchIcon className="absolute left-3.5 top-1/2 transform -translate-y-1/2 h-4 w-4 text-muted-foreground" />
|
|
</div>
|
|
<button
|
|
type="submit"
|
|
disabled={isLoading || !query.trim()}
|
|
className="px-6 py-3 bg-nvidia-green text-white hover:bg-nvidia-green/90 rounded-xl font-medium transition-colors disabled:opacity-50 disabled:pointer-events-none shadow-sm"
|
|
>
|
|
{isLoading ? "Searching..." : "Search"}
|
|
</button>
|
|
<button
|
|
type="button"
|
|
onClick={handleReset}
|
|
disabled={isLoading}
|
|
className="px-4 py-3 border border-border/60 bg-background hover:bg-muted/30 text-foreground rounded-xl transition-colors disabled:opacity-50 disabled:pointer-events-none"
|
|
>
|
|
Clear
|
|
</button>
|
|
</div>
|
|
|
|
{error && (
|
|
<div className="p-4 bg-destructive/10 border border-destructive/30 rounded-xl text-destructive text-sm">
|
|
<div className="flex items-start gap-2">
|
|
<div className="w-4 h-4 rounded-full bg-destructive/20 flex items-center justify-center mt-0.5 flex-shrink-0">
|
|
<span className="text-xs">!</span>
|
|
</div>
|
|
<div>
|
|
<strong>Error:</strong> {error}
|
|
</div>
|
|
</div>
|
|
</div>
|
|
)}
|
|
|
|
<div>
|
|
<button
|
|
type="button"
|
|
onClick={() => setShowAdvanced(!showAdvanced)}
|
|
className="text-sm text-muted-foreground hover:text-foreground transition-colors flex items-center gap-2 font-medium"
|
|
>
|
|
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className={`h-4 w-4 transition-transform duration-200 ${showAdvanced ? 'rotate-180' : ''}`}>
|
|
<path d="m6 9 6 6 6-6"/>
|
|
</svg>
|
|
{showAdvanced ? "Hide" : "Show"} Advanced Parameters
|
|
</button>
|
|
|
|
{showAdvanced && (
|
|
<div className="mt-6 p-6 border border-border/40 rounded-xl bg-muted/10 space-y-6">
|
|
{/* Only show graph-related parameters for traditional and vector search modes */}
|
|
{queryMode !== 'pure-rag' && (
|
|
<>
|
|
<div>
|
|
<div className="flex justify-between items-center mb-2">
|
|
<label className="text-sm font-medium text-foreground">KNN Neighbors</label>
|
|
<span className="text-sm font-semibold text-nvidia-green bg-nvidia-green/10 px-2 py-1 rounded-md">{params.kNeighbors}</span>
|
|
</div>
|
|
<input
|
|
type="range"
|
|
min="256"
|
|
max="8192"
|
|
step="256"
|
|
value={params.kNeighbors}
|
|
onChange={(e) => updateParam('kNeighbors', parseInt(e.target.value))}
|
|
className="w-full h-2 bg-muted/30 rounded-lg appearance-none cursor-pointer slider-thumb"
|
|
/>
|
|
</div>
|
|
|
|
<div>
|
|
<div className="flex justify-between items-center mb-2">
|
|
<label className="text-sm font-medium text-foreground">Fanout</label>
|
|
<span className="text-sm font-semibold text-nvidia-green bg-nvidia-green/10 px-2 py-1 rounded-md">{params.fanout}</span>
|
|
</div>
|
|
<input
|
|
type="range"
|
|
min="50"
|
|
max="1000"
|
|
step="50"
|
|
value={params.fanout}
|
|
onChange={(e) => updateParam('fanout', parseInt(e.target.value))}
|
|
className="w-full h-2 bg-muted/30 rounded-lg appearance-none cursor-pointer slider-thumb"
|
|
/>
|
|
</div>
|
|
|
|
<div>
|
|
<div className="flex justify-between items-center mb-2">
|
|
<label className="text-sm font-medium text-foreground">Number of Hops</label>
|
|
<span className="text-sm font-semibold text-nvidia-green bg-nvidia-green/10 px-2 py-1 rounded-md">{params.numHops}</span>
|
|
</div>
|
|
<input
|
|
type="range"
|
|
min="1"
|
|
max="4"
|
|
step="1"
|
|
value={params.numHops}
|
|
onChange={(e) => updateParam('numHops', parseInt(e.target.value))}
|
|
className="w-full h-2 bg-muted/30 rounded-lg appearance-none cursor-pointer slider-thumb"
|
|
/>
|
|
</div>
|
|
</>
|
|
)}
|
|
|
|
<div>
|
|
<div className="flex justify-between items-center mb-2">
|
|
<label className="text-sm font-medium text-foreground">Top K Results</label>
|
|
<span className="text-sm font-semibold text-nvidia-green bg-nvidia-green/10 px-2 py-1 rounded-md">{params.topK}</span>
|
|
</div>
|
|
<input
|
|
type="range"
|
|
min="1"
|
|
max="50"
|
|
step="1"
|
|
value={params.topK}
|
|
onChange={(e) => updateParam('topK', parseInt(e.target.value))}
|
|
className="w-full h-2 bg-muted/30 rounded-lg appearance-none cursor-pointer slider-thumb"
|
|
/>
|
|
</div>
|
|
</div>
|
|
)}
|
|
</div>
|
|
</form>
|
|
</div>
|
|
);
|
|
}
|