From e5f67183f1a5bdd8e7966be168f2790b23872b9b Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 15 Jul 2024 15:22:13 +0700
Subject: [PATCH] fix: remote engines should accept normalized headers only

---
 .../src/domain/abstracts/oai.abstract.ts      |  4 +++-
 .../controllers/chat.controller.ts            | 13 +++++------
 cortex-js/src/utils/request.ts                | 22 +++++++++++++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)
 create mode 100644 cortex-js/src/utils/request.ts

diff --git a/cortex-js/src/domain/abstracts/oai.abstract.ts b/cortex-js/src/domain/abstracts/oai.abstract.ts
index 646afc921..3b75edfbf 100644
--- a/cortex-js/src/domain/abstracts/oai.abstract.ts
+++ b/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',
diff --git a/cortex-js/src/infrastructure/controllers/chat.controller.ts b/cortex-js/src/infrastructure/controllers/chat.controller.ts
index f61daca54..3a0aed579 100644
--- a/cortex-js/src/infrastructure/controllers/chat.controller.ts
+++ b/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);
         })
diff --git a/cortex-js/src/utils/request.ts b/cortex-js/src/utils/request.ts
new file mode 100644
index 000000000..e81c1b0de
--- /dev/null
+++ b/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;
+}