Skip to content

Commit 9c2203b

Browse files
Feature: Add Jina AI Rerank Retriever (#3898)
1 parent 5c9f178 commit 9c2203b

File tree

3 files changed

+185
-0
lines changed

3 files changed

+185
-0
lines changed
Loading
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import { Callbacks } from '@langchain/core/callbacks/manager'
2+
import { Document } from '@langchain/core/documents'
3+
import axios from 'axios'
4+
import { BaseDocumentCompressor } from 'langchain/retrievers/document_compressors'
5+
6+
export class JinaRerank extends BaseDocumentCompressor {
7+
private jinaAPIKey: string
8+
private readonly JINA_RERANK_API_URL = 'https://api.jina.ai/v1/rerank'
9+
private model: string = 'jina-reranker-v2-base-multilingual'
10+
private readonly topN: number
11+
12+
constructor(jinaAPIKey: string, model: string, topN: number) {
13+
super()
14+
this.jinaAPIKey = jinaAPIKey
15+
this.model = model
16+
this.topN = topN
17+
}
18+
async compressDocuments(
19+
documents: Document<Record<string, any>>[],
20+
query: string,
21+
_?: Callbacks | undefined
22+
): Promise<Document<Record<string, any>>[]> {
23+
if (documents.length === 0) {
24+
return []
25+
}
26+
const config = {
27+
headers: {
28+
Authorization: `Bearer ${this.jinaAPIKey}`,
29+
'Content-Type': 'application/json'
30+
}
31+
}
32+
const data = {
33+
model: this.model,
34+
query: query,
35+
documents: documents.map((doc) => doc.pageContent),
36+
top_n: this.topN
37+
}
38+
try {
39+
let returnedDocs = await axios.post(this.JINA_RERANK_API_URL, data, config)
40+
const finalResults: Document<Record<string, any>>[] = []
41+
returnedDocs.data.results.forEach((result: any) => {
42+
const doc = documents[result.index]
43+
doc.metadata.relevance_score = result.relevance_score
44+
finalResults.push(doc)
45+
})
46+
return finalResults
47+
} catch (error) {
48+
return documents
49+
}
50+
}
51+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import { BaseRetriever } from '@langchain/core/retrievers'
2+
import { ContextualCompressionRetriever } from 'langchain/retrievers/contextual_compression'
3+
import { getCredentialData, getCredentialParam, handleEscapeCharacters } from '../../../src'
4+
import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
5+
import { JinaRerank } from './JinaRerank'
6+
7+
class JinaRerankRetriever_Retrievers implements INode {
8+
label: string
9+
name: string
10+
version: number
11+
description: string
12+
type: string
13+
icon: string
14+
category: string
15+
baseClasses: string[]
16+
inputs: INodeParams[]
17+
credential: INodeParams
18+
badge: string
19+
outputs: INodeOutputsValue[]
20+
21+
constructor() {
22+
this.label = 'Jina AI Rerank Retriever'
23+
this.name = 'JinaRerankRetriever'
24+
this.version = 1.0
25+
this.type = 'JinaRerankRetriever'
26+
this.icon = 'JinaAI.svg'
27+
this.category = 'Retrievers'
28+
this.description = 'Jina AI Rerank indexes the documents from most to least semantically relevant to the query.'
29+
this.baseClasses = [this.type, 'BaseRetriever']
30+
this.credential = {
31+
label: 'Connect Credential',
32+
name: 'credential',
33+
type: 'credential',
34+
credentialNames: ['jinaAIApi']
35+
}
36+
this.inputs = [
37+
{
38+
label: 'Vector Store Retriever',
39+
name: 'baseRetriever',
40+
type: 'VectorStoreRetriever'
41+
},
42+
{
43+
label: 'Model Name',
44+
name: 'model',
45+
type: 'options',
46+
options: [
47+
{
48+
label: 'jina-reranker-v2-base-multilingual',
49+
name: 'jina-reranker-v2-base-multilingual'
50+
},
51+
{
52+
label: 'jina-colbert-v2',
53+
name: 'jina-colbert-v2'
54+
}
55+
],
56+
default: 'jina-reranker-v2-base-multilingual',
57+
optional: true
58+
},
59+
{
60+
label: 'Query',
61+
name: 'query',
62+
type: 'string',
63+
description: 'Query to retrieve documents from retriever. If not specified, user question will be used',
64+
optional: true,
65+
acceptVariable: true
66+
},
67+
{
68+
label: 'Top N',
69+
name: 'topN',
70+
description: 'Number of top results to fetch. Default to 4',
71+
placeholder: '4',
72+
default: 4,
73+
type: 'number',
74+
additionalParams: true,
75+
optional: true
76+
}
77+
]
78+
this.outputs = [
79+
{
80+
label: 'Jina AI Rerank Retriever',
81+
name: 'retriever',
82+
baseClasses: this.baseClasses
83+
},
84+
{
85+
label: 'Document',
86+
name: 'document',
87+
description: 'Array of document objects containing metadata and pageContent',
88+
baseClasses: ['Document', 'json']
89+
},
90+
{
91+
label: 'Text',
92+
name: 'text',
93+
description: 'Concatenated string from pageContent of documents',
94+
baseClasses: ['string', 'json']
95+
}
96+
]
97+
}
98+
99+
async init(nodeData: INodeData, input: string, options: ICommonObject): Promise<any> {
100+
const baseRetriever = nodeData.inputs?.baseRetriever as BaseRetriever
101+
const model = nodeData.inputs?.model as string
102+
const query = nodeData.inputs?.query as string
103+
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
104+
const jinaApiKey = getCredentialParam('jinaAIAPIKey', credentialData, nodeData)
105+
const topN = nodeData.inputs?.topN ? parseFloat(nodeData.inputs?.topN as string) : 4
106+
const output = nodeData.outputs?.output as string
107+
108+
const jinaCompressor = new JinaRerank(jinaApiKey, model, topN)
109+
110+
const retriever = new ContextualCompressionRetriever({
111+
baseCompressor: jinaCompressor,
112+
baseRetriever: baseRetriever
113+
})
114+
115+
if (output === 'retriever') return retriever
116+
else if (output === 'document') return await retriever.invoke(query ? query : input)
117+
else if (output === 'text') {
118+
const docs = await retriever.invoke(query ? query : input)
119+
let finaltext = ''
120+
for (const doc of docs) finaltext += `${doc.pageContent}\n`
121+
122+
return handleEscapeCharacters(finaltext, false)
123+
}
124+
125+
return retriever
126+
}
127+
}
128+
129+
module.exports = { nodeClass: JinaRerankRetriever_Retrievers }

0 commit comments

Comments
 (0)