feat: two-phase PDF follow-up message handling
Add two-phase message handling for PDF files sent without text (common
in Telegram reply-with-attachment flow):
Phase 1 (pre-OCR, 5s wait):
- Send hint message when bare PDF arrives
- Wait up to 5 seconds for follow-up message with OCR keywords
(e.g. "figures"/"図版") to set OCR flags before processing starts
Phase 2 (during OCR):
- Run processMessage/OCR concurrently while draining the message queue
- Buffer same-chat messages as user instructions for post-OCR LLM turn
- Support cancel keywords ("cancel"/"中止"/"キャンセル") to abort OCR
- After OCR + initial LLM response, merge buffered messages into a
follow-up LLM turn so user instructions are processed with context
Refactors llmWorker into llmWorkerNormal/llmWorkerPDF with shared
helpers (sendResponseIfNeeded, resetMessageTool).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dj-oyu
parent
f1705d8ce4
commit
d7a2e3cb1c
4 changed files with 461 additions and 28 deletions
|
|
@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
266
pkg/agent/loop_pdf_wait.go
Normal file
266
pkg/agent/loop_pdf_wait.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
87
pkg/agent/loop_pdf_wait_test.go
Normal file
87
pkg/agent/loop_pdf_wait_test.go
Normal file
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue