ollie/cmd/olliesrv/internal/agent/chatlog.go

package agent

import (
	"ollie/cmd/olliesrv/internal/backend"
	"ollie/format"
)

// initChatHandler sets up the agent's output handler to write to its own chat log.
// Called once at construction.
func (ag *Agent) initChatHandler() {
	var streamingRole string
	var streamingResponseID string
	var streamingFence string

	ag.output = func(ev Event) {
		switch ev.Role {
		case "assistant":
			if streamingRole != ev.Role || (ev.Role == "assistant" && ev.ResponseID != streamingResponseID) {
				if streamingRole != "" {
					if streamingFence != "" {
						ag.AppendChat([]byte("\n" + streamingFence + "\n"))
						streamingFence = ""
					}
					ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				}
				if ev.Role == "assistant" && ev.ResponseID != "" {
					ag.AppendChat([]byte(format.AssistantHeaderPrefix + ev.ResponseID + "]]]\n"))
					streamingResponseID = ev.ResponseID
				} else {
					ag.AppendChat([]byte(format.RoleDelim(ev.Role) + "\n"))
				}
				streamingRole = ev.Role
			}
			ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		case "reasoning":
			if streamingRole != ev.Role {
				if streamingRole != "" {
					if streamingFence != "" {
						ag.AppendChat([]byte("\n" + streamingFence + "\n"))
						streamingFence = ""
					}
					ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				}
				ag.AppendChat([]byte(format.RoleDelim(ev.Role) + "\n"))
				streamingRole = ev.Role
			}
			ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		case "tool":
			if streamingRole == "tool" {
				if ev.Content == "" {
					if streamingFence != "" {
						ag.AppendChat([]byte("\n" + streamingFence + "\n"))
						streamingFence = ""
					}
					ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
					streamingRole = ""
				} else {
					ag.AppendChat([]byte(ev.Content))
				}
				return
			}
			if streamingRole != "" {
				if streamingFence != "" {
					ag.AppendChat([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			if ev.Content == "" {
				ag.AppendChat([]byte(format.ToolDelim(ev.Name) + "```" + ev.OutputFormat + "\n"))
				streamingFence = format.ToolFence
				streamingRole = "tool"
			} else {
				ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))
			}

		case "retry":
			if streamingRole == "retry" {
				ag.AppendChat([]byte("\n" + ev.Content))
				return
			}
			if streamingRole != "" {
				ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			ag.AppendChat([]byte(format.RetryTag + "\n"))
			streamingRole = "retry"
			ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		case "call":
			if streamingRole != "" {
				if streamingFence != "" {
					ag.AppendChat([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))

		default:
			if streamingRole != "" {
				if streamingFence != "" {
					ag.AppendChat([]byte("\n" + streamingFence + "\n"))
					streamingFence = ""
				}
				ag.AppendChat([]byte("\n" + format.EndTag + "\n"))
				streamingRole = ""
			}
			ag.AppendChat(format.FormatEvent(ev.Role, ev.Name, ev.Content, ev.OutputFormat))
		}
	}
}

// ReplayMessages renders persisted messages into the agent's chat log.
func (ag *Agent) ReplayMessages(messages []backend.Message) {
	const maxReplay = 20
	start := len(messages) - maxReplay
	if start < 0 {
		start = 0
	}
	for start < len(messages) && messages[start].Role == "system" {
		start++
	}

	// Build a lookup from ToolCallID → tool name.
	callNames := make(map[string]string)
	for _, m := range messages[start:] {
		for _, tc := range m.ToolCalls {
			if tc.ID != "" {
				callNames[tc.ID] = tc.Name
			}
		}
	}

	for _, m := range messages[start:] {
		switch m.Role {
		case "system":
		case "user":
			ag.AppendChat(format.FormatEvent("user", "", m.Content, ""))
		case "assistant":
			if m.Reasoning != "" {
				ag.AppendChat([]byte(format.RoleDelim("reasoning") + "\n" + m.Reasoning + "\n" + format.EndTag + "\n"))
			}
			if m.ID != "" {
				ag.AppendChat([]byte(format.AssistantHeaderPrefix + m.ID + "]]]\n"))
			} else {
				ag.AppendChat([]byte(format.AssistantTag + "\n"))
			}
			if m.Content != "" {
				ag.AppendChat([]byte(m.Content + "\n"))
			}
			ag.AppendChat([]byte(format.EndTag + "\n"))
			for _, tc := range m.ToolCalls {
				ag.AppendChat(format.FormatEvent("call", tc.Name, string(tc.Arguments), ""))
			}
		case "tool":
			name := callNames[m.ToolCallID]
			if name == "" {
				name = m.ToolCallID
			}
			ag.AppendChat(format.FormatEvent("tool", name, m.Content, ""))
		}
	}
}