fix: remote engines should accept normalized headers only (#885)

Author: louis-jan

cortex-js/src/domain/abstracts/oai.abstract.ts

@@ -18,7 +18,9 @@ export abstract class OAIEngineExtension extends EngineExtension {
     createChatDto: any,
     headers: Record<string, string>,
   ): Promise<stream.Readable | any> {
-    const payload = this.transformPayload ? this.transformPayload(createChatDto) : createChatDto;
+    const payload = this.transformPayload
+      ? this.transformPayload(createChatDto)
+      : createChatDto;
     const { stream: isStream } = payload;
     const additionalHeaders = _.omit(headers, [
       'content-type',

cortex-js/src/infrastructure/controllers/chat.controller.ts

@@ -5,10 +5,8 @@ import { Response } from 'express';
 import { ApiOperation, ApiTags, ApiResponse } from '@nestjs/swagger';
 import { ChatCompletionResponseDto } from '../dtos/chat/chat-completion-response.dto';
 import { TelemetryUsecases } from '@/usecases/telemetry/telemetry.usecases';
-import {
-  EventName,
-  TelemetrySource,
-} from '@/domain/telemetry/telemetry.interface';
+import { EventName } from '@/domain/telemetry/telemetry.interface';
+import { extractCommonHeaders } from '@/utils/request';
 
 @ApiTags('Inference')
 @Controller('chat')
@@ -20,7 +18,8 @@ export class ChatController {
 
   @ApiOperation({
     summary: 'Create chat completion',
-    description: 'Creates a model response for the given conversation. The following parameters are not working for the `TensorRT-LLM` engine:\n- `frequency_penalty`\n- `presence_penalty`\n- `top_p`',
+    description:
+      'Creates a model response for the given conversation. The following parameters are not working for the `TensorRT-LLM` engine:\n- `frequency_penalty`\n- `presence_penalty`\n- `top_p`',
   })
   @HttpCode(200)
   @ApiResponse({
@@ -38,7 +37,7 @@ export class ChatController {
 
     if (stream) {
       this.chatService
-        .inference(createChatDto, headers)
+        .inference(createChatDto, extractCommonHeaders(headers))
         .then((stream) => {
           res.header('Content-Type', 'text/event-stream');
           stream.pipe(res);
@@ -49,7 +48,7 @@ export class ChatController {
     } else {
       res.header('Content-Type', 'application/json');
       this.chatService
-        .inference(createChatDto, headers)
+        .inference(createChatDto, extractCommonHeaders(headers))
         .then((response) => {
           res.json(response);
         })

cortex-js/src/utils/request.ts

@@ -0,0 +1,22 @@
+export function extractCommonHeaders(headers: any) {
+  const commonHeaders = [
+    'Content-Type',
+    'User-Agent',
+    'Accept',
+    'Authorization',
+    'Origin',
+    'Referer',
+    'Connection',
+  ];
+
+  const extractedHeaders: Record<string, string> = {};
+
+  for (const header of commonHeaders) {
+    const value = headers[header];
+    if (value) {
+      extractedHeaders[header] = value;
+    }
+  }
+
+  return extractedHeaders;
+}