-
Notifications
You must be signed in to change notification settings - Fork 2k
/
retrievers.ts
88 lines (76 loc) 路 2.16 KB
/
retrievers.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
import Exa, {
ContentsOptions,
RegularSearchOptions,
SearchResponse,
SearchResult,
} from "exa-js";
import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";
/**
* Interface for the fields required during the initialization of a
* `ExaRetriever` instance. It extends the `BaseRetrieverInput`
* interface and adds a `client` field of type `Exa`.
*/
export interface ExaRetrieverFields<T extends ContentsOptions = { text: true }>
extends BaseRetrieverInput {
client: Exa.default;
searchArgs?: RegularSearchOptions & T;
}
export function _getMetadata<T extends ContentsOptions = { text: true }>(
result: SearchResult<T>
): Record<string, unknown> {
const newMetadata: Record<string, unknown> = { ...result };
delete newMetadata.text;
return newMetadata;
}
/**
* @example
* ```typescript
* const retriever = new ExaRetriever({
* client: new Exa(
* process.env.EXA_API_KEY,
* process.env.EXA_BASE_URL,
* ),
* });
* const docs = await retriever.getRelevantDocuments("hello");
* ```
*/
export class ExaRetriever<
T extends ContentsOptions = { text: true }
> extends BaseRetriever {
static lc_name() {
return "ExaRetriever";
}
lc_namespace = ["langchain", "retrievers", "exa"];
private client: Exa.default;
searchArgs?: RegularSearchOptions & T;
constructor(fields: ExaRetrieverFields<T>) {
super(fields);
this.client = fields.client;
this.searchArgs = fields.searchArgs;
}
async _getRelevantDocuments(query: string): Promise<Document[]> {
const res: SearchResponse<T> = await this.client.searchAndContents<T>(
query,
this.searchArgs
);
const documents: Document[] = [];
for (const result of res.results) {
let pageContent;
if ("text" in result) {
pageContent = result.text;
} else if ("highlights" in result) {
pageContent = result.highlights.join("\n\n");
} else {
pageContent = "No results found.";
}
documents.push(
new Document({
pageContent,
metadata: _getMetadata<T>(result),
})
);
}
return documents;
}
}