Merge pull request #1163 from Portkey-AI/feat/circuit_breaker

feat: base integration for circuit breaker

Author: VisargD

src/handlers/handlerUtils.ts

@@ -598,6 +598,7 @@ export async function tryTargetsRecursively(
 
   // start: merge inherited config with current target config (preference given to current)
   const currentInheritedConfig: Record<string, any> = {
+    id: inheritedConfig.id || currentTarget.id,
     overrideParams: {
       ...inheritedConfig.overrideParams,
       ...currentTarget.overrideParams,
@@ -751,19 +752,34 @@ export async function tryTargetsRecursively(
   ];
   // end: merge inherited config with current target config (preference given to current)
 
+  const isHandlingCircuitBreaker = currentInheritedConfig.id;
+  if (isHandlingCircuitBreaker) {
+    const healthyTargets = (currentTarget.targets || [])
+      .map((t: any, index: number) => ({
+        ...t,
+        originalIndex: index,
+      }))
+      .filter((t: any) => !t.isOpen);
+
+    if (healthyTargets.length) {
+      currentTarget.targets = healthyTargets;
+    }
+  }
+
   let response;
 
   switch (strategyMode) {
     case StrategyModes.FALLBACK:
       for (const [index, target] of currentTarget.targets.entries()) {
+        const originalIndex = target.originalIndex || index;
         response = await tryTargetsRecursively(
           c,
           target,
           request,
           requestHeaders,
           fn,
           method,
-          `${currentJsonPath}.targets[${index}]`,
+          `${currentJsonPath}.targets[${originalIndex}]`,
           currentInheritedConfig
         );
         if (response?.headers.get('x-portkey-gateway-exception') === 'true') {
@@ -791,8 +807,9 @@ export async function tryTargetsRecursively(
 
       let randomWeight = Math.random() * totalWeight;
       for (const [index, provider] of currentTarget.targets.entries()) {
+        const originalIndex = provider.originalIndex || index;
         if (randomWeight < provider.weight) {
-          currentJsonPath = currentJsonPath + `.targets[${index}]`;
+          currentJsonPath = currentJsonPath + `.targets[${originalIndex}]`;
           response = await tryTargetsRecursively(
             c,
             provider,
@@ -836,28 +853,30 @@ export async function tryTargetsRecursively(
         throw new RouterError(conditionalRouter.message);
       }
 
+      const originalIndex = finalTarget.originalIndex || finalTarget.index;
       response = await tryTargetsRecursively(
         c,
         finalTarget,
         request,
         requestHeaders,
         fn,
         method,
-        `${currentJsonPath}.targets[${finalTarget.index}]`,
+        `${currentJsonPath}.targets[${originalIndex}]`,
         currentInheritedConfig
       );
       break;
     }
 
     case StrategyModes.SINGLE:
+      const originalIndex = currentTarget.targets[0].originalIndex || 0;
       response = await tryTargetsRecursively(
         c,
         currentTarget.targets[0],
         request,
         requestHeaders,
         fn,
         method,
-        `${currentJsonPath}.targets[0]`,
+        `${currentJsonPath}.targets[${originalIndex}]`,
         currentInheritedConfig
       );
       break;
@@ -873,6 +892,15 @@ export async function tryTargetsRecursively(
           currentJsonPath,
           method
         );
+        if (isHandlingCircuitBreaker) {
+          await c.get('handleCircuitBreakerResponse')?.(
+            response,
+            currentInheritedConfig.id,
+            currentTarget.cbConfig,
+            currentJsonPath,
+            c
+          );
+        }
       } catch (error: any) {
         // tryPost always returns a Response.
         // TypeError will check for all unhandled exceptions.
@@ -898,6 +926,15 @@ export async function tryTargetsRecursively(
           );
         } else {
           response = error.response;
+          if (isHandlingCircuitBreaker) {
+            await c.get('recordCircuitBreakerFailure')?.(
+              env(c),
+              currentInheritedConfig.id,
+              currentTarget.cbConfig,
+              currentJsonPath,
+              response.status
+            );
+          }
         }
       }
       break;
@@ -1207,6 +1244,7 @@ export function constructConfigFromRequestHeaders(
       'output_guardrails',
       'default_input_guardrails',
       'default_output_guardrails',
+      'cb_config',
     ]) as any;
   }
 

src/types/requestBody.ts

@@ -197,6 +197,7 @@ export interface Targets {
 
   defaultInputGuardrails?: HookObject[];
   defaultOutputGuardrails?: HookObject[];
+  originalIndex?: number;
 }
 
 /**