extensions.ts72 lines · main
1import { codeBlock } from 'common-tags'
2import OpenAI from 'openai'
3import { expect } from 'vitest'
4
5expect.extend({
6 async toMatchCriteria(received: string, criteria: string) {
7 const openAiKey = process.env.OPENAI_API_KEY
8 const openai = new OpenAI({ apiKey: openAiKey })
9
10 const model = 'gpt-4o-2024-05-13'
11
12 const completionResponse = await openai.chat.completions.create({
13 model,
14 messages: [
15 {
16 role: 'system',
17 content: codeBlock`
18 You are a test runner. Your job is to evaluate whether 'Received' adheres to the test 'Criteria'.
19
20 You must output JSON, specifically an object containing a "pass" boolean and "reason" string:
21 - \`{ "pass": true, "reason": "<reason>" }\` if 'Received' adheres to the test 'Criteria'
22 - \`{ "pass": false, "reason": "<reason>" }\` if 'Received' does not adhere to the test 'Criteria'
23
24 The "reason" must explain exactly which part of 'Received' did or did not pass the test 'Criteria'.
25 `,
26 },
27 {
28 role: 'user',
29 content: codeBlock`
30 Received:
31 ${received}
32
33 Criteria:
34 ${criteria}
35 `,
36 },
37 ],
38 max_tokens: 256,
39 temperature: 0,
40 response_format: {
41 type: 'json_object',
42 },
43 stream: false,
44 })
45
46 const [choice] = completionResponse.choices
47
48 if (!choice.message.content) {
49 throw new Error('LLM evaluator returned invalid response')
50 }
51
52 const { pass, reason }: { pass?: boolean; reason?: string } = JSON.parse(choice.message.content)
53
54 if (pass === undefined) {
55 throw new Error('LLM evaluator returned invalid response')
56 }
57
58 return {
59 message: () =>
60 codeBlock`
61 ${this.utils.matcherHint('toMatchCriteria', received, criteria, {
62 comment: `evaluated by LLM '${model}'`,
63 isNot: this.isNot,
64 promise: this.promise,
65 })}
66
67 ${reason}
68 `,
69 pass,
70 }
71 },
72})