From 53f2ca158124dca891080beec03f91d480bfc860 Mon Sep 17 00:00:00 2001
From: vishwasnavadak <vishwas@antstack.io>
Date: Mon, 18 Nov 2024 16:27:07 +0530
Subject: [PATCH] feat: add ai gateway config to openai

---
 .../Resources/ChatGPTRegularLambda.yaml       |  2 +
 .../Resources/ChatGPTStreamingFunction.yaml   |  2 +
 backend/src/ChatGPTRegularLambda/index.mjs    |  5 +-
 backend/src/ChatGPTStreamingLambda/index.mjs  | 53 +++++++++----------
 4 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/backend/infra_resources/Resources/ChatGPTRegularLambda.yaml b/backend/infra_resources/Resources/ChatGPTRegularLambda.yaml
index e3a7a48..161bf0f 100644
--- a/backend/infra_resources/Resources/ChatGPTRegularLambda.yaml
+++ b/backend/infra_resources/Resources/ChatGPTRegularLambda.yaml
@@ -19,3 +19,5 @@ Properties:
     Variables:
       API_KEY: "{{resolve:secretsmanager:CHAT_GPT:SecretString:API_KEY}}"
       ORG_ID: "{{resolve:secretsmanager:CHAT_GPT:SecretString:ORG_ID}}"
+      ACCOUNT_ID: "{{resolve:secretsmanager:CLOUDFLARE_AI_GATEWAY:SecretString:ACCOUNT_ID}}"
+      GATEWAY_ID: "{{resolve:secretsmanager:CLOUDFLARE_AI_GATEWAY:SecretString:GATEWAY_ID}}"
diff --git a/backend/infra_resources/Resources/ChatGPTStreamingFunction.yaml b/backend/infra_resources/Resources/ChatGPTStreamingFunction.yaml
index 51b1c68..7eb9175 100644
--- a/backend/infra_resources/Resources/ChatGPTStreamingFunction.yaml
+++ b/backend/infra_resources/Resources/ChatGPTStreamingFunction.yaml
@@ -12,3 +12,5 @@ Properties:
     Variables:
       API_KEY: "{{resolve:secretsmanager:CHAT_GPT:SecretString:API_KEY}}"
       ORG_ID: "{{resolve:secretsmanager:CHAT_GPT:SecretString:ORG_ID}}"
+      ACCOUNT_ID: "{{resolve:secretsmanager:CLOUDFLARE_AI_GATEWAY:SecretString:ACCOUNT_ID}}"
+      GATEWAY_ID: "{{resolve:secretsmanager:CLOUDFLARE_AI_GATEWAY:SecretString:GATEWAY_ID}}"
diff --git a/backend/src/ChatGPTRegularLambda/index.mjs b/backend/src/ChatGPTRegularLambda/index.mjs
index a4c50cb..778cfbc 100644
--- a/backend/src/ChatGPTRegularLambda/index.mjs
+++ b/backend/src/ChatGPTRegularLambda/index.mjs
@@ -1,9 +1,10 @@
 // Open AI Initialization
-import OpenAI from 'openai';
+import OpenAI from "openai";
 
 const openai = new OpenAI({
   organization: process.env.ORG_ID,
   apiKey: process.env.API_KEY,
+  baseURL: `https://gateway.ai.cloudflare.com/v1/${process.env.ACCOUNT_ID}/${process.env.GATEWAY_ID}/openai/`,
 });
 
 // Handler Function - Regular Lambda
@@ -11,7 +12,7 @@ export const handler = async (event) => {
   // Call ChatGPT API
   const completion = await openai.chat.completions.create({
     model: "gpt-3.5-turbo",
-    messages: [{"role": "user", "content": event.body}]
+    messages: [{ role: "user", content: event.body }],
   });
 
   // Send Response
diff --git a/backend/src/ChatGPTStreamingLambda/index.mjs b/backend/src/ChatGPTStreamingLambda/index.mjs
index b42e6f7..4ffb10f 100644
--- a/backend/src/ChatGPTStreamingLambda/index.mjs
+++ b/backend/src/ChatGPTStreamingLambda/index.mjs
@@ -1,41 +1,36 @@
 // Open AI Initialization
-import OpenAI from 'openai';
+import OpenAI from "openai";
 
 const openai = new OpenAI({
   organization: process.env.ORG_ID,
   apiKey: process.env.API_KEY,
+  baseURL: `https://gateway.ai.cloudflare.com/v1/${process.env.ACCOUNT_ID}/${process.env.GATEWAY_ID}/openai/`,
 });
 
 // Handler Function - Streaming Lambda
-export const handler = awslambda.streamifyResponse(
-  async (event, responseStream, context) => {
-    const httpResponseMetadata = {
-      statusCode: 200,
-      headers: {
-        "Content-Type": "text/html",
-        "Access-Control-Allow-Origin": "*",
-      },
-    };
+export const handler = awslambda.streamifyResponse(async (event, responseStream, context) => {
+  const httpResponseMetadata = {
+    statusCode: 200,
+    headers: {
+      "Content-Type": "text/html",
+      "Access-Control-Allow-Origin": "*",
+    },
+  };
 
-    // Stream Object
-    responseStream = awslambda.HttpResponseStream.from(
-      responseStream,
-      httpResponseMetadata
-    );
+  // Stream Object
+  responseStream = awslambda.HttpResponseStream.from(responseStream, httpResponseMetadata);
 
-    // Call ChatGPT API
-    const completion = await openai.chat.completions.create({
-      model: "gpt-3.5-turbo",
-      messages: [{"role": "user", "content": event.body}],
-      stream: true,
-    });
-    for await (const part of completion) {
-      if (Object.keys(part.choices[0].delta).includes("content")){
-        responseStream.write(part.choices[0].delta.content)
-      }
-      else{
-        responseStream.end()
-      }
+  // Call ChatGPT API
+  const completion = await openai.chat.completions.create({
+    model: "gpt-3.5-turbo",
+    messages: [{ role: "user", content: event.body }],
+    stream: true,
+  });
+  for await (const part of completion) {
+    if (Object.keys(part.choices[0].delta).includes("content")) {
+      responseStream.write(part.choices[0].delta.content);
+    } else {
+      responseStream.end();
     }
   }
-);
+});