package chatpb import ( "fmt" "os" "path/filepath" "time" "github.com/google/uuid" "go.wit.com/log" "google.golang.org/protobuf/proto" timestamppb "google.golang.org/protobuf/types/known/timestamppb" ) func (c *Chats) AddGeminiComment(s string) *ChatEntry { chat := new(ChatEntry) chat.From = Who_GEMINI chat.Content = s chat.Ctime = timestamppb.New(time.Now()) c.AppendNew(chat) return chat } func (c *Chats) AddUserComment(s string) *ChatEntry { chat := new(ChatEntry) chat.From = Who_USER chat.Content = s c.AppendNew(chat) return chat } func UnmarshalChats(data []byte) (*Chats, error) { c := new(Chats) err := c.Unmarshal(data) return c, err } func UnmarshalChatsTEXT(data []byte) (*Chats, error) { c := new(Chats) err := c.UnmarshalTEXT(data) return c, err } func (all *Chats) AddFile(filename string) error { // Nil checks for safety. if all == nil { return fmt.Errorf("cannot call AddFile on a nil *Chats object") } if all.Chats == nil { all.Chats = make([]*Chat, 0) } data, err := os.ReadFile(filename) if err != nil { log.Fatalf("Error reading file %s: %v", filename, err) return err } logData, err := UnmarshalChatsTEXT(data) if err != nil { log.Fatalf("Error unmarshaling log file %s: %v", filename, err) return err } // New logic to handle both flat and nested formats for _, chatOrGroup := range logData.GetChats() { if len(chatOrGroup.GetEntries()) > 0 { // This is a new-style Chat group with entries for _, entry := range chatOrGroup.GetEntries() { // Convert the ChatEntry into a new Chat object for the flat list. newChat := convertEntryToChat(entry, filename) if newChat != nil { newChat.VerifyUuid() all.AppendByUuid(newChat) } } } else { // This is an old-style flat Chat entry. // We still process it to handle its external content file correctly. newChat := convertChatToChat(chatOrGroup, filename) if newChat != nil { newChat.VerifyUuid() all.AppendByUuid(newChat) } } } return nil } // convertChatToChat handles an old-style Chat message. It creates a clean // copy and resolves its external content file. func convertChatToChat(chat *Chat, filename string) *Chat { if chat == nil { return nil } // Manually create a ChatEntry from the Chat fields to reuse the logic. entry := &ChatEntry{ From: chat.GetFrom(), Ctime: chat.GetCtime(), Content: chat.GetContent(), Table: chat.GetTable(), ToolCalls: chat.GetToolCalls(), ContentFile: chat.GetContentFile(), Uuid: chat.GetUuid(), Snippets: chat.GetSnippets(), } return convertEntryToChat(entry, filename) } // convertEntryToChat creates a new Chat object from a ChatEntry's data // and resolves its external content file. func convertEntryToChat(entry *ChatEntry, filename string) *Chat { if entry == nil { return nil } // Create a new Chat object and copy the fields. newChat := &Chat{ From: entry.GetFrom(), Ctime: entry.GetCtime(), Table: entry.GetTable(), ToolCalls: entry.GetToolCalls(), Snippets: entry.GetSnippets(), Uuid: entry.GetUuid(), } // Handle content: prefer content_file, fallback to content. var content string if contentFile := entry.GetContentFile(); contentFile != "" { logDir := filepath.Dir(filename) contentPath := filepath.Join(logDir, contentFile) contentBytes, err := os.ReadFile(contentPath) if err != nil { content = fmt.Sprintf("--- ERROR: Could not read content file %s: %v ---", contentPath, err) } else { content = string(contentBytes) } } else { content = entry.GetContent() } newChat.Content = content return newChat } func (chats *Chats) VerifyUuids() bool { var changed bool all := chats.SortByUuid() for all.Scan() { chat := all.Next() if chat.Uuid == "" { chat.Uuid = uuid.New().String() changed = true } } return changed } func (c *Chat) VerifyUuid() bool { if c.Uuid == "" { c.Uuid = uuid.New().String() return true } return false } func (x *Chats) AppendNew(y *ChatEntry) { x.Lock() defer x.Unlock() var chat *Chat chat = proto.Clone(y).(*ChatEntry) x.Chats = append(x.Chats, chat) }