Skip to content

Commit

Permalink
autoformat
Browse files Browse the repository at this point in the history
  • Loading branch information
franktip committed Jul 1, 2024
1 parent 6004c8d commit c19dc12
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 12 deletions.
10 changes: 8 additions & 2 deletions benchmark/run.ts
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ if (require.main === module) {
demandOption: false,
description:
"use custom rate-limiting for benchmarking (if specified, this supercedes the rateLimit option)",
}
},
});
const argv = await parser.argv;

Expand All @@ -191,7 +191,13 @@ if (require.main === module) {
"Warning: --strictResponses has no effect when not using --responses"
);
}
model = new ChatModel(argv.model, argv.nrAttempts, argv.rateLimit, argv.benchmark, { max_tokens: argv.maxTokens } );
model = new ChatModel(
argv.model,
argv.nrAttempts,
argv.rateLimit,
argv.benchmark,
{ max_tokens: argv.maxTokens }
);
} else {
model = MockCompletionModel.fromFile(
argv.responses,
Expand Down
34 changes: 25 additions & 9 deletions src/chatmodel.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
import axios from "axios";
import { performance } from "perf_hooks";
import { ICompletionModel } from "./completionModel";
import { retry, RateLimiter, BenchmarkRateLimiter, FixedRateLimiter } from "./promise-utils";
import {
retry,
RateLimiter,
BenchmarkRateLimiter,
FixedRateLimiter,
} from "./promise-utils";

const defaultPostOptions = {
max_tokens: 1000, // maximum number of tokens to return
Expand Down Expand Up @@ -32,21 +37,26 @@ export class ChatModel implements ICompletionModel {
private readonly nrAttempts: number,
private readonly rateLimit: number,
private readonly benchmark: boolean,
private readonly instanceOptions: PostOptions = {},
private readonly instanceOptions: PostOptions = {}
) {
this.apiEndpoint = getEnv("TESTPILOT_LLM_API_ENDPOINT");
this.authHeaders = getEnv("TESTPILOT_LLM_AUTH_HEADERS");
if (this.benchmark) {
this.rateLimiter = new BenchmarkRateLimiter();
console.log(`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and benchmark rate limit.`);
console.log(
`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and benchmark rate limit.`
);
} else if (this.rateLimit > 0) {
this.rateLimiter = new FixedRateLimiter(this.rateLimit);
console.log(`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and fixed rate of ${this.rateLimit} ms.`);
console.log(
`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and fixed rate of ${this.rateLimit} ms.`
);
} else {
this.rateLimiter = undefined;
console.log(`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and no rate limit.`);
console.log(
`Using ${this.model} at ${this.apiEndpoint} with ${this.nrAttempts} attempts and no rate limit.`
);
}

}

/**
Expand Down Expand Up @@ -92,12 +102,18 @@ export class ChatModel implements ICompletionModel {

let res;
if (this.rateLimiter) {
res = await retry( () =>
this.rateLimiter!.next(() => axios.post(this.apiEndpoint, postOptions, { headers })),
res = await retry(
() =>
this.rateLimiter!.next(() =>
axios.post(this.apiEndpoint, postOptions, { headers })
),
this.nrAttempts
);
} else {
res = await retry( () => axios.post(this.apiEndpoint, postOptions, { headers }), this.nrAttempts);
res = await retry(
() => axios.post(this.apiEndpoint, postOptions, { headers }),
this.nrAttempts
);
}

performance.measure(
Expand Down
2 changes: 1 addition & 1 deletion src/promise-utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -112,4 +112,4 @@ export class BenchmarkRateLimiter extends RateLimiter {
}
return super.next(p);
}
}
}

0 comments on commit c19dc12

Please sign in to comment.