feat: add unofficial midjourney client

bin/scratch.ts

@@ -8,7 +8,7 @@ import restoreCursor from 'restore-cursor'
 // import { ProxycurlClient } from '../src/services/proxycurl-client.js'
 // import { WikipediaClient } from '../src/index.js'
 // import { PerigonClient } from '../src/index.js'
-import { FirecrawlClient } from '../src/index.js'
+// import { FirecrawlClient } from '../src/index.js'
 // import { ExaClient } from '../src/index.js'
 // import { DiffbotClient } from '../src/index.js'
 // import { WolframClient } from '../src/index.js'
@@ -16,6 +16,7 @@ import { FirecrawlClient } from '../src/index.js'
 //   createTwitterV2Client,
 //   TwitterClient
 // } from '../src/services/twitter/index.js'
+import { MidjourneyClient } from '../src/index.js'
 
 /**
  * Scratch pad for testing.
@@ -56,13 +57,13 @@ async function main() {
   // })
   // console.log(JSON.stringify(res, null, 2))
 
-  const firecrawl = new FirecrawlClient()
-  const res = await firecrawl.scrapeUrl({
-    url: 'https://www.bbc.com/news/articles/cp4475gwny1o'
-    // url: 'https://www.theguardian.com/technology/article/2024/jun/04/openai-google-ai-risks-letter'
-    // url: 'https://www.firecrawl.dev'
-  })
-  console.log(JSON.stringify(res, null, 2))
+  // const firecrawl = new FirecrawlClient()
+  // const res = await firecrawl.scrapeUrl({
+  //   url: 'https://www.bbc.com/news/articles/cp4475gwny1o'
+  //   // url: 'https://www.theguardian.com/technology/article/2024/jun/04/openai-google-ai-risks-letter'
+  //   // url: 'https://www.firecrawl.dev'
+  // })
+  // console.log(JSON.stringify(res, null, 2))
 
   // const exa = new ExaClient()
   // const res = await exa.search({
@@ -96,6 +97,12 @@ async function main() {
   //   query: 'open source AI agents'
   // })
   // console.log(res)
+
+  const midjourney = new MidjourneyClient()
+  const res = await midjourney.imagine(
+    'tiny lil baby kittens playing with an inquisitive AI robot, kawaii, anime'
+  )
+  console.log(JSON.stringify(res, null, 2))
 }
 
 try {

src/errors.ts

@@ -1,3 +1,5 @@
 export class RetryableError extends Error {}
 
 export class ParseError extends RetryableError {}
+
+export class TimeoutError extends Error {}

src/services/index.ts

@@ -3,6 +3,7 @@ export * from './dexa-client.js'
 export * from './diffbot-client.js'
 export * from './exa-client.js'
 export * from './firecrawl-client.js'
+export * from './midjourney-client.js'
 export * from './people-data-labs-client.js'
 export * from './perigon-client.js'
 export * from './predict-leads-client.js'

src/services/midjourney-client.ts

@@ -0,0 +1,139 @@
+import defaultKy, { type KyInstance } from 'ky'
+import { z } from 'zod'
+
+import { TimeoutError } from '../errors.js'
+import { aiFunction, AIFunctionsProvider } from '../fns.js'
+import { assert, delay, getEnv, pruneNullOrUndefined } from '../utils.js'
+
+export namespace midjourney {
+  export const API_BASE_URL = 'https://cl.imagineapi.dev'
+
+  export type JobStatus = 'pending' | 'in-progress' | 'completed' | 'failed'
+
+  export interface ImagineResponse {
+    data: Job
+  }
+
+  export interface Job {
+    id: string
+    prompt: string
+    status: JobStatus
+    user_created: string
+    date_created: string
+    results?: string
+    progress?: string
+    url?: string
+    error?: string
+    upscaled_urls?: string[]
+    ref?: string
+    upscaled?: string[]
+  }
+}
+
+/**
+ * Unofficial Midjourney API client.
+ *
+ * @see https://www.imagineapi.dev
+ */
+export class MidjourneyClient extends AIFunctionsProvider {
+  readonly ky: KyInstance
+  readonly apiKey: string
+  readonly apiBaseUrl: string
+
+  constructor({
+    apiKey = getEnv('MIDJOURNEY_IMAGINE_API_KEY'),
+    apiBaseUrl = midjourney.API_BASE_URL,
+    ky = defaultKy
+  }: {
+    apiKey?: string
+    apiBaseUrl?: string
+    ky?: KyInstance
+  } = {}) {
+    assert(
+      apiKey,
+      'MidjourneyClient missing required "apiKey" (defaults to "MIDJOURNEY_IMAGINE_API_KEY")'
+    )
+    super()
+
+    this.apiKey = apiKey
+    this.apiBaseUrl = apiBaseUrl
+
+    this.ky = ky.extend({
+      prefixUrl: apiBaseUrl,
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`
+      }
+    })
+  }
+
+  @aiFunction({
+    name: 'midjourney_create_images',
+    description:
+      'Creates 4 images from a prompt using the Midjourney API. Useful for generating images on the fly.',
+    inputSchema: z.object({
+      prompt: z
+        .string()
+        .describe(
+          'Simple, short, comma-separated list of phrases which describe the image you want to generate'
+        )
+    })
+  })
+  async imagine(
+    promptOrOptions: string | { prompt: string }
+  ): Promise<midjourney.Job> {
+    const options =
+      typeof promptOrOptions === 'string'
+        ? { prompt: promptOrOptions }
+        : promptOrOptions
+
+    const res = await this.ky
+      .post('items/images', {
+        json: { ...options }
+      })
+      .json<midjourney.ImagineResponse>()
+
+    return pruneNullOrUndefined(res.data)
+  }
+
+  async getJobById(jobId: string): Promise<midjourney.Job> {
+    const res = await this.ky
+      .get(`items/images/${jobId}`)
+      .json<midjourney.ImagineResponse>()
+
+    return pruneNullOrUndefined(res.data)
+  }
+
+  async waitForJobById(
+    jobId: string,
+    {
+      timeoutMs = 5 * 60 * 1000, // 5 minutes
+      intervalMs = 1000
+    }: {
+      timeoutMs?: number
+      intervalMs?: number
+    } = {}
+  ) {
+    const startTimeMs = Date.now()
+
+    function checkForTimeout() {
+      const elapsedTimeMs = Date.now() - startTimeMs
+      if (elapsedTimeMs >= timeoutMs) {
+        throw new TimeoutError(
+          `MidjourneyClient timeout waiting for job "${jobId}"`
+        )
+      }
+    }
+
+    do {
+      checkForTimeout()
+
+      const job = await this.getJobById(jobId)
+      if (job.status === 'completed' || job.status === 'failed') {
+        return job
+      }
+
+      checkForTimeout()
+      await delay(intervalMs)
+    } while (true)
+  }
+}

src/url-utils.test.ts

@@ -0,0 +1,31 @@
+import { describe, expect, test } from 'vitest'
+
+import { normalizeUrl } from './url-utils.js'
+
+describe('normalizeUrl', () => {
+  test('valid urls', async () => {
+    expect(normalizeUrl('https://www.google.com')).toBe(
+      'https://www.google.com'
+    )
+    expect(normalizeUrl('//www.google.com')).toBe('https://www.google.com')
+    expect(normalizeUrl('https://www.google.com/foo?')).toBe(
+      'https://www.google.com/foo'
+    )
+    expect(normalizeUrl('https://www.google.com/?foo=bar&dog=cat')).toBe(
+      'https://www.google.com/?dog=cat&foo=bar'
+    )
+    expect(normalizeUrl('https://google.com/abc/123//')).toBe(
+      'https://google.com/abc/123'
+    )
+  })
+
+  test('invalid urls', async () => {
+    expect(normalizeUrl('/foo')).toBe(null)
+    expect(normalizeUrl('/foo/bar/baz')).toBe(null)
+    expect(normalizeUrl('://foo.com')).toBe(null)
+    expect(normalizeUrl('foo')).toBe(null)
+    expect(normalizeUrl('')).toBe(null)
+    expect(normalizeUrl(undefined as unknown as string)).toBe(null)
+    expect(normalizeUrl(null as unknown as string)).toBe(null)
+  })
+})

src/utils.ts

@@ -57,6 +57,16 @@ export function pruneUndefined<T extends Record<string, any>>(
   ) as NonNullable<T>
 }
 
+export function pruneNullOrUndefined<T extends Record<string, any>>(
+  obj: T
+): NonNullable<{ [K in keyof T]: Exclude<T[K], undefined | null> }> {
+  return Object.fromEntries(
+    Object.entries(obj).filter(
+      ([, value]) => value !== undefined && value !== null
+    )
+  ) as NonNullable<T>
+}
+
 export function getEnv(name: string): string | undefined {
   try {
     return typeof process !== 'undefined'