-
-
Notifications
You must be signed in to change notification settings - Fork 23.4k
feat: add Rerankers from Azure #5576
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
ba6ab63
54a6d91
a675e90
e887f53
fd781b3
1ce5c79
94593f4
ceee2cf
d799420
14313de
1e947da
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <a target="_blank" href="https://docs.microsoft.com/en-us/azure/ai-foundry/">Azure AI Foundry documentation</a> 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 } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<Record<string, any>>[], | ||
| query: string, | ||
| _?: Callbacks | undefined | ||
| ): Promise<Document<Record<string, any>>[]> { | ||
| // 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<Record<string, any>>[] = [] | ||
| 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 | ||
| } | ||
| } | ||
|
Comment on lines
+44
to
+56
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This part of the code can be improved for robustness and clarity:
try {
const returnedDocs = await axios.post(this.AZURE_API_URL, data, config)
const finalResults: Document<Record<string, any>>[] = []
interface RerankResult {
index: number
relevance_score: number
}
returnedDocs.data.results.forEach((result: RerankResult) => {
const doc = documents[result.index]
doc.metadata.relevance_score = result.relevance_score
finalResults.push(doc)
})
return finalResults
} catch (error) {
throw new Error(`Azure Rerank API call failed: ${error.message}`)
}
} |
||
| } | ||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -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<any> { | ||||||||||||||||||||||
| 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) | ||||||||||||||||||||||
|
Comment on lines
+125
to
+126
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The
Suggested change
|
||||||||||||||||||||||
| 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) | ||||||||||||||||||||||
|
Comment on lines
+130
to
+133
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The variable
Suggested change
|
||||||||||||||||||||||
|
|
||||||||||||||||||||||
| 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 } | ||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are a few opportunities to improve naming, typing, and immutability in the class properties and constructor, following common TypeScript best practices:
AzureAPIKeyandAZURE_API_URLshould be incamelCase(e.g.,azureApiKey,azureApiUrl).AzureAPIKeyshould bestring, notany.readonly.Applying these changes will require updating the property access in the
compressDocumentsmethod as well.