diff --git a/src/api-types.ts b/src/api-types.ts
index dca8c4b..0e10247 100644
--- a/src/api-types.ts
+++ b/src/api-types.ts
@@ -442,16 +442,21 @@ export const ChatInputSchema = z.object({
   parameters: ParametersSchema.nullish(),
 });
 export type ChatInput = z.input<typeof ChatInputSchema>;
-export const ChatOutputSchema = z.object({
-  conversation_id: z.string(),
-  results: z.array(
-    z
-      .object({
-        generated_text: z.string(),
-      })
-      .partial(),
-  ),
-});
+export const ChatOutputSchema = z
+  .object({
+    id: z.string(),
+    model_id: z.string(),
+    created_at: z.coerce.date(),
+    conversation_id: z.string(),
+    results: z.array(
+      z
+        .object({
+          generated_text: z.string(),
+        })
+        .partial(),
+    ),
+  })
+  .passthrough();
 export type ChatOutput = z.output<typeof ChatOutputSchema>;
 
 export const ChatStreamInputSchema = ChatInputSchema;
diff --git a/src/tests/mocks/handlers.ts b/src/tests/mocks/handlers.ts
index 2be4c44..e0d9fce 100644
--- a/src/tests/mocks/handlers.ts
+++ b/src/tests/mocks/handlers.ts
@@ -426,6 +426,9 @@ export const handlers: RestHandler<MockedRequest<DefaultBodyType>>[] = [
     return res(
       ctx.status(200),
       ctx.json({
+        id: randomUUID(),
+        model_id: body.model_id,
+        created_at: new Date('2022-12-19T22:53:22.000Z'),
         conversation_id,
         results: conversation
           ?.slice(-1)