tmp

platformatic · Apr 3, 2024 · 4672742 · 4672742
1 parent 3716f9b
commit 4672742
Show file tree

Hide file tree

Showing 14 changed files with 395 additions and 108 deletions.
diff --git a/ai-providers/open-ai.ts b/ai-providers/open-ai.ts
@@ -120,6 +120,6 @@ export class OpenAiProvider implements AiProvider {
       ],
       stream: true
     })
-    return new ReadableStream(new OpenAiByteSource(response.toReadableStream()))
+    return new ReadableStream(new OpenAiByteSource(response.toReadableStream(), chunkCallback))
   }
 }
diff --git a/index.d.ts b/index.d.ts
@@ -1,5 +1,6 @@
 import { ReadableStream } from 'node:stream/web'
 import { PlatformaticApp } from '@platformatic/service'
+import { errorResponseBuilderContext } from '@fastify/rate-limit'
 import { AiWarpConfig } from './config'
 
 declare module 'fastify' {
@@ -21,8 +22,9 @@ declare module 'fastify' {
         ) => object
         onExceeding?: (req: FastifyRequest, key: string) => void
         onExceeded?: (req: FastifyRequest, key: string) => void
-      }
-    }
+      },
+    },
+    rateLimitMax?: ((req: FastifyRequest, key: string) => number) | ((req: FastifyRequest, key: string) => Promise<number>)
   }
 }
 

