From 82394ac7d5d4c9cb65066971336fdfaa458204f6 Mon Sep 17 00:00:00 2001
From: Adam Howard <91115+codeincontext@users.noreply.github.com>
Date: Thu, 6 Mar 2025 10:07:37 +0100
Subject: [PATCH] fix: use llmMessage JSON schema in prompt

---
 .../aila/src/protocol/jsonPatchProtocol.ts    | 22 ++++++++-----------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/packages/aila/src/protocol/jsonPatchProtocol.ts b/packages/aila/src/protocol/jsonPatchProtocol.ts
index ce0536ece..24932fb27 100644
--- a/packages/aila/src/protocol/jsonPatchProtocol.ts
+++ b/packages/aila/src/protocol/jsonPatchProtocol.ts
@@ -472,19 +472,6 @@ export const JsonPatchDocumentJsonSchema = zodToJsonSchema(
   "patchDocumentSchema",
 );
 
-export const LLMResponseSchema = z.discriminatedUnion("type", [
-  PatchDocumentSchema,
-  PromptDocumentSchema,
-  StateDocumentSchema,
-  CommentDocumentSchema,
-  ErrorDocumentSchema,
-]);
-
-export const LLMResponseJsonSchema = zodToJsonSchema(
-  LLMResponseSchema,
-  "llmResponseSchema",
-);
-
 export const MessagePartDocumentSchema = z.discriminatedUnion("type", [
   ModerationDocumentSchema,
   ErrorDocumentSchema,
@@ -575,6 +562,15 @@ const LLMMessageSchemaWhileStreaming = z.object({
   status: z.literal("complete").optional(),
 });
 
+export const LLMResponseSchema = z.discriminatedUnion("type", [
+  LLMMessageSchema,
+]);
+
+export const LLMResponseJsonSchema = zodToJsonSchema(
+  LLMResponseSchema,
+  "llmResponseSchema",
+);
+
 function tryParseJson(str: string): {
   parsed: { type: string } | null;
   isPartial: boolean;