feat: Add Cohere reranking capability to vector stores (#16014)

Co-authored-by: Yiorgis Gozadinos <yiorgis@n8n.io>
Co-authored-by: Mutasem Aldmour <mutasem@n8n.io>
This commit is contained in:
Benjamin Schroth
2025-06-05 10:16:22 +02:00
committed by GitHub
parent 47ad74d137
commit 8a1cabe62a
14 changed files with 585 additions and 20 deletions

View File

@@ -6,6 +6,7 @@ import { Embeddings } from '@langchain/core/embeddings';
import type { InputValues, MemoryVariables, OutputValues } from '@langchain/core/memory';
import type { BaseMessage } from '@langchain/core/messages';
import { BaseRetriever } from '@langchain/core/retrievers';
import { BaseDocumentCompressor } from '@langchain/core/retrievers/document_compressors';
import type { StructuredTool, Tool } from '@langchain/core/tools';
import { VectorStore } from '@langchain/core/vectorstores';
import { TextSplitter } from '@langchain/textsplitters';
@@ -18,7 +19,12 @@ import type {
ITaskMetadata,
NodeConnectionType,
} from 'n8n-workflow';
import { NodeOperationError, NodeConnectionTypes, parseErrorMetadata } from 'n8n-workflow';
import {
NodeOperationError,
NodeConnectionTypes,
parseErrorMetadata,
deepCopy,
} from 'n8n-workflow';
import { logAiEvent, isToolsInstance, isBaseChatMemory, isBaseChatMessageHistory } from './helpers';
import { N8nBinaryLoader } from './N8nBinaryLoader';
@@ -102,6 +108,7 @@ export function logWrapper<
| BaseChatMemory
| BaseChatMessageHistory
| BaseRetriever
| BaseDocumentCompressor
| Embeddings
| Document[]
| Document
@@ -297,6 +304,32 @@ export function logWrapper<
}
}
// ========== Rerankers ==========
if (originalInstance instanceof BaseDocumentCompressor) {
if (prop === 'compressDocuments' && 'compressDocuments' in target) {
return async (documents: Document[], query: string): Promise<Document[]> => {
connectionType = NodeConnectionTypes.AiReranker;
const { index } = executeFunctions.addInputData(connectionType, [
[{ json: { query, documents } }],
]);
const response = (await callMethodAsync.call(target, {
executeFunctions,
connectionType,
currentNodeRunIndex: index,
method: target[prop],
// compressDocuments mutates the original object
// messing up the input data logging
arguments: [deepCopy(documents), query],
})) as Document[];
logAiEvent(executeFunctions, 'ai-document-reranked', { query });
executeFunctions.addOutputData(connectionType, index, [[{ json: { response } }]]);
return response;
};
}
}
// ========== N8n Loaders Process All ==========
if (
originalInstance instanceof N8nJsonLoader ||