(EAI-649): Draft code example classifier type (#595)

* migrate classifier to core

* start creating docs team classifier

* configurable model

* working e2e

* fix build err

---------

Co-authored-by: Ben Perlmutter <mongodben@mongodb.com>

Author: mongodben

packages/datasets/.env.example

@@ -1,4 +1,6 @@
 OPENAI_API_KEY=<YOUR_API_KEY>
-OPENAI_DEPLOYMENT_NAME=<YOUR_DEPLOYMENT_NAME>
+OPENAI_CHAT_COMPLETION_DEPLOYMENT="gpt-4o-mini"
+OPENAI_CHAT_COMPLETION_MODEL_VERSION="placeholder...must include but value doesn't matter"
 OPENAI_ENDPOINT=<YOUR_ENDPOINT>
-OPENAI_GPT_4o_COMPLETION_DEPLOYMENT="<some deployment>"
+OPENAI_API_VERSION="2024-06-01"
+BRAINTRUST_API_KEY="<some key>"

packages/datasets/src/classifyCodeExampleDocsTeam.eval.ts

@@ -0,0 +1,79 @@
+import "dotenv/config";
+import { EvalCase, Eval, EvalScorer } from "mongodb-rag-core/braintrust";
+import {
+  CodeExampleClassification,
+  makeClassifyCodeExampleDocsTeam,
+} from "./classifyCodeExampleDocsTeam.js";
+import {
+  assertEnvVars,
+  Classification,
+  CORE_OPENAI_CHAT_COMPLETION_ENV_VARS,
+} from "mongodb-rag-core";
+import { AzureOpenAI } from "mongodb-rag-core/openai";
+
+const {
+  OPENAI_API_KEY,
+  OPENAI_API_VERSION,
+  OPENAI_CHAT_COMPLETION_DEPLOYMENT,
+  OPENAI_ENDPOINT,
+} = assertEnvVars(CORE_OPENAI_CHAT_COMPLETION_ENV_VARS);
+
+type ClassifyCodeExampleEvalCaseInput = {
+  text: string;
+};
+
+type ClassifyCodeExampleEvalCaseOutput = Classification;
+
+type ClassifyCodeExampleEvalCaseExpected = CodeExampleClassification;
+
+type ClassifyCodeExampleEvalCase = EvalCase<
+  ClassifyCodeExampleEvalCaseInput,
+  ClassifyCodeExampleEvalCaseExpected,
+  Record<string, unknown> | undefined
+>;
+
+// TODO: add eval cases
+const evalCases = [
+  {
+    input: {
+      text: `\`\`\`ruby
+puts "Hello, World!"
+\`\`\``,
+    },
+    expected: "unknown",
+  },
+] satisfies ClassifyCodeExampleEvalCase[];
+
+const isCorrectClassification: EvalScorer<
+  ClassifyCodeExampleEvalCaseInput,
+  ClassifyCodeExampleEvalCaseOutput,
+  ClassifyCodeExampleEvalCaseExpected
+> = function ({ output, expected }) {
+  return {
+    name: "CorrectClassification",
+    score: output.type === expected ? 1 : 0,
+  };
+};
+
+const classifyCodeExample = makeClassifyCodeExampleDocsTeam({
+  model: OPENAI_CHAT_COMPLETION_DEPLOYMENT,
+  openAiClient: new AzureOpenAI({
+    endpoint: OPENAI_ENDPOINT,
+    apiKey: OPENAI_API_KEY,
+    apiVersion: OPENAI_API_VERSION,
+  }),
+});
+
+Eval("classify-code-example-docs-team", {
+  maxConcurrency: 10,
+  experimentName: "classify-code-example",
+  data: evalCases,
+  async task(input) {
+    const { text } = input;
+
+    const { classification } = await classifyCodeExample({ input: text });
+
+    return classification;
+  },
+  scores: [isCorrectClassification],
+});

packages/datasets/src/classifyCodeExampleDocsTeam.ts

@@ -0,0 +1,84 @@
+/**
+   @fileoverview Contains classification as defined by MongoDB Docs team (https://docs.google.com/document/d/199fTaNlvMXb1GmloQRkh8MYPfYyuQakaYIsAuZvgKvU/)
+ */
+import { ClassificationType, makeClassifier } from "mongodb-rag-core";
+
+const classificationTypes = [
+  {
+    type: "api_method_signature",
+    description:
+      "API method signature code block showing an API method name and arguments",
+    examples: [
+      {
+        // From https://www.mongodb.com/docs/languages/python/pymongo-driver/current/write/insert/#insert-one-document
+        text: `\`\`\`\`python
+sample_restaurants.restaurants.insert_one({"name" : "Mongo's Burgers"})
+        \`\`\``,
+        reason:
+          'Shows the minimum required arguments for the "insert_one" method. Therefore this is an API method signature.',
+      },
+      {
+        // From https://www.mongodb.com/docs/manual/reference/method/db.collection.insertOne/#syntax
+        text: `\`\`\`\`javascript
+db.collection.insertOne(
+    <document>,
+    {
+      writeConcern: <document>
+    }
+)
+\`\`\``,
+        reason:
+          "Complete API method signature, showing all possible arguments in an abstract form. Therefore this is an API method signature.",
+      },
+    ],
+  },
+  {
+    type: "return_example",
+    description:
+      "A JSON blob, example document, or other return object type demonstrating what a user might expect from executing a piece of code",
+    // TODO: add other examples
+    examples: [],
+  },
+  {
+    type: "example_configuration_object",
+    description:
+      "Example object, often represented in YAML or JSON, enumerating parameters and their types",
+    // TODO: add other examples
+    examples: [],
+  },
+  {
+    type: "usage_example",
+    description:
+      "A longer code snippet that establishes parameters, performs basic set-up code, and includes the larger context to demonstrate how to accomplish a task",
+    // TODO: add other examples
+    examples: [],
+  },
+  {
+    type: "sample_application",
+    description:
+      "Runnable applications that connect more discrete pieces of code, and may include error handling, framework integrations, or User Interface elements",
+    // TODO: add other examples
+    examples: [],
+  },
+  // Note: adding this b/c there will certainly be types of code examples that don't cleanly fit into these buckets.
+  // For example, if there was a file tree in a code block.
+  {
+    type: "unknown",
+    description:
+      "Unknown classification type. The code example doesn't easily fit into any of the other categories.",
+  },
+] as const satisfies ClassificationType[];
+
+export type CodeExampleClassification =
+  (typeof classificationTypes)[number]["type"];
+
+export const makeClassifyCodeExampleDocsTeam = ({
+  openAiClient,
+  model,
+}: Pick<Parameters<typeof makeClassifier>[0], "openAiClient" | "model">) =>
+  makeClassifier({
+    classificationTypes,
+    chainOfThought: true,
+    openAiClient,
+    model,
+  });

packages/mongodb-artifact-generator/src/chat/makeClassifierWithLogger.ts

@@ -0,0 +1,28 @@
+import { Classifier, makeClassifier } from "mongodb-rag-core";
+import { RunLogger } from "../runlogger";
+import { stripIndents } from "common-tags";
+
+export function makeClassifierWithLogger(
+  args: Parameters<typeof makeClassifier>[0] & { logger?: RunLogger }
+) {
+  const { logger, ...classifierBuilder } = args;
+  const classifier = makeClassifier(classifierBuilder);
+  const classifierWithLogger: Classifier = async (
+    args: Parameters<typeof classifier>[0]
+  ) => {
+    const result = await classifier(args);
+    logger?.appendArtifact(
+      `chatTemplates/classifier-${Date.now()}.json`,
+      stripIndents`
+        <SystemMessage>
+          ${result.inputMessages[0].content}
+        </SystemMessage>
+        <Classification>
+          ${JSON.stringify(result.classification)}
+        </Classification>
+      `
+    );
+    return result;
+  };
+  return classifierWithLogger;
+}

packages/mongodb-artifact-generator/src/commands/generateReleaseNotes.ts

@@ -20,6 +20,7 @@ import {
   ClassifiedChangelog,
   makeClassifyChangelogs,
 } from "../release-notes/classifyChangelog";
+import { assertEnvVars } from "mongodb-rag-core";
 
 let logger: RunLogger;
 
@@ -29,6 +30,9 @@ type GenerateReleaseNotesCommandArgs = {
   llmMaxConcurrency: number;
 };
 
+const { OPENAI_CHAT_COMPLETION_DEPLOYMENT } = assertEnvVars({
+  OPENAI_CHAT_COMPLETION_DEPLOYMENT: "",
+});
 export default createCommand<GenerateReleaseNotesCommandArgs>({
   command: "generateReleaseNotes",
   builder(args) {
@@ -175,6 +179,7 @@ export const action = createConfiguredAction<GenerateReleaseNotesCommandArgs>(
 
     const classifyChangelogs = makeClassifyChangelogs({
       openAiClient,
+      model: OPENAI_CHAT_COMPLETION_DEPLOYMENT,
       logger,
     });
 

packages/mongodb-artifact-generator/src/release-notes/classifyChangelog.ts

@@ -4,7 +4,7 @@ import { makeClassifyChangelogAudience } from "./classifyChangelogAudience";
 import { makeClassifyChangelogScope } from "./classifyChangelogScope";
 import { iOfN } from "../utils";
 import { RunLogger } from "../runlogger";
-import { Classification } from "../chat/makeClassifier";
+import { Classification } from "mongodb-rag-core";
 
 export type ClassifiedChangelog = {
   audience: Classification;
@@ -14,19 +14,23 @@ export type ClassifiedChangelog = {
 
 export type MakeClassifyChangelogScope = {
   openAiClient: OpenAI;
+  model: string;
   logger?: RunLogger;
 };
 
 export function makeClassifyChangelog({
   openAiClient,
+  model,
   logger,
 }: MakeClassifyChangelogScope) {
   const classifyChangelogAudience = makeClassifyChangelogAudience({
     openAiClient,
+    model,
     logger,
   });
   const classifyChangelogScope = makeClassifyChangelogScope({
     openAiClient,
+    model,
     logger,
   });
 
@@ -35,8 +39,12 @@ export function makeClassifyChangelog({
   }: {
     changelog: string;
   }): Promise<ClassifiedChangelog> {
-    const audience = await classifyChangelogAudience({ input: changelog });
-    const scope = await classifyChangelogScope({ input: changelog });
+    const { classification: audience } = await classifyChangelogAudience({
+      input: changelog,
+    });
+    const { classification: scope } = await classifyChangelogScope({
+      input: changelog,
+    });
 
     return {
       audience,
@@ -49,8 +57,13 @@ export function makeClassifyChangelog({
 export function makeClassifyChangelogs({
   openAiClient,
   logger,
+  model,
 }: MakeClassifyChangelogScope) {
-  const classifyChangelog = makeClassifyChangelog({ openAiClient, logger });
+  const classifyChangelog = makeClassifyChangelog({
+    openAiClient,
+    logger,
+    model,
+  });
 
   return async function classifyChangelogs({
     changelogs,

packages/mongodb-artifact-generator/src/release-notes/classifyChangelogAudience.ts

@@ -1,5 +1,5 @@
 import { OpenAI } from "mongodb-rag-core/openai";
-import { makeClassifier } from "../chat/makeClassifier";
+import { makeClassifierWithLogger } from "../chat/makeClassifierWithLogger";
 import { RunLogger } from "../runlogger";
 
 const classificationTypes = [
@@ -41,14 +41,17 @@ const classificationTypes = [
 
 export type MakeClassifyChangelogAudienceArgs = {
   openAiClient: OpenAI;
+  model: string;
   logger?: RunLogger;
 };
 
 export function makeClassifyChangelogAudience({
   openAiClient,
+  model,
 }: MakeClassifyChangelogAudienceArgs) {
-  return makeClassifier({
+  return makeClassifierWithLogger({
     openAiClient,
+    model,
     classificationTypes,
     chainOfThought: true,
   });

packages/mongodb-artifact-generator/src/release-notes/classifyChangelogScope.ts

@@ -1,5 +1,5 @@
 import { OpenAI } from "mongodb-rag-core/openai";
-import { makeClassifier } from "../chat/makeClassifier";
+import { makeClassifierWithLogger } from "../chat/makeClassifierWithLogger";
 import { RunLogger } from "../runlogger";
 
 const classificationTypes = [
@@ -86,17 +86,20 @@ const classificationTypes = [
 
 export type MakeClassifyChangelogScope = {
   openAiClient: OpenAI;
+  model: string;
   logger?: RunLogger;
 };
 
 export function makeClassifyChangelogScope({
   openAiClient,
+  model,
   logger,
 }: MakeClassifyChangelogScope) {
-  return makeClassifier({
+  return makeClassifierWithLogger({
     openAiClient,
     logger,
     classificationTypes,
     chainOfThought: true,
+    model,
   });
 }

packages/mongodb-rag-core/src/index.ts

@@ -20,6 +20,7 @@ export * from "./DatabaseConnection";
 export * from "./DataStreamer";
 export * from "./logger";
 export * from "./conversations/MongoDbConversations";
+export * from "./makeClassifier";
 export * from "./References";
 export * from "./VectorStore";
 export * from "./arrayFilters";