-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
retrieval_qa.ts
155 lines (140 loc) Β· 4.29 KB
/
retrieval_qa.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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import type { BaseLanguageModelInterface } from "@langchain/core/language_models/base";
import type { BaseRetrieverInterface } from "@langchain/core/retrievers";
import { ChainValues } from "@langchain/core/utils/types";
import { CallbackManagerForChainRun } from "@langchain/core/callbacks/manager";
import { BaseChain, ChainInputs } from "./base.js";
import { SerializedVectorDBQAChain } from "./serde.js";
import {
StuffQAChainParams,
loadQAStuffChain,
} from "./question_answering/load.js";
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type LoadValues = Record<string, any>;
/**
* Interface for the input parameters of the RetrievalQAChain class.
*/
export interface RetrievalQAChainInput extends Omit<ChainInputs, "memory"> {
retriever: BaseRetrieverInterface;
combineDocumentsChain: BaseChain;
inputKey?: string;
returnSourceDocuments?: boolean;
}
/**
* Class representing a chain for performing question-answering tasks with
* a retrieval component.
* @example
* ```typescript
* // Initialize the OpenAI model and the remote retriever with the specified configuration
* const model = new ChatOpenAI({});
* const retriever = new RemoteLangChainRetriever({
* url: "http://example.com/api",
* auth: { bearer: "foo" },
* inputKey: "message",
* responseKey: "response",
* });
*
* // Create a RetrievalQAChain using the model and retriever
* const chain = RetrievalQAChain.fromLLM(model, retriever);
*
* // Execute the chain with a query and log the result
* const res = await chain.call({
* query: "What did the president say about Justice Breyer?",
* });
* console.log({ res });
*
* ```
*/
export class RetrievalQAChain
extends BaseChain
implements RetrievalQAChainInput
{
static lc_name() {
return "RetrievalQAChain";
}
inputKey = "query";
get inputKeys() {
return [this.inputKey];
}
get outputKeys() {
return this.combineDocumentsChain.outputKeys.concat(
this.returnSourceDocuments ? ["sourceDocuments"] : []
);
}
retriever: BaseRetrieverInterface;
combineDocumentsChain: BaseChain;
returnSourceDocuments = false;
constructor(fields: RetrievalQAChainInput) {
super(fields);
this.retriever = fields.retriever;
this.combineDocumentsChain = fields.combineDocumentsChain;
this.inputKey = fields.inputKey ?? this.inputKey;
this.returnSourceDocuments =
fields.returnSourceDocuments ?? this.returnSourceDocuments;
}
/** @ignore */
async _call(
values: ChainValues,
runManager?: CallbackManagerForChainRun
): Promise<ChainValues> {
if (!(this.inputKey in values)) {
throw new Error(`Question key "${this.inputKey}" not found.`);
}
const question: string = values[this.inputKey];
const docs = await this.retriever.getRelevantDocuments(
question,
runManager?.getChild("retriever")
);
const inputs = { question, input_documents: docs, ...values };
const result = await this.combineDocumentsChain.call(
inputs,
runManager?.getChild("combine_documents")
);
if (this.returnSourceDocuments) {
return {
...result,
sourceDocuments: docs,
};
}
return result;
}
_chainType() {
return "retrieval_qa" as const;
}
static async deserialize(
_data: SerializedVectorDBQAChain,
_values: LoadValues
): Promise<RetrievalQAChain> {
throw new Error("Not implemented");
}
serialize(): SerializedVectorDBQAChain {
throw new Error("Not implemented");
}
/**
* Creates a new instance of RetrievalQAChain using a BaseLanguageModel
* and a BaseRetriever.
* @param llm The BaseLanguageModel used to generate a new question.
* @param retriever The BaseRetriever used to retrieve relevant documents.
* @param options Optional parameters for the RetrievalQAChain.
* @returns A new instance of RetrievalQAChain.
*/
static fromLLM(
llm: BaseLanguageModelInterface,
retriever: BaseRetrieverInterface,
options?: Partial<
Omit<
RetrievalQAChainInput,
"retriever" | "combineDocumentsChain" | "index"
>
> &
StuffQAChainParams
): RetrievalQAChain {
const qaChain = loadQAStuffChain(llm, {
prompt: options?.prompt,
});
return new this({
...options,
retriever,
combineDocumentsChain: qaChain,
});
}
}