Upgrade OpenAI image generation to GPT Image

README.md

@@ -77,6 +77,7 @@ Optional tools:
    S3_BUCKET_NAME=<S3 bucket name for uploaded pins>
 
    OPENAI_API_KEY=<Get from OpenAI>
+   OPENAI_IMAGE_MODEL=gpt-image-1
 
    NODE_ENV=development
 

package.json

@@ -35,7 +35,7 @@
     "lodash": "^4.17.21",
     "mongoose": "^5.12.12",
     "morgan": "~1.9.0",
-    "openai": "^4.16.1",
+    "openai": "^4.52.1",
     "passportbundle": "^1.1.1",
     "react": "^17.0.0",
     "react-dom": "^17.0.0",

server/controllers/post.ts

@@ -98,6 +98,7 @@ export const addPin = async (req: Request, res: genericResponseType) => {
 export const generateAIimage = async (req: Request, res: genericResponseType) => {
   const { userId } = getUserProfile(req.user as UserType);
   const { description } = req.body;
+  const OPENAI_IMAGE_MODEL = process.env.OPENAI_IMAGE_MODEL || 'gpt-image-1';
   try {
     const aiGeneratedByUser = await aiGenerated.find({ userId });
     if (!description.trim().length || aiGeneratedByUser.length >= 5) {
@@ -109,7 +110,7 @@ export const generateAIimage = async (req: Request, res: genericResponseType) =>
     });
     debug(`UserId -> ${userId} Generating AI image and title for -> ${description} from openAI`);
     const imageResponse = await openai.images.generate({
-      model: 'dall-e-3',
+      model: OPENAI_IMAGE_MODEL,
       prompt: description,
       n: 1,
       size: '1024x1024',
@@ -122,15 +123,29 @@ export const generateAIimage = async (req: Request, res: genericResponseType) =>
       }],
       max_tokens: 10,
     });
+    const [imageData] = imageResponse.data;
+    const imgURL = imageData?.b64_json
+      ? `data:image/png;base64,${imageData.b64_json}`
+      : imageData?.url;
+    if (!imgURL) {
+      throw new Error('OpenAI returned no image data');
+    }
     const { _id } = await aiGenerated.create({
       userId,
       description,
-      response: { imageResponse, titleResponse },
+      response: {
+        imageResponse: {
+          created: imageResponse.created,
+          hasUrl: Boolean(imageData?.url),
+          hasB64Json: Boolean(imageData?.b64_json),
+          model: OPENAI_IMAGE_MODEL,
+        },
+        titleResponse,
+      },
     });
-    const [linkObject] = imageResponse.data;
     const [titleObject] = titleResponse.choices;
 
