Changes (#133)

Co-authored-by: David Montague <35119617+dmontagu@users.noreply.github.com>

Author: Kludex

deploy/example.config.ts

@@ -47,35 +47,25 @@ export const config: Config<ProviderKeys> = {
       injectCost: true,
       // credentials are used by the ProviderProxy to authenticate the forwarded request
       credentials: env.OPENAI_API_KEY,
-      apiTypes: ['chat', 'responses'],
-    },
-    b: {
-      providerId: 'groq',
-      baseUrl: 'https://api.groq.com',
-      injectCost: true,
-      credentials: env.GROQ_API_KEY,
-      apiTypes: ['groq'],
     },
+    b: { providerId: 'groq', baseUrl: 'https://api.groq.com', injectCost: true, credentials: env.GROQ_API_KEY },
     c: {
       providerId: 'google-vertex',
       baseUrl: 'https://us-central1-aiplatform.googleapis.com',
       injectCost: true,
       credentials: env.GOOGLE_SERVICE_ACCOUNT_KEY,
-      apiTypes: ['gemini', 'anthropic'],
     },
     d: {
       providerId: 'anthropic',
       baseUrl: 'https://api.anthropic.com',
       injectCost: true,
       credentials: env.ANTHROPIC_API_KEY,
-      apiTypes: ['anthropic'],
     },
     e: {
       providerId: 'bedrock',
       baseUrl: 'https://bedrock-runtime.us-east-1.amazonaws.com',
       injectCost: true,
       credentials: env.AWS_BEARER_TOKEN_BEDROCK,
-      apiTypes: ['anthropic', 'converse'],
     },
   },
   // individual apiKeys

deploy/src/db.ts

@@ -10,11 +10,19 @@ export class ConfigDB extends KeysDbD1 {
     const project = config.projects[keyInfo.project]!
     const user = keyInfo.user ? project.users[keyInfo.user] : undefined
 
-    let providers: ProviderProxy[]
+    let providersWithKeys: (ProviderProxy & { key: string })[]
     if (keyInfo.providers === '__all__') {
-      providers = Object.values(config.providers)
+      providersWithKeys = Object.entries(config.providers).map(([key, provider]) => ({ ...provider, key }))
     } else {
-      providers = keyInfo.providers.map((name) => config.providers[name])
+      providersWithKeys = keyInfo.providers.map((key) => ({ ...config.providers[key], key }))
+    }
+
+    // Transform routes to routingGroups
+    const routingGroups: Record<string, { key: string }[]> = {}
+    if (config.routes) {
+      for (const [routeName, routeProviderKeys] of Object.entries(config.routes)) {
+        routingGroups[routeName] = routeProviderKeys.map((providerKey) => ({ key: providerKey }))
+      }
     }
 
     // if keyInfo.id is unset, hash the API key to give something unique without explicitly using the key directly
@@ -49,7 +57,8 @@ export class ConfigDB extends KeysDbD1 {
       userSpendingLimitDaily: user?.spendingLimitDaily,
       userSpendingLimitWeekly: user?.spendingLimitWeekly,
       userSpendingLimitMonthly: user?.spendingLimitMonthly,
-      providers,
+      providers: providersWithKeys,
+      routingGroups,
       otelSettings: user?.otel ?? project.otel,
     }
   }

deploy/src/types.ts

@@ -3,6 +3,7 @@ import type { OtelSettings, ProviderProxy } from '@pydantic/ai-gateway'
 export interface Config<ProviderKey extends string = string> {
   /** @param project: record keys are the project ids */
   projects: Record<number, Project>
+  routes?: Record<string, ProviderKey[]>
   providers: Record<ProviderKey, ProviderProxy>
   apiKeys: Record<string, ApiKey<ProviderKey>>
 }

deploy/test.config.ts

@@ -32,14 +32,12 @@ export const config: Config<ProviderKeys> = {
       injectCost: true,
       // credentials are used by the ProviderProxy to authenticate the forwarded request
       credentials: env.OPENAI_API_KEY,
-      apiTypes: ['chat'],
     },
     groq: {
       baseUrl: 'http://localhost:8005/groq',
       providerId: 'groq',
       injectCost: true,
       credentials: env.GROQ_API_KEY,
-      apiTypes: ['groq'],
     },
     // google: {
     //     baseUrl:
