-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
contextual_compression.ts
54 lines (46 loc) · 1.71 KB
/
contextual_compression.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import { BaseDocumentCompressor } from "./document_compressors/index.js";
import { Document } from "../document.js";
import { BaseRetriever, BaseRetrieverInput } from "../schema/retriever.js";
import { CallbackManagerForRetrieverRun } from "../callbacks/manager.js";
/**
* Interface for the arguments required to construct a
* ContextualCompressionRetriever. It extends the BaseRetrieverInput
* interface with two additional fields: baseCompressor and baseRetriever.
*/
export interface ContextualCompressionRetrieverArgs extends BaseRetrieverInput {
baseCompressor: BaseDocumentCompressor;
baseRetriever: BaseRetriever;
}
/**
* A retriever that wraps a base retriever and compresses the results. It
* retrieves relevant documents based on a given query and then compresses
* these documents using a specified document compressor.
*/
export class ContextualCompressionRetriever extends BaseRetriever {
static lc_name() {
return "ContextualCompressionRetriever";
}
lc_namespace = ["langchain", "retrievers", "contextual_compression"];
baseCompressor: BaseDocumentCompressor;
baseRetriever: BaseRetriever;
constructor(fields: ContextualCompressionRetrieverArgs) {
super(fields);
this.baseCompressor = fields.baseCompressor;
this.baseRetriever = fields.baseRetriever;
}
async _getRelevantDocuments(
query: string,
runManager?: CallbackManagerForRetrieverRun
): Promise<Document[]> {
const docs = await this.baseRetriever.getRelevantDocuments(
query,
runManager?.getChild("base_retriever")
);
const compressedDocs = await this.baseCompressor.compressDocuments(
docs,
query,
runManager?.getChild("base_compressor")
);
return compressedDocs;
}
}