From f12eb4a5cf0147fe364c87ef04e26c204a245e1c Mon Sep 17 00:00:00 2001
From: Chris Weaver <25087905+Weves@users.noreply.github.com>
Date: Fri, 29 Nov 2024 20:45:40 -0800
Subject: [PATCH] Fix assistant prompt zero-ing (#3293)

---
 backend/danswer/db/persona.py       |  4 ++-
 web/src/app/admin/assistants/lib.ts | 47 +++++++++++++++--------------
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/backend/danswer/db/persona.py b/backend/danswer/db/persona.py
index b71df22181e..3f3dead3d80 100644
--- a/backend/danswer/db/persona.py
+++ b/backend/danswer/db/persona.py
@@ -185,7 +185,7 @@ def create_update_persona(
             "persona_id": persona_id,
             "user": user,
             "db_session": db_session,
-            **create_persona_request.dict(exclude={"users", "groups"}),
+            **create_persona_request.model_dump(exclude={"users", "groups"}),
         }
 
         persona = upsert_persona(**persona_data)
@@ -758,6 +758,8 @@ def get_prompt_by_name(
     if user and user.role != UserRole.ADMIN:
         stmt = stmt.where(Prompt.user_id == user.id)
 
+    # Order by ID to ensure consistent result when multiple prompts exist
+    stmt = stmt.order_by(Prompt.id).limit(1)
     result = db_session.execute(stmt).scalar_one_or_none()
     return result
 
diff --git a/web/src/app/admin/assistants/lib.ts b/web/src/app/admin/assistants/lib.ts
index 2ba5057b298..f7a11f8ced3 100644
--- a/web/src/app/admin/assistants/lib.ts
+++ b/web/src/app/admin/assistants/lib.ts
@@ -259,29 +259,8 @@ export async function updatePersona(
 ): Promise<[Response, Response | null]> {
   const { id, existingPromptId } = personaUpdateRequest;
 
-  let fileId = null;
-  if (personaUpdateRequest.uploaded_image) {
-    fileId = await uploadFile(personaUpdateRequest.uploaded_image);
-    if (!fileId) {
-      return [new Response(null, { status: 400 }), null];
-    }
-  }
-
-  const updatePersonaResponse = await fetch(`/api/persona/${id}`, {
-    method: "PATCH",
-    headers: {
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify(
-      buildPersonaAPIBody(personaUpdateRequest, existingPromptId ?? 0, fileId)
-    ),
-  });
-
-  if (!updatePersonaResponse.ok) {
-    return [updatePersonaResponse, null];
-  }
-
   let promptResponse;
+  let promptId: number | null = null;
   if (existingPromptId !== undefined) {
     promptResponse = await updatePrompt({
       promptId: existingPromptId,
@@ -290,6 +269,7 @@ export async function updatePersona(
       taskPrompt: personaUpdateRequest.task_prompt,
       includeCitations: personaUpdateRequest.include_citations,
     });
+    promptId = existingPromptId;
   } else {
     promptResponse = await createPrompt({
       personaName: personaUpdateRequest.name,
@@ -297,8 +277,31 @@ export async function updatePersona(
       taskPrompt: personaUpdateRequest.task_prompt,
       includeCitations: personaUpdateRequest.include_citations,
     });
+    promptId = promptResponse.ok
+      ? ((await promptResponse.json()).id as number)
+      : null;
+  }
+  let fileId = null;
+  if (personaUpdateRequest.uploaded_image) {
+    fileId = await uploadFile(personaUpdateRequest.uploaded_image);
+    if (!fileId) {
+      return [promptResponse, null];
+    }
   }
 
+  const updatePersonaResponse =
+    promptResponse.ok && promptId !== null
+      ? await fetch(`/api/persona/${id}`, {
+          method: "PATCH",
+          headers: {
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify(
+            buildPersonaAPIBody(personaUpdateRequest, promptId, fileId)
+          ),
+        })
+      : null;
+
   return [promptResponse, updatePersonaResponse];
 }