extensions.ts72 lines · main
| 1 | import { codeBlock } from 'common-tags' |
| 2 | import OpenAI from 'openai' |
| 3 | import { expect } from 'vitest' |
| 4 | |
| 5 | expect.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 | }) |