Merge pull request #1053 from Portkey-AI/feat/passthough-oai-embed-without-opening-body

feat: passthrough oai embed response without opening body

Author: VisargD

src/handlers/handlerUtils.ts

@@ -344,7 +344,7 @@ export async function tryPost(
 
   let mappedResponse: Response;
   let retryCount: number | undefined;
-  let originalResponseJson: Record<string, any> | undefined;
+  let originalResponseJson: Record<string, any> | null | undefined;
 
   let cacheKey: string | undefined;
   let { cacheMode, cacheMaxAge, cacheStatus } = getCacheOptions(
@@ -355,6 +355,12 @@ export async function tryPost(
   const requestOptions = c.get('requestOptions') ?? [];
   let transformedRequestBody: ReadableStream | FormData | Params = {};
   let fetchOptions: RequestInit = {};
+  const areSyncHooksAvailable = Boolean(
+    hooksManager.getHooksToExecute(hookSpan, [
+      'syncBeforeRequestHook',
+      'syncAfterRequestHook',
+    ]).length
+  );
 
   // before_request_hooks handler
   ({
@@ -471,7 +477,8 @@ export async function tryPost(
           isCacheHit,
           params,
           strictOpenAiCompliance,
-          c.req.url
+          c.req.url,
+          areSyncHooksAvailable
         ));
     }
 
@@ -1247,7 +1254,7 @@ export async function recursiveAfterRequestHookHandler(
   mappedResponse: Response;
   retryCount: number;
   createdAt: Date;
-  originalResponseJson?: Record<string, any>;
+  originalResponseJson?: Record<string, any> | null;
 }> {
   let response, retryCount, createdAt, executionTime, retrySkipped;
   const requestTimeout =
@@ -1287,6 +1294,17 @@ export async function recursiveAfterRequestHookHandler(
     retry?.useRetryAfterHeader || false
   ));
 
+  const hooksManager = c.get('hooksManager') as HooksManager;
+  const hookSpan = hooksManager.getSpan(hookSpanId) as HookSpan;
+  // Check if sync hooks are available
+  // This will be used to determine if we need to parse the response body or simply passthrough the response as is
+  const areSyncHooksAvailable = Boolean(
+    hooksManager.getHooksToExecute(hookSpan, [
+      'syncBeforeRequestHook',
+      'syncAfterRequestHook',
+    ]).length
+  );
+
   const {
     response: mappedResponse,
     responseJson: mappedResponseJson,
@@ -1300,7 +1318,8 @@ export async function recursiveAfterRequestHookHandler(
     false,
     gatewayParams,
     strictOpenAiCompliance,
-    c.req.url
+    c.req.url,
+    areSyncHooksAvailable
   );
 
   const arhResponse = await afterRequestHookHandler(

src/handlers/responseHandlers.ts

@@ -42,11 +42,12 @@ export async function responseHandler(
   isCacheHit: boolean = false,
   gatewayRequest: Params,
   strictOpenAiCompliance: boolean,
-  gatewayRequestUrl: string
+  gatewayRequestUrl: string,
+  areSyncHooksAvailable: boolean
 ): Promise<{
   response: Response;
   responseJson: Record<string, any> | null;
-  originalResponseJson?: Record<string, any>;
+  originalResponseJson?: Record<string, any> | null;
 }> {
   let responseTransformerFunction: Function | undefined;
   const responseContentType = response.headers?.get('content-type');
@@ -160,7 +161,8 @@ export async function responseHandler(
     responseTransformerFunction,
     strictOpenAiCompliance,
     gatewayRequestUrl,
-    gatewayRequest
+    gatewayRequest,
+    areSyncHooksAvailable
   );
 
   return {

src/handlers/streamHandler.ts

@@ -226,11 +226,12 @@ export async function handleNonStreamingMode(
   responseTransformer: Function | undefined,
   strictOpenAiCompliance: boolean,
   gatewayRequestUrl: string,
-  gatewayRequest: Params
+  gatewayRequest: Params,
+  areSyncHooksAvailable: boolean
 ): Promise<{
   response: Response;
-  json: Record<string, any>;
-  originalResponseBodyJson?: Record<string, any>;
+  json: Record<string, any> | null;
+  originalResponseBodyJson?: Record<string, any> | null;
 }> {
   // 408 is thrown whenever a request takes more than request_timeout to respond.
   // In that case, response thrown by gateway is already in OpenAI format.
@@ -244,7 +245,9 @@ export async function handleNonStreamingMode(
     return { response, json: await response.clone().json() };
   }
 
-  const originalResponseBodyJson: Record<string, any> = await response.json();
+  const isJsonParsingRequired = responseTransformer || areSyncHooksAvailable;
+  const originalResponseBodyJson: Record<string, any> | null =
+    isJsonParsingRequired ? await response.json() : null;
   let responseBodyJson = originalResponseBodyJson;
   if (responseTransformer) {
     responseBodyJson = responseTransformer(
@@ -255,6 +258,12 @@ export async function handleNonStreamingMode(
       gatewayRequestUrl,
       gatewayRequest
     );
+  } else if (!areSyncHooksAvailable) {
+    return {
+      response: new Response(response.body, response),
+      json: null,
+      originalResponseBodyJson,
+    };
   }
 
   return {

src/middlewares/hooks/index.ts

@@ -465,7 +465,7 @@ export class HooksManager {
       .join(', ');
   }
 
-  private getHooksToExecute(
+  public getHooksToExecute(
     span: HookSpan,
     eventTypePresets: string[]
   ): HookObject[] {

src/providers/openai/index.ts

@@ -76,7 +76,6 @@ const OpenAIConfig: ProviderConfigs = {
     // 'stream-complete': OpenAICompleteResponseTransform,
     chatComplete: OpenAIChatCompleteResponseTransform,
     // 'stream-chatComplete': OpenAIChatCompleteResponseTransform,
-    embed: OpenAIEmbedResponseTransform,
     imageGenerate: OpenAIImageGenerateResponseTransform,
     createSpeech: OpenAICreateSpeechResponseTransform,
     createTranscription: OpenAICreateTranscriptionResponseTransform,