-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
index.ts
123 lines (113 loc) Β· 3.79 KB
/
index.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
import {
BaseCallbackConfig,
CallbackManager,
CallbackManagerForRetrieverRun,
Callbacks,
parseCallbackConfigArg,
} from "../callbacks/manager.js";
import type { DocumentInterface } from "../documents/document.js";
import { Runnable, type RunnableInterface } from "../runnables/base.js";
import { RunnableConfig, ensureConfig } from "../runnables/config.js";
/**
* Base Retriever class. All indexes should extend this class.
*/
export interface BaseRetrieverInput {
callbacks?: Callbacks;
tags?: string[];
metadata?: Record<string, unknown>;
verbose?: boolean;
}
export interface BaseRetrieverInterface<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
Metadata extends Record<string, any> = Record<string, any>
> extends RunnableInterface<string, DocumentInterface<Metadata>[]> {
getRelevantDocuments(
query: string,
config?: Callbacks | BaseCallbackConfig
): Promise<DocumentInterface<Metadata>[]>;
}
/**
* Abstract base class for a Document retrieval system. A retrieval system
* is defined as something that can take string queries and return the
* most 'relevant' Documents from some source.
*/
export abstract class BaseRetriever<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
Metadata extends Record<string, any> = Record<string, any>
>
extends Runnable<string, DocumentInterface<Metadata>[]>
implements BaseRetrieverInterface
{
callbacks?: Callbacks;
tags?: string[];
metadata?: Record<string, unknown>;
verbose?: boolean;
constructor(fields?: BaseRetrieverInput) {
super(fields);
this.callbacks = fields?.callbacks;
this.tags = fields?.tags ?? [];
this.metadata = fields?.metadata ?? {};
this.verbose = fields?.verbose ?? false;
}
/**
* TODO: This should be an abstract method, but we'd like to avoid breaking
* changes to people currently using subclassed custom retrievers.
* Change it on next major release.
*/
_getRelevantDocuments(
_query: string,
_callbacks?: CallbackManagerForRetrieverRun
): Promise<DocumentInterface<Metadata>[]> {
throw new Error("Not implemented!");
}
async invoke(
input: string,
options?: RunnableConfig
): Promise<DocumentInterface<Metadata>[]> {
return this.getRelevantDocuments(input, ensureConfig(options));
}
/**
* @deprecated Use .invoke() instead. Will be removed in 0.3.0.
*
* Main method used to retrieve relevant documents. It takes a query
* string and an optional configuration object, and returns a promise that
* resolves to an array of `Document` objects. This method handles the
* retrieval process, including starting and ending callbacks, and error
* handling.
* @param query The query string to retrieve relevant documents for.
* @param config Optional configuration object for the retrieval process.
* @returns A promise that resolves to an array of `Document` objects.
*/
async getRelevantDocuments(
query: string,
config?: Callbacks | BaseCallbackConfig
): Promise<DocumentInterface<Metadata>[]> {
const parsedConfig = ensureConfig(parseCallbackConfigArg(config));
const callbackManager_ = await CallbackManager.configure(
parsedConfig.callbacks,
this.callbacks,
parsedConfig.tags,
this.tags,
parsedConfig.metadata,
this.metadata,
{ verbose: this.verbose }
);
const runManager = await callbackManager_?.handleRetrieverStart(
this.toJSON(),
query,
parsedConfig.runId,
undefined,
undefined,
undefined,
parsedConfig.runName
);
try {
const results = await this._getRelevantDocuments(query, runManager);
await runManager?.handleRetrieverEnd(results);
return results;
} catch (error) {
await runManager?.handleRetrieverError(error);
throw error;
}
}
}