From b41be465ad45e3854e79f8a56c5fb92f1e3ec769 Mon Sep 17 00:00:00 2001
From: staneyffer <staneyffer@gmail.com>
Date: Mon, 26 Jun 2023 18:13:31 +0800
Subject: [PATCH] fix: documents chat not working with chatglm2-6b

---
 pilot/model/llm_out/chatglm_llm.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/pilot/model/llm_out/chatglm_llm.py b/pilot/model/llm_out/chatglm_llm.py
index 690be0f06..299adbf2a 100644
--- a/pilot/model/llm_out/chatglm_llm.py
+++ b/pilot/model/llm_out/chatglm_llm.py
@@ -54,7 +54,12 @@ def chatglm_generate_stream(
     try:
         query = messages[-2].split("human:")[1]
     except IndexError:
-        query = messages[-3].split("human:")[1]
+        # fix doc qa: https://github.com/csunny/DB-GPT/issues/274
+        doc_qa_message = messages[-2]
+        if "system:" in doc_qa_message:
+            query = doc_qa_message.split("system:")[1]
+        else:
+            query = messages[-3].split("human:")[1]
     print("Query Message: ", query)
     # output = ""
     # i = 0