From 98a83543e858ec4486957edb0b4b7a920912aba2 Mon Sep 17 00:00:00 2001
From: sinopec <jayleezju@hotmail.com>
Date: Tue, 5 Aug 2025 14:32:39 +0800
Subject: [PATCH] Fix: fix mismatch of assitant message and its reference
 (#9233)

### What problem does this PR solve?

#9232

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

1. When creating a new session, initialize an empty reference that
includes both the app api and sdk API.
2. Fix the logic when retrieving references for historical messages: the
number of dialogue messages and reference messages may differ, but it
should match the number of assistant messages.

Co-authored-by: Li Ye <liye@unittec.com>
---
 api/apps/conversation_app.py |  3 ++-
 api/apps/sdk/session.py      | 16 +++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/api/apps/conversation_app.py b/api/apps/conversation_app.py
index 3799d01bc..5e5b44a47 100644
--- a/api/apps/conversation_app.py
+++ b/api/apps/conversation_app.py
@@ -66,7 +66,8 @@ def set_conversation():
         e, dia = DialogService.get_by_id(req["dialog_id"])
         if not e:
             return get_data_error_result(message="Dialog not found")
-        conv = {"id": conv_id, "dialog_id": req["dialog_id"], "name": name, "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}],"user_id": current_user.id}
+        conv = {"id": conv_id, "dialog_id": req["dialog_id"], "name": name, "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}],"user_id": current_user.id,
+                "reference":[{}],}
         ConversationService.save(**conv)
         return get_json_result(data=conv)
     except Exception as e:
diff --git a/api/apps/sdk/session.py b/api/apps/sdk/session.py
index b7e9f109c..281f88855 100644
--- a/api/apps/sdk/session.py
+++ b/api/apps/sdk/session.py
@@ -51,6 +51,7 @@ def create(tenant_id, chat_id):
         "name": req.get("name", "New session"),
         "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue")}],
         "user_id": req.get("user_id", ""),
+        "reference":[{}],
     }
     if not conv.get("name"):
         return get_error_data_result(message="`name` can not be empty.")
@@ -512,16 +513,16 @@ def list_session(tenant_id, chat_id):
             if "prompt" in info:
                 info.pop("prompt")
         conv["chat_id"] = conv.pop("dialog_id")
-        if conv["reference"]:
+        ref_messages = conv["reference"]
+        if ref_messages:
             messages = conv["messages"]
             message_num = 0
-            while message_num < len(messages) and message_num < len(conv["reference"]):
-                if message_num != 0 and messages[message_num]["role"] != "user":
-                    if message_num >= len(conv["reference"]):
-                        break
+            ref_num = 0
+            while message_num < len(messages) and ref_num < len(ref_messages):
+                if messages[message_num]["role"] != "user":
                     chunk_list = []
-                    if "chunks" in conv["reference"][message_num]:
-                        chunks = conv["reference"][message_num]["chunks"]
+                    if "chunks" in ref_messages[ref_num]:
+                        chunks = ref_messages[ref_num]["chunks"]
                         for chunk in chunks:
                             new_chunk = {
                                 "id": chunk.get("chunk_id", chunk.get("id")),
@@ -535,6 +536,7 @@ def list_session(tenant_id, chat_id):
 
                             chunk_list.append(new_chunk)
                     messages[message_num]["reference"] = chunk_list
+                    ref_num += 1
                 message_num += 1
         del conv["reference"]
     return get_result(data=convs)