-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
apify_dataset.ts
143 lines (125 loc) Β· 5.37 KB
/
apify_dataset.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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
/* eslint-disable @typescript-eslint/no-explicit-any */
import {
ActorCallOptions,
ApifyClient,
ApifyClientOptions,
TaskCallOptions,
} from "apify-client";
import { AsyncCaller, AsyncCallerParams } from "../../util/async_caller.js";
import { BaseDocumentLoader, DocumentLoader } from "../base.js";
import { Document } from "../../document.js";
import { getEnvironmentVariable } from "../../util/env.js";
/**
* A type that represents a function that takes a single object (an Apify
* dataset item) and converts it to an instance of the Document class.
*/
export type ApifyDatasetMappingFunction<Metadata extends Record<string, any>> =
(
item: Record<string | number, unknown>
) => Document<Metadata> | Promise<Document<Metadata>>;
export interface ApifyDatasetLoaderConfig<Metadata extends Record<string, any>>
extends AsyncCallerParams {
datasetMappingFunction: ApifyDatasetMappingFunction<Metadata>;
clientOptions?: ApifyClientOptions;
}
/**
* A class that extends the BaseDocumentLoader and implements the
* DocumentLoader interface. It represents a document loader that loads
* documents from an Apify dataset.
*/
export class ApifyDatasetLoader<Metadata extends Record<string, any>>
extends BaseDocumentLoader
implements DocumentLoader
{
protected apifyClient: ApifyClient;
protected datasetId: string;
protected datasetMappingFunction: ApifyDatasetMappingFunction<Metadata>;
protected caller: AsyncCaller;
constructor(datasetId: string, config: ApifyDatasetLoaderConfig<Metadata>) {
super();
const { clientOptions, datasetMappingFunction, ...asyncCallerParams } =
config;
const token = ApifyDatasetLoader._getApifyApiToken(clientOptions);
this.apifyClient = new ApifyClient({ ...clientOptions, token });
this.datasetId = datasetId;
this.datasetMappingFunction = datasetMappingFunction;
this.caller = new AsyncCaller(asyncCallerParams);
}
private static _getApifyApiToken(config?: { token?: string }) {
return config?.token ?? getEnvironmentVariable("APIFY_API_TOKEN");
}
/**
* Retrieves the dataset items from the Apify platform and applies the
* datasetMappingFunction to each item to create an array of Document
* instances.
* @returns An array of Document instances.
*/
async load(): Promise<Document<Metadata>[]> {
const datasetItems = (
await this.apifyClient.dataset(this.datasetId).listItems({ clean: true })
).items;
return await Promise.all(
datasetItems.map((item) =>
this.caller.call(async () => this.datasetMappingFunction(item))
)
);
}
/**
* Create an ApifyDatasetLoader by calling an Actor on the Apify platform and waiting for its results to be ready.
* @param actorId The ID or name of the Actor on the Apify platform.
* @param input The input object of the Actor that you're trying to run.
* @param options Options specifying settings for the Actor run.
* @param options.datasetMappingFunction A function that takes a single object (an Apify dataset item) and converts it to an instance of the Document class.
* @returns An instance of `ApifyDatasetLoader` with the results from the Actor run.
*/
static async fromActorCall<Metadata extends Record<string, any>>(
actorId: string,
input: Record<string | number, unknown>,
config: {
callOptions?: ActorCallOptions;
clientOptions?: ApifyClientOptions;
datasetMappingFunction: ApifyDatasetMappingFunction<Metadata>;
}
): Promise<ApifyDatasetLoader<Metadata>> {
const apifyApiToken = ApifyDatasetLoader._getApifyApiToken(
config.clientOptions
);
const apifyClient = new ApifyClient({ token: apifyApiToken });
const actorCall = await apifyClient
.actor(actorId)
.call(input, config.callOptions ?? {});
return new ApifyDatasetLoader(actorCall.defaultDatasetId, {
datasetMappingFunction: config.datasetMappingFunction,
clientOptions: { ...config.clientOptions, token: apifyApiToken },
});
}
/**
* Create an ApifyDatasetLoader by calling a saved Actor task on the Apify platform and waiting for its results to be ready.
* @param taskId The ID or name of the task on the Apify platform.
* @param input The input object of the task that you're trying to run. Overrides the task's saved input.
* @param options Options specifying settings for the task run.
* @param options.datasetMappingFunction A function that takes a single object (an Apify dataset item) and converts it to an instance of the Document class.
* @returns An instance of `ApifyDatasetLoader` with the results from the task's run.
*/
static async fromActorTaskCall<Metadata extends Record<string, any>>(
taskId: string,
input: Record<string | number, unknown>,
config: {
callOptions?: TaskCallOptions;
clientOptions?: ApifyClientOptions;
datasetMappingFunction: ApifyDatasetMappingFunction<Metadata>;
}
): Promise<ApifyDatasetLoader<Metadata>> {
const apifyApiToken = ApifyDatasetLoader._getApifyApiToken(
config.clientOptions
);
const apifyClient = new ApifyClient({ token: apifyApiToken });
const taskCall = await apifyClient
.task(taskId)
.call(input, config.callOptions ?? {});
return new ApifyDatasetLoader(taskCall.defaultDatasetId, {
datasetMappingFunction: config.datasetMappingFunction,
clientOptions: { ...config.clientOptions, token: apifyApiToken },
});
}
}