@@ -53,22 +51,14 @@ export const config: Config<ProviderKeys> = {
       providerId: 'anthropic',
       injectCost: true,
       credentials: env.ANTHROPIC_API_KEY,
-      apiTypes: ['anthropic'],
     },
     bedrock: {
       baseUrl: 'http://localhost:8005/bedrock',
       providerId: 'bedrock',
       injectCost: true,
       credentials: env.AWS_BEARER_TOKEN_BEDROCK,
-      apiTypes: ['anthropic', 'converse'],
-    },
-    test: {
-      baseUrl: 'http://test.example.com/test',
-      providerId: 'test',
-      injectCost: true,
-      credentials: 'test',
-      apiTypes: ['test'],
     },
+    test: { baseUrl: 'http://test.example.com/test', providerId: 'test', injectCost: true, credentials: 'test' },
   },
   // individual apiKeys
   apiKeys: {

deploy/test/index.spec.ts

@@ -72,7 +72,7 @@ describe('deploy', () => {
 
     const client = new OpenAI({
       apiKey: 'healthy-key',
-      baseURL: 'https://example.com/chat',
+      baseURL: 'https://example.com/openai',
       fetch: SELF.fetch.bind(SELF),
     })
 

deploy/test/index.spec.ts.snap

@@ -14,8 +14,8 @@ exports[`deploy > should call openai via gateway > llm 1`] = `
       },
     },
   ],
-  "created": 1762272055,
-  "id": "chatcmpl-CYDklwaN7x9okuWTnABMCrZykoiRj",
+  "created": 1762861820,
+  "id": "chatcmpl-CahB69PCs04fSZmb69YXvx65zQ0XE",
   "model": "gpt-5-2025-08-07",
   "object": "chat.completion",
   "service_tier": "default",
@@ -118,7 +118,7 @@ exports[`deploy > should call openai via gateway > span 1`] = `
   {
     "key": "gen_ai.response.id",
     "value": {
-      "stringValue": "chatcmpl-CYDklwaN7x9okuWTnABMCrZykoiRj",
+      "stringValue": "chatcmpl-CahB69PCs04fSZmb69YXvx65zQ0XE",
     },
   },
   {
@@ -293,7 +293,7 @@ exports[`deploy > should call openai via gateway > span 1`] = `
   {
     "key": "http.response.body.text",
     "value": {
-      "stringValue": "{"id":"chatcmpl-CYDklwaN7x9okuWTnABMCrZykoiRj","object":"chat.completion","created":1762272055,"model":"gpt-5-2025-08-07","choices":[{"index":0,"message":{"role":"assistant","content":"Paris.","refusal":null,"annotations":[]},"finish_reason":"stop"}],"usage":{"prompt_tokens":23,"completion_tokens":75,"total_tokens":98,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":64,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0},"pydantic_ai_gateway":{"cost_estimate":0.00077875}},"service_tier":"default","system_fingerprint":null}",
+      "stringValue": "{"id":"chatcmpl-CahB69PCs04fSZmb69YXvx65zQ0XE","object":"chat.completion","created":1762861820,"model":"gpt-5-2025-08-07","choices":[{"index":0,"message":{"role":"assistant","content":"Paris.","refusal":null,"annotations":[]},"finish_reason":"stop"}],"usage":{"prompt_tokens":23,"completion_tokens":75,"total_tokens":98,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":64,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0},"pydantic_ai_gateway":{"cost_estimate":0.00077875}},"service_tier":"default","system_fingerprint":null}",
     },
   },
   {

examples/pai_openai.py

@@ -26,7 +26,7 @@ def validate_dob(cls, v: date) -> date:
 
 
 person_agent = Agent(
-    'gateway/openai-chat:gpt-4.1-mini',
+    'gateway/openai:gpt-4.1-mini',
     output_type=Person,
     instructions='Extract information about the person',
 )

examples/pyproject.toml

@@ -11,5 +11,10 @@ dependencies = [
     "mypy-boto3-bedrock-runtime",
 ]
 
+# Install pydantic ai from git
+[tool.uv.sources]
+# Please don't remove this line, it's useful to test branches.
+pydantic-ai = { git = "https://github.com/pydantic/pydantic-ai.git", rev = '395cbe73' }
+
 [tool.uv]
 package = false

gateway/src/gateway.ts

@@ -5,8 +5,7 @@ import { currentScopeIntervals, type ExceededScope, endOfMonth, endOfWeek, type
 import { OtelTrace } from './otel'
 import { genAiOtelAttributes } from './otel/attributes'
 import { getProvider } from './providers'
-import type { APIType } from './types'
-import { type ApiKeyInfo, apiTypesArray, guardAPIType } from './types'
+import type { ApiKeyInfo, ProviderProxy } from './types'
 import { runAfter, textResponse } from './utils'
 
 export async function gateway(
@@ -15,23 +14,19 @@ export async function gateway(
   ctx: ExecutionContext,
   options: GatewayOptions,
 ): Promise<Response> {
-  const apiTypeMatch = /^\/([^/]+)\/(.*)$/.exec(proxyPath)
-  if (!apiTypeMatch) {
+  const routeMatch = /^\/([^/]+)\/(.*)$/.exec(proxyPath)
+  if (!routeMatch) {
     return textResponse(404, 'Path not found')
   }
-  let [, apiType, restOfPath] = apiTypeMatch as unknown as [string, string, string]
-
-  // support for other common names for openai api types
-  if (apiType === 'openai' || apiType === 'openai-chat') {
-    apiType = 'chat'
-  } else if (apiType === 'openai-responses') {
-    apiType = 'responses'
-  } else if (apiType === 'google-vertex') {
-    apiType = 'gemini'
-  }
-
-  if (!guardAPIType(apiType)) {
-    return textResponse(400, `Invalid API type '${apiType}', should be one of ${apiTypesArray.join(', ')}`)
+  let [, route, restOfPath] = routeMatch as unknown as [string, string, string]
+
+  // Backwards compatibility with the old route format.
+  if (route === 'openai-responses' || route === 'openai-chat' || route === 'chat' || route === 'responses') {
+    route = 'openai'
+  } else if (route === 'gemini') {
+    route = 'google-vertex'
+  } else if (route === 'converse') {
+    route = 'bedrock'
   }
 
   const rateLimiter = options.rateLimiter ?? noopLimiter
@@ -41,16 +36,43 @@ export async function gateway(
   }
   const apiKeyInfo = authResult
   try {
-    return await gatewayWithLimiter(request, restOfPath, apiType, apiKeyInfo, ctx, options)
+    return await gatewayWithLimiter(request, restOfPath, route, apiKeyInfo, ctx, options)
   } finally {
     runAfter(ctx, 'options.rateLimiter.requestFinish', rateLimiter.requestFinish())
   }
 }
 
+const getProviderProxies = (
+  route: string,
+  providerProxyMapping: Record<string, ProviderProxy>,
+  routingGroups: ApiKeyInfo['routingGroups'],
+): ProviderProxy[] | { status: number; message: string } => {
+  if (route in providerProxyMapping) {
+    return [providerProxyMapping[route]!]
+  }
+  const routingGroup = routingGroups?.[route]
+  if (!routingGroup) {
+    const supportedValues = [...new Set([...Object.keys(providerProxyMapping), ...Object.keys(routingGroups ?? {})])]
+      .sort()
+      .join(', ')
+    return { status: 404, message: `Route not found: ${route}. Supported values: ${supportedValues}` }
+  }
+  const providerProxies = routingGroup
+    .map(({ key }) => providerProxyMapping[key])
+    .filter((x): x is ProviderProxy & { key: string } => !!x)
+  if (providerProxies.length === 0) {
+    return {
+      status: 400,
+      message: `No providers included in routing group '${route}'. Add one or more providers to this routing group in the Pydantic AI Gateway console.`,
+    }
+  }
+  return providerProxies
+}
+
 export async function gatewayWithLimiter(
   request: Request,
   restOfPath: string,
-  apiType: APIType,
+  route: string,
   apiKeyInfo: ApiKeyInfo,
   ctx: ExecutionContext,
   options: GatewayOptions,
@@ -59,23 +81,13 @@ export async function gatewayWithLimiter(
     return textResponse(403, `Unauthorized - Key ${apiKeyInfo.status}`)
   }
 
-  let providerProxies = apiKeyInfo.providers.filter((p) => p.apiTypes.includes(apiType))
-
-  const routingGroup = request.headers.get('pydantic-ai-gateway-routing-group')
-  if (routingGroup !== null) {
-    providerProxies = providerProxies.filter((p) => p.routingGroup === routingGroup)
-  }
-
-  const profile = request.headers.get('pydantic-ai-gateway-profile')
-  if (profile !== null) {
-    providerProxies = providerProxies.filter((p) => p.profile === profile)
-  }
-
-  // sort providers on priority, highest first
-  providerProxies.sort((a, b) => (b.priority ?? 0) - (a.priority ?? 0))
-
-  if (providerProxies.length === 0) {
-    return textResponse(403, 'Forbidden - Provider not supported by this API Key')
+  const { routingGroups } = apiKeyInfo
+  const providerProxyMapping: Record<string, ProviderProxy> = Object.fromEntries(
+    apiKeyInfo.providers.map((p) => [p.key, p]),
+  )
+  const providerProxies = getProviderProxies(route, providerProxyMapping, routingGroups)
+  if (!Array.isArray(providerProxies)) {
+    return textResponse(providerProxies.status, providerProxies.message)
   }
 
   const otel = new OtelTrace(request, apiKeyInfo.otelSettings, options)
@@ -102,10 +114,7 @@ export async function gatewayWithLimiter(
     try {
       result = await proxy.dispatch()
     } catch (error) {
-      logfire.reportError('Connection error', error as Error, {
-        providerId: providerProxy.providerId,
-        routingGroup: providerProxy.routingGroup,
-      })
+      logfire.reportError('Connection error', error as Error, { providerId: providerProxy.providerId, route })
       continue
     }
 
@@ -120,7 +129,7 @@ export async function gatewayWithLimiter(
       logfire.info('Provider failed with retryable error, trying next provider', {
         providerId: providerProxy.providerId,
         status: result.unexpectedStatus,
-        routingGroup: providerProxy.routingGroup,
+        route,
       })
       continue
     }

gateway/src/types.ts

@@ -8,7 +8,7 @@ export type KeyStatus =
   | 'blocked' // when we got a valid response that we couldn't calculate the cost for
 
 // Info about an API key for a particular provider returned by the DB during a request
-export interface ApiKeyInfo {
+export interface ApiKeyInfo<ProviderKey extends string = string> {
   id: number
   user?: number
   project: number
@@ -30,29 +30,27 @@ export interface ApiKeyInfo {
   userSpendingLimitDaily?: number
   userSpendingLimitWeekly?: number
   userSpendingLimitMonthly?: number
-  providers: ProviderProxy[]
+  providers: (ProviderProxy & { key: ProviderKey })[]
+  routingGroups: Record<string, { key: ProviderKey }[]>
   otelSettings?: OtelSettings
 }
 
 export type ProviderID = 'groq' | 'openai' | 'google-vertex' | 'anthropic' | 'test' | 'bedrock'
 // TODO | 'azure' | 'fireworks' | 'mistral' | 'cohere'
 
-export type APIType = 'chat' | 'responses' | 'converse' | 'anthropic' | 'gemini' | 'groq' | 'test'
-
-const apiTypes: Record<APIType, boolean> = {
-  chat: true,
-  responses: true,
-  converse: true,
-  anthropic: true,
-  gemini: true,
+const providerIds: Record<ProviderID, boolean> = {
   groq: true,
+  openai: true,
+  'google-vertex': true,
+  anthropic: true,
   test: true,
+  bedrock: true,
 }
 
-export const apiTypesArray = Object.keys(apiTypes) as APIType[]
+export const providerIdsArray = Object.keys(providerIds) as ProviderID[]
 
-export function guardAPIType(type: string): type is APIType {
-  return type in apiTypes
+export function guardProviderID(id: string): id is ProviderID {
+  return id in providerIds
 }
 
 export interface ProviderProxy {
@@ -73,18 +71,12 @@ export interface ProviderProxy {
   profile?: string
 
   /** Higher priority providers will be used first */
+  // TODO(Marcelo): Remove now - this should live in the routingGroups.
   priority?: number
 
   /** Weather to disable the key in case of error, if missing defaults to True. */
   disableKey?: boolean
 
-  /** The APIs that the provider supports. Example: ['chat', 'responses'] */
-  apiTypes: APIType[]
-
-  /** A grouping of APIs that serve the same models.
-   * @example: 'anthropic' would route the requests to Anthropic, Bedrock and Vertex AI. */
-  routingGroup?: string
-
   /** Whether the provider is managed by the platform and not by the user. */
   isBuiltIn?: boolean
 }