diff --git a/index.ts b/index.ts
@@ -26,9 +26,12 @@ const stackable: Stackable<AiWarpConfig> = async function (fastify, opts) {
   const { rateLimiting: rateLimitingConfig } = config
   await fastify.register(fastifyRateLimit, {
     max: async (req, key) => {
+      console.log('stackable max callback')
       if (rateLimiting.max !== undefined) {
+        console.log(' user land callback defined')
         return await rateLimiting.max(req, key)
       } else {
+        console.log(' user land callback undefined')
         return rateLimitingConfig?.max ?? 1000
       }
     },
@@ -79,6 +82,8 @@ const stackable: Stackable<AiWarpConfig> = async function (fastify, opts) {
   if (!dontRegisterPlatformaticService) {
     await fastify.register(platformaticService, opts)
   }
+  console.log('stackable registered')
+
 }
 
 stackable.configType = 'ai-warp-app'

diff --git a/tests/e2e/api.test.ts b/tests/e2e/api.test.ts
@@ -1,10 +1,11 @@
-import { before, after, beforeEach, describe, it } from 'node:test'
+/* eslint-disable @typescript-eslint/no-floating-promises */
+import { before, after, describe, it } from 'node:test'
 import assert from 'node:assert'
-import { buildServer } from '@platformatic/service'
-import stackable, { setDontRegisterPlatformaticService } from '../../index'
 import { FastifyInstance } from 'fastify'
+import fastifyPlugin from 'fastify-plugin'
 import { MOCK_CONTENT_RESPONSE, buildExpectedStreamBodyString, mockMistralApi, mockOpenAiApi } from '../utils/mocks'
 import { AiWarpConfig } from '../../config'
+import { buildAiWarpApp } from '../utils/stackable'
 
 const expectedStreamBody = buildExpectedStreamBodyString()
 
@@ -34,45 +35,32 @@ const providers: Provider[] = [
   }
 ]
 
-let i = 0
+mockOpenAiApi()
+mockMistralApi()
+
+// Test the prompt and stream endpoint for each provider
 for (const { name, config } of providers) {
   describe(name, () => {
-    const port = 3042 + i
-    i++
-
     let app: FastifyInstance
+    let port: number
+    let chunkCallbackCalled = false
     before(async () => {
-      app = await buildServer({
-        server: {
-          port,
-          forceCloseConnections: true,
-          healthCheck: {
-            enabled: false
-          },
-          logger: {
-            level: 'error'
-          }
-        }
-      })
+      [app, port] = await buildAiWarpApp({ aiProvider: config })
 
-      app.platformatic.config.aiProvider = config
-
-      setDontRegisterPlatformaticService(true)
-      app.register(stackable)
+      await app.register(fastifyPlugin(async () => {
+        app.ai.preResponseChunkCallback = (_, response) => {
+          chunkCallbackCalled = true
+          return response
+        }
+      }))
 
-      // @ts-expect-error
       await app.start()
     })
 
     after(async () => {
       await app.close()
     })
 
-    beforeEach(async () => {
-      mockOpenAiApi()
-      mockMistralApi()
-    })
-
     it('/api/v1/prompt returns expected response', async () => {
       const res = await fetch(`http://localhost:${port}/api/v1/prompt`, {
         method: 'POST',
@@ -90,6 +78,8 @@ for (const { name, config } of providers) {
     })
 
     it('/api/v1/stream returns expected response', async () => {
+      assert.strictEqual(chunkCallbackCalled, false)
+
       const res = await fetch(`http://localhost:${port}/api/v1/stream`, {
         method: 'POST',
         headers: {
@@ -101,6 +91,8 @@ for (const { name, config } of providers) {
       })
       assert.strictEqual(res.headers.get('content-type'), 'text/event-stream')
 
+      assert.strictEqual(chunkCallbackCalled, true)
+
       assert.notStrictEqual(res.body, undefined)
 
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
@@ -120,5 +112,42 @@ for (const { name, config } of providers) {
       assert.strictEqual(body, expectedStreamBody)
     })
   })
-  i++
 }
+
+it('calls the preResponseCallback', async () => {
+  const [app, port] = await buildAiWarpApp({
+    aiProvider: {
+      openai: {
+        model: 'gpt-3.5-turbo',
+        apiKey: ''
+      }
+    }
+  })
+
+  let callbackCalled = false
+  await app.register(fastifyPlugin(async () => {
+    app.ai.preResponseCallback = (_, response) => {
+      callbackCalled = true
+      return response + ' modified'
+    }
+  }))
+
+  await app.start()
+
+  const res = await fetch(`http://localhost:${port}/api/v1/prompt`, {
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json'
+    },
+    body: JSON.stringify({
+      prompt: 'asd'
+    })
+  })
+
+  assert.strictEqual(callbackCalled, true)
+
+  const body = await res.json()
+  assert.strictEqual(body.response, `${MOCK_CONTENT_RESPONSE} modified`)
+
+  await app.close()
+})
diff --git a/tests/e2e/auth.test.ts b/tests/e2e/auth.test.ts
@@ -0,0 +1,30 @@
+/* eslint-disable @typescript-eslint/no-floating-promises */
+import { before, after, beforeEach, describe, it } from 'node:test'
+import assert from 'node:assert'
+import { buildAiWarpApp } from '../utils/stackable';
+import { AiWarpConfig } from '../../config';
+
+const aiProvider: AiWarpConfig['aiProvider'] = {
+    openai: {
+        model: 'gpt-3.5-turbo',
+        apiKey: ''
+    }
+}
+
+it('returns 401 for unauthorized user', async () => {
+  const [app, port] = await buildAiWarpApp({
+    aiProvider,
+    auth: {
+      required: true
+    }
+  })
+
+  try {
+    await app.start()
+
+    const response = await fetch(`http://localhost:${port}`);
+    assert.strictEqual(response.status, 401)
+  } finally {
+    await app.close()
+  }
+})
diff --git a/tests/e2e/index.ts b/tests/e2e/index.ts
@@ -1 +1,3 @@
 import './api.test'
+import './rate-limiting.test'
+import './auth.test'
diff --git a/tests/e2e/rate-limiting.test.ts b/tests/e2e/rate-limiting.test.ts
@@ -0,0 +1,121 @@
+/* eslint-disable @typescript-eslint/no-floating-promises */
+import { it } from 'node:test'
+import assert from 'node:assert'
+import fastifyPlugin from 'fastify-plugin'
+import { AiWarpConfig } from '../../config'
+import { buildAiWarpApp } from '../utils/stackable'
+
+const aiProvider: AiWarpConfig['aiProvider'] = {
+  openai: {
+    model: 'gpt-3.5-turbo',
+    apiKey: ''
+  }
+}
+
+it('calls ai.rateLimiting.max callback', async () => {
+  const [app, port] = await buildAiWarpApp({ aiProvider })
+
+  try {
+    const expectedMax = 100
+    let callbackCalled = false
+    await app.register(fastifyPlugin(async () => {
+      app.ai.rateLimiting.max = () => {
+        console.log('user land rate limit max callback')
+        callbackCalled = true
+        return expectedMax
+      }
+    }))
+
+    await app.start()
+
+    const res = await fetch(`http://localhost:${port}`)
+    assert.strictEqual(callbackCalled, true)
+    assert.strictEqual(res.headers.get('x-ratelimit-limit'), expectedMax)
+  } finally {
+    await app.close()
+  }
+})
+
+// it('calls ai.rateLimiting.allowList callback', async () => {
+//   const [app, port] = await buildAiWarpApp({ aiProvider })
+
+//   let callbackCalled = false
+//   app.register(fastifyPlugin(async () => {
+//     app.ai.rateLimiting.allowList = () => {
+//       callbackCalled = true
+//       return true
+//     }
+//   }))
+
+//   await app.start()
+
+//   await fetch(`http://localhost:${port}`)
+//   assert.strictEqual(callbackCalled, true)
+
+//   await app.close()
+// })
+
+// it('calls ai.rateLimiting.onBanReach callback', async () => {
+//   const [app, port] = await buildAiWarpApp({ aiProvider })
+
+//   let onBanReachCalled = false
+//   let errorResponseBuilderCalled = false
+//   app.register(fastifyPlugin(async () => {
+//     app.ai.rateLimiting.max = () => 0
+
+//     app.ai.rateLimiting.onBanReach = () => {
+//       onBanReachCalled = true
+//     }
+
+//     app.ai.rateLimiting.errorResponseBuilder = () => {
+//       errorResponseBuilderCalled = true
+//       return { error: 'rate limited' }
+//     }
+//   }))
+
+//   await app.start()
+
+//   await fetch(`http://localhost:${port}`)
+//   assert.strictEqual(onBanReachCalled, true)
+//   assert.strictEqual(errorResponseBuilderCalled, true)
+
+//   await app.close()
+// })
+
+// it('calls ai.rateLimiting.keyGenerator callback', async () => {
+//   const [app, port] = await buildAiWarpApp({ aiProvider })
+
+//   let callbackCalled = false
+//   app.register(fastifyPlugin(async () => {
+//     app.ai.rateLimiting.keyGenerator = (req) => {
+//       callbackCalled = true
+//       return req.ip
+//     }
+//   }))
+
+//   await app.start()
+
+//   await fetch(`http://localhost:${port}`)
+//   assert.strictEqual(callbackCalled, true)
+
+//   await app.close()
+// })
+
+// // it('calls ai.rateLimiting.errorResponseBuilder callback', async () => {
+// //     const [app, port] = await buildAiWarpApp({ aiProvider })
+
+// //     let callbackCalled = false
+// //     app.register(fastifyPlugin(async () => {
+// //         app.ai.rateLimiting.errorResponseBuilder = () => {
+// //             callbackCalled = true
+// //             return { error: 'rate limited' }
+// //         }
+// //     }))
+
+// //     await app.start()
+
+// //     await fetch(`http://localhost:${port}`)
+// //     assert.strictEqual(callbackCalled, true)
+
+// //     await app.close()
+// // })
diff --git a/tests/types/fastify.test-d.ts b/tests/types/fastify.test-d.ts
@@ -0,0 +1,37 @@
+import { ReadableStream } from 'node:stream/web'
+import { FastifyInstance, FastifyRequest } from 'fastify'
+import { expectAssignable } from 'tsd'
+import '../../index.d'
+import { errorResponseBuilderContext } from '@fastify/rate-limit'
+
+expectAssignable<FastifyInstance['ai']['warp']>(async (_: FastifyRequest, _2: string) => '')
+
+expectAssignable<FastifyInstance['ai']['warpStream']>(async (_: FastifyRequest, _2: string) => new ReadableStream())
+
+expectAssignable<FastifyInstance['ai']['preResponseCallback']>((_: FastifyRequest, _2: string) => '')
+expectAssignable<FastifyInstance['ai']['preResponseCallback']>(async (_: FastifyRequest, _2: string) => '')
+
+expectAssignable<FastifyInstance['ai']['preResponseChunkCallback']>((_: FastifyRequest, _2: string) => '')
+expectAssignable<FastifyInstance['ai']['preResponseChunkCallback']>(async (_: FastifyRequest, _2: string) => '')
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['max']>((_: FastifyRequest, _2: string) => 0)
+expectAssignable<FastifyInstance['ai']['rateLimiting']['max']>(async (_: FastifyRequest, _2: string) => 0)
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['allowList']>((_: FastifyRequest, _2: string) => true)
+expectAssignable<FastifyInstance['ai']['rateLimiting']['allowList']>(async (_: FastifyRequest, _2: string) => true)
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onBanReach']>((_: FastifyRequest, _2: string) => {})
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onBanReach']>(async (_: FastifyRequest, _2: string) => {})
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['keyGenerator']>((_: FastifyRequest) => '')
+expectAssignable<FastifyInstance['ai']['rateLimiting']['keyGenerator']>((_: FastifyRequest) => 0)
+expectAssignable<FastifyInstance['ai']['rateLimiting']['keyGenerator']>(async (_: FastifyRequest) => '')
+expectAssignable<FastifyInstance['ai']['rateLimiting']['keyGenerator']>(async (_: FastifyRequest) => 0)
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['errorResponseBuilder']>((_: FastifyRequest, _2: errorResponseBuilderContext) => { return {} })
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onExceeding']>((_: FastifyRequest, _2: string) => {})
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onExceeding']>(async (_: FastifyRequest, _2: string) => {})
+
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onExceeded']>((_: FastifyRequest, _2: string) => {})
+expectAssignable<FastifyInstance['ai']['rateLimiting']['onExceeded']>(async (_: FastifyRequest, _2: string) => {})
diff --git a/tests/types/schema.test-d.ts b/tests/types/schema.test-d.ts
@@ -16,13 +16,6 @@ expectAssignable<AiWarpConfig['aiProvider']>({
   }
 })
 
-expectError<AiWarpConfig['aiProvider']>({
-  openai: {
-    model: 'mistral-tiny',
-    apiKey: ''
-  }
-})
-
 expectAssignable<AiWarpConfig['aiProvider']>({
   mistral: {
     model: 'mistral-tiny',