Feat: divide continue-gen and next-msg-gen

4 files changed, 56 insertions, 29 deletions

diff --git a/bot.go b/bot.go
index cbf1de3..c3ce273 100644
--- a/bot.go
+++ b/bot.go
@@ -83,12 +83,12 @@ func sendMsgToLLM(body io.Reader) {
 	reader := bufio.NewReader(resp.Body)
 	counter := uint32(0)
 	for {
+		var (
+			answerText string
+			content    string
+			stop       bool
+		)
 		counter++
-		if interruptResp {
-			interruptResp = false
-			logger.Info("interrupted bot response", "chunk_counter", counter)
-			break
-		}
 		// to stop from spiriling in infinity read of bad bytes that happens with poor connection
 		if cfg.ChunkLimit > 0 && counter > cfg.ChunkLimit {
 			logger.Warn("response hit chunk limit", "limit", cfg.ChunkLimit)
@@ -105,12 +105,15 @@ func sendMsgToLLM(body io.Reader) {
 			continue
 		}
 		if len(line) <= 1 {
+			if interruptResp {
+				goto interrupt // get unstuck from bad connection
+			}
 			continue // skip \n
 		}
 		// starts with -> data:
 		line = line[6:]
 		logger.Debug("debugging resp", "line", string(line))
-		content, stop, err := chunkParser.ParseChunk(line)
+		content, stop, err = chunkParser.ParseChunk(line)
 		if err != nil {
 			logger.Error("error parsing response body", "error", err, "line", string(line), "url", cfg.CurrentAPI)
 			streamDone <- true
@@ -127,8 +130,15 @@ func sendMsgToLLM(body io.Reader) {
 			content = strings.TrimPrefix(content, " ")
 		}
 		// bot sends way too many \n
-		answerText := strings.ReplaceAll(content, "\n\n", "\n")
+		answerText = strings.ReplaceAll(content, "\n\n", "\n")
 		chunkChan <- answerText
+	interrupt:
+		if interruptResp { // read bytes, so it would not get into beginning of the next req
+			interruptResp = false
+			logger.Info("interrupted bot response", "chunk_counter", counter)
+			streamDone <- true
+			break
+		}
 	}
 }
 
@@ -173,20 +183,21 @@ func roleToIcon(role string) string {
 	return "<" + role + ">: "
 }
 
-func chatRound(userMsg, role string, tv *tview.TextView, regen bool) {
+func chatRound(userMsg, role string, tv *tview.TextView, regen, resume bool) {
 	botRespMode = true
 	// reader := formMsg(chatBody, userMsg, role)
-	reader, err := chunkParser.FormMsg(userMsg, role)
+	reader, err := chunkParser.FormMsg(userMsg, role, resume)
 	if reader == nil || err != nil {
 		logger.Error("empty reader from msgs", "role", role, "error", err)
 		return
 	}
 	go sendMsgToLLM(reader)
-	// if userMsg != "" && !regen { // no need to write assistant icon since we continue old message
-	if userMsg != "" || regen {
-		fmt.Fprintf(tv, "(%d) ", len(chatBody.Messages))
+	logger.Debug("looking at vars in chatRound", "msg", userMsg, "regen", regen, "resume", resume)
+	// TODO: consider case where user msg is regened (not assistant one)
+	if !resume {
+		fmt.Fprintf(tv, "[-:-:b](%d) ", len(chatBody.Messages))
 		fmt.Fprint(tv, roleToIcon(cfg.AssistantRole))
-		fmt.Fprint(tv, "\n")
+		fmt.Fprint(tv, "[-:-:-]\n")
 		if cfg.ThinkUse && !strings.Contains(cfg.CurrentAPI, "v1") {
 			// fmt.Fprint(tv, "<think>")
 			chunkChan <- "<think>"
@@ -197,7 +208,6 @@ out:
 	for {
 		select {
 		case chunk := <-chunkChan:
-			// fmt.Printf(chunk)
 			fmt.Fprint(tv, chunk)
 			respText.WriteString(chunk)
 			tv.ScrollToEnd()
@@ -207,10 +217,15 @@ out:
 		}
 	}
 	botRespMode = false
-	// how can previous messages be affected?
-	chatBody.Messages = append(chatBody.Messages, models.RoleMsg{
-		Role: cfg.AssistantRole, Content: respText.String(),
-	})
+	// numbers in chatbody and displayed must be the same
+	if resume {
+		chatBody.Messages[len(chatBody.Messages)-1].Content += respText.String()
+		// lastM.Content = lastM.Content + respText.String()
+	} else {
+		chatBody.Messages = append(chatBody.Messages, models.RoleMsg{
+			Role: cfg.AssistantRole, Content: respText.String(),
+		})
+	}
 	colorText()
 	updateStatusLine()
 	// bot msg is done;
@@ -239,12 +254,12 @@ func findCall(msg string, tv *tview.TextView) {
 	f, ok := fnMap[fc.Name]
 	if !ok {
 		m := fc.Name + "%s is not implemented"
-		chatRound(m, cfg.ToolRole, tv, false)
+		chatRound(m, cfg.ToolRole, tv, false, false)
 		return
 	}
 	resp := f(fc.Args...)
 	toolMsg := fmt.Sprintf("tool response: %+v", string(resp))
-	chatRound(toolMsg, cfg.ToolRole, tv, false)
+	chatRound(toolMsg, cfg.ToolRole, tv, false, false)
 }
 
 func chatToTextSlice(showSys bool) []string {
diff --git a/llm.go b/llm.go
index 4e18d6d..1289552 100644
--- a/llm.go
+++ b/llm.go
@@ -10,7 +10,7 @@ import (
 
 type ChunkParser interface {
 	ParseChunk([]byte) (string, bool, error)
-	FormMsg(msg, role string) (io.Reader, error)
+	FormMsg(msg, role string, cont bool) (io.Reader, error)
 }
 
 func initChunkParser() {
@@ -28,7 +28,7 @@ type LlamaCPPeer struct {
 type OpenAIer struct {
 }
 
-func (lcp LlamaCPPeer) FormMsg(msg, role string) (io.Reader, error) {
+func (lcp LlamaCPPeer) FormMsg(msg, role string, cont bool) (io.Reader, error) {
 	if msg != "" { // otherwise let the bot continue
 		newMsg := models.RoleMsg{Role: role, Content: msg}
 		chatBody.Messages = append(chatBody.Messages, newMsg)
@@ -49,11 +49,13 @@ func (lcp LlamaCPPeer) FormMsg(msg, role string) (io.Reader, error) {
 	}
 	prompt := strings.Join(messages, "\n")
 	// strings builder?
-	if cfg.ToolUse && msg != "" {
+	if cfg.ToolUse && msg != "" && !cont {
 		prompt += "\n" + cfg.ToolRole + ":\n" + toolSysMsg
 	}
-	botMsgStart := "\n" + cfg.AssistantRole + ":\n"
-	prompt += botMsgStart
+	if !cont {
+		botMsgStart := "\n" + cfg.AssistantRole + ":\n"
+		prompt += botMsgStart
+	}
 	// if cfg.ThinkUse && msg != "" && !cfg.ToolUse {
 	if cfg.ThinkUse && !cfg.ToolUse {
 		prompt += "<think>"
@@ -98,7 +100,7 @@ func (op OpenAIer) ParseChunk(data []byte) (string, bool, error) {
 	return content, false, nil
 }
 
-func (op OpenAIer) FormMsg(msg, role string) (io.Reader, error) {
+func (op OpenAIer) FormMsg(msg, role string, resume bool) (io.Reader, error) {
 	if msg != "" { // otherwise let the bot continue
 		newMsg := models.RoleMsg{Role: role, Content: msg}
 		chatBody.Messages = append(chatBody.Messages, newMsg)
diff --git a/models/models.go b/models/models.go
index ceb98fd..bb61abf 100644
--- a/models/models.go
+++ b/models/models.go
@@ -181,6 +181,7 @@ func NewLCPReq(prompt string, cfg *config.Config, props map[string]float32) Llam
 		Stop: []string{
 			cfg.UserRole + ":\n", "<|im_end|>",
 			cfg.ToolRole + ":\n",
+			cfg.AssistantRole + ":\n",
 		},
 	}
 }
diff --git a/tui.go b/tui.go
index 8cbdf0b..16d63e5 100644
--- a/tui.go
+++ b/tui.go
@@ -55,6 +55,7 @@ var (
 [yellow]F10[white]: manage loaded rag files (that already in vector db)
 [yellow]F11[white]: switch RAGEnabled boolean
 [yellow]F12[white]: show this help page
+[yellow]Ctrl+w[white]: resume generation on the last msg
 [yellow]Ctrl+s[white]: load new char/agent
 [yellow]Ctrl+e[white]: export chat to json file
 [yellow]Ctrl+n[white]: start a new chat
@@ -450,7 +451,7 @@ func init() {
 			// regen last msg
 			chatBody.Messages = chatBody.Messages[:len(chatBody.Messages)-1]
 			textView.SetText(chatToText(cfg.ShowSys))
-			go chatRound("", cfg.UserRole, textView, true)
+			go chatRound("", cfg.UserRole, textView, true, false)
 			return nil
 		}
 		if event.Key() == tcell.KeyF3 && !botRespMode {
@@ -649,6 +650,13 @@ func init() {
 			pages.AddPage(RAGPage, chatRAGTable, true, true)
 			return nil
 		}
+		if event.Key() == tcell.KeyCtrlW {
+			// INFO: continue bot/text message
+			// without new role
+			lastRole := chatBody.Messages[len(chatBody.Messages)-1].Role
+			go chatRound("", lastRole, textView, false, true)
+			return nil
+		}
 		// cannot send msg in editMode or botRespMode
 		if event.Key() == tcell.KeyEscape && !editMode && !botRespMode {
 			// read all text into buffer
@@ -660,7 +668,8 @@ func init() {
 			if strings.HasSuffix(prevText, nl) {
 				nl = ""
 			}
-			if msgText != "" { // continue
+			if msgText != "" {
+				// add user icon before user msg
 				fmt.Fprintf(textView, "%s[-:-:b](%d) <%s>: [-:-:-]\n%s\n",
 					nl, len(chatBody.Messages), cfg.UserRole, msgText)
 				textArea.SetText("", true)
@@ -668,7 +677,7 @@ func init() {
 				colorText()
 			}
 			// update statue line
-			go chatRound(msgText, cfg.UserRole, textView, false)
+			go chatRound(msgText, cfg.UserRole, textView, false, false)
 			return nil
 		}
 		if event.Key() == tcell.KeyPgUp || event.Key() == tcell.KeyPgDn {