diff --git a/apis/cloudflare/src/env.ts b/apis/cloudflare/src/env.ts index 6039a7f..7bf1d09 100644 --- a/apis/cloudflare/src/env.ts +++ b/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; diff --git a/apis/cloudflare/src/proxy.ts b/apis/cloudflare/src/proxy.ts index 84e8f89..5171a9a 100644 --- a/apis/cloudflare/src/proxy.ts +++ b/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( diff --git a/packages/proxy/edge/index.ts b/packages/proxy/edge/index.ts index 5a99230..b9c9a15 100644 --- a/packages/proxy/edge/index.ts +++ b/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; + } | { + 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 => { + 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, diff --git a/packages/proxy/schema/models.ts b/packages/proxy/schema/models.ts index d0292cf..ddc88b0 100644 --- a/packages/proxy/schema/models.ts +++ b/packages/proxy/schema/models.ts @@ -26,6 +26,130 @@ export const ModelSchema = z.object({ export type ModelSpec = z.infer; +// 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; + +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": {