feat: support cloudflare secrets based auth for self hosting

apis/cloudflare/src/env.ts

@@ -1,5 +1,8 @@
+import { Secrets } from "@braintrust/proxy/edge";
+
 declare global {
-  interface Env {
+  interface Env extends Secrets {
+    API_KEY: string;
     ai_proxy: KVNamespace;
     BRAINTRUST_APP_URL: string;
     DISABLE_METRICS?: boolean;

apis/cloudflare/src/proxy.ts

@@ -4,6 +4,7 @@ import {
   ProxyOpts,
   makeFetchApiSecrets,
   encryptedGet,
+  getApiSecret,
 } from "@braintrust/proxy/edge";
 import { NOOP_METER_PROVIDER, initMetrics } from "@braintrust/proxy";
 import { PrometheusMetricAggregator } from "./metric-aggregator";
@@ -111,11 +112,27 @@ export async function handleProxyV1(
         cacheSetLatency.record(end - start);
       },
     },
-    braintrustApiUrl: braintrustAppUrl(env).toString(),
     meterProvider,
     whitelist,
   };
 
+  const optsWithCloudflareAuth: ProxyOpts = {
+    ...opts,
+    authConfig: {
+      type: "cloudflare",
+      apiKey: env.API_KEY,
+      getSecret: (model: string) => getApiSecret(model, env),
+    },
+  };
+
+  const optsWithBraintrustAuth: ProxyOpts = {
+    ...opts,
+    authConfig: {
+      type: "braintrust",
+      braintrustApiUrl: braintrustAppUrl(env).toString(),
+    },
+  };
+
   const url = new URL(request.url);
   if (url.pathname === `${proxyV1Prefix}/realtime`) {
     return await handleRealtimeProxy({
@@ -135,7 +152,9 @@ export async function handleProxyV1(
     });
   }
 
-  return EdgeProxyV1(opts)(request, ctx);
+  // decide if you want to use braintrust or cloudflare auth
+  return EdgeProxyV1(optsWithCloudflareAuth)(request, ctx);
+  // return EdgeProxyV1(optsWithBraintrustAuth)(request, ctx);
 }
 
 export async function handlePrometheusScrape(

packages/proxy/edge/index.ts

@@ -4,7 +4,7 @@ import { proxyV1 } from "@lib/proxy";
 import { isEmpty } from "@lib/util";
 import { MeterProvider } from "@opentelemetry/sdk-metrics";
 
-import { APISecret, getModelEndpointTypes } from "@schema";
+import { APISecret, getModelEndpointTypes, isFireworksModel, isAnthropicModel, isBedrockModel, isGroqModel, isOpenAIModel, isGoogleModel, isXAIModel, isMistralModel, isPerplexityModel } from "@schema";
 import { verifyTempCredentials, isTempCredential } from "utils";
 import {
   decryptMessage,
@@ -31,7 +31,20 @@ export interface ProxyOpts {
   cors?: boolean;
   credentialsCache?: Cache;
   completionsCache?: Cache;
-  braintrustApiUrl?: string;
+  authConfig?: {
+    type: "cloudflare";
+    /**
+     * The API key to use for proxy authentication
+     */
+    apiKey: string;
+    /**
+     * A function that returns the API secret for a given model
+     */
+    getSecret: (model: string) => Promise<APISecret> | APISecret;
+  } | {
+    type: "braintrust";
+    braintrustApiUrl?: string;
+  }
   meterProvider?: MeterProvider;
   whitelist?: (string | RegExp)[];
 }
@@ -67,9 +80,9 @@ export function getCorsHeaders(
 
   return origin
     ? {
-        "access-control-allow-origin": origin,
-        ...baseCorsHeaders,
-      }
+      "access-control-allow-origin": origin,
+      ...baseCorsHeaders,
+    }
     : {};
 }
 
@@ -122,6 +135,19 @@ export function makeFetchApiSecrets({
     model: string | null,
     org_name?: string,
   ): Promise<APISecret[]> => {
+    if (opts.authConfig?.type === "cloudflare") {
+      if (authToken !== opts.authConfig.apiKey) {
+        throw new Error("Forbidden");
+      }
+
+      if (!model) {
+        throw new Error("no model provided");
+      }
+
+      const secret = await opts.authConfig.getSecret(model);
+      return [secret];
+    }
+
     // First try to decode & verify as JWT. We gate this on Braintrust JWT
     // format, not just any JWT, in case a future model provider uses JWT as
     // the auth token.
@@ -167,7 +193,7 @@ export function makeFetchApiSecrets({
     let ttl = 60;
     try {
       const response = await fetch(
-        `${opts.braintrustApiUrl || DEFAULT_BRAINTRUST_APP_URL}/api/secret`,
+        `${opts.authConfig?.braintrustApiUrl || DEFAULT_BRAINTRUST_APP_URL}/api/secret`,
         {
           method: "POST",
           headers: {
@@ -330,6 +356,83 @@ export function EdgeProxyV1(opts: ProxyOpts) {
   };
 }
 
+export interface Secrets {
+  OPENAI_API_KEY: string;
+  ANTHROPIC_API_KEY: string;
+  PERPLEXITY_API_KEY: string;
+  REPLICATE_API_KEY: string;
+  FIREWORKS_API_KEY: string;
+  GOOGLE_API_KEY: string;
+  XAI_API_KEY: string;
+
+  TOGETHER_API_KEY: string;
+  LEPTON_API_KEY: string;
+  MISTRAL_API_KEY: string;
+  OLLAMA_API_KEY: string;
+  GROQ_API_KEY: string;
+  CEREBRAS_API_KEY: string;
+
+  BEDROCK_SECRET_KEY: string;
+  BEDROCK_ACCESS_KEY: string;
+  BEDROCK_REGION: string;
+}
+
+export function getApiSecret(model: string, secrets: Secrets): APISecret {
+  if (isOpenAIModel(model)) {
+    return {
+      secret: secrets.OPENAI_API_KEY,
+      type: "openai",
+    };
+  } else if (isAnthropicModel(model)) {
+    return {
+      secret: secrets.ANTHROPIC_API_KEY,
+      type: "anthropic",
+    };
+  } else if (isBedrockModel(model)) {
+    return {
+      secret: secrets.BEDROCK_SECRET_KEY,
+      type: "bedrock",
+      metadata: {
+        "region": secrets.BEDROCK_REGION,
+        "access_key": secrets.BEDROCK_ACCESS_KEY,
+        supportsStreaming: true,
+      },
+    };
+  } else if (isGroqModel(model)) {
+    return {
+      secret: secrets.GROQ_API_KEY,
+      type: "groq",
+    };
+  } else if (isFireworksModel(model)) {
+    return {
+      secret: secrets.FIREWORKS_API_KEY,
+      type: "fireworks",
+    };
+  } else if (isGoogleModel(model)) {
+    return {
+      secret: secrets.GOOGLE_API_KEY,
+      type: "google",
+    };
+  } else if (isXAIModel(model)) {
+    return {
+      secret: secrets.XAI_API_KEY,
+      type: "xAI",
+    };
+  } else if (isMistralModel(model)) {
+    return {
+      secret: secrets.MISTRAL_API_KEY,
+      type: "mistral",
+    };
+  } else if (isPerplexityModel(model)) {
+    return {
+      secret: secrets.PERPLEXITY_API_KEY,
+      type: "perplexity",
+    };
+  }
+
+  throw new Error(`could not find secret for model ${model}`);
+}
+
 // We rely on the fact that Upstash will automatically serialize and deserialize things for us
 export async function encryptedGet(
   cache: Cache,

packages/proxy/schema/models.ts

@@ -26,6 +26,130 @@ export const ModelSchema = z.object({
 
 export type ModelSpec = z.infer<typeof ModelSchema>;
 
+// OpenAI/Azure Models
+export const OpenAIModelSchema = z.enum([
+  'gpt-4o', 'gpt-4o-mini', 'gpt-4o-2024-11-20', 'gpt-4o-2024-08-06', 'gpt-4o-2024-05-13',
+  'gpt-4o-mini-2024-07-18', 'o1', 'o1-preview', 'o1-mini', 'o1-2024-12-17',
+  'o1-preview-2024-09-12', 'o1-mini-2024-09-12', 'gpt-4-turbo', 'gpt-4-turbo-2024-04-09',
+  'gpt-4-turbo-preview', 'gpt-4-0125-preview', 'gpt-4-1106-preview', 'gpt-4', 'gpt-4-0613',
+  'gpt-4-0314', 'gpt-3.5-turbo-0125', 'gpt-3.5-turbo', 'gpt-35-turbo', 'gpt-3.5-turbo-1106',
+  'gpt-3.5-turbo-instruct', 'gpt-3.5-turbo-instruct-0914', 'gpt-4-32k', 'gpt-4-32k-0613',
+  'gpt-4-32k-0314', 'gpt-4-vision-preview', 'gpt-4-1106-vision-preview', 'gpt-3.5-turbo-16k',
+  'gpt-35-turbo-16k', 'gpt-3.5-turbo-16k-0613', 'gpt-3.5-turbo-0613', 'gpt-3.5-turbo-0301',
+  'text-embedding-3-large', 'text-embedding-3-small', 'text-embedding-ada-002',
+]);
+
+export const AnthropicModelSchema = z.enum([
+  'claude-3-5-sonnet-latest', 'claude-3-5-sonnet-20241022', 'claude-3-5-sonnet-20240620',
+  'claude-3-5-haiku-20241022', 'claude-3-haiku-20240307', 'claude-3-sonnet-20240229',
+  'claude-3-opus-20240229', 'claude-instant-1.2', 'claude-instant-1', 'claude-2.1',
+  'claude-2.0', 'claude-2',
+]);
+
+export const BedrockModelSchema = z.enum([
+  'anthropic.claude-3-5-sonnet-20241022-v2:0', 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+  'anthropic.claude-3-haiku-20240307-v1:0', 'anthropic.claude-3-sonnet-20240229-v1:0',
+  'anthropic.claude-3-opus-20240229-v1:0', 'amazon.nova-pro-v1:0', 'amazon.nova-lite-v1:0',
+  'amazon.nova-micro-v1:0', 'amazon.titan-embed-text-v2:0',
+]);
+
+export const PerplexityModelSchema = z.enum([
+  'pplx-7b-chat', 'pplx-7b-online', 'pplx-70b-chat', 'pplx-70b-online',
+  'codellama-34b-instruct', 'codellama-70b-instruct', 'llama-3-8b-instruct',
+  'llama-3-70b-instruct', 'llama-2-13b-chat', 'llama-2-70b-chat', 'mistral-7b-instruct',
+  'mixtral-8x7b-instruct', 'mixtral-8x22b-instruct', 'openhermes-2-mistral-7b',
+  'openhermes-2.5-mistral-7b',
+]);
+
+export const TogetherModelSchema = z.enum([
+  'meta-llama/Llama-3.3-70B-Instruct-Turbo', 'meta-llama/Llama-3.2-3B-Instruct-Turbo',
+  'meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo', 'meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo',
+  'meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo', 'meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo',
+  'meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo', 'meta-llama/Meta-Llama-3-70B',
+  'meta-llama/Llama-3-8b-hf', 'meta-llama/Llama-3-8b-chat-hf', 'meta-llama/Llama-3-70b-chat-hf',
+  'mistralai/Mistral-7B-Instruct-v0.1', 'mistralai/mixtral-8x7b-32kseqlen', 'mistralai/Mixtral-8x7B-Instruct-v0.1',
+  'mistralai/Mixtral-8x7B-Instruct-v0.1-json', 'mistralai/Mixtral-8x22B', 'mistralai/Mixtral-8x22B-Instruct-v0.1',
+  'NousResearch/Nous-Hermes-2-Yi-34B', 'deepseek-ai/DeepSeek-V3', 'deepseek-ai/deepseek-coder-33b-instruct',
+]);
+
+export const MistralModelSchema = z.enum([
+  'mistral-large-latest', 'pixtral-12b-2409', 'open-mistral-nemo', 'codestral-latest',
+  'open-mixtral-8x22b', 'open-codestral-mamba', 'mistral-tiny', 'mistral-small',
+  'mistral-medium',
+]);
+
+export const GroqModelSchema = z.enum([
+  'llama-3.3-70b-versatile', 'llama-3.1-8b-instant', 'llama-3.1-70b-versatile',
+  'llama-3.1-405b-reasoning', 'gemma-7b-it', 'llama3-8b-8192', 'llama3-70b-8192',
+  'llama2-70b-4096', 'mixtral-8x7b-32768',
+]);
+
+export const FireworksModelSchema = z.enum([
+  'accounts/fireworks/models/llama-v3p3-70b-instruct',
+  'accounts/fireworks/models/llama-v3p2-3b-instruct',
+  'accounts/fireworks/models/llama-v3p1-8b-instruct',
+  'accounts/fireworks/models/llama-v3p2-11b-vision-instruct',
+  'accounts/fireworks/models/llama-v3p1-70b-instruct',
+  'accounts/fireworks/models/llama-v3p2-90b-vision-instruct',
+  'accounts/fireworks/models/llama-v3p1-405b-instruct',
+  'accounts/fireworks/models/deepseek-v3',
+]);
+
+export const GoogleModelSchema = z.enum([
+  'gemini-1.5-pro', 'gemini-1.5-flash', 'gemini-2.0-flash-exp', 'gemini-exp-1206',
+  'gemini-exp-1114', 'gemini-exp-1121', 'gemini-1.5-pro-002', 'gemini-1.5-flash-002',
+  'gemini-1.5-pro-latest', 'gemini-1.5-flash-latest', 'gemini-1.5-flash-8b',
+  'gemini-1.0-pro', 'gemini-pro',
+]);
+
+export const XAIModelSchema = z.enum([
+  'grok-beta',
+]);
+
+export const ModelNameSchema = z.union([OpenAIModelSchema, AnthropicModelSchema, BedrockModelSchema, PerplexityModelSchema, TogetherModelSchema, MistralModelSchema, GroqModelSchema, FireworksModelSchema, GoogleModelSchema, XAIModelSchema]);
+
+export type ModelName = z.infer<typeof ModelNameSchema>;
+
+export function isOpenAIModel(model: string) {
+  return OpenAIModelSchema.safeParse(model).success;
+}
+
+export function isAnthropicModel(model: string) {
+  return AnthropicModelSchema.safeParse(model).success;
+}
+
+export function isBedrockModel(model: string) {
+  return BedrockModelSchema.safeParse(model).success;
+}
+
+export function isPerplexityModel(model: string) {
+  return PerplexityModelSchema.safeParse(model).success;
+}
+
+export function isTogetherModel(model: string) {
+  return TogetherModelSchema.safeParse(model).success;
+}
+
+export function isMistralModel(model: string) {
+  return MistralModelSchema.safeParse(model).success;
+}
+
+export function isGroqModel(model: string) {
+  return GroqModelSchema.safeParse(model).success;
+}
+
+export function isFireworksModel(model: string) {
+  return FireworksModelSchema.safeParse(model).success;
+}
+
+export function isGoogleModel(model: string) {
+  return GoogleModelSchema.safeParse(model).success;
+}
+
+export function isXAIModel(model: string) {
+  return XAIModelSchema.safeParse(model).success;
+}
+
 export const AvailableModels: { [name: string]: ModelSpec } = {
   // OPENAI / AZURE MODELS
   "gpt-4o": {