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
         }
       }