Add starts-with assertion type (promptfoo#64)

sparticleinc · Jul 13, 2023 · ed2589e · ed2589e
1 parent e921803
commit ed2589e
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -96,6 +96,7 @@ See [Test assertions](https://promptfoo.dev/docs/configuration/expected-outputs)
 | `contains`      | output contains substring                                                 |
 | `icontains`     | output contains substring, case insensitive                               |
 | `regex`         | output matches regex                                                      |
+| `starts-with`   | output starts with string                                                 |
 | `contains-any ` | output contains any of the listed substrings                              |
 | `contains-all`  | output contains all list of substrings                                    |
 | `is-json`       | output is valid json                                                      |

diff --git a/src/assertions.ts b/src/assertions.ts
@@ -207,6 +207,22 @@ export async function runAssertion(
     };
   }
 
+  if (baseType === 'starts-with') {
+    invariant(assertion.value, '"starts-with" assertion type must have a string value');
+    invariant(
+      typeof assertion.value === 'string',
+      '"starts-with" assertion type must have a string value',
+    );
+    pass = output.startsWith(String(assertion.value)) !== inverse;
+    return {
+      pass,
+      score: pass ? 1 : 0,
+      reason: pass
+        ? 'Assertion passed'
+        : `Expected output to ${inverse ? 'not ' : ''}start with "${assertion.value}"`,
+    };
+  }
+
   if (baseType === 'contains-json') {
     pass = containsJSON(output) !== inverse;
     return {
@@ -480,7 +496,7 @@ export function assertionFromString(expected: string): Assertion {
 
   // New options
   const assertionRegex =
-    /^(not-)?(equals|contains-any|contains-all|contains-json|is-json|regex|icontains|contains|webhook|rouge-n|similar)(?::|\((\d+(\.\d+)?)\):)?(.*)$/;
+    /^(not-)?(equals|contains-any|contains-all|contains-json|is-json|regex|icontains|contains|webhook|rouge-n|similar|starts-with)(?::|\((\d+(\.\d+)?)\):)?(.*)$/;
   const regexMatch = expected.match(assertionRegex);
 
   if (regexMatch) {
@@ -497,7 +513,7 @@ export function assertionFromString(expected: string): Assertion {
       return {
         type: fullType as AssertionType,
       };
-    } else if (type === 'rouge-n' || type === 'similar') {
+    } else if (type === 'rouge-n' || type === 'similar' || type === 'starts-with') {
       return {
         type: fullType as AssertionType,
         value,

diff --git a/src/types.ts b/src/types.ts
@@ -136,6 +136,7 @@ type BaseAssertionTypes =
   | 'icontains'
   | 'contains-all'
   | 'contains-any'
+  | 'starts-with'
   | 'regex'
   | 'is-json'
   | 'contains-json'

diff --git a/test/assertions.test.ts b/test/assertions.test.ts
@@ -493,6 +493,36 @@ describe('runAssertion', () => {
     expect(result.pass).toBeFalsy();
     expect(result.reason).toBe('ROUGE-N score 0.2 is less than threshold 0.75');
   });
+
+  // Test for starts-with assertion
+  const startsWithAssertion: Assertion = {
+    type: 'starts-with',
+    value: 'Expected',
+  };
+
+  it('should pass when the starts-with assertion passes', async () => {
+    const output = 'Expected output';
+
+    const result: GradingResult = await runAssertion(
+      startsWithAssertion,
+      {} as AtomicTestCase,
+      output,
+    );
+    expect(result.pass).toBeTruthy();
+    expect(result.reason).toBe('Assertion passed');
+  });
+
+  it('should fail when the starts-with assertion fails', async () => {
+    const output = 'Different output';
+
+    const result: GradingResult = await runAssertion(
+      startsWithAssertion,
+      {} as AtomicTestCase,
+      output,
+    );
+    expect(result.pass).toBeFalsy();
+    expect(result.reason).toBe('Expected output to start with "Expected"');
+  });
 });
 
 describe('assertionFromString', () => {
@@ -632,6 +662,14 @@ describe('assertionFromString', () => {
     expect(result.value).toBe('foo');
     expect(result.threshold).toBeCloseTo(0.225);
   });
+
+  it('should create a starts-with assertion', () => {
+    const expected = 'starts-with:Expected';
+
+    const result: Assertion = assertionFromString(expected);
+    expect(result.type).toBe('starts-with');
+    expect(result.value).toBe('Expected');
+  });
 });
 
 describe('matchesSimilarity', () => {