fix(chat.go): using a pointer to the updated chat message instead of …

…the value. Updating the whole currentChatMessages (user & assistant) instead of just the content. Fixed a bug on order where I frogot to add +1 to the length of the array
MohammadBnei · Mar 13, 2024 · b8731b2 · b8731b2
1 parent e2903a1
commit b8731b2
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 11 deletions.
diff --git a/cmd/prompt.go b/cmd/prompt.go
@@ -56,7 +56,7 @@ var promptCmd = &cobra.Command{
 			}
 		}
 
-		updateChan := make(chan service.ChatMessage)
+		updateChan := make(chan *service.ChatMessage)
 		defer close(updateChan)
 		promptConfig.UpdateChan = updateChan
 

diff --git a/mission.png b/mission.png
diff --git a/prince.png b/prince.png
diff --git a/service/chatMessages.go b/service/chatMessages.go
@@ -100,7 +100,7 @@ func (c *ChatMessages) SaveChatInModelfileFormat(filename string) error {
 			continue
 		}
 
-		builder.WriteString(fmt.Sprintf("%s\n\n", m.Content))
+		builder.WriteString(fmt.Sprintf("%s\n\n", strings.ReplaceAll(m.Content, "\"", "\\\"")))
 	}
 
 	return tool.SaveToFile([]byte(builder.String()), filename, false)
@@ -206,7 +206,7 @@ func (c *ChatMessages) AddMessage(content string, role ROLES) (*ChatMessage, err
 			ApiType: viper.GetString(config.AI_API_TYPE),
 			Model:   viper.GetString(config.AI_MODEL_NAME),
 		},
-		Order: uint(len(c.Messages)),
+		Order: uint(len(c.Messages)) + 1,
 	}
 
 	msg.Tokens = tokenCount

diff --git a/service/config.go b/service/config.go
@@ -19,7 +19,7 @@ type PromptConfig struct {
 	ChatMessages   *ChatMessages
 	PreviousPrompt string
 	UserPrompt     string
-	UpdateChan     chan ChatMessage
+	UpdateChan     chan *ChatMessage
 	Contexts       []ContextHold
 	*FileService
 }

diff --git a/ui/chat/chat.go b/ui/chat/chat.go
@@ -276,13 +276,15 @@ func (m chatModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.transition = msg.Title != ""
 		m.transitionModel.Title = msg.Title
 
-	case service.ChatMessage:
+	case *service.ChatMessage:
 		if m.currentChatMessages.user != nil && msg.Id == m.currentChatMessages.user.Id {
 			m.userPrompt = msg.Content
+			m.currentChatMessages.user = msg
 		}
 
 		if m.currentChatMessages.assistant != nil && msg.Id == m.currentChatMessages.assistant.Id {
 			m.aiResponse = msg.Content
+			m.currentChatMessages.assistant = msg
 		}
 
 		cmds = append(cmds, tea.Sequence(event.UpdateChatContent(m.userPrompt, m.aiResponse), waitForUpdate(m.promptConfig.UpdateChan)))
@@ -391,7 +393,7 @@ func (m chatModel) GetTitleView() string {
 	return style.TitleStyle.Render(wordwrap.String(userPrompt, m.size.Width-8))
 }
 
-func waitForUpdate(updateChan chan service.ChatMessage) tea.Cmd {
+func waitForUpdate(updateChan chan *service.ChatMessage) tea.Cmd {
 	return func() tea.Msg {
 		return <-updateChan
 	}

diff --git a/ui/chat/chatUpdate.go b/ui/chat/chatUpdate.go
@@ -211,7 +211,7 @@ func sendPrompt(pc *service.PromptConfig, currentChatMsgs currentChatMessages) e
 			previous.Content += string(chunk)
 			pc.ChatMessages.UpdateMessage(*previous)
 			if pc.UpdateChan != nil {
-				pc.UpdateChan <- *previous
+				pc.UpdateChan <- previous
 			}
 			return nil
 		}),
@@ -229,7 +229,7 @@ func sendPrompt(pc *service.PromptConfig, currentChatMsgs currentChatMessages) e
 	}
 
 	if pc.UpdateChan != nil {
-		pc.UpdateChan <- *pc.ChatMessages.FindById(currentChatMsgs.assistant.Id.Int64())
+		pc.UpdateChan <- pc.ChatMessages.FindById(currentChatMsgs.assistant.Id.Int64())
 	}
 
 	if viper.GetBool(config.C_COMPLETION_MODE) {
@@ -256,7 +256,7 @@ func sendAgentPrompt(m chatModel, currentChatMsgs currentChatMessages) error {
 	defer m.promptConfig.CloseContext(ctx)
 
 	if m.promptConfig.UpdateChan != nil {
-		m.promptConfig.UpdateChan <- *m.promptConfig.ChatMessages.FindById(currentChatMsgs.assistant.Id.Int64())
+		m.promptConfig.UpdateChan <- m.promptConfig.ChatMessages.FindById(currentChatMsgs.assistant.Id.Int64())
 	}
 
 	userMessages, _ := m.promptConfig.ChatMessages.FilterMessages(service.RoleUser)
@@ -281,7 +281,7 @@ func sendAgentPrompt(m chatModel, currentChatMsgs currentChatMessages) error {
 		previous.Content += string(chunk)
 		m.promptConfig.ChatMessages.UpdateMessage(*previous)
 		if m.promptConfig.UpdateChan != nil {
-			m.promptConfig.UpdateChan <- *previous
+			m.promptConfig.UpdateChan <- previous
 		}
 		return nil
 	}))
@@ -296,7 +296,7 @@ func sendAgentPrompt(m chatModel, currentChatMsgs currentChatMessages) error {
 	currentChatMsgs.user.Meta.Agent = m.chainName
 	m.promptConfig.ChatMessages.UpdateMessage(*currentChatMsgs.assistant)
 	if m.promptConfig.UpdateChan != nil {
-		m.promptConfig.UpdateChan <- *currentChatMsgs.assistant
+		m.promptConfig.UpdateChan <- currentChatMsgs.assistant
 	}
 
 	return nil