feat: Improve agent mode message formatting

internal/slack/agentCallbackHandler.go

@@ -2,7 +2,11 @@ package slackbot
 
 import (
 	"context"
+	"encoding/json"
 	"github.com/tmc/langchaingo/callbacks"
+	"regexp"
+
+	"github.com/slack-go/slack"
 )
 
 type sendMessageFunc func(message string)
@@ -15,7 +19,49 @@ type agentCallbackHandler struct {
 func (handler *agentCallbackHandler) HandleChainEnd(_ context.Context, outputs map[string]any) {
 	if text, ok := outputs["text"]; ok {
 		if textStr, ok := text.(string); ok {
+			if isThinkingMessage(textStr) {
+				textStr = formatContextMessageBlock(textStr)
+			} else {
+				textStr = formatFinalResponse(textStr)
+			}
 			handler.sendMessage(textStr)
 		}
 	}
 }
+
+var (
+	thinkingPattern = regexp.MustCompile(`Do I need to use a tool\? Yes`)
+	cleanupPattern  = regexp.MustCompile(`(Do I need to use a tool\? No|AI:)`)
+)
+
+func isThinkingMessage(msg string) bool {
+	return thinkingPattern.MatchString(msg)
+}
+
+// formatFinalResponse removes LLM agent response prefixes.
+// The agent response format is defined in internal/llm/langchain.go
+// > Thought: Do I need to use a tool? No
+// > AI: [your response here]
+func formatFinalResponse(msg string) string {
+	maxRemoves := 2
+	return cleanupPattern.ReplaceAllStringFunc(msg, func(s string) string {
+		if maxRemoves > 0 {
+			maxRemoves--
+			return ""
+		}
+		return s
+	})
+}
+
+func formatContextMessageBlock(message string) string {
+	mrkdwnBlock := slack.NewTextBlockObject("mrkdwn", message, false, false)
+	contextBlock := slack.NewContextBlock("", []slack.MixedElement{mrkdwnBlock}...)
+	blockMessage := slack.NewBlockMessage(contextBlock)
+
+	jsonByte, err := json.Marshal(blockMessage)
+	if err != nil {
+		// Fallback to plain message if marshaling fails
+		return message
+	}
+	return string(jsonByte)
+}

internal/slack/formatter/formatter.go

@@ -113,9 +113,12 @@ func FormatMessage(text string, options FormatOptions) []slack.MsgOption {
 				case "divider":
 					slackBlock = slack.NewDividerBlock()
 				case "context":
-					var context slack.ContextBlock
-					if err := json.Unmarshal(blockJSON, &context); err == nil {
-						slackBlock = context
+					var contextBlock struct {
+						Elements slack.ContextElements `json:"elements"`
+					}
+
+					if err := json.Unmarshal(blockJSON, &contextBlock); err == nil {
+						slackBlock = slack.NewContextBlock("", contextBlock.Elements.Elements...)
 					}
 					// Add more block types as needed
 				}