diff --git a/packages/components/credentials/AzureRerankerApi.credential.ts b/packages/components/credentials/AzureRerankerApi.credential.ts new file mode 100644 index 00000000000..73d676368e4 --- /dev/null +++ b/packages/components/credentials/AzureRerankerApi.credential.ts @@ -0,0 +1,34 @@ +import { INodeParams, INodeCredential } from '../src/Interface' + +class AzureRerankerApi implements INodeCredential { + label: string + name: string + version: number + description: string + inputs: INodeParams[] + + constructor() { + this.label = 'Azure Foundry API' + this.name = 'azureFoundryApi' + this.version = 1.0 + this.description = + 'Refer to Azure AI Foundry documentation for setup instructions' + this.inputs = [ + { + label: 'Azure Foundry API Key', + name: 'azureFoundryApiKey', + type: 'password', + description: 'Your Azure AI Foundry API key' + }, + { + label: 'Azure Foundry Endpoint', + name: 'azureFoundryEndpoint', + type: 'string', + placeholder: 'https://your-foundry-instance.services.ai.azure.com/providers/cohere/v2/rerank', + description: 'Your Azure AI Foundry endpoint URL' + } + ] + } +} + +module.exports = { credClass: AzureRerankerApi } diff --git a/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerank.ts b/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerank.ts new file mode 100644 index 00000000000..31f45520525 --- /dev/null +++ b/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerank.ts @@ -0,0 +1,57 @@ +import axios from 'axios' +import { Callbacks } from '@langchain/core/callbacks/manager' +import { Document } from '@langchain/core/documents' +import { BaseDocumentCompressor } from 'langchain/retrievers/document_compressors' + +export class AzureRerank extends BaseDocumentCompressor { + private AzureAPIKey: any + private AZURE_API_URL: string + private readonly model: string + private readonly k: number + private readonly maxChunksPerDoc: number + constructor(AzureAPIKey: string, AZURE_API_URL: string, model: string, k: number, maxChunksPerDoc: number) { + super() + this.AzureAPIKey = AzureAPIKey + this.AZURE_API_URL = AZURE_API_URL + this.model = model + this.k = k + this.maxChunksPerDoc = maxChunksPerDoc + } + async compressDocuments( + documents: Document>[], + query: string, + _?: Callbacks | undefined + ): Promise>[]> { + // avoid empty api call + if (documents.length === 0) { + return [] + } + const config = { + headers: { + 'api-key': `${this.AzureAPIKey}`, + 'Content-Type': 'application/json', + Accept: 'application/json' + } + } + const data = { + model: this.model, + top_n: this.k, + max_chunks_per_doc: this.maxChunksPerDoc, + query: query, + return_documents: false, + documents: documents.map((doc) => doc.pageContent) + } + try { + let returnedDocs = await axios.post(this.AZURE_API_URL, data, config) + const finalResults: Document>[] = [] + returnedDocs.data.results.forEach((result: any) => { + const doc = documents[result.index] + doc.metadata.relevance_score = result.relevance_score + finalResults.push(doc) + }) + return finalResults.splice(0, this.k) + } catch (error) { + return documents + } + } +} diff --git a/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerankRetriever.ts b/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerankRetriever.ts new file mode 100644 index 00000000000..7f61b3ceea0 --- /dev/null +++ b/packages/components/nodes/retrievers/AzureRerankRetriever/AzureRerankRetriever.ts @@ -0,0 +1,156 @@ +import { BaseRetriever } from '@langchain/core/retrievers' +import { VectorStoreRetriever } from '@langchain/core/vectorstores' +import { ContextualCompressionRetriever } from 'langchain/retrievers/contextual_compression' +import { AzureRerank } from './AzureRerank' +import { getCredentialData, getCredentialParam, handleEscapeCharacters } from '../../../src' +import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface' + +class AzureRerankRetriever_Retrievers implements INode { + label: string + name: string + version: number + description: string + type: string + icon: string + category: string + baseClasses: string[] + inputs: INodeParams[] + credential: INodeParams + badge: string + outputs: INodeOutputsValue[] + + constructor() { + this.label = 'Azure Rerank Retriever' + this.name = 'AzureRerankRetriever' + this.version = 1.0 + this.type = 'Azure Rerank Retriever' + this.icon = 'azurefoundry.svg' + this.category = 'Retrievers' + this.description = 'Azure Rerank indexes the documents from most to least semantically relevant to the query.' + this.baseClasses = [this.type, 'BaseRetriever'] + this.credential = { + label: 'Connect Credential', + name: 'credential', + type: 'credential', + credentialNames: ['azureFoundryApi'] + } + this.inputs = [ + { + label: 'Vector Store Retriever', + name: 'baseRetriever', + type: 'VectorStoreRetriever' + }, + { + label: 'Model Name', + name: 'model', + type: 'options', + options: [ + { + label: 'rerank-v3.5', + name: 'rerank-v3.5' + }, + { + label: 'rerank-english-v3.0', + name: 'rerank-english-v3.0' + }, + { + label: 'rerank-multilingual-v3.0', + name: 'rerank-multilingual-v3.0' + }, + { + label: 'Cohere-rerank-v4.0-fast', + name: 'Cohere-rerank-v4.0-fast' + }, + { + label: 'Cohere-rerank-v4.0-pro', + name: 'Cohere-rerank-v4.0-pro' + } + ], + default: 'Cohere-rerank-v4.0-fast', + optional: true + }, + { + label: 'Query', + name: 'query', + type: 'string', + description: 'Query to retrieve documents from retriever. If not specified, user question will be used', + optional: true, + acceptVariable: true + }, + { + label: 'Top K', + name: 'topK', + description: 'Number of top results to fetch. Default to the TopK of the Base Retriever', + placeholder: '4', + type: 'number', + additionalParams: true, + optional: true + }, + { + label: 'Max Chunks Per Doc', + name: 'maxChunksPerDoc', + description: 'The maximum number of chunks to produce internally from a document. Default to 10', + placeholder: '10', + type: 'number', + additionalParams: true, + optional: true + } + ] + this.outputs = [ + { + label: 'Azure Rerank Retriever', + name: 'retriever', + baseClasses: this.baseClasses + }, + { + label: 'Document', + name: 'document', + description: 'Array of document objects containing metadata and pageContent', + baseClasses: ['Document', 'json'] + }, + { + label: 'Text', + name: 'text', + description: 'Concatenated string from pageContent of documents', + baseClasses: ['string', 'json'] + } + ] + } + + async init(nodeData: INodeData, input: string, options: ICommonObject): Promise { + const baseRetriever = nodeData.inputs?.baseRetriever as BaseRetriever + const model = nodeData.inputs?.model as string + const query = nodeData.inputs?.query as string + const credentialData = await getCredentialData(nodeData.credential ?? '', options) + const azureApiKey = getCredentialParam('azureFoundryApiKey', credentialData, nodeData) + const azureEndpoint = getCredentialParam('azureFoundryEndpoint', credentialData, nodeData) + const topK = nodeData.inputs?.topK as string + const k = topK ? parseFloat(topK) : (baseRetriever as VectorStoreRetriever).k ?? 4 + const maxChunksPerDoc = nodeData.inputs?.maxChunksPerDoc as string + const max_chunks_per_doc = maxChunksPerDoc ? parseFloat(maxChunksPerDoc) : 10 + const output = nodeData.outputs?.output as string + + const azureCompressor = new AzureRerank(azureApiKey, azureEndpoint, model, k, max_chunks_per_doc) + + const retriever = new ContextualCompressionRetriever({ + baseCompressor: azureCompressor, + baseRetriever: baseRetriever + }) + + if (output === 'retriever') return retriever + else if (output === 'document') return await retriever.getRelevantDocuments(query ? query : input) + else if (output === 'text') { + let finaltext = '' + + const docs = await retriever.getRelevantDocuments(query ? query : input) + + for (const doc of docs) finaltext += `${doc.pageContent}\n` + + return handleEscapeCharacters(finaltext, false) + } + + return retriever + } +} + +module.exports = { nodeClass: AzureRerankRetriever_Retrievers } diff --git a/packages/components/nodes/retrievers/AzureRerankRetriever/azurefoundry.svg b/packages/components/nodes/retrievers/AzureRerankRetriever/azurefoundry.svg new file mode 100644 index 00000000000..41f503f589a --- /dev/null +++ b/packages/components/nodes/retrievers/AzureRerankRetriever/azurefoundry.svg @@ -0,0 +1 @@ +AzureAI \ No newline at end of file