Update Image Gen example with mock photoes (#48)

* fix: put proper mock images

* fix: add mock api

before, we were simply using context.sleep, which can be surprising for users. Now we call a mock endpoint.

examples/image-gen-with-workflow/app/api/mock-image-gen-endpoint/route.ts

@@ -0,0 +1,22 @@
+import { NextRequest } from "next/server"
+import { IMAGES, MOCK_WAIT_MS } from "utils/constants"
+import { ImageResponse, Prompt } from "utils/types"
+
+
+export const POST = async (request: NextRequest) => {
+  const params = (await request.json()) as { prompt: Prompt }
+
+  const prompt = params.prompt
+
+  await new Promise((r) => setTimeout(r, MOCK_WAIT_MS))
+  const response: ImageResponse = {
+    created: "mock",
+    data: [
+      {
+        prompt,
+        url: IMAGES[prompt]
+      }
+    ]
+  }
+  return new Response(JSON.stringify(response))
+}

examples/image-gen-with-workflow/app/api/regular/route.ts

@@ -11,8 +11,8 @@ export const maxDuration = 30
 import { NextRequest, NextResponse } from 'next/server'
 import { ratelimit, validateRequest } from 'utils/redis'
 import { getFetchParameters } from 'utils/request'
-import { CallPayload, ImageResponse, RedisEntry } from 'utils/types'
-import { PLACEHOLDER_IMAGE, PROMPTS, RATELIMIT_CODE } from 'utils/constants'
+import { CallPayload, ImageResponse, Prompt, RedisEntry } from 'utils/types'
+import { PROMPTS, RATELIMIT_CODE } from 'utils/constants'
 
 export const POST = async (request: NextRequest) => {
   // check the ratelimit
@@ -26,7 +26,7 @@ export const POST = async (request: NextRequest) => {
   const prompt = PROMPTS[promptIndex]
 
   // call Ideogram and record the time
-  const url = await makeRequest(prompt)
+  const url = await makeRequest(prompt, request.url)
   const time = performance.now() - t1
 
   // return the results in the same format as how Worklow saves
@@ -43,19 +43,12 @@ export const POST = async (request: NextRequest) => {
  * Calls Ideogram to get an image and returns its URL
  *
  * @param prompt prompt to use
+ * @param requestUrl url of the request. passed to getFetchParameters 
  * @returns image url
  */
-const makeRequest = async (prompt: string) => {
+const makeRequest = async (prompt: Prompt, requestUrl: string) => {
   // get parameters for fetch
-  const parameters = getFetchParameters(prompt)
-
-  if (!parameters) {
-    // Exists for development purposes.
-    // if the parameters are not present, return a mock image
-    // after waiting for 3 seconds.
-    await new Promise((r) => setTimeout(r, 3000))
-    return PLACEHOLDER_IMAGE
-  }
+  const parameters = getFetchParameters(prompt, requestUrl)
 
   // make the fetch request
   const response = await fetch(parameters.url, {
@@ -66,5 +59,6 @@ const makeRequest = async (prompt: string) => {
 
   // return the response
   const data = (await response.json()) as ImageResponse
+
   return data.data[0].url
 }

examples/image-gen-with-workflow/app/api/workflow-simple/route.ts

@@ -10,7 +10,7 @@ export const { POST } = serve<{ prompt: string }>(async (context) => {
 
   // make the call to Idogram through QStash
   const { body: result } = await context.call(
-    'call Ideogram',
+    'call image generation API',
     {
       url: 'https://api.ideogram.ai/generate',
       method: 'POST',

examples/image-gen-with-workflow/app/api/workflow/route.ts

@@ -14,7 +14,7 @@ import { waitUntil } from '@vercel/functions'
 import { ratelimit, redis, validateRequest } from 'utils/redis'
 import { getFetchParameters } from 'utils/request'
 import { CallPayload, ImageResponse, RedisEntry } from 'utils/types'
-import { PLACEHOLDER_IMAGE, PROMPTS, RATELIMIT_CODE } from 'utils/constants'
+import { PROMPTS, RATELIMIT_CODE } from 'utils/constants'
 
 // get key to store the time for each workflow run
 const getTimeKey = (key: string) => `time-${key}`
@@ -26,6 +26,7 @@ export const POST = async (request: NextRequest) => {
 
   // record the start time and run the workflow serve method
   const t1 = performance.now()
+
   const result = await serveMethod(request)
 
   // get the workflow run identifier header
@@ -68,32 +69,14 @@ const { POST: serveMethod } = serve<CallPayload>(async (context) => {
   const prompt = PROMPTS[payload.promptIndex]
 
   // get parameters for context.call
-  const parameters = getFetchParameters(prompt)
-  let result: ImageResponse
+  const parameters = getFetchParameters(prompt, context.url)
 
-  if (parameters) {
-    // if the parameters are present, make context.call request
-    // to call Ideogram through QStash
-    const { body } = await context.call(
-      'call Ideogram',
-      parameters
-    );
-    result = body as ImageResponse
-  } else {
-    // Exists for development purposes.
-    // if the parameters are not present, return a mock image
-    // after waiting for 2 seconds.
-    await context.sleep('mock call', 2)
-    result = {
-      created: '',
-      data: [
-        {
-          prompt,
-          url: PLACEHOLDER_IMAGE,
-        },
-      ],
-    }
-  }
+  // if the parameters are present, make context.call request
+  // to call Ideogram through QStash
+  const { body } = await context.call<ImageResponse>(
+    'call image generation API',
+    parameters
+  );
 
   await context.run('save results in redis', async () => {
     // get callKey from headers
@@ -108,7 +91,7 @@ const { POST: serveMethod } = serve<CallPayload>(async (context) => {
       callKey,
       {
         time: (await redis.get(getTimeKey(callKey))) ?? 0,
-        url: result.data[0].url,
+        url: body.data[0].url,
       },
       { ex: 120 },
     ) // expire in 120 seconds

examples/image-gen-with-workflow/components/call-workflow.tsx

@@ -120,7 +120,7 @@ export const { POST } = serve<{ prompt: string }>(async (context) => {
 
   // make the call to Idogram through QStash
   const { body: result } = await context.call(
-    'call Ideogram',
+    'call image generation API',
     {
       url: 'https://api.ideogram.ai/generate',
       method: 'POST',

examples/image-gen-with-workflow/package.json

@@ -14,7 +14,7 @@
     "@upstash/qstash": "^2.7.8",
     "@upstash/ratelimit": "^2.0.3",
     "@upstash/redis": "^1.34.0",
-    "@upstash/workflow": "latest",
+    "@upstash/workflow": "^0.2.2",
     "@vercel/functions": "^1.4.1",
     "clsx": "^2.1.1",
     "next": "14.2.13",

examples/image-gen-with-workflow/utils/constants.ts

@@ -1,11 +1,21 @@
+import { Prompt } from "./types"
+
 export const RATELIMIT_CODE = 429
 export const REDIS_PREFIX = 'llm-call'
-export const PLACEHOLDER_IMAGE =
-  'https://mintlify.s3-us-west-1.amazonaws.com/upstash/img/qstash/qstash-benefits.png'
+
 export const PROMPTS = [
   'A supersonic jet rising to the stars in 1980s propaganda posters style. For coloring, use a contrast between a calm white/blue and a striking red',
   'A futuristic city skyline at dusk, with towering skyscrapers and flying vehicles in the style of retro sci-fi art. Colors should feature deep purples, bright neon pinks, and glowing electric blues.',
   'A high-speed train racing through a futuristic city, inspired by cyberpunk aesthetics. Use a mix of metallic greys and dark purples, with neon accents lighting up the scene.',
   'A tranquil mountain village under a starry night sky, painted in the style of traditional Japanese woodblock prints with a modern touch. Use soft blues and greens for the landscape, with glowing golden stars in the sky.',
   'A group of astronauts exploring a distant planet, depicted in the vibrant, surreal style of 1970s space art.'
-]
+] as const
+
+export const IMAGES: Record<Prompt, string> = {
+  "A futuristic city skyline at dusk, with towering skyscrapers and flying vehicles in the style of retro sci-fi art. Colors should feature deep purples, bright neon pinks, and glowing electric blues.": "futuristic-city.png",
+  "A group of astronauts exploring a distant planet, depicted in the vibrant, surreal style of 1970s space art.": "astronauts.png",
+  "A high-speed train racing through a futuristic city, inspired by cyberpunk aesthetics. Use a mix of metallic greys and dark purples, with neon accents lighting up the scene.": "train.png",
+  "A supersonic jet rising to the stars in 1980s propaganda posters style. For coloring, use a contrast between a calm white/blue and a striking red": "supersonic-jet.png",
+  "A tranquil mountain village under a starry night sky, painted in the style of traditional Japanese woodblock prints with a modern touch. Use soft blues and greens for the landscape, with glowing golden stars in the sky.": "mountain-village.png"
+}
+export const MOCK_WAIT_MS = 7_000

examples/image-gen-with-workflow/utils/request.ts

@@ -1,8 +1,20 @@
 import { FetchParameters } from './types'
 
+/**
+ * Creates the call parameters for the iamge generation endpoint.
+ * 
+ * If the OpenAI credentials are set, returns OpenAI info.
+ * Otherwise, returns Ideogram credentials if they are set.
+ * Finally, returns the mock endpoint of not env vars are set.
+ * 
+ * @param prompt 
+ * @param requestUrl 
+ * @returns 
+ */
 export const getFetchParameters = (
   prompt: string,
-): FetchParameters | undefined => {
+  requestUrl: string,
+): FetchParameters => {
   if (process.env.OPENAI_API_KEY) {
     return {
       url: 'https://api.openai.com/v1/images/generations',
@@ -40,5 +52,21 @@ export const getFetchParameters = (
   }
 
   console.warn('No credential env var is set. Using placeholder.')
-  return
+  const mockRoute = "mock-image-gen-endpoint"
+  return {
+    url: process.env.UPSTASH_WORKFLOW_URL
+      ? `${process.env.UPSTASH_WORKFLOW_URL}/api/${mockRoute}`
+      : `${requestUrl.split("/").slice(0,-1).join("/")}/${mockRoute}`,
+    method: 'POST',
+    body: {
+      model: 'dall-e-2',
+      prompt,
+      n: 1,
+      size: '512x512',
+    },
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${process.env.OPENAI_API_KEY}`,
+    },
+  }
 }

examples/image-gen-with-workflow/utils/types.ts

@@ -1,3 +1,5 @@
+import { PROMPTS } from "./constants"
+
 export type CallInfo = {
   duration: number
   result: string
@@ -36,4 +38,6 @@ export type ImageResponse = IdeogramResponse | OpenAIResponse
 
 export type CallPayload = {
   promptIndex: number  
-}
+}
+
+export type Prompt = typeof PROMPTS[number]