-
Notifications
You must be signed in to change notification settings - Fork 39
/
rag.ts
69 lines (59 loc) 路 1.99 KB
/
rag.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 {
dedup,
Generate,
type GenerateOptions,
Signature
} from '../dsp/index.js';
import { Program, type ProgramForwardOptions } from '../dsp/program.js';
import type { AIService } from '../text/types.js';
import { ChainOfThought } from './cot.js';
export class RAG extends Program<{ question: string }, { answer: string }> {
private genQuery: Generate<
{ context: string[]; question: string },
{ query: string }
>;
private genAnswer: Generate<
{ context: string[]; question: string },
{ answer: string }
>;
private queryFn: (query: string) => Promise<string>;
private maxHops: number;
constructor(
ai: AIService,
queryFn: (query: string) => Promise<string>,
options: Readonly<GenerateOptions & { maxHops?: number }>
) {
super();
this.maxHops = options?.maxHops ?? 3;
const qsig = new Signature(
'"Write a simple search query that will help answer a complex question." context?:string[] "may contain relevant facts", question -> query "question to further our understanding"'
);
this.genQuery = new Generate(ai, qsig);
const asig = new Signature(
'"Answer questions with short factoid answers." context:string[] "may contain relevant facts", question -> answer'
);
this.genAnswer = new ChainOfThought(ai, asig);
this.queryFn = queryFn;
this.register(this.genQuery);
this.register(this.genAnswer);
}
override getSignature = (): Signature => this.genAnswer.getSignature();
override forward = async (
{ question }: Readonly<{ question: string }>,
options?: Readonly<ProgramForwardOptions>
): Promise<{ answer: string }> => {
let context: string[] = [];
for (let i = 0; i < this.maxHops; i++) {
const { query } = await this.genQuery.forward(
{
context,
question
},
options
);
const val = await this.queryFn(query);
context = dedup([...context, val]);
}
return this.genAnswer.forward({ context, question }, options);
};
}