diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 70f53a098..dafecf3b3 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -482,42 +482,121 @@ func (al *AgentLoop) llmWorker(ctx context.Context, queue <-chan bus.InboundMess return } - // Reset per-round message-tool state so a previous round's - // tool-sent flag does not suppress this round's response. - if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { - if tool, ok := defaultAgent.Tools.Get("message"); ok { - if mt, ok := tool.(*tools.MessageTool); ok { - mt.ResetSentInRound() - } - } + // PDF two-phase handling: + // Phase 1: wait briefly for OCR keyword follow-up ("figures"/"図版") + // Phase 2: buffer messages during OCR, support cancel + if messageHasBareFile(msg) { + al.llmWorkerPDF(ctx, msg, queue) + continue } - response, err := al.processMessage(ctx, msg) - if err != nil { - response = fmt.Sprintf("Error processing message: %v", err) + al.llmWorkerNormal(ctx, msg) + } +} + +// llmWorkerNormal processes a single non-PDF message. +func (al *AgentLoop) llmWorkerNormal(ctx context.Context, msg bus.InboundMessage) { + // Reset per-round message-tool state so a previous round's + // tool-sent flag does not suppress this round's response. + if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + mt.ResetSentInRound() + } + } + } + + response, err := al.processMessage(ctx, msg) + if err != nil { + response = fmt.Sprintf("Error processing message: %v", err) + } + + al.sendResponseIfNeeded(ctx, msg, response) +} + +// llmWorkerPDF handles a bare-PDF message with two-phase follow-up collection. +func (al *AgentLoop) llmWorkerPDF(ctx context.Context, msg bus.InboundMessage, queue <-chan bus.InboundMessage) { + // Phase 1: wait for OCR keywords (figures/図版) — up to 5 seconds. + var overflow []bus.InboundMessage + msg, overflow = al.waitForPDFFollowUp(ctx, msg, queue) + + // Phase 2: run processMessage (OCR) concurrently while buffering + // messages from the same chat. Cancel on "中止"/"cancel". + al.resetMessageTool() + + response, err, buffered := al.processPDFWithBuffering(ctx, msg, queue, overflow) + if err != nil { + response = fmt.Sprintf("Error processing message: %v", err) + } + + al.sendResponseIfNeeded(ctx, msg, response) + + // Process buffered same-chat messages as a single follow-up turn + // so the LLM sees user instructions alongside the OCR result. + followUpText := mergeBufferedMessages(buffered, msg.ChatID) + if followUpText != "" { + notice := formatBufferedNotice(len(buffered)) + if notice != "" { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: msg.Channel, + ChatID: msg.ChatID, + Content: notice, + SkipPlaceholder: true, + IsStatus: true, + }) } - if response != "" { - alreadySent := false + followUpMsg := bus.InboundMessage{ + Channel: msg.Channel, + SenderID: msg.SenderID, + Sender: msg.Sender, + ChatID: msg.ChatID, + Peer: msg.Peer, + SessionKey: msg.SessionKey, + Content: followUpText, + Metadata: msg.Metadata, + } + al.llmWorkerNormal(ctx, followUpMsg) + } - defaultAgent := al.registry.GetDefaultAgent() + // Re-queue messages from other chats that were buffered. + otherMsgs := extractNonChatMessages(buffered, msg.ChatID) + for _, other := range otherMsgs { + al.llmWorkerNormal(ctx, other) + } +} - if defaultAgent != nil { - if tool, ok := defaultAgent.Tools.Get("message"); ok { - if mt, ok := tool.(*tools.MessageTool); ok { - alreadySent = mt.HasSentInRound() - } - } +// sendResponseIfNeeded sends the LLM response unless the message tool +// already sent it during this round. +func (al *AgentLoop) sendResponseIfNeeded(ctx context.Context, msg bus.InboundMessage, response string) { + if response == "" { + return + } + + alreadySent := false + if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + alreadySent = mt.HasSentInRound() } + } + } - if !alreadySent { - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: msg.Channel, + if !alreadySent { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: msg.Channel, + ChatID: msg.ChatID, + Content: response, + }) + } +} - ChatID: msg.ChatID, - - Content: response, - }) +// resetMessageTool resets the per-round message-tool state. +func (al *AgentLoop) resetMessageTool() { + if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + mt.ResetSentInRound() } } } diff --git a/pkg/agent/loop_media.go b/pkg/agent/loop_media.go index 766dc5815..6b0a211f4 100644 --- a/pkg/agent/loop_media.go +++ b/pkg/agent/loop_media.go @@ -537,13 +537,14 @@ func (al *AgentLoop) ocrPDF( totalPages := mediacache.PDFPageCount(pdfPath) totalStr := mediacache.FormatPageCount(totalPages) - // Send hint message and start progress indicator + // Send hint message (only if not already sent by Phase 1 waitForPDFFollowUp) if al.bus != nil && channel != "" && chatID != "" { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: channel, ChatID: chatID, Content: pdfHintMessage, SkipPlaceholder: true, + IsStatus: true, }) } diff --git a/pkg/agent/loop_pdf_wait.go b/pkg/agent/loop_pdf_wait.go new file mode 100644 index 000000000..f4c4372db --- /dev/null +++ b/pkg/agent/loop_pdf_wait.go @@ -0,0 +1,266 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/logger" +) + +// pdfFollowUpWait is how long to wait for a follow-up message after +// receiving a PDF file without accompanying text. During this window +// the user can send OCR keywords like "figures" / "図版". +const pdfFollowUpWait = 5 * time.Second + +// pdfFollowUpHint is sent when a PDF arrives without text. +const pdfFollowUpHint = "PDF received. You can send OCR options " + + "(e.g. \"figures\" / \"\u56f3\u7248\") within a few seconds, or processing will start automatically." + +// pdfCancelKeywords triggers OCR cancellation when found in a message +// received during Phase 2 (OCR in progress). +var pdfCancelKeywords = []string{ + "cancel", "abort", "stop", + "\u4e2d\u6b62", "\u30ad\u30e3\u30f3\u30bb\u30eb", "\u3084\u3081", +} + +// isCancelKeyword returns true if content contains a cancel keyword. +func isCancelKeyword(content string) bool { + lower := strings.ToLower(content) + for _, kw := range pdfCancelKeywords { + if strings.Contains(lower, kw) { + return true + } + } + return false +} + +// ── Phase 1: Pre-OCR keyword wait ────────────────────────────────────── + +// waitForPDFFollowUp checks whether msg contains a bare PDF file (no text). +// If so, it sends a hint and waits up to pdfFollowUpWait for a follow-up +// message that may contain OCR keywords (e.g. "figures"). The follow-up +// text is merged into msg.Content so that wantFigures() picks it up. +// +// If the message already has text (Caption) or is not a file, it returns +// the message unchanged. +func (al *AgentLoop) waitForPDFFollowUp( + ctx context.Context, msg bus.InboundMessage, queue <-chan bus.InboundMessage, +) (bus.InboundMessage, []bus.InboundMessage) { + if !messageHasBareFile(msg) { + return msg, nil + } + + // Send hint so the user knows they can add instructions + if al.bus != nil && msg.Channel != "" && msg.ChatID != "" { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: msg.Channel, + ChatID: msg.ChatID, + Content: pdfFollowUpHint, + SkipPlaceholder: true, + }) + } + + logger.DebugCF("agent", "Phase 1: waiting for PDF follow-up keywords", map[string]any{ + "chat_id": msg.ChatID, + "wait": pdfFollowUpWait.String(), + }) + + timer := time.NewTimer(pdfFollowUpWait) + defer timer.Stop() + + // Collect messages that arrived for other chats during the wait; + // they must not be lost. + var overflow []bus.InboundMessage + + for { + select { + case <-ctx.Done(): + return msg, overflow + case <-timer.C: + return msg, overflow + case followUp, ok := <-queue: + if !ok { + return msg, overflow + } + + // Different chat or slash command → don't absorb + if followUp.ChatID != msg.ChatID || + strings.HasPrefix(strings.TrimSpace(followUp.Content), "/") { + overflow = append(overflow, followUp) + continue + } + + followUpText := strings.TrimSpace(followUp.Content) + if followUpText == "" { + continue + } + + // Merge follow-up text into the PDF message so wantFigures() + // and the LLM both see it. + msg.Content = msg.Content + "\n" + followUpText + + logger.InfoCF("agent", "Phase 1: merged PDF follow-up keywords", map[string]any{ + "chat_id": msg.ChatID, + "follow_up": followUpText, + }) + + return msg, overflow + } + } +} + +// ── Phase 2: Buffer during OCR ───────────────────────────────────────── + +// processResult holds the return values of processMessage. +type processResult struct { + response string + err error +} + +// processPDFWithBuffering runs processMessage (which triggers OCR) in a +// goroutine while draining the llmQueue for the same chat. Messages +// arriving during OCR are buffered and returned so the caller can +// merge them into a follow-up LLM turn. +// +// If a cancel keyword is detected, the context is canceled (killing the +// OCR process) and the cancel message is returned as the response. +func (al *AgentLoop) processPDFWithBuffering( + ctx context.Context, + msg bus.InboundMessage, + queue <-chan bus.InboundMessage, + overflow []bus.InboundMessage, +) (response string, err error, buffered []bus.InboundMessage) { + // Run processMessage concurrently so we can drain the queue. + ocrCtx, ocrCancel := context.WithCancel(ctx) + defer ocrCancel() + + resultCh := make(chan processResult, 1) + go func() { + resp, e := al.processMessage(ocrCtx, msg) + resultCh <- processResult{resp, e} + }() + + // Notify the user that messages will be collected. + if al.bus != nil && msg.Channel != "" && msg.ChatID != "" { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: msg.Channel, + ChatID: msg.ChatID, + Content: "OCR starting. You can send instructions — they will be included after processing. Send \"cancel\" to abort.", + SkipPlaceholder: true, + IsStatus: true, + }) + } + + // Check overflow messages from Phase 1 first (they may contain cancel + // or instructions for same chat). + for _, o := range overflow { + if o.ChatID == msg.ChatID && + !strings.HasPrefix(strings.TrimSpace(o.Content), "/") { + text := strings.TrimSpace(o.Content) + if isCancelKeyword(text) { + ocrCancel() + <-resultCh // wait for processMessage to finish + return "PDF processing canceled.", nil, nil + } + buffered = append(buffered, o) + } else { + buffered = append(buffered, o) + } + } + + // Drain queue while OCR is running. + for { + select { + case result := <-resultCh: + // processMessage finished (OCR done + LLM responded) + return result.response, result.err, buffered + + case followUp, ok := <-queue: + if !ok { + result := <-resultCh + return result.response, result.err, buffered + } + + // Same chat, not a command → buffer or cancel + if followUp.ChatID == msg.ChatID && + !strings.HasPrefix(strings.TrimSpace(followUp.Content), "/") { + text := strings.TrimSpace(followUp.Content) + if isCancelKeyword(text) { + ocrCancel() + <-resultCh + return "PDF processing canceled.", nil, nil + } + if text != "" { + buffered = append(buffered, followUp) + logger.InfoCF("agent", "Phase 2: buffered message during OCR", map[string]any{ + "chat_id": msg.ChatID, + "content": text, + }) + } + } else { + // Different chat or command → keep for later processing + buffered = append(buffered, followUp) + } + } + } +} + +// ── Helpers ──────────────────────────────────────────────────────────── + +// messageHasBareFile returns true if the message contains a [file] tag +// (indicating a document attachment) but has no meaningful user text +// alongside it. A message like "[file]" or "\n[file]" is considered bare. +func messageHasBareFile(msg bus.InboundMessage) bool { + content := msg.Content + if !strings.Contains(content, "[file]") { + return false + } + // Strip all media tags to see if there's any remaining user text + stripped := content + for _, tag := range []string{"[file]", "[image: photo]", "[voice]", "[audio]"} { + stripped = strings.ReplaceAll(stripped, tag, "") + } + stripped = strings.TrimSpace(stripped) + return stripped == "" +} + +// mergeBufferedMessages creates a combined content string from buffered +// messages for use as a follow-up user message after OCR. +func mergeBufferedMessages(buffered []bus.InboundMessage, chatID string) string { + var parts []string + for _, b := range buffered { + if b.ChatID == chatID { + text := strings.TrimSpace(b.Content) + if text != "" { + parts = append(parts, text) + } + } + } + return strings.Join(parts, "\n") +} + +// extractNonChatMessages returns messages not belonging to the given chatID. +func extractNonChatMessages(buffered []bus.InboundMessage, chatID string) []bus.InboundMessage { + var result []bus.InboundMessage + for _, b := range buffered { + if b.ChatID != chatID { + result = append(result, b) + } + } + return result +} + +// formatBufferedNotice creates a status message listing how many messages +// were collected during OCR processing. +func formatBufferedNotice(count int) string { + if count == 0 { + return "" + } + if count == 1 { + return "(1 message received during processing — included below)" + } + return fmt.Sprintf("(%d messages received during processing — included below)", count) +} diff --git a/pkg/agent/loop_pdf_wait_test.go b/pkg/agent/loop_pdf_wait_test.go new file mode 100644 index 000000000..86d4198f4 --- /dev/null +++ b/pkg/agent/loop_pdf_wait_test.go @@ -0,0 +1,87 @@ +package agent + +import ( + "testing" + + "github.com/sipeed/picoclaw/pkg/bus" +) + +func TestMessageHasBareFile(t *testing.T) { + tests := []struct { + name string + content string + want bool + }{ + {"bare file only", "[file]", true}, + {"file with newline", "\n[file]", true}, + {"file with whitespace", " [file] ", true}, + {"file with caption", "please analyze this\n[file]", false}, + {"file with Japanese", "\u56f3\u7248\u4ed8\u304d\n[file]", false}, + {"no file", "hello world", false}, + {"image only", "[image: photo]", false}, + {"file and image bare", "[image: photo]\n[file]", true}, + {"empty", "", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + msg := bus.InboundMessage{Content: tt.content} + if got := messageHasBareFile(msg); got != tt.want { + t.Errorf("messageHasBareFile(%q) = %v, want %v", tt.content, got, tt.want) + } + }) + } +} + +func TestIsCancelKeyword(t *testing.T) { + tests := []struct { + content string + want bool + }{ + {"cancel", true}, + {"Cancel this", true}, + {"ABORT", true}, + {"\u4e2d\u6b62\u3057\u3066", true}, + {"キャンセル", true}, + {"やめ", true}, + {"やめて", true}, + {"please continue", false}, + {"summarize this", false}, + {"", false}, + } + + for _, tt := range tests { + t.Run(tt.content, func(t *testing.T) { + if got := isCancelKeyword(tt.content); got != tt.want { + t.Errorf("isCancelKeyword(%q) = %v, want %v", tt.content, got, tt.want) + } + }) + } +} + +func TestMergeBufferedMessages(t *testing.T) { + buffered := []bus.InboundMessage{ + {ChatID: "chat1", Content: "summarize this"}, + {ChatID: "chat2", Content: "other chat"}, + {ChatID: "chat1", Content: "in Japanese"}, + } + + got := mergeBufferedMessages(buffered, "chat1") + want := "summarize this\nin Japanese" + if got != want { + t.Errorf("mergeBufferedMessages = %q, want %q", got, want) + } +} + +func TestExtractNonChatMessages(t *testing.T) { + buffered := []bus.InboundMessage{ + {ChatID: "chat1", Content: "same"}, + {ChatID: "chat2", Content: "other"}, + {ChatID: "chat1", Content: "same2"}, + } + + got := extractNonChatMessages(buffered, "chat1") + if len(got) != 1 || got[0].Content != "other" { + t.Errorf("extractNonChatMessages: got %d messages, want 1", len(got)) + } +}