From b2d69b20b73347b8f0508c05e4ff2f605f3be546 Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Tue, 7 Jan 2025 00:13:33 -0800 Subject: [PATCH 1/4] Add API key support for various models in Cloudflare Worker configuration - Updated `Env` interface to include new API keys for OpenAI, Anthropic, Perplexity, Replicate, Fireworks, Google, XAI, Together, Lepton, Mistral, Ollama, GROQ, and Cerebras. - Enhanced `handleProxyV1` function to retrieve secrets based on model type, supporting new models and their respective API keys. - Modified `ProxyOpts` interface to include an `authConfig` for Cloudflare authentication. - Expanded model schemas to include new models for OpenAI, Anthropic, Bedrock, Perplexity, Together, Mistral, Groq, Fireworks, Google, and XAI. - Added support for embedding models in the AvailableModels list. --- apis/cloudflare/src/proxy.ts | 65 +++++++++++- apis/cloudflare/worker-configuration.d.ts | 19 ++++ packages/proxy/edge/index.ts | 23 +++- packages/proxy/schema/models.ts | 124 ++++++++++++++++++++++ 4 files changed, 228 insertions(+), 3 deletions(-) diff --git a/apis/cloudflare/src/proxy.ts b/apis/cloudflare/src/proxy.ts index 84e8f89..bb015f3 100644 --- a/apis/cloudflare/src/proxy.ts +++ b/apis/cloudflare/src/proxy.ts @@ -9,6 +9,8 @@ import { NOOP_METER_PROVIDER, initMetrics } from "@braintrust/proxy"; import { PrometheusMetricAggregator } from "./metric-aggregator"; import { handleRealtimeProxy } from "./realtime"; import { braintrustAppUrl } from "./env"; +import { isFireworksModel, isGoogleModel, isGroqModel, isMistralModel, isPerplexityModel, isXAIModel } from "@braintrust/proxy/schema"; +import { isAnthropicModel, isBedrockModel, isOpenAIModel } from "@braintrust/proxy/schema"; export const proxyV1Prefixes = ["/v1/proxy", "/v1"]; @@ -66,6 +68,68 @@ export async function handleProxyV1( const cache = await caches.open("apikey:cache"); const opts: ProxyOpts = { + authConfig: { + type: "cloudflare", + getSecret: async (model: string) => { + if (isOpenAIModel(model)) { + return { + secret: env.OPENAI_API_KEY, + type: "openai", + }; + } else if (isAnthropicModel(model)) { + return { + secret: env.ANTHROPIC_API_KEY, + type: "anthropic", + }; + } else if (isBedrockModel(model)) { + return { + secret: env.BEDROCK_SECRET_KEY, + type: "bedrock", + metadata: { + "region": env.BEDROCK_REGION, + "access_key": env.BEDROCK_ACCESS_KEY, + supportsStreaming: true, + }, + }; + } else if (isGroqModel(model)) { + return { + secret: env.GROQ_API_KEY, + type: "groq", + }; + } else if (isFireworksModel(model)) { + return { + secret: env.FIREWORKS_API_KEY, + type: "fireworks", + }; + } else if (isGoogleModel(model)) { + return { + secret: env.GOOGLE_API_KEY, + type: "google", + }; + } else if (isXAIModel(model)) { + return { + secret: env.XAI_API_KEY, + type: "xAI", + }; + } else if (isMistralModel(model)) { + return { + secret: env.MISTRAL_API_KEY, + type: "mistral", + }; + } else if (isPerplexityModel(model)) { + return { + secret: env.PERPLEXITY_API_KEY, + type: "perplexity", + }; + } + + throw new Error(`could not find secret for model ${model}`); + } + }, + // authConfig: { + // type: "braintrust", + // braintrustApiUrl: braintrustAppUrl(env).toString(), + // }, getRelativeURL(request: Request): string { return new URL(request.url).pathname.slice(proxyV1Prefix.length); }, @@ -111,7 +175,6 @@ export async function handleProxyV1( cacheSetLatency.record(end - start); }, }, - braintrustApiUrl: braintrustAppUrl(env).toString(), meterProvider, whitelist, }; diff --git a/apis/cloudflare/worker-configuration.d.ts b/apis/cloudflare/worker-configuration.d.ts index 43f5e68..f97c034 100644 --- a/apis/cloudflare/worker-configuration.d.ts +++ b/apis/cloudflare/worker-configuration.d.ts @@ -1,4 +1,23 @@ interface Env { + 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; + // Example binding to KV. Learn more at https://developers.cloudflare.com/workers/runtime-apis/kv/ // MY_KV_NAMESPACE: KVNamespace; // diff --git a/packages/proxy/edge/index.ts b/packages/proxy/edge/index.ts index 5a99230..30b6550 100644 --- a/packages/proxy/edge/index.ts +++ b/packages/proxy/edge/index.ts @@ -31,7 +31,13 @@ export interface ProxyOpts { cors?: boolean; credentialsCache?: Cache; completionsCache?: Cache; - braintrustApiUrl?: string; + authConfig?: { + type: "cloudflare"; + getSecret?: (model: string) => Promise; + } | { + type: "braintrust"; + braintrustApiUrl?: string; + } meterProvider?: MeterProvider; whitelist?: (string | RegExp)[]; } @@ -122,6 +128,19 @@ export function makeFetchApiSecrets({ model: string | null, org_name?: string, ): Promise => { + if (opts.authConfig?.type === "cloudflare") { + if (!opts.authConfig.getSecret) { + throw new Error("cloudflare auth config requires getSecret"); + } + + if (!model) { + throw new Error("cloudflare auth config requires model"); + } + + 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 +186,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: { 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": { From 0286892fa1f4712db94bbf0e1a81873596982486 Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Tue, 7 Jan 2025 00:37:40 -0800 Subject: [PATCH 2/4] master auth token --- apis/cloudflare/src/proxy.ts | 1 + apis/cloudflare/worker-configuration.d.ts | 2 ++ packages/proxy/edge/index.ts | 5 +++++ 3 files changed, 8 insertions(+) diff --git a/apis/cloudflare/src/proxy.ts b/apis/cloudflare/src/proxy.ts index bb015f3..19f7e4a 100644 --- a/apis/cloudflare/src/proxy.ts +++ b/apis/cloudflare/src/proxy.ts @@ -70,6 +70,7 @@ export async function handleProxyV1( const opts: ProxyOpts = { authConfig: { type: "cloudflare", + authToken: env.AUTH_TOKEN, getSecret: async (model: string) => { if (isOpenAIModel(model)) { return { diff --git a/apis/cloudflare/worker-configuration.d.ts b/apis/cloudflare/worker-configuration.d.ts index f97c034..35d7537 100644 --- a/apis/cloudflare/worker-configuration.d.ts +++ b/apis/cloudflare/worker-configuration.d.ts @@ -1,4 +1,6 @@ interface Env { + AUTH_TOKEN: string; + OPENAI_API_KEY: string; ANTHROPIC_API_KEY: string; PERPLEXITY_API_KEY: string; diff --git a/packages/proxy/edge/index.ts b/packages/proxy/edge/index.ts index 30b6550..8348260 100644 --- a/packages/proxy/edge/index.ts +++ b/packages/proxy/edge/index.ts @@ -33,6 +33,7 @@ export interface ProxyOpts { completionsCache?: Cache; authConfig?: { type: "cloudflare"; + authToken: string; getSecret?: (model: string) => Promise; } | { type: "braintrust"; @@ -129,6 +130,10 @@ export function makeFetchApiSecrets({ org_name?: string, ): Promise => { if (opts.authConfig?.type === "cloudflare") { + if (authToken !== opts.authConfig.authToken) { + throw new Error("Forbidden"); + } + if (!opts.authConfig.getSecret) { throw new Error("cloudflare auth config requires getSecret"); } From c1fc32a467459c8bda8f96e5d5749a79502e41b9 Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Tue, 7 Jan 2025 00:45:46 -0800 Subject: [PATCH 3/4] show both braintrust and cloudflare based auth --- apis/cloudflare/src/proxy.ts | 115 +++++++++++++++++++---------------- 1 file changed, 62 insertions(+), 53 deletions(-) diff --git a/apis/cloudflare/src/proxy.ts b/apis/cloudflare/src/proxy.ts index 19f7e4a..6b6175d 100644 --- a/apis/cloudflare/src/proxy.ts +++ b/apis/cloudflare/src/proxy.ts @@ -9,8 +9,7 @@ import { NOOP_METER_PROVIDER, initMetrics } from "@braintrust/proxy"; import { PrometheusMetricAggregator } from "./metric-aggregator"; import { handleRealtimeProxy } from "./realtime"; import { braintrustAppUrl } from "./env"; -import { isFireworksModel, isGoogleModel, isGroqModel, isMistralModel, isPerplexityModel, isXAIModel } from "@braintrust/proxy/schema"; -import { isAnthropicModel, isBedrockModel, isOpenAIModel } from "@braintrust/proxy/schema"; +import { isFireworksModel, isGoogleModel, isGroqModel, isMistralModel, isPerplexityModel, isXAIModel, isAnthropicModel, isBedrockModel, isOpenAIModel } from "@braintrust/proxy/schema"; export const proxyV1Prefixes = ["/v1/proxy", "/v1"]; @@ -68,6 +67,57 @@ export async function handleProxyV1( const cache = await caches.open("apikey:cache"); const opts: ProxyOpts = { + getRelativeURL(request: Request): string { + return new URL(request.url).pathname.slice(proxyV1Prefix.length); + }, + cors: true, + credentialsCache: { + async get(key: string): Promise { + const response = await cache.match(apiCacheKey(key)); + if (response) { + return (await response.json()) as T; + } else { + return null; + } + }, + async set(key: string, value: T, { ttl }: { ttl?: number }) { + await cache.put( + apiCacheKey(key), + new Response(JSON.stringify(value), { + headers: { + "Cache-Control": `public${ttl ? `, max-age=${ttl}}` : ""}`, + }, + }), + ); + }, + }, + completionsCache: { + get: async (key) => { + const start = performance.now(); + const ret = await env.ai_proxy.get(key); + const end = performance.now(); + cacheGetLatency.record(end - start); + if (ret) { + return JSON.parse(ret); + } else { + return null; + } + }, + set: async (key, value, { ttl }: { ttl?: number }) => { + const start = performance.now(); + await env.ai_proxy.put(key, JSON.stringify(value), { + expirationTtl: ttl, + }); + const end = performance.now(); + cacheSetLatency.record(end - start); + }, + }, + meterProvider, + whitelist, + }; + + const optsWithCloudflareAuth: ProxyOpts = { + ...opts, authConfig: { type: "cloudflare", authToken: env.AUTH_TOKEN, @@ -127,57 +177,14 @@ export async function handleProxyV1( throw new Error(`could not find secret for model ${model}`); } }, - // authConfig: { - // type: "braintrust", - // braintrustApiUrl: braintrustAppUrl(env).toString(), - // }, - getRelativeURL(request: Request): string { - return new URL(request.url).pathname.slice(proxyV1Prefix.length); - }, - cors: true, - credentialsCache: { - async get(key: string): Promise { - const response = await cache.match(apiCacheKey(key)); - if (response) { - return (await response.json()) as T; - } else { - return null; - } - }, - async set(key: string, value: T, { ttl }: { ttl?: number }) { - await cache.put( - apiCacheKey(key), - new Response(JSON.stringify(value), { - headers: { - "Cache-Control": `public${ttl ? `, max-age=${ttl}}` : ""}`, - }, - }), - ); - }, - }, - completionsCache: { - get: async (key) => { - const start = performance.now(); - const ret = await env.ai_proxy.get(key); - const end = performance.now(); - cacheGetLatency.record(end - start); - if (ret) { - return JSON.parse(ret); - } else { - return null; - } - }, - set: async (key, value, { ttl }: { ttl?: number }) => { - const start = performance.now(); - await env.ai_proxy.put(key, JSON.stringify(value), { - expirationTtl: ttl, - }); - const end = performance.now(); - cacheSetLatency.record(end - start); - }, + }; + + const optsWithBraintrustAuth: ProxyOpts = { + ...opts, + authConfig: { + type: "braintrust", + braintrustApiUrl: braintrustAppUrl(env).toString(), }, - meterProvider, - whitelist, }; const url = new URL(request.url); @@ -199,7 +206,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( From 9b24639c8411c987d4f66d523d3d2d1a73d207f2 Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Tue, 7 Jan 2025 02:45:16 -0800 Subject: [PATCH 4/4] refactor: use apiKey instead of authToken in proxyOptions --- apis/cloudflare/src/env.ts | 5 +- apis/cloudflare/src/proxy.ts | 60 +------------ apis/cloudflare/worker-configuration.d.ts | 21 ----- packages/proxy/edge/index.ts | 103 +++++++++++++++++++--- 4 files changed, 98 insertions(+), 91 deletions(-) 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 6b6175d..5171a9a 100644 --- a/apis/cloudflare/src/proxy.ts +++ b/apis/cloudflare/src/proxy.ts @@ -4,12 +4,12 @@ import { ProxyOpts, makeFetchApiSecrets, encryptedGet, + getApiSecret, } from "@braintrust/proxy/edge"; import { NOOP_METER_PROVIDER, initMetrics } from "@braintrust/proxy"; import { PrometheusMetricAggregator } from "./metric-aggregator"; import { handleRealtimeProxy } from "./realtime"; import { braintrustAppUrl } from "./env"; -import { isFireworksModel, isGoogleModel, isGroqModel, isMistralModel, isPerplexityModel, isXAIModel, isAnthropicModel, isBedrockModel, isOpenAIModel } from "@braintrust/proxy/schema"; export const proxyV1Prefixes = ["/v1/proxy", "/v1"]; @@ -120,62 +120,8 @@ export async function handleProxyV1( ...opts, authConfig: { type: "cloudflare", - authToken: env.AUTH_TOKEN, - getSecret: async (model: string) => { - if (isOpenAIModel(model)) { - return { - secret: env.OPENAI_API_KEY, - type: "openai", - }; - } else if (isAnthropicModel(model)) { - return { - secret: env.ANTHROPIC_API_KEY, - type: "anthropic", - }; - } else if (isBedrockModel(model)) { - return { - secret: env.BEDROCK_SECRET_KEY, - type: "bedrock", - metadata: { - "region": env.BEDROCK_REGION, - "access_key": env.BEDROCK_ACCESS_KEY, - supportsStreaming: true, - }, - }; - } else if (isGroqModel(model)) { - return { - secret: env.GROQ_API_KEY, - type: "groq", - }; - } else if (isFireworksModel(model)) { - return { - secret: env.FIREWORKS_API_KEY, - type: "fireworks", - }; - } else if (isGoogleModel(model)) { - return { - secret: env.GOOGLE_API_KEY, - type: "google", - }; - } else if (isXAIModel(model)) { - return { - secret: env.XAI_API_KEY, - type: "xAI", - }; - } else if (isMistralModel(model)) { - return { - secret: env.MISTRAL_API_KEY, - type: "mistral", - }; - } else if (isPerplexityModel(model)) { - return { - secret: env.PERPLEXITY_API_KEY, - type: "perplexity", - }; - } - - throw new Error(`could not find secret for model ${model}`); - } + apiKey: env.API_KEY, + getSecret: (model: string) => getApiSecret(model, env), }, }; diff --git a/apis/cloudflare/worker-configuration.d.ts b/apis/cloudflare/worker-configuration.d.ts index 35d7537..43f5e68 100644 --- a/apis/cloudflare/worker-configuration.d.ts +++ b/apis/cloudflare/worker-configuration.d.ts @@ -1,25 +1,4 @@ interface Env { - AUTH_TOKEN: string; - - 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; - // Example binding to KV. Learn more at https://developers.cloudflare.com/workers/runtime-apis/kv/ // MY_KV_NAMESPACE: KVNamespace; // diff --git a/packages/proxy/edge/index.ts b/packages/proxy/edge/index.ts index 8348260..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, @@ -33,8 +33,14 @@ export interface ProxyOpts { completionsCache?: Cache; authConfig?: { type: "cloudflare"; - authToken: string; - getSecret?: (model: string) => Promise; + /** + * 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; @@ -74,9 +80,9 @@ export function getCorsHeaders( return origin ? { - "access-control-allow-origin": origin, - ...baseCorsHeaders, - } + "access-control-allow-origin": origin, + ...baseCorsHeaders, + } : {}; } @@ -130,16 +136,12 @@ export function makeFetchApiSecrets({ org_name?: string, ): Promise => { if (opts.authConfig?.type === "cloudflare") { - if (authToken !== opts.authConfig.authToken) { + if (authToken !== opts.authConfig.apiKey) { throw new Error("Forbidden"); } - if (!opts.authConfig.getSecret) { - throw new Error("cloudflare auth config requires getSecret"); - } - if (!model) { - throw new Error("cloudflare auth config requires model"); + throw new Error("no model provided"); } const secret = await opts.authConfig.getSecret(model); @@ -354,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,