-
Notifications
You must be signed in to change notification settings - Fork 2k
/
base.ts
116 lines (103 loc) 路 3.19 KB
/
base.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
import {
BaseRetriever,
type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import {
AsyncCaller,
type AsyncCallerParams,
} from "@langchain/core/utils/async_caller";
import type { DocumentInterface } from "@langchain/core/documents";
/**
* Type for the authentication method used by the RemoteRetriever. It can
* either be false (no authentication) or an object with a bearer token.
*/
export type RemoteRetrieverAuth = false | { bearer: string };
/**
* Type for the JSON response values from the remote server.
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type RemoteRetrieverValues = Record<string, any>;
/**
* Interface for the parameters required to initialize a RemoteRetriever
* instance.
*/
export interface RemoteRetrieverParams
extends AsyncCallerParams,
BaseRetrieverInput {
/**
* The URL of the remote retriever server
*/
url: string;
/**
* The authentication method to use, currently implemented is
* - false: no authentication
* - { bearer: string }: Bearer token authentication
*/
auth: RemoteRetrieverAuth;
}
/**
* Abstract class for interacting with a remote server to retrieve
* relevant documents based on a given query.
*/
export abstract class RemoteRetriever
extends BaseRetriever
implements RemoteRetrieverParams
{
get lc_secrets(): { [key: string]: string } | undefined {
return {
"auth.bearer": "REMOTE_RETRIEVER_AUTH_BEARER",
};
}
url: string;
auth: RemoteRetrieverAuth;
headers: Record<string, string>;
asyncCaller: AsyncCaller;
constructor(fields: RemoteRetrieverParams) {
super(fields);
const { url, auth, ...rest } = fields;
this.url = url;
this.auth = auth;
this.headers = {
Accept: "application/json",
"Content-Type": "application/json",
...(this.auth && this.auth.bearer
? { Authorization: `Bearer ${this.auth.bearer}` }
: {}),
};
this.asyncCaller = new AsyncCaller(rest);
}
/**
* Abstract method that should be implemented by subclasses to create the
* JSON body of the request based on the given query.
* @param query The query based on which the JSON body of the request is created.
* @returns The JSON body of the request.
*/
abstract createJsonBody(query: string): RemoteRetrieverValues;
/**
* Abstract method that should be implemented by subclasses to process the
* JSON response from the server and convert it into an array of Document
* instances.
* @param json The JSON response from the server.
* @returns An array of Document instances.
*/
abstract processJsonResponse(
json: RemoteRetrieverValues
): DocumentInterface[];
async _getRelevantDocuments(query: string): Promise<DocumentInterface[]> {
const body = this.createJsonBody(query);
const response = await this.asyncCaller.call(() =>
fetch(this.url, {
method: "POST",
headers: this.headers,
body: JSON.stringify(body),
})
);
if (!response.ok) {
throw new Error(
`Failed to retrieve documents from ${this.url}: ${response.status} ${response.statusText}`
);
}
const json = await response.json();
return this.processJsonResponse(json);
}
}