-    res.json({ imgURL: linkObject.url, title: titleObject.message?.content?.trim().replace(/[".]/g, ''), _id });
+    res.json({ imgURL, title: titleObject.message?.content?.trim().replace(/[".]/g, ''), _id });
   } catch (error) {
     debug(`Error Generating AI image and title for -> UserId -> ${userId} and description: ${description}, Error: ${error}`);
     res.json({ imgURL: '', title: '', _id: null });

tests/server/controllers/post.test.ts

@@ -28,11 +28,11 @@ jest.mock('@aws-sdk/client-s3', () => {
 /* Mock open ai api */
 const mockOpenAiInstance = {
   images: {
-    generate: jest.fn(() => Promise.resolve({ data: [{ url: 'http:/stub-ai-image-url' }] })),
+    generate: jest.fn<Promise<unknown>, unknown[]>(() => Promise.resolve({ data: [{ url: 'http:/stub-ai-image-url' }] })),
   },
   chat: {
     completions: {
-      create: jest.fn(() => Promise.resolve({ choices: [{ message: { content: '["TEST-LABEL-A", "TEST-LABEL-B"]' } }] })),
+      create: jest.fn<Promise<unknown>, unknown[]>(() => Promise.resolve({ choices: [{ message: { content: '["TEST-LABEL-A", "TEST-LABEL-B"]' } }] })),
     },
   },
 };
@@ -382,14 +382,18 @@ describe('generating an AI image', () => {
   };
   beforeEach(() => {
     res = { json: jest.fn(), end: jest.fn() };
+    mockOpenAiInstance.images.generate.mockResolvedValue({ data: [{ url: 'http:/stub-ai-image-url' }] });
+    mockOpenAiInstance.chat.completions.create.mockResolvedValue({ choices: [{ message: { content: '["TEST-LABEL-A", "TEST-LABEL-B"]' } }] });
     aiGenerated.create = jest.fn().mockResolvedValue({ _id: 'stub_ai_mongoose_storage_ID' });
     aiGenerated.find = jest.fn().mockResolvedValue([1, 2, 3]);
+    delete process.env.OPENAI_IMAGE_MODEL;
   });
   afterEach(() => {
     jest.restoreAllMocks();
   });
 
-  test('will make request to openAI to generate a new image', async () => {
+  test('will make request to openAI to generate a new image with the configured GPT Image model', async () => {
+    process.env.OPENAI_IMAGE_MODEL = 'stub-image-model';
     const req = {
       user,
       body: {
@@ -398,7 +402,7 @@ describe('generating an AI image', () => {
     };
     await generateAIimage(req as genericRequest, res as unknown as Response);
     expect(mockOpenAiInstance.images.generate).toHaveBeenCalledWith({
-      model: 'dall-e-3',
+      model: 'stub-image-model',
       n: 1,
       prompt: 'open ai image creation prompt',
       size: '1024x1024',
@@ -418,6 +422,76 @@ describe('generating an AI image', () => {
     });
   });
 
+  test('will return a PNG data URI for a GPT Image base64 response without persisting the payload', async () => {
+    const stubImage = Buffer.from('stub-image').toString('base64');
+    mockOpenAiInstance.images.generate.mockResolvedValue({ data: [{ b64_json: stubImage }] });
+    const req = {
+      user,
+      body: {
+        description: 'open ai image creation prompt',
+      },
+    };
+    await generateAIimage(req as genericRequest, res as unknown as Response);
+    expect(mockOpenAiInstance.images.generate).toHaveBeenCalledWith({
+      model: 'gpt-image-1',
+      n: 1,
+      prompt: 'open ai image creation prompt',
+      size: '1024x1024',
+    });
+    expect(res.json).toHaveBeenCalledWith({
+      imgURL: `data:image/png;base64,${stubImage}`,
+      title: '[TEST-LABEL-A, TEST-LABEL-B]',
+      _id: 'stub_ai_mongoose_storage_ID',
+    });
+    expect(aiGenerated.create).toHaveBeenCalledWith({
+      userId: user._id,
+      description: 'open ai image creation prompt',
+      response: {
+        imageResponse: {
+          created: undefined,
+          hasUrl: false,
+          hasB64Json: true,
+          model: 'gpt-image-1',
+        },
+        titleResponse: { choices: [{ message: { content: '["TEST-LABEL-A", "TEST-LABEL-B"]' } }] },
+      },
+    });
+    const persistedAIRecord = JSON.stringify((aiGenerated.create as jest.Mock).mock.calls[0][0]);
+    expect(persistedAIRecord).not.toContain(stubImage);
+  });
+
+  test('will return an existing image URL if OpenAI provides a URL fallback', async () => {
+    const req = {
+      user,
+      body: {
+        description: 'open ai image creation prompt',
+      },
+    };
+    await generateAIimage(req as genericRequest, res as unknown as Response);
+    expect(res.json).toHaveBeenCalledWith({
+      imgURL: 'http:/stub-ai-image-url',
+      title: '[TEST-LABEL-A, TEST-LABEL-B]',
+      _id: 'stub_ai_mongoose_storage_ID',
+    });
+  });
+
+  test('will respond with error if OpenAI returns no image payload', async () => {
+    mockOpenAiInstance.images.generate.mockResolvedValue({ data: [{}] });
+    const req = {
+      user,
+      body: {
+        description: 'open ai image creation prompt',
+      },
+    };
+    await generateAIimage(req as genericRequest, res as unknown as Response);
+    expect(aiGenerated.create).not.toHaveBeenCalled();
+    expect(res.json).toHaveBeenCalledWith({
+      imgURL: '',
+      title: '',
+      _id: null,
+    });
+  });
+
   test('will end response if no prompt provided', async () => {
     const req = {
       user,