check via braintrust secret type

Author: knjiang

apis/cloudflare/src/billing.ts

@@ -3,17 +3,12 @@ import { type BillingEvent } from "@braintrust/proxy";
 const DEFAULT_BILLING_TELEMETRY_URL =
   "https://api.braintrust.dev/billing/telemetry/ingest";
 
-function isBrainModel(model: string): boolean {
-  return model.startsWith("brain-");
-}
-
 function buildPayloadEvent(event: BillingEvent) {
   if (!event.model) {
     console.warn("billing event skipped: missing model");
     return null;
   }
-  // Skip non-brain models since braintrust only hosts brain models.
-  if (!isBrainModel(event.model)) {
+  if (!event.isNativeInference) {
     return null;
   }
 

packages/proxy/src/proxy.ts

@@ -102,6 +102,7 @@ import {
   flattenChunksArray,
   getRandomInt,
   isEmpty,
+  isNativeInferenceSecret,
   isObject,
   ModelResponse,
   parseAuthHeader,
@@ -194,6 +195,7 @@ export type BillingEvent = {
   event_name: "NativeInferenceTokenUsageEvent";
   auth_token: string;
   org_id?: string;
+  isNativeInference?: boolean;
   model?: string | null;
   resolved_model?: string | null;
   org_name?: string;
@@ -1151,6 +1153,7 @@ export async function proxyV1({
               event_name: "NativeInferenceTokenUsageEvent",
               auth_token: authToken,
               org_id: billingOrgId,
+              isNativeInference: nativeInferenceSecret !== undefined,
               model,
               resolved_model: resolvedModel,
               org_name: resolvedOrgName,
@@ -1330,6 +1333,10 @@ async function fetchModelLoop(
   }
 
   const initialIdx = getRandomInt(secrets.length);
+  const nativeInferenceSecret =
+    model === null || model === undefined
+      ? undefined
+      : secrets.find((secret) => isNativeInferenceSecret(secret, model));
   let proxyResponse: ModelResponse | null = null;
   let secretName: string | null | undefined = null;
   let lastException = null;

packages/proxy/src/util.test.ts

@@ -1,5 +1,10 @@
 import { describe, expect, test } from "vitest";
-import { parseFileMetadataFromUrl, _urljoin } from "./util";
+import { type APISecret } from "@schema";
+import {
+  isNativeInferenceSecret,
+  parseFileMetadataFromUrl,
+  _urljoin,
+} from "./util";
 
 describe("parseFileMetadataFromUrl", () => {
   test("handles basic URLs", () => {
@@ -222,3 +227,58 @@ test("_urljoin", () => {
   expect(_urljoin()).toBe("");
   expect(_urljoin("a")).toBe("a");
 });
+
+describe("isNativeInferenceSecret", () => {
+  test("treats braintrust secrets without custom models as native inference", () => {
+    const secret: APISecret = {
+      type: "braintrust",
+      secret: "test-secret",
+      metadata: {
+        api_base: "http://localhost:8001/native-inference-test",
+      },
+    };
+
+    expect(isNativeInferenceSecret(secret, "brain-test-native-1")).toBe(true);
+  });
+
+  test("treats custom model secrets as native inference only for matching models", () => {
+    const secret: APISecret = {
+      type: "openai",
+      secret: "test-secret",
+      metadata: {
+        customModels: {
+          "brain-test-native-1": {
+            format: "openai",
+            flavor: "chat",
+          },
+        },
+      },
+    };
+
+    expect(isNativeInferenceSecret(secret, "brain-test-native-1")).toBe(true);
+    expect(isNativeInferenceSecret(secret, "brain-test-native-2")).toBe(false);
+  });
+
+  test("treats non-braintrust secrets without matching custom models as non-native", () => {
+    const secret: APISecret = {
+      type: "openai",
+      secret: "test-secret",
+      metadata: {
+        api_base: "https://api.openai.com/v1",
+      },
+    };
+
+    expect(isNativeInferenceSecret(secret, "brain-test-native-1")).toBe(false);
+  });
+
+  test("returns false when the model is missing", () => {
+    const secret: APISecret = {
+      type: "braintrust",
+      secret: "test-secret",
+      metadata: null,
+    };
+
+    expect(isNativeInferenceSecret(secret, null)).toBe(false);
+    expect(isNativeInferenceSecret(secret, undefined)).toBe(false);
+  });
+});

packages/proxy/src/util.ts

@@ -1,4 +1,5 @@
 import contentDisposition from "content-disposition";
+import { type APISecret } from "@schema";
 export interface ModelResponse {
   stream: ReadableStream<Uint8Array> | null;
   response: Response;
@@ -79,6 +80,27 @@ export function isEmpty(a: any): a is null | undefined {
   return a === undefined || a === null;
 }
 
+export function isNativeInferenceSecret(
+  secret: APISecret,
+  model: string | null | undefined,
+): boolean {
+  if (model === null || model === undefined) {
+    return false;
+  }
+
+  const customModels = secret.metadata?.customModels;
+  if (
+    customModels === null ||
+    customModels === undefined ||
+    typeof customModels !== "object" ||
+    Array.isArray(customModels)
+  ) {
+    return secret.type === "braintrust";
+  }
+
+  return Object.prototype.hasOwnProperty.call(customModels, model);
+}
+
 export function getRandomInt(max: number) {
   return Math.floor(Math.random() * max);
 }