Merge pull request #1407 from narengogi/chore/send-hook-results-in-streaming-chunks

send hook results when streaming response

Author: VisargD

src/handlers/handlerUtils.ts

@@ -1191,6 +1191,7 @@ export async function recursiveAfterRequestHookHandler(
     responseJson: mappedResponseJson,
     originalResponseJson,
   } = await responseHandler(
+    c,
     response,
     isStreamingMode,
     providerOption,
@@ -1200,7 +1201,8 @@ export async function recursiveAfterRequestHookHandler(
     gatewayParams,
     strictOpenAiCompliance,
     c.req.url,
-    areSyncHooksAvailable
+    areSyncHooksAvailable,
+    hookSpanId
   );
 
   const arhResponse = await afterRequestHookHandler(

src/handlers/responseHandlers.ts

@@ -18,6 +18,7 @@ import { HookSpan } from '../middlewares/hooks';
 import { env } from 'hono/adapter';
 import { OpenAIModelResponseJSONToStreamGenerator } from '../providers/open-ai-base/createModelResponse';
 import { anthropicMessagesJsonToStreamGenerator } from '../providers/anthropic-base/utils/streamGenerator';
+import { endpointStrings } from '../providers/types';
 
 /**
  * Handles various types of responses based on the specified parameters
@@ -35,6 +36,7 @@ import { anthropicMessagesJsonToStreamGenerator } from '../providers/anthropic-b
  * @returns {Promise<{response: Response, json?: any}>} - The mapped response.
  */
 export async function responseHandler(
+  c: Context,
   response: Response,
   streamingMode: boolean,
   providerOptions: Options,
@@ -44,7 +46,8 @@ export async function responseHandler(
   gatewayRequest: Params,
   strictOpenAiCompliance: boolean,
   gatewayRequestUrl: string,
-  areSyncHooksAvailable: boolean
+  areSyncHooksAvailable: boolean,
+  hookSpanId: string
 ): Promise<{
   response: Response;
   responseJson: Record<string, any> | null;
@@ -96,28 +99,31 @@ export async function responseHandler(
     responseTransformerFunction = undefined;
   }
 
-  if (
-    streamingMode &&
-    isSuccessStatusCode &&
-    isCacheHit &&
-    responseTransformerFunction
-  ) {
-    const streamingResponse = await handleJSONToStreamResponse(
-      response,
-      provider,
-      responseTransformerFunction
-    );
-    return { response: streamingResponse, responseJson: null };
-  }
   if (streamingMode && isSuccessStatusCode) {
+    const hooksManager = c.get('hooksManager');
+    const span = hooksManager.getSpan(hookSpanId) as HookSpan;
+    const hooksResult = span.getHooksResult();
+    if (isCacheHit && responseTransformerFunction) {
+      const streamingResponse = await handleJSONToStreamResponse(
+        response,
+        provider,
+        responseTransformerFunction,
+        strictOpenAiCompliance,
+        responseTransformer as endpointStrings,
+        hooksResult
+      );
+      return { response: streamingResponse, responseJson: null };
+    }
     return {
       response: handleStreamingMode(
         response,
         provider,
         responseTransformerFunction,
         requestURL,
         strictOpenAiCompliance,
-        gatewayRequest
+        gatewayRequest,
+        responseTransformer as endpointStrings,
+        hooksResult
       ),
       responseJson: null,
     };
@@ -264,7 +270,6 @@ export async function afterRequestHookHandler(
           ...response,
           status: 246,
           statusText: 'Hooks failed',
-          headers: response.headers,
         });
       }
       return response;

src/handlers/services/responseService.ts

@@ -81,6 +81,7 @@ export class ResponseService {
   }> {
     const url = this.context.requestURL;
     return await responseHandler(
+      this.context.honoContext,
       response,
       this.context.isStreaming,
       this.context.providerOption,
@@ -90,7 +91,8 @@ export class ResponseService {
       this.context.params,
       this.context.strictOpenAiCompliance,
       this.context.honoContext.req.url,
-      this.hooksService.areSyncHooksAvailable
+      this.hooksService.areSyncHooksAvailable,
+      this.hooksService.hookSpan?.id as string
     );
   }
 

src/handlers/streamHandler.ts

@@ -8,9 +8,11 @@ import {
   PRECONDITION_CHECK_FAILED_STATUS_CODE,
   GOOGLE_VERTEX_AI,
 } from '../globals';
+import { HookSpan } from '../middlewares/hooks';
 import { VertexLlamaChatCompleteStreamChunkTransform } from '../providers/google-vertex-ai/chatComplete';
 import { OpenAIChatCompleteResponse } from '../providers/openai/chatComplete';
 import { OpenAICompleteResponse } from '../providers/openai/complete';
+import { endpointStrings } from '../providers/types';
 import { Params } from '../types/requestBody';
 import { getStreamModeSplitPattern, type SplitPatternType } from '../utils';
 
@@ -24,6 +26,15 @@ function readUInt32BE(buffer: Uint8Array, offset: number) {
   ); // Ensure the result is an unsigned integer
 }
 
+const shouldSendHookResultChunk = (
+  strictOpenAiCompliance: boolean,
+  hooksResult: HookSpan['hooksResult']
+) => {
+  return (
+    !strictOpenAiCompliance && hooksResult?.beforeRequestHooksResult?.length > 0
+  );
+};
+
 function getPayloadFromAWSChunk(chunk: Uint8Array): string {
   const decoder = new TextDecoder();
   const chunkLength = readUInt32BE(chunk, 0);
@@ -292,7 +303,9 @@ export function handleStreamingMode(
   responseTransformer: Function | undefined,
   requestURL: string,
   strictOpenAiCompliance: boolean,
-  gatewayRequest: Params
+  gatewayRequest: Params,
+  fn: endpointStrings,
+  hooksResult: HookSpan['hooksResult']
 ): Response {
   const splitPattern = getStreamModeSplitPattern(proxyProvider, requestURL);
   // If the provider doesn't supply completion id,
@@ -311,6 +324,12 @@ export function handleStreamingMode(
   if (proxyProvider === BEDROCK) {
     (async () => {
       try {
+        if (shouldSendHookResultChunk(strictOpenAiCompliance, hooksResult)) {
+          const hookResultChunk = constructHookResultChunk(hooksResult, fn);
+          if (hookResultChunk) {
+            await writer.write(encoder.encode(hookResultChunk));
+          }
+        }
         for await (const chunk of readAWSStream(
           reader,
           responseTransformer,
@@ -337,6 +356,12 @@ export function handleStreamingMode(
   } else {
     (async () => {
       try {
+        if (shouldSendHookResultChunk(strictOpenAiCompliance, hooksResult)) {
+          const hookResultChunk = constructHookResultChunk(hooksResult, fn);
+          if (hookResultChunk) {
+            await writer.write(encoder.encode(hookResultChunk));
+          }
+        }
         for await (const chunk of readStream(
           reader,
           splitPattern,
@@ -389,7 +414,10 @@ export function handleStreamingMode(
 export async function handleJSONToStreamResponse(
   response: Response,
   provider: string,
-  responseTransformerFunction: Function
+  responseTransformerFunction: Function,
+  strictOpenAiCompliance: boolean,
+  fn: endpointStrings,
+  hooksResult: HookSpan['hooksResult']
 ): Promise<Response> {
   const { readable, writable } = new TransformStream();
   const writer = writable.getWriter();
@@ -403,6 +431,12 @@ export async function handleJSONToStreamResponse(
   ) {
     const generator = responseTransformerFunction(responseJSON, provider);
     (async () => {
+      if (shouldSendHookResultChunk(strictOpenAiCompliance, hooksResult)) {
+        const hookResultChunk = constructHookResultChunk(hooksResult, fn);
+        if (hookResultChunk) {
+          await writer.write(encoder.encode(hookResultChunk));
+        }
+      }
       while (true) {
         const chunk = generator.next();
         if (chunk.done) {
@@ -418,6 +452,12 @@ export async function handleJSONToStreamResponse(
       provider
     );
     (async () => {
+      if (shouldSendHookResultChunk(strictOpenAiCompliance, hooksResult)) {
+        const hookResultChunk = constructHookResultChunk(hooksResult, fn);
+        if (hookResultChunk) {
+          await writer.write(encoder.encode(hookResultChunk));
+        }
+      }
       for (const chunk of streamChunkArray) {
         await writer.write(encoder.encode(chunk));
       }
@@ -434,3 +474,21 @@ export async function handleJSONToStreamResponse(
     statusText: response.statusText,
   });
 }
+
+const constructHookResultChunk = (
+  hooksResult: HookSpan['hooksResult'],
+  fn: endpointStrings
+) => {
+  if (fn === 'messages') {
+    return `event: hook_results\ndata: ${JSON.stringify({
+      hook_results: {
+        before_request_hooks: hooksResult.beforeRequestHooksResult,
+      },
+    })}\n\n`;
+  }
+  return `data: ${JSON.stringify({
+    hook_results: {
+      before_request_hooks: hooksResult.beforeRequestHooksResult,
+    },
+  })}\n\n`;
+};