🍀

readme.md

@@ -128,6 +128,7 @@ The SDK-specific imports are all isolated to keep the main `@agentic/stdlib` as
 - diffbot
 - exa
 - firecrawl (WIP)
+- midjourney
 - people data labs (WIP)
 - perigon
 - predict leads
@@ -169,7 +170,6 @@ The SDK-specific imports are all isolated to keep the main `@agentic/stdlib` as
   - replicate
   - huggingface
   - [skyvern](https://github.com/Skyvern-AI/skyvern)
-  - midjourney
   - unstructured
   - pull from [langchain](https://github.com/langchain-ai/langchainjs/tree/main/langchain)
     - provide a converter for langchain `DynamicStructuredTool`

src/services/midjourney-client.ts

@@ -5,6 +5,8 @@ import { TimeoutError } from '../errors.js'
 import { aiFunction, AIFunctionsProvider } from '../fns.js'
 import { assert, delay, getEnv, pruneNullOrUndefined } from '../utils.js'
 
+// TODO: add additional methods for upscaling, variations, etc.
+
 export namespace midjourney {
   export const API_BASE_URL = 'https://cl.imagineapi.dev'
 
@@ -28,6 +30,12 @@ export namespace midjourney {
     ref?: string
     upscaled?: string[]
   }
+
+  export interface JobOptions {
+    wait?: boolean
+    timeoutMs?: number
+    intervalMs?: number
+  }
 }
 
 /**
@@ -79,12 +87,22 @@ export class MidjourneyClient extends AIFunctionsProvider {
     })
   })
   async imagine(
-    promptOrOptions: string | { prompt: string }
+    promptOrOptions:
+      | string
+      | ({
+          prompt: string
+        } & midjourney.JobOptions)
   ): Promise<midjourney.Job> {
-    const options =
-      typeof promptOrOptions === 'string'
-        ? { prompt: promptOrOptions }
-        : promptOrOptions
+    const {
+      wait = true,
+      timeoutMs,
+      intervalMs,
+      ...options
+    } = typeof promptOrOptions === 'string'
+      ? ({ prompt: promptOrOptions } as {
+          prompt: string
+        } & midjourney.JobOptions)
+      : promptOrOptions
 
     const res = await this.ky
       .post('items/images', {
@@ -92,15 +110,56 @@ export class MidjourneyClient extends AIFunctionsProvider {
       })
       .json<midjourney.ImagineResponse>()
 
-    return pruneNullOrUndefined(res.data)
+    const job = pruneNullOrUndefined(res.data)
+    if (!wait) {
+      return job
+    }
+
+    if (job.status === 'completed' || job.status === 'failed') {
+      return job
+    }
+
+    return this.waitForJobById(job.id, {
+      timeoutMs,
+      intervalMs
+    })
   }
 
-  async getJobById(jobId: string): Promise<midjourney.Job> {
+  async getJobById(
+    jobIdOrOptions:
+      | string
+      | ({
+          jobId: string
+        } & midjourney.JobOptions)
+  ): Promise<midjourney.Job> {
+    const {
+      jobId,
+      wait = true,
+      timeoutMs,
+      intervalMs
+    } = typeof jobIdOrOptions === 'string'
+      ? ({ jobId: jobIdOrOptions } as {
+          jobId: string
+        } & midjourney.JobOptions)
+      : jobIdOrOptions
+
     const res = await this.ky
       .get(`items/images/${jobId}`)
       .json<midjourney.ImagineResponse>()
 
-    return pruneNullOrUndefined(res.data)
+    const job = pruneNullOrUndefined(res.data)
+    if (!wait) {
+      return job
+    }
+
+    if (job.status === 'completed' || job.status === 'failed') {
+      return job
+    }
+
+    return this.waitForJobById(job.id, {
+      timeoutMs,
+      intervalMs
+    })
   }
 
   async waitForJobById(
@@ -108,11 +167,8 @@ export class MidjourneyClient extends AIFunctionsProvider {
     {
       timeoutMs = 5 * 60 * 1000, // 5 minutes
       intervalMs = 1000
-    }: {
-      timeoutMs?: number
-      intervalMs?: number
-    } = {}
-  ) {
+    }: Omit<midjourney.JobOptions, 'wait'> = {}
+  ): Promise<midjourney.Job> {
     const startTimeMs = Date.now()
 
     function checkForTimeout() {