From b14c8e3d3ce78919ab2b8e3a204b001f658675a3 Mon Sep 17 00:00:00 2001
From: flakey5 <73616808+flakey5@users.noreply.github.com>
Date: Mon, 22 Apr 2024 05:18:40 -0700
Subject: [PATCH] Allow preResponseCallback and preResponseChunkCallback to
 return void (#21)

---
 index.d.ts      | 10 ++++++++--
 plugins/warp.ts |  4 ++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/index.d.ts b/index.d.ts
index eccbaf6..3e6a952 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -9,8 +9,14 @@ declare module 'fastify' {
     ai: {
       warp: (request: FastifyRequest, prompt: string) => Promise<string>
       warpStream: (request: FastifyRequest, prompt: string) => Promise<ReadableStream>
-      preResponseCallback?: ((request: FastifyRequest, response: string) => string) | ((request: FastifyRequest, response: string) => Promise<string>)
-      preResponseChunkCallback?: ((request: FastifyRequest, response: string) => string) | ((request: FastifyRequest, response: string) => Promise<string>)
+      preResponseCallback?: ((request: FastifyRequest, response: string) => void) |
+      ((request: FastifyRequest, response: string) => string) |
+      ((request: FastifyRequest, response: string) => Promise<void>) |
+      ((request: FastifyRequest, response: string) => Promise<string>)
+      preResponseChunkCallback?: ((request: FastifyRequest, response: string) => void) |
+      ((request: FastifyRequest, response: string) => string) |
+      ((request: FastifyRequest, response: string) => Promise<void>) |
+      ((request: FastifyRequest, response: string) => Promise<string>)
       rateLimiting: {
         max?: ((req: FastifyRequest, key: string) => number) | ((req: FastifyRequest, key: string) => Promise<number>)
         allowList?: (req: FastifyRequest, key: string) => boolean | Promise<boolean>
diff --git a/plugins/warp.ts b/plugins/warp.ts
index ed99e09..0bda413 100644
--- a/plugins/warp.ts
+++ b/plugins/warp.ts
@@ -39,7 +39,7 @@ export default fastifyPlugin(async (fastify) => {
 
       let response = await provider.ask(decoratedPrompt)
       if (fastify.ai.preResponseCallback !== undefined) {
-        response = await fastify.ai.preResponseCallback(request, response)
+        response = await fastify.ai.preResponseCallback(request, response) ?? response
       }
 
       return response
@@ -57,7 +57,7 @@ export default fastifyPlugin(async (fastify) => {
           if (fastify.ai.preResponseChunkCallback === undefined) {
             return response
           }
-          return await fastify.ai.preResponseChunkCallback(request, response)
+          return await fastify.ai.preResponseChunkCallback(request, response) ?? response
         }
       }