feat: add azure openai

ctrf-io · Oct 6, 2024 · 14519b1 · 14519b1
1 parent fc98e68
commit 14519b1
Show file tree

Hide file tree

Showing 7 changed files with 139 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -62,6 +62,38 @@ You will be responsible for any charges incurred from using your selected OpenAI
 
 `--log`: Whether to log the AI responses to the console (default: true).
 
+## Azure OpenAI
+
+Run the following command:
+
+```bash
+npx ai-ctrf azure-openai <path-to-ctrf-report>
+```
+
+An AI summary for each failed test will be added to your test report.
+
+The package interacts with the Azure OpenAI API, you must set `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_ENDPOINT`, and `AZURE_OPENAI_DEPLOYMENT_NAME` environment variable or provide them as arguments.
+
+You will be responsible for any charges incurred from using your selected Azure OpenAI model. Make sure you are aware of the associated cost.
+
+### Options
+
+`--model`: OpenAI model to use (default: gpt-3.5-turbo).
+
+`--systemPrompt`: Custom system prompt to guide the AI response.
+
+`--frequencyPenalty`: OpenAI frequency penalty parameter (default: 0).
+
+`--maxTokens`: Maximum number of tokens for the response.
+
+`--presencePenalty`: OpenAI presence penalty parameter (default: 0).
+
+`--temperature`: Sampling temperature (conflicts with topP).
+
+`--topP`: Top-p sampling parameter (conflicts with temperature).
+
+`--log`: Whether to log the AI responses to the console (default: true).
+
 ## Claude
 
 Run the following command:

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "ai-ctrf",
-  "version": "0.0.3",
-  "description": "",
+  "version": "0.0.4",
+  "description": "AI Test Reporter - Create human-readable summaries of test results with LLMs like OpenAI GPT",
   "main": "index.js",
   "scripts": {
     "test": "echo \"Error: no test specified\" && exit 1",

diff --git a/src/index.ts b/src/index.ts
@@ -2,6 +2,7 @@
 import yargs from 'yargs/yargs';
 import { hideBin } from 'yargs/helpers';
 import { openAISummary } from './models/openai';
+import { azureOpenAISummary } from './models/azure-openai';
 import { validateCtrfFile } from './common';
 import { claudeSummary } from './models/claude';
 
@@ -16,6 +17,7 @@ export interface Arguments {
     temperature?: number;
     topP?: number;
     log?: boolean;
+    deploymentId?: string;
 }
 
 const argv: Arguments = yargs(hideBin(process.argv))
@@ -31,7 +33,7 @@ const argv: Arguments = yargs(hideBin(process.argv))
                 describe: 'OpenAI model to use',
                 type: 'string',
                 default: 'gpt-4o', 
-            });;
+            });
         }
     )
     .command(
@@ -45,10 +47,29 @@ const argv: Arguments = yargs(hideBin(process.argv))
             .option('model', {
                 describe: 'Claude model to use',
                 type: 'string',
-                default: 'claude-3-5-sonnet-20240620',  
+                default: 'claude-3-5-sonnet-20240620', 
             });
         }
     )
