Modify to add a custom prompt in RFF retriever

Author: Amrrx

packages/components/nodes/retrievers/RRFRetriever/RRFRetriever.ts

@@ -77,6 +77,26 @@ class RRFRetriever_Retrievers implements INode {
                 default: 60,
                 additionalParams: true,
                 optional: true
+            },
+            {
+                label: 'System Message',
+                name: 'systemMessage',
+                description: 'System message for query generation. Leave empty to use default.',
+                type: 'string',
+                rows: 2,
+                placeholder: 'You are a helpful assistant that generates multiple search queries based on a single input query.',
+                additionalParams: true,
+                optional: true
+            },
+            {
+                label: 'Query Generation Prompt',
+                name: 'queryPrompt',
+                description: 'Prompt template for generating query variations. Use {input} to refer to the original query.',
+                type: 'string',
+                rows: 4,
+                placeholder: 'Generate multiple search queries related to: {input}. Provide these alternative questions separated by newlines, do not add any numbers.',
+                additionalParams: true,
+                optional: true
             }
         ]
         this.outputs = [
@@ -110,9 +130,11 @@ class RRFRetriever_Retrievers implements INode {
         const k = topK ? parseFloat(topK) : (baseRetriever as VectorStoreRetriever).k ?? 4
         const constantC = nodeData.inputs?.c as string
         const c = topK ? parseFloat(constantC) : 60
+        const systemMessage = nodeData.inputs?.systemMessage as string
+        const queryPrompt = nodeData.inputs?.queryPrompt as string
         const output = nodeData.outputs?.output as string
 
-        const ragFusion = new ReciprocalRankFusion(llm, baseRetriever as VectorStoreRetriever, q, k, c)
+        const ragFusion = new ReciprocalRankFusion(llm, baseRetriever as VectorStoreRetriever, q, k, c, systemMessage, queryPrompt)
         const retriever = new ContextualCompressionRetriever({
             baseCompressor: ragFusion,
             baseRetriever: baseRetriever

packages/components/nodes/retrievers/RRFRetriever/ReciprocalRankFusion.ts

@@ -11,14 +11,26 @@ export class ReciprocalRankFusion extends BaseDocumentCompressor {
     private readonly queryCount: number
     private readonly topK: number
     private readonly c: number
+    private readonly systemMessage?: string
+    private readonly queryPrompt?: string
     private baseRetriever: VectorStoreRetriever
-    constructor(llm: BaseLanguageModel, baseRetriever: VectorStoreRetriever, queryCount: number, topK: number, c: number) {
+    constructor(
+        llm: BaseLanguageModel,
+        baseRetriever: VectorStoreRetriever,
+        queryCount: number,
+        topK: number,
+        c: number,
+        systemMessage?: string,
+        queryPrompt?: string
+    ) {
         super()
         this.queryCount = queryCount
         this.llm = llm
         this.baseRetriever = baseRetriever
         this.topK = topK
         this.c = c
+        this.systemMessage = systemMessage
+        this.queryPrompt = queryPrompt
     }
     async compressDocuments(
         documents: Document<Record<string, any>>[],
@@ -29,13 +41,17 @@ export class ReciprocalRankFusion extends BaseDocumentCompressor {
         if (documents.length === 0) {
             return []
         }
+        // Use custom prompts if provided, otherwise use defaults
+        const defaultSystemMessage = 'You are a helpful assistant that generates multiple search queries based on a single input query.'
+        const defaultQueryPrompt =
+            'Generate multiple search queries related to: {input}. Provide these alternative questions separated by newlines, do not add any numbers.'
+
+        const systemMsg = this.systemMessage || defaultSystemMessage
+        const queryMsg = this.queryPrompt || defaultQueryPrompt
+
         const chatPrompt = ChatPromptTemplate.fromMessages([
-            SystemMessagePromptTemplate.fromTemplate(
-                'You are a helpful assistant that generates multiple search queries based on a single input query.'
-            ),
-            HumanMessagePromptTemplate.fromTemplate(
-                'Generate multiple search queries related to: {input}. Provide these alternative questions separated by newlines, do not add any numbers.'
-            ),
+            SystemMessagePromptTemplate.fromTemplate(systemMsg),
+            HumanMessagePromptTemplate.fromTemplate(queryMsg),
             HumanMessagePromptTemplate.fromTemplate('OUTPUT (' + this.queryCount + ' queries):')
         ])
         const llmChain = new LLMChain({
@@ -50,7 +66,7 @@ export class ReciprocalRankFusion extends BaseDocumentCompressor {
         })
         const docList: Document<Record<string, any>>[][] = []
         for (let i = 0; i < queries.length; i++) {
-            const resultOne = await this.baseRetriever.vectorStore.similaritySearch(queries[i], 5, this.baseRetriever.filter)
+            const resultOne = await this.baseRetriever.vectorStore.similaritySearch(queries[i], this.topK, this.baseRetriever.filter)
             const docs: any[] = []
             resultOne.forEach((doc) => {
                 docs.push(doc)

packages/ui/src/ui-component/dialog/HistoryDialog.jsx

@@ -1,4 +1,4 @@
-import { useState, useEffect } from 'react'
+import { useState, useEffect, useCallback } from 'react'
 import { createPortal } from 'react-dom'
 import {
     Dialog,
@@ -43,7 +43,7 @@ const HistoryDialog = ({ show, dialogProps, onCancel, onRestore }) => {
     const dispatch = useDispatch()
     useNotifier() // Side effect hook
 
-    const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
+    const enqueueSnackbar = useCallback((...args) => dispatch(enqueueSnackbarAction(...args)), [dispatch])
 
     const [historyItems, setHistoryItems] = useState([])
     const [loading, setLoading] = useState(false)
@@ -85,7 +85,7 @@ const HistoryDialog = ({ show, dialogProps, onCancel, onRestore }) => {
         }
 
         loadHistory()
-    }, [show, entityType, entityId, currentPage, reloadTrigger, latestVersion, enqueueSnackbar])
+    }, [show, entityType, entityId, currentPage, reloadTrigger, enqueueSnackbar])
 
     const handleRestore = async (historyItem) => {
         const confirmed = await confirm({