-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
chaindesk.ts
103 lines (89 loc) · 2.17 KB
/
chaindesk.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
import {
BaseRetriever,
type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";
import {
AsyncCaller,
type AsyncCallerParams,
} from "@langchain/core/utils/async_caller";
export interface ChaindeskRetrieverArgs
extends AsyncCallerParams,
BaseRetrieverInput {
datastoreId: string;
topK?: number;
filter?: Record<string, unknown>;
apiKey?: string;
}
interface Berry {
text: string;
score: number;
source?: string;
[key: string]: unknown;
}
/**
* @example
* ```typescript
* const retriever = new ChaindeskRetriever({
* datastoreId: "DATASTORE_ID",
* apiKey: "CHAINDESK_API_KEY",
* topK: 8,
* });
* const docs = await retriever.getRelevantDocuments("hello");
* ```
*/
export class ChaindeskRetriever extends BaseRetriever {
static lc_name() {
return "ChaindeskRetriever";
}
lc_namespace = ["langchain", "retrievers", "chaindesk"];
caller: AsyncCaller;
datastoreId: string;
topK?: number;
filter?: Record<string, unknown>;
apiKey?: string;
constructor({
datastoreId,
apiKey,
topK,
filter,
...rest
}: ChaindeskRetrieverArgs) {
super();
this.caller = new AsyncCaller(rest);
this.datastoreId = datastoreId;
this.apiKey = apiKey;
this.topK = topK;
this.filter = filter;
}
async getRelevantDocuments(query: string): Promise<Document[]> {
const r = await this.caller.call(
fetch,
`https://app.chaindesk.ai/api/datastores/${this.datastoreId}/query`,
{
method: "POST",
body: JSON.stringify({
query,
...(this.topK ? { topK: this.topK } : {}),
...(this.filter ? { filters: this.filter } : {}),
}),
headers: {
"Content-Type": "application/json",
...(this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {}),
},
}
);
const { results } = (await r.json()) as { results: Berry[] };
return results.map(
({ text, score, source, ...rest }) =>
new Document({
pageContent: text,
metadata: {
score,
source,
...rest,
},
})
);
}
}