Skip to content

Commit c23cc44

Browse files
nshiabnshiab
committed
Concurrent requests
1 parent 476c1a4 commit c23cc44

File tree

4 files changed

+387
-173
lines changed

4 files changed

+387
-173
lines changed

‎src/class/SimpleTable.ts

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -520,18 +520,20 @@ export default class SimpleTable extends Simple {
520520
*
521521
* This method currently supports Google Gemini and Vertex AI. It retrieves credentials and the model from environment variables (`AI_KEY`, `AI_PROJECT`, `AI_LOCATION`, `AI_MODEL`) or accepts them as options. Options take precedence over environment variables.
522522
*
523-
* This method can be slow for large tables. To avoid exceeding rate limits, you can process multiple rows at once with the `batchSize` option. You can also use the `rateLimitPerMinute` option to automatically add a delay between requests to comply with the rate limit.
523+
* To avoid exceeding rate limits, you can process multiple rows at once with the `batchSize` option. You can also use the `rateLimitPerMinute` option to automatically add a delay between requests to comply with the rate limit.
524+
*
525+
* On the other hand, if you have a business or professional account with high rate limits, you can set the `concurrent` option to process multiple requests concurrently and speed up the process.
524526
*
525527
* The `cache` option allows you to cache locally the results of each request, saving resources and time. The data is cached in the local hidden folder `.journalism` (because this method uses the `askAI` function from the [journalism library](https://github.com/nshiab/journalism)). So don't forget to add `.journalism` to your `.gitignore` file!
526528
*
527-
* Sometimes, the AI returns less items than the batch size, which throws an error. If you want to automatically retry the request, you can use the `retry` option. The method will retry the request up to the specified number of times.
529+
* Sometimes, the AI returns fewer items than the batch size, which throws an error. If you want to automatically retry the request, you can use the `retry` option. The method will retry the request up to the specified number of times.
528530
*
529531
* The temperature is set to 0 to ensure reproducible results. However, consistent results cannot be guaranteed.
530532
*
531533
* This method won't work if you have geometries in your table.
532534
*
533535
* @example
534-
* Basic usage with cache, batchSize and retry options
536+
* Basic usage with cache, batchSize and rate limit
535537
* ```ts
536538
* // New table with column "city".
537539
* await table.loadArray([
@@ -547,7 +549,7 @@ export default class SimpleTable extends Simple {
547549
* "country",
548550
* `Give me the country of the city.`,
549551
* // Don't forget to add .journalism to your .gitignore file!
550-
* { cache: true, batchSize: 10, retry: 3, verbose: true },
552+
* { cache: true, batchSize: 10, rateLimitPerMinute: 15, verbose: true },
551553
* );
552554
*
553555
* // Result:
@@ -563,6 +565,7 @@ export default class SimpleTable extends Simple {
563565
* @param prompt - The input string to guide the AI's response.
564566
* @param options - Configuration options for the AI request.
565567
* @param options.batchSize - The number of rows to process in each batch. By default, it is 1.
568+
* @param options.concurrent - The number of concurrent requests to send. By default, it is 1.
566569
* @param options.cache - If true, the results will be cached locally. By default, it is false.
567570
* @param options.retry - The number of times to retry the request in case of failure. By default, it is 0.
568571
* @param options.rateLimitPerMinute - The rate limit for the AI requests in requests per minute. If necessary, the method will wait between requests. By default, there is no limit.
@@ -579,6 +582,7 @@ export default class SimpleTable extends Simple {
579582
prompt: string,
580583
options: {
581584
batchSize?: number;
585+
concurrent?: number;
582586
cache?: boolean;
583587
retry?: number;
584588
model?: string;

‎src/helpers/tryAI.ts

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import { formatNumber } from "@nshiab/journalism";
2+
import { askAI } from "@nshiab/journalism";
3+
4+
export default async function tryAI(
5+
i: number,
6+
batchSize: number,
7+
rows: {
8+
[key: string]: string | number | boolean | Date | null;
9+
}[],
10+
column: string,
11+
newColumn: string,
12+
prompt: string,
13+
options: {
14+
batchSize?: number;
15+
concurrent?: number;
16+
cache?: boolean;
17+
retry?: number;
18+
model?: string;
19+
apiKey?: string;
20+
vertex?: boolean;
21+
project?: string;
22+
location?: string;
23+
verbose?: boolean;
24+
rateLimitPerMinute?: number;
25+
} = {},
26+
) {
27+
options.verbose &&
28+
console.log(
29+
`\n${Math.min(i + batchSize, rows.length)}/${rows.length} | ${
30+
formatNumber(
31+
(Math.min(i + batchSize, rows.length)) / rows.length * 100,
32+
{
33+
significantDigits: 3,
34+
suffix: "%",
35+
},
36+
)
37+
}`,
38+
);
39+
const batch = rows.slice(i, i + batchSize);
40+
const fullPrompt = `${prompt}\nHere are the ${column} values as a list: ${
41+
JSON.stringify(batch.map((d) => d[column]))
42+
}\nReturn the results in a list as well. It's critical you return the same number of items, which is ${batch.length}, exactly in the same order.`;
43+
44+
if (options.verbose) {
45+
console.log("\nPrompt:");
46+
console.log(fullPrompt);
47+
}
48+
49+
const retry = options.retry ?? 1;
50+
51+
let testPassed = false;
52+
let iterations = 1;
53+
let newValues: (string | number | boolean | Date | null)[] = [];
54+
while (!testPassed && iterations <= retry) {
55+
try {
56+
// Types could be improved
57+
newValues = await askAI(
58+
fullPrompt,
59+
{
60+
...options,
61+
returnJson: true,
62+
test: (response: unknown) => {
63+
if (!Array.isArray(response)) {
64+
throw new Error(
65+
`The AI returned a non-array value: ${
66+
JSON.stringify(response)
67+
}`,
68+
);
69+
}
70+
if (response.length !== batch.length) {
71+
throw new Error(
72+
`The AI returned ${response.length} values, but the batch size is ${batch.length}.`,
73+
);
74+
}
75+
},
76+
},
77+
) as (string | number | boolean | Date | null)[];
78+
79+
testPassed = true;
80+
} catch (e: unknown) {
81+
if (iterations < retry) {
82+
console.log(
83+
`Error: the AI didn't return the expected number of items.\nRetrying... (${iterations}/${retry})`,
84+
);
85+
iterations++;
86+
} else {
87+
console.log(
88+
`Error: the AI didn't return the expected number of items.\nNo more retries left. (${iterations}/${retry}).`,
89+
);
90+
throw e;
91+
}
92+
}
93+
}
94+
95+
if (options.verbose) {
96+
console.log("\nResponse:", newValues);
97+
}
98+
99+
for (let j = 0; j < newValues.length; j++) {
100+
rows[i + j][newColumn] = newValues[j];
101+
}
102+
}

‎src/methods/aiRowByRow.ts

Lines changed: 70 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
import { askAI, formatNumber, prettyDuration, sleep } from "@nshiab/journalism";
1+
import { prettyDuration, sleep } from "@nshiab/journalism";
22
import type { SimpleTable } from "../index.ts";
3+
import tryAI from "../helpers/tryAI.ts";
34

45
export default async function aiRowByRow(
56
simpleTable: SimpleTable,
@@ -8,6 +9,7 @@ export default async function aiRowByRow(
89
prompt: string,
910
options: {
1011
batchSize?: number;
12+
concurrent?: number;
1113
cache?: boolean;
1214
retry?: number;
1315
model?: string;
@@ -25,99 +27,82 @@ export default async function aiRowByRow(
2527
}
2628

2729
const batchSize = options.batchSize ?? 1;
30+
const concurrent = options.concurrent ?? 1;
2831

32+
let requests = [];
2933
for (let i = 0; i < rows.length; i += batchSize) {
30-
options.verbose &&
31-
console.log(
32-
`\n${Math.min(i + batchSize, rows.length)}/${rows.length} | ${
33-
formatNumber(
34-
(Math.min(i + batchSize, rows.length)) / rows.length * 100,
35-
{
36-
significantDigits: 3,
37-
suffix: "%",
38-
},
39-
)
40-
}`,
34+
if (concurrent === 1) {
35+
const start = new Date();
36+
await tryAI(
37+
i,
38+
batchSize,
39+
rows,
40+
column,
41+
newColumn,
42+
prompt,
43+
options,
4144
);
42-
const batch = rows.slice(i, i + batchSize);
43-
const fullPrompt = `${prompt}\nHere are the ${column} values as a list: ${
44-
JSON.stringify(batch.map((d) => d[column]))
45-
}\nReturn the results in a list as well. It's critical you return the same number of items, which is ${batch.length}, exactly in the same order.`;
45+
const end = new Date();
4646

47-
if (options.verbose) {
48-
console.log("\nPrompt:");
49-
console.log(fullPrompt);
50-
}
51-
52-
const start = new Date();
53-
54-
const retry = options.retry ?? 1;
55-
let testPassed = false;
56-
let iterations = 1;
57-
let newValues: (string | number | boolean | Date | null)[] = [];
58-
while (!testPassed && iterations <= retry) {
59-
try {
60-
// Types could be improved
61-
newValues = await askAI(
62-
fullPrompt,
63-
{
64-
...options,
65-
returnJson: true,
66-
test: (response: unknown) => {
67-
if (!Array.isArray(response)) {
68-
throw new Error(
69-
`The AI returned a non-array value: ${
70-
JSON.stringify(response)
71-
}`,
72-
);
73-
}
74-
if (response.length !== batch.length) {
75-
throw new Error(
76-
`The AI returned ${response.length} values, but the batch size is ${batchSize}.`,
77-
);
78-
}
79-
},
80-
},
81-
) as (string | number | boolean | Date | null)[];
82-
83-
testPassed = true;
84-
} catch (e: unknown) {
85-
if (iterations < retry) {
86-
console.log(
87-
`Error: the AI didn't return the expected number of items.\nRetrying... (${iterations}/${retry})`,
88-
);
89-
iterations++;
90-
} else {
91-
console.log(
92-
`Error: the AI didn't return the expected number of items.\nNo more retries left. (${iterations}/${retry}).`,
93-
);
94-
throw e;
47+
const duration = end.getTime() - start.getTime();
48+
// If duration is less than 50ms, it means data comes from cache and we don't need to wait
49+
if (
50+
typeof options.rateLimitPerMinute === "number" && duration > 50
51+
) {
52+
const delay = Math.round((60 / options.rateLimitPerMinute) * 1000) -
53+
duration;
54+
if (delay > 0) {
55+
if (options.verbose) {
56+
console.log(
57+
`Waiting ${
58+
prettyDuration(0, { end: delay })
59+
} to respect rate limit...`,
60+
);
61+
}
62+
await sleep(delay);
9563
}
9664
}
97-
}
98-
99-
const end = new Date();
100-
101-
if (options.verbose) {
102-
console.log("\nResponse:", newValues);
103-
}
65+
} else if (concurrent) {
66+
if (requests.length < concurrent) {
67+
requests.push(
68+
tryAI(
69+
i,
70+
batchSize,
71+
rows,
72+
column,
73+
newColumn,
74+
prompt,
75+
options,
76+
),
77+
);
78+
}
79+
if (requests.length === concurrent || i + batchSize >= rows.length) {
80+
const start = new Date();
81+
await Promise.all(requests);
82+
const end = new Date();
10483

105-
for (let j = 0; j < newValues.length; j++) {
106-
rows[i + j][newColumn] = newValues[j];
107-
}
84+
requests = [];
10885

109-
if (typeof options.rateLimitPerMinute === "number") {
110-
const delay = Math.round((60 / options.rateLimitPerMinute) * 1000) -
111-
(end.getTime() - start.getTime());
112-
if (delay > 0) {
113-
if (options.verbose) {
114-
console.log(
115-
`Waiting ${
116-
prettyDuration(0, { end: delay })
117-
} to respect rate limit...`,
118-
);
86+
const duration = end.getTime() - start.getTime();
87+
// If duration is less than 50ms, it means data comes from cache and we don't need to wait
88+
if (
89+
typeof options.rateLimitPerMinute === "number" && duration > 50
90+
) {
91+
const delay = Math.round(
92+
(60 / (options.rateLimitPerMinute / concurrent)) * 1000,
93+
) -
94+
(end.getTime() - start.getTime());
95+
if (delay > 0) {
96+
if (options.verbose) {
97+
console.log(
98+
`Waiting ${
99+
prettyDuration(0, { end: delay })
100+
} to respect rate limit...`,
101+
);
102+
}
103+
await sleep(delay);
104+
}
119105
}
120-
await sleep(delay);
121106
}
122107
}
123108
}

0 commit comments

Comments
 (0)