How Zapier uses assertions to evaluate tool usage in chatbots

import { wrapOpenAI } from "braintrust";
import pick from "lodash/pick";
import { ChatCompletionTool } from "openai/resources/chat/completions";
import OpenAI from "openai";
import { z } from "zod";
import zodToJsonSchema from "zod-to-json-schema";
 
// This wrap function adds some useful tracing in Braintrust
const openai = wrapOpenAI(new OpenAI());
 
// Convenience function for defining an OpenAI function call
const makeFunctionDefinition = (
  name: string,
  description: string,
  schema: z.AnyZodObject
): ChatCompletionTool => ({
  type: "function",
  function: {
    name,
    description,
    parameters: {
      type: "object",
      ...pick(
        zodToJsonSchema(schema, {
          name: "root",
          $refStrategy: "none",
        }).definitions?.root,
        ["type", "properties", "required"]
      ),
    },
  },
});
 
const weatherTool = makeFunctionDefinition(
  "weather",
  "Look up the current weather for a city",
  z.object({
    city: z.string().describe("The city to look up the weather for"),
    date: z.string().optional().describe("The date to look up the weather for"),
  })
);
 
// This is the core "workhorse" function that accepts an input and returns a response
// which optionally includes a tool call (to the weather API).
async function task(input: string) {
  const completion = await openai.chat.completions.create({
    model: "gpt-3.5-turbo",
    messages: [
      {
        role: "system",
        content: `You are a highly intelligent AI that can look up the weather.`,
      },
      { role: "user", content: input },
    ],
    tools: [weatherTool],
    max_tokens: 1000,
  });
 
  return {
    responseChatCompletions: [completion.choices[0].message],
  };
}

JSON.stringify(await task("What's the weather in San Francisco?"), null, 2);

JSON.stringify(await task("What is my bank balance?"), null, 2);

JSON.stringify(await task("What is the weather?"), null, 2);

type AssertionTypes =
  | "equals"
  | "exists"
  | "not_exists"
  | "llm_criteria_met"
  | "semantic_contains";
 
type Assertion = {
  path: string;
  assertion_type: AssertionTypes;
  value: string;
};

import { ClosedQA } from "autoevals";
import get from "lodash/get";
import every from "lodash/every";
 
/**
 * Uses an LLM call to classify if a substring is semantically contained in a text.
 * @param text The full text you want to check against
 * @param needle The string you want to check if it is contained in the text
 */
async function semanticContains({
  text1,
  text2,
}: {
  text1: string;
  text2: string;
}): Promise<boolean> {
  const system = `
  You are a highly intelligent AI. You will be given two texts, TEXT_1 and TEXT_2. Your job is to tell me if TEXT_2 is semantically present in TEXT_1.
  Examples:
  \`\`\`
  TEXT_1: "I've just sent “hello world” to the #testing channel on Slack as you requested. Can I assist you with anything else?"
  TEXT_2: "Can I help you with something else?"
  Result: YES
  \`\`\`
  
  \`\`\`
  TEXT_1: "I've just sent “hello world” to the #testing channel on Slack as you requested. Can I assist you with anything else?"
  TEXT_2: "Sorry, something went wrong."
  Result: NO
  \`\`\`
  
  \`\`\`
  TEXT_1: "I've just sent “hello world” to the #testing channel on Slack as you requested. Can I assist you with anything else?"
  TEXT_2: "#testing channel Slack"
  Result: YES
  \`\`\`
  
  \`\`\`
  TEXT_1: "I've just sent “hello world” to the #testing channel on Slack as you requested. Can I assist you with anything else?"
  TEXT_2: "#general channel Slack"
  Result: NO
  \`\`\`
  `;
 
  const toolSchema = z.object({
    rationale: z
      .string()
      .describe(
        "A string that explains the reasoning behind your answer. It's a step-by-step explanation of how you determined that TEXT_2 is or isn't semantically present in TEXT_1."
      ),
    answer: z.boolean().describe("Your answer"),
  });
 
  const completion = await openai.chat.completions.create({
    model: "gpt-3.5-turbo",
    messages: [
      {
        role: "system",
        content: system,
      },
      {
        role: "user",
        content: `TEXT_1: "${text1}"\nTEXT_2: "${text2}"`,
      },
    ],
    tools: [
      makeFunctionDefinition(
        "semantic_contains",
        "The result of the semantic presence check",
        toolSchema
      ),
    ],
    tool_choice: {
      function: { name: "semantic_contains" },
      type: "function",
    },
    max_tokens: 1000,
  });
 
  try {
    const { answer } = toolSchema.parse(
      JSON.parse(
        completion.choices[0].message.tool_calls![0].function.arguments
      )
    );
    return answer;
  } catch (e) {
    console.error(e, "Error parsing semanticContains response");
    return false;
  }
}
 
