From 3cd405e61301620ec1bedb0296e3392a819fc342 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Wed, 10 Jan 2024 11:11:45 +0100
Subject: [PATCH] fix(chat): add missing chat response properties

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 src/api-types.ts            | 25 +++++++++++++++----------
 src/tests/mocks/handlers.ts |  3 +++
 2 files changed, 18 insertions(+), 10 deletions(-)

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)