Fix: export chat save in db; botPersona switching back to assistant

3 files changed, 30 insertions, 19 deletions

diff --git a/bot.go b/bot.go
index 8373edf..1f7a15b 100644
--- a/bot.go
+++ b/bot.go
@@ -335,6 +335,10 @@ func checkGame(role string, tv *tview.TextView) {
 
 func chatRound(userMsg, role string, tv *tview.TextView, regen, resume bool) {
 	botRespMode = true
+	botPersona := cfg.AssistantRole
+	if cfg.WriteNextMsgAsCompletionAgent != "" {
+		botPersona = cfg.WriteNextMsgAsCompletionAgent
+	}
 	defer func() { botRespMode = false }()
 	// check that there is a model set to use if is not local
 	if cfg.CurrentAPI == cfg.DeepSeekChatAPI || cfg.CurrentAPI == cfg.DeepSeekCompletionAPI {
@@ -362,7 +366,7 @@ func chatRound(userMsg, role string, tv *tview.TextView, regen, resume bool) {
 	logger.Debug("looking at vars in chatRound", "msg", userMsg, "regen", regen, "resume", resume)
 	if !resume {
 		fmt.Fprintf(tv, "[-:-:b](%d) ", len(chatBody.Messages))
-		fmt.Fprint(tv, roleToIcon(cfg.AssistantRole))
+		fmt.Fprint(tv, roleToIcon(botPersona))
 		fmt.Fprint(tv, "[-:-:-]\n")
 		if cfg.ThinkUse && !strings.Contains(cfg.CurrentAPI, "v1") {
 			// fmt.Fprint(tv, "<think>")
@@ -404,7 +408,7 @@ out:
 		// lastM.Content = lastM.Content + respText.String()
 	} else {
 		chatBody.Messages = append(chatBody.Messages, models.RoleMsg{
-			Role: cfg.AssistantRole, Content: respText.String(),
+			Role: botPersona, Content: respText.String(),
 		})
 	}
 	colorText()
@@ -490,6 +494,26 @@ func removeThinking(chatBody *models.ChatBody) {
 	chatBody.Messages = msgs
 }
 
+func addNewChat(chatName string) {
+	id, err := store.ChatGetMaxID()
+	if err != nil {
+		logger.Error("failed to get max chat id from db;", "id:", id)
+		// INFO: will rewrite first chat
+	}
+	chat := &models.Chat{
+		ID:        id + 1,
+		CreatedAt: time.Now(),
+		UpdatedAt: time.Now(),
+		Agent:     cfg.AssistantRole,
+	}
+	if chatName == "" {
+		chatName = fmt.Sprintf("%d_%s", chat.ID, cfg.AssistantRole)
+	}
+	chat.Name = chatName
+	chatMap[chat.Name] = chat
+	activeChatName = chat.Name
+}
+
 func applyCharCard(cc *models.CharCard) {
 	cfg.AssistantRole = cc.Role
 	// FIXME: remove
@@ -506,20 +530,7 @@ func applyCharCard(cc *models.CharCard) {
 			{Role: "system", Content: cc.SysPrompt},
 			{Role: cfg.AssistantRole, Content: cc.FirstMsg},
 		}
-		id, err := store.ChatGetMaxID()
-		if err != nil {
-			logger.Error("failed to get max chat id from db;", "id:", id)
-			// INFO: will rewrite first chat
-		}
-		chat := &models.Chat{
-			ID:        id + 1,
-			CreatedAt: time.Now(),
-			UpdatedAt: time.Now(),
-			Agent:     cfg.AssistantRole,
-		}
-		chat.Name = fmt.Sprintf("%d_%s", chat.ID, cfg.AssistantRole)
-		chatMap[chat.Name] = chat
-		activeChatName = chat.Name
+		addNewChat("")
 	}
 	chatBody.Messages = history
 }
diff --git a/session.go b/session.go
index ee03397..c978f53 100644
--- a/session.go
+++ b/session.go
@@ -47,6 +47,9 @@ func importChat(filename string) error {
 		return err
 	}
 	activeChatName = filepath.Base(filename)
+	if _, ok := chatMap[activeChatName]; !ok {
+		addNewChat(activeChatName)
+	}
 	chatBody.Messages = messages
 	cfg.AssistantRole = messages[1].Role
 	if cfg.AssistantRole == cfg.UserRole {
diff --git a/tui.go b/tui.go
index 5f6ceb1..644bb53 100644
--- a/tui.go
+++ b/tui.go
@@ -880,9 +880,6 @@ func init() {
 				colorText()
 			}
 			go chatRound(msgText, persona, textView, false, false)
-			// if !cfg.SkipLLMResp {
-			// 	// update statue line
-			// }
 			return nil
 		}
 		if event.Key() == tcell.KeyPgUp || event.Key() == tcell.KeyPgDn {