const AssertionScorer = async ({
  input,
  output,
  expected: assertions,
}: {
  input: string;
  output: any;
  expected: Assertion[];
}) => {
  // for each assertion, perform the comparison
  const assertionResults: {
    status: string;
    path: string;
    assertion_type: string;
    value: string;
    actualValue: string;
  }[] = [];
  for (const assertion of assertions) {
    const { assertion_type, path, value } = assertion;
    const actualValue = get(output, path);
    let passedTest = false;
 
    try {
      switch (assertion_type) {
        case "equals":
          passedTest = actualValue === value;
          break;
        case "exists":
          passedTest = actualValue !== undefined;
          break;
        case "not_exists":
          passedTest = actualValue === undefined;
          break;
        case "llm_criteria_met":
          const closedQA = await ClosedQA({
            input:
              "According to the provided criterion is the submission correct?",
            criteria: value,
            output: actualValue,
          });
          passedTest = !!closedQA.score && closedQA.score > 0.5;
          break;
        case "semantic_contains":
          passedTest = await semanticContains({
            text1: actualValue,
            text2: value,
          });
          break;
        default:
          assertion_type satisfies never; // if you see a ts error here, its because your switch is not exhaustive
          throw new Error(`unknown assertion type ${assertion_type}`);
      }
    } catch (e) {
      passedTest = false;
    }
    assertionResults.push({
      status: passedTest ? "passed" : "failed",
      path,
      assertion_type,
      value,
      actualValue,
    });
  }
 
  const allPassed = every(assertionResults, (r) => r.status === "passed");
 
  return {
    name: "Assertions Score",
    score: allPassed ? 1 : 0,
    metadata: {
      assertionResults,
    },
  };
};

const data = [
  {
    input: "What's the weather like in San Francisco?",
    expected: [
      {
        path: "responseChatCompletions[0].tool_calls[0].function.name",
        assertion_type: "equals",
        value: "weather",
      },
    ],
  },
  {
    input: "What's the weather like?",
    expected: [
      {
        path: "responseChatCompletions[0].tool_calls[0].function.name",
        assertion_type: "not_exists",
        value: "",
      },
      {
        path: "responseChatCompletions[0].content",
        assertion_type: "llm_criteria_met",
        value:
          "Response reflecting the bot does not have enough information to look up the weather",
      },
    ],
  },
  {
    input: "How much is AAPL stock today?",
    expected: [
      {
        path: "responseChatCompletions[0].tool_calls[0].function.name",
        assertion_type: "not_exists",
        value: "",
      },
      {
        path: "responseChatCompletions[0].content",
        assertion_type: "llm_criteria_met",
        value:
          "Response reflecting the bot does not have access to the ability or tool to look up stock prices.",
      },
    ],
  },
  {
    input: "What can you do?",
    expected: [
      {
        path: "responseChatCompletions[0].content",
        assertion_type: "semantic_contains",
        value: "look up the weather",
      },
    ],
  },
];

import { Eval } from "braintrust";
 
await Eval("Weather Bot", {
  data,
  task: async (input) => {
    const result = await task(input);
    return result;
  },
  scores: [AssertionScorer],
});

async function task(input: string) {
  const completion = await openai.chat.completions.create({
    model: "gpt-3.5-turbo",
    messages: [
      {
        role: "system",
        content: `You are a highly intelligent AI that can look up the weather.
        
Do not try to use tools other than those provided to you. If you do not have the tools needed to solve a problem, just say so.
 
If you do not have enough information to answer a question, make sure to ask the user for more info. Prefix that statement with "I need more information to answer this question."
        `,
      },
      { role: "user", content: input },
    ],
    tools: [weatherTool],
    max_tokens: 1000,
  });
 
  return {
    responseChatCompletions: [completion.choices[0].message],
  };
}

JSON.stringify(await task("How much is AAPL stock today?"), null, 2);

await Eval("Weather Bot", {
  data: data,
  task: async (input) => {
    const result = await task(input);
    return result;
  },
  scores: [AssertionScorer],
});