support moderations

braintrustdata · Jan 14, 2025 · 528ec7f · 528ec7f
1 parent d3977ad
commit 528ec7f
Showing 1 changed file with 22 additions and 5 deletions.
diff --git a/packages/proxy/src/proxy.ts b/packages/proxy/src/proxy.ts
@@ -47,6 +47,7 @@ import {
   ChatCompletionCreateParams,
   CompletionUsage,
   CreateEmbeddingResponse,
+  ModerationCreateResponse,
 } from "openai/resources";
 import { fetchBedrockAnthropic, fetchBedrockOpenAI } from "./providers/bedrock";
 import { Buffer } from "node:buffer";
@@ -210,7 +211,8 @@ export async function proxyV1({
     url === "/auto" ||
     url === "/embeddings" ||
     url === "/chat/completions" ||
-    url === "/completions";
+    url === "/completions" ||
+    url === "/moderations";
 
   let bodyData = null;
   if (
@@ -474,7 +476,7 @@ export async function proxyV1({
         async flush(controller) {
           const data = flattenChunksArray(allChunks);
           const dataB64 = Buffer.from(data).toString("base64");
-          cachePut(
+          await cachePut(
             encryptionKey,
             cacheKey,
             JSON.stringify({ headers: proxyResponseHeaders, data: dataB64 }),
@@ -637,6 +639,14 @@ export async function proxyV1({
                 });
               }
               break;
+            case "moderation":
+              {
+                const data = dataRaw as ModerationCreateResponse;
+                spanLogger.log({
+                  output: data.results,
+                });
+              }
+              break;
           }
         }
 
@@ -1598,7 +1608,7 @@ function tryParseRateLimitReset(headers: Headers): number | null {
   return null;
 }
 
-export type SpanType = "chat" | "completion" | "embedding";
+export type SpanType = "chat" | "completion" | "embedding" | "moderation";
 
 function spanTypeToName(spanType: SpanType): string {
   switch (spanType) {
@@ -1608,6 +1618,8 @@ function spanTypeToName(spanType: SpanType): string {
       return "Completion";
     case "embedding":
       return "Embedding";
+    case "moderation":
+      return "Moderation";
   }
 }
 
@@ -1622,7 +1634,9 @@ export function guessSpanType(
         ? "completion"
         : url === "/embeddings"
           ? "embedding"
-          : undefined;
+          : url === "/moderations"
+            ? "moderation"
+            : undefined;
   if (spanName) {
     return spanName;
   }
@@ -1632,6 +1646,8 @@ export function guessSpanType(
     return "chat";
   } else if (flavor === "completion") {
     return "completion";
+  } else if (url === "/moderations") {
+    return "moderation";
   } else {
     return undefined;
   }
@@ -1660,7 +1676,8 @@ function logSpanInputs(
       });
       break;
     }
-    case "embedding": {
+    case "embedding":
+    case "moderation": {
       const { input, ...rest } = bodyData;
       spanLogger.log({
         input: bodyData,