chore: move groundTruthCompletion into adapters

ubq-testing · Oct 25, 2024 · f7fa7b5 · f7fa7b5
1 parent 359451e
commit f7fa7b5
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 41 deletions.
diff --git a/src/adapters/openai/helpers/completions.ts b/src/adapters/openai/helpers/completions.ts
@@ -1,6 +1,7 @@
 import OpenAI from "openai";
 import { Context } from "../../../types";
 import { SuperOpenAi } from "./openai";
+import { CompletionsModelHelper, ModelApplications } from "../../../types/llm";
 const MAX_TOKENS = 7000;
 
 export interface CompletionsType {
@@ -76,4 +77,39 @@ export class Completions extends SuperOpenAi {
     }
     return { answer: "", tokenUsage: { input: 0, output: 0, total: 0 } };
   }
+
+  async createGroundTruthCompletion<TApp extends ModelApplications>(
+    context: Context,
+    groundTruthSource: string,
+    systemMsg: string,
+    model: CompletionsModelHelper<TApp>
+  ): Promise<string | null> {
+    const {
+      env: { OPENAI_API_KEY },
+      config: { openAiBaseUrl },
+    } = context;
+
+    const openAi = new OpenAI({
+      apiKey: OPENAI_API_KEY,
+      ...(openAiBaseUrl && { baseURL: openAiBaseUrl }),
+    });
+
+    const msgs = [
+      {
+        role: "system",
+        content: systemMsg,
+      },
+      {
+        role: "user",
+        content: groundTruthSource,
+      },
+    ] as OpenAI.Chat.Completions.ChatCompletionMessageParam[];
+
+    const res = await openAi.chat.completions.create({
+      messages: msgs,
+      model: model,
+    });
+
+    return res.choices[0].message.content;
+  }
 }
diff --git a/src/handlers/ground-truths/create-ground-truth-completion.ts b/src/handlers/ground-truths/create-ground-truth-completion.ts
diff --git a/src/handlers/ground-truths/find-ground-truths.ts b/src/handlers/ground-truths/find-ground-truths.ts
@@ -4,7 +4,6 @@ import { GROUND_TRUTHS_SYSTEM_MESSAGES } from "./prompts";
 import { chatBotPayloadTypeguard, codeReviewPayloadTypeguard } from "../../types/typeguards";
 import { validateGroundTruths } from "./validate";
 import { logger } from "../../helpers/errors";
-import { createGroundTruthCompletion } from "./create-ground-truth-completion";
 import { createGroundTruthSysMsg } from "./create-system-message";
 
 export async function findGroundTruths<TApp extends ModelApplications = ModelApplications>(
@@ -32,8 +31,9 @@ async function findChatBotTruths(
   params: AppParamsHelper<"chat-bot">,
   systemMsgObj: GroundTruthsSystemMessage<"chat-bot">
 ): Promise<string[]> {
+  const { adapters: { openai: { completions } } } = context;
   const systemMsg = createGroundTruthSysMsg(systemMsgObj);
-  const truths = await createGroundTruthCompletion<"chat-bot">(context, JSON.stringify(params), systemMsg, "o1-mini");
+  const truths = await completions.createGroundTruthCompletion<"chat-bot">(context, JSON.stringify(params), systemMsg, "o1-mini");
   return validateGroundTruths(truths);
 }
 
@@ -42,7 +42,8 @@ async function findCodeReviewTruths(
   params: AppParamsHelper<"code-review">,
   systemMsgObj: GroundTruthsSystemMessage<"code-review">
 ): Promise<string[]> {
+  const { adapters: { openai: { completions } } } = context;
   const systemMsg = createGroundTruthSysMsg(systemMsgObj);
-  const truths = await createGroundTruthCompletion<"code-review">(context, params.taskSpecification, systemMsg, "gpt-4o");
+  const truths = await completions.createGroundTruthCompletion<"code-review">(context, params.taskSpecification, systemMsg, "gpt-4o");
   return validateGroundTruths(truths);
 }