+    .command(
+        'azure-openai <file>',
+        'Generate test summary from a CTRF report using Azure OpenAI',
+        (yargs) => {
+            return yargs.positional('file', {
+                describe: 'Path to the CTRF file',
+                type: 'string',
+            })
+                .option('deploymentId', {
+                    describe: 'Deployment ID for Azure OpenAI',
+                    type: 'string',
+                })
+                .option('model', {
+                    describe: 'Model to use',
+                    type: 'string',
+                    default: 'gpt-4o',
+                });
+        }
+    )
     .option('systemPrompt', {
         describe: 'System prompt to guide the AI',
         type: 'string',
@@ -116,4 +137,13 @@ if (argv._.includes('openai') && argv.file) {
     } catch (error) {
         console.error('Failed to read file:', error);
     }
+} else if (argv._.includes('azure-openai') && argv.file) {
+    try {
+        const report = validateCtrfFile(argv.file);
+        if (report !== null) {
+            azureOpenAISummary(report, file, argv);
+        }
+    } catch (error) {
+        console.error('Failed to read file:', error);
+    }
 }
diff --git a/src/models/azure-openai.ts b/src/models/azure-openai.ts
@@ -0,0 +1,61 @@
+import { AzureOpenAI } from "openai";
+import { CtrfReport } from "../../types/ctrf";
+import { Arguments } from "../index";
+import { saveUpdatedReport, stripAnsi } from "../common";
+
+export async function azureOpenAISummary(report: CtrfReport, file: string, args: Arguments) {
+    const apiKey = process.env.AZURE_OPENAI_API_KEY;
+    const endpoint = process.env.AZURE_OPENAI_ENDPOINT;
+    const deployment = args.deploymentId || process.env.AZURE_OPENAI_DEPLOYMENT_NAME;
+
+    if (!apiKey || !endpoint || !deployment) {
+        console.error('Missing Azure OpenAI configuration. Please set AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_DEPLOYMENT_NAME environment variables or provide them as arguments.');
+        return;
+    }
+
+    const client = new AzureOpenAI({ endpoint, apiKey, deployment });
+
+    const failedTests = report.results.tests.filter(test => test.status === 'failed');
+
+    for (const test of failedTests) {
+
+        const systemPrompt = args.systemPrompt || "";
+
+        const prompt = `Report:\n${JSON.stringify(test, null, 2)}.\n\nTool: ${report.results.tool.name}.\n\nPlease provide a human-readable failure summary that explains why you think the test might have failed and ways to fix it.`;
+
+        try {
+            const response = await client.chat.completions.create({
+                model: "gpt-4o",
+                messages: [
+                    { role: "system", content: systemPrompt },
+                    {
+                        role: "user",
+                        content: stripAnsi(prompt),
+                    },
+                ],
+                max_tokens: args.maxTokens || null,
+                frequency_penalty: args.frequencyPenalty,
+                presence_penalty: args.presencePenalty,
+                ...(args.temperature !== undefined ? { temperature: args.temperature } : {}),
+                ...(args.topP !== undefined ? { top_p: args.topP } : {}),
+            });
+
+            const aiResponse = response.choices[0]?.message?.content;
+
+            if (aiResponse) {
+                test.ai = aiResponse;
+                if (args.log) {
+                    console.log(`\n─────────────────────────────────────────────────────────────────────────────────────────────────────────────`);
+                    console.log(`✨ AI Test Reporter Summary`);
+                    console.log(`─────────────────────────────────────────────────────────────────────────────────────────────────────────────\n`);
+                    console.log(`❌ Failed Test: ${test.name}\n`)
+                    console.log(`${aiResponse}\n`);                }
+            }
+        } catch (error) {
+            console.error(`Error generating summary for test ${test.name}:`, error);
+            test.ai = 'Failed to generate summary due to an error.';
+        }
+    }
+
+    saveUpdatedReport(file, report);
+}
diff --git a/src/models/claude.ts b/src/models/claude.ts
@@ -36,8 +36,11 @@ export async function claudeSummary(report: CtrfReport, file: string, args: Argu
             if (aiResponse) {
                 test.ai = aiResponse;
                 if (args.log) {
-                    console.log(`AI summary for test: ${test.name}\n`, aiResponse);
-                }
+                    console.log(`\n─────────────────────────────────────────────────────────────────────────────────────────────────────────────`);
+                    console.log(`✨ AI Test Reporter Summary`);
+                    console.log(`─────────────────────────────────────────────────────────────────────────────────────────────────────────────\n`);
+                    console.log(`❌ Failed Test: ${test.name}\n`)
+                    console.log(`${aiResponse}\n`);                }
             }
         } catch (error) {
             console.error(`Error generating summary for test ${test.name}:`, error);

diff --git a/src/models/openai.ts b/src/models/openai.ts
@@ -35,7 +35,11 @@ export async function openAISummary(report: CtrfReport, file: string, args: Argu
             if (aiResponse) {
                 test.ai = aiResponse;
                 if (args.log) {
-                    console.log(`AI summary for test: ${test.name}\n`, aiResponse);
+                        console.log(`\n─────────────────────────────────────────────────────────────────────────────────────────────────────────────`);
+                        console.log(`✨ AI Test Reporter Summary`);
+                        console.log(`─────────────────────────────────────────────────────────────────────────────────────────────────────────────\n`);
+                        console.log(`❌ Failed Test: ${test.name}\n`)
+                        console.log(`${aiResponse}\n`);
                 }
             }
         } catch (error) {