-
Notifications
You must be signed in to change notification settings - Fork 2k
/
vectara_summary.ts
69 lines (60 loc) · 1.68 KB
/
vectara_summary.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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import { Document } from "@langchain/core/documents";
import {
BaseRetriever,
type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import { CallbackManagerForRetrieverRun } from "@langchain/core/callbacks/manager";
import {
VectaraStore,
type VectaraSummary,
type VectaraFilter,
DEFAULT_FILTER,
} from "../vectorstores/vectara.js";
export interface VectaraRetrieverInput extends BaseRetrieverInput {
vectara: VectaraStore;
filter?: VectaraFilter;
topK?: number;
summaryConfig?: VectaraSummary;
}
export class VectaraSummaryRetriever extends BaseRetriever {
static lc_name() {
return "VectaraSummaryRetriever";
}
lc_namespace = ["langchain", "retrievers"];
private filter = DEFAULT_FILTER;
private vectara: VectaraStore;
private topK: number;
private summaryConfig: VectaraSummary;
constructor(fields: VectaraRetrieverInput) {
super(fields);
this.vectara = fields.vectara;
this.topK = fields.topK ?? 10;
this.filter = fields.filter ?? DEFAULT_FILTER;
this.summaryConfig = fields.summaryConfig ?? {
enabled: false,
maxSummarizedResults: 0,
responseLang: "eng",
};
}
async _getRelevantDocuments(
query: string,
_callbacks?: CallbackManagerForRetrieverRun
): Promise<Document[]> {
const summaryResult = await this.vectara.vectaraQuery(
query,
this.topK,
this.filter,
this.summaryConfig ? this.summaryConfig : undefined
);
const docs = summaryResult.documents;
if (this.summaryConfig.enabled) {
docs.push(
new Document({
pageContent: summaryResult.summary,
metadata: { summary: true },
})
);
}
return docs;
}
}