diff --git a/pkg/agent/loop_media.go b/pkg/agent/loop_media.go index 768f7ef42..3086b44cb 100644 --- a/pkg/agent/loop_media.go +++ b/pkg/agent/loop_media.go @@ -7,18 +7,23 @@ package agent import ( + "bufio" "bytes" "context" "encoding/base64" "fmt" "io" "os" + "os/exec" + "path/filepath" "strings" + "sync/atomic" "time" "github.com/h2non/filetype" "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/mediacache" @@ -216,8 +221,8 @@ func (al *AgentLoop) describeImagesInMessages( if imageCount > 1 { label = fmt.Sprintf("Processing %d images...", imageCount) } - stopIndicator := al.processingIndicator(ctx, channel, chatID, label) - defer stopIndicator() + indicator := al.processingIndicator(ctx, channel, chatID, label) + defer indicator.Stop() for i, m := range result { if len(m.Media) == 0 { @@ -326,15 +331,38 @@ func (al *AgentLoop) describeImage( // a smooth rotating animation when displayed sequentially. var brailleSpinnerFrames = [...]string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"} +// progressIndicator manages a braille spinner with a dynamically updatable label. +// Call UpdateLabel to change the displayed text during processing. +type progressIndicator struct { + label atomic.Value // string + done chan struct{} +} + +// UpdateLabel changes the label shown alongside the spinner. +func (p *progressIndicator) UpdateLabel(label string) { + p.label.Store(label) +} + +// Stop terminates the spinner goroutine. +func (p *progressIndicator) Stop() { + select { + case <-p.done: + default: + close(p.done) + } +} + // processingIndicator publishes draft status messages with a braille spinner -// animation to indicate active processing. It runs until the returned stop -// function is called. The label describes what is being processed. -func (al *AgentLoop) processingIndicator(ctx context.Context, channel, chatID, label string) (stop func()) { +// animation to indicate active processing. It runs until Stop is called. +// Use UpdateLabel to change the displayed text during long operations. +func (al *AgentLoop) processingIndicator(ctx context.Context, channel, chatID, label string) *progressIndicator { + p := &progressIndicator{done: make(chan struct{})} + p.label.Store(label) + if al.bus == nil || channel == "" || chatID == "" { - return func() {} + return p } - done := make(chan struct{}) go func() { ticker := time.NewTicker(150 * time.Millisecond) defer ticker.Stop() @@ -342,12 +370,13 @@ func (al *AgentLoop) processingIndicator(ctx context.Context, channel, chatID, l frame := 0 for { select { - case <-done: + case <-p.done: return case <-ctx.Done(): return case <-ticker.C: - content := fmt.Sprintf("%s %s", brailleSpinnerFrames[frame%len(brailleSpinnerFrames)], label) + lbl, _ := p.label.Load().(string) + content := fmt.Sprintf("%s %s", brailleSpinnerFrames[frame%len(brailleSpinnerFrames)], lbl) _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: channel, ChatID: chatID, @@ -359,9 +388,7 @@ func (al *AgentLoop) processingIndicator(ctx context.Context, channel, chatID, l } }() - return func() { - close(done) - } + return p } // injectImageDescriptions replaces "[image: photo]" tags in content with @@ -379,3 +406,258 @@ func injectImageDescriptions(content string, descriptions []string) string { } return content } + +// maxPreviewRunes is the maximum number of runes to store as preview +// in the media cache for PDF OCR results. +const maxPreviewRunes = 500 + +// processPDFsInMessages finds [file:/path.pdf] tags in messages and replaces +// them with [document: preview... (full: /path/to.md, N pages)] tags after +// running OCR. A braille spinner with page progress is shown during processing. +func (al *AgentLoop) processPDFsInMessages( + ctx context.Context, messages []providers.Message, ocrCfg *config.OCRConfig, + channel, chatID string, +) []providers.Message { + result := make([]providers.Message, len(messages)) + copy(result, messages) + + for i, m := range result { + if !strings.Contains(m.Content, "[file:") { + continue + } + result[i].Content = al.replacePDFTags(ctx, m.Content, ocrCfg, channel, chatID) + } + + return result +} + +// pdfTagPrefix is the file tag pattern for PDF files injected by resolveMediaRefs. +const pdfTagPrefix = "[file:" + +// replacePDFTags finds [file:*.pdf] tags and replaces them with OCR results. +func (al *AgentLoop) replacePDFTags( + ctx context.Context, content string, ocrCfg *config.OCRConfig, + channel, chatID string, +) string { + var out strings.Builder + rest := content + + for { + idx := strings.Index(rest, pdfTagPrefix) + if idx < 0 { + out.WriteString(rest) + break + } + endRel := strings.Index(rest[idx:], "]") + if endRel < 0 { + out.WriteString(rest) + break + } + end := idx + endRel + 1 + + tag := rest[idx:end] + path := tag[len(pdfTagPrefix) : len(tag)-1] + + out.WriteString(rest[:idx]) + + if strings.HasSuffix(strings.ToLower(path), ".pdf") { + out.WriteString(al.ocrPDF(ctx, path, ocrCfg, channel, chatID)) + } else { + out.WriteString(tag) + } + + rest = rest[end:] + } + + return out.String() +} + +// ocrPDF runs OCR on a PDF file and returns a document tag with preview. +// Uses the media cache to avoid redundant OCR runs. +func (al *AgentLoop) ocrPDF( + ctx context.Context, pdfPath string, ocrCfg *config.OCRConfig, + channel, chatID string, +) string { + // Hash the file content for cache lookup + pdfData, err := os.ReadFile(pdfPath) + if err != nil { + logger.WarnCF("agent", "Failed to read PDF", map[string]any{"path": pdfPath, "error": err.Error()}) + return fmt.Sprintf("[file:%s]", pdfPath) + } + hash := mediacache.HashData(pdfData) + + // Check cache + if al.mediaCache != nil { + if entry, ok := al.mediaCache.GetEntry(hash, mediacache.TypePDFOCR); ok { + logger.DebugCF("agent", "PDF OCR cache hit", map[string]any{"hash": hash}) + return formatDocumentTag(entry.Result, entry.FilePath, entry.Pages) + } + } + + // Get page count for progress display + totalPages := mediacache.PDFPageCount(pdfPath) + totalStr := mediacache.FormatPageCount(totalPages) + + // Start progress indicator + indicator := al.processingIndicator(ctx, channel, chatID, + fmt.Sprintf("Processing PDF (0/%s)...", totalStr)) + defer indicator.Stop() + + // Determine output directory for OCR results + outputDir := al.ocrOutputDir() + os.MkdirAll(outputDir, 0o755) + + // Build command + timeout := time.Duration(ocrCfg.GetOCRTimeout()) * time.Second + cmdCtx, cmdCancel := context.WithTimeout(ctx, timeout) + defer cmdCancel() + + args := make([]string, 0, len(ocrCfg.Args)+4) + args = append(args, ocrCfg.Args...) + args = append(args, pdfPath, "-o", outputDir) + + cmd := exec.CommandContext(cmdCtx, ocrCfg.Command, args...) + + // Set environment + if len(ocrCfg.Env) > 0 { + cmd.Env = append(os.Environ(), ocrEnvSlice(ocrCfg.Env)...) + } + + // Pipe stderr for progress tracking + stderrPipe, err := cmd.StderrPipe() + if err != nil { + logger.WarnCF("agent", "Failed to create stderr pipe", map[string]any{"error": err.Error()}) + return fmt.Sprintf("[file:%s]", pdfPath) + } + + logger.InfoCF("agent", "Starting PDF OCR", map[string]any{ + "path": pdfPath, + "pages": totalStr, + "cmd": ocrCfg.Command, + }) + + if err := cmd.Start(); err != nil { + logger.WarnCF("agent", "Failed to start OCR command", map[string]any{"error": err.Error()}) + return fmt.Sprintf("[file:%s]", pdfPath) + } + + // Track progress via stderr + page := 0 + scanner := bufio.NewScanner(stderrPipe) + for scanner.Scan() { + line := scanner.Text() + if strings.Contains(line, "TextDetector __call__") { + page++ + indicator.UpdateLabel(fmt.Sprintf("Processing PDF (%d/%s)...", page, totalStr)) + } + } + + if err := cmd.Wait(); err != nil { + logger.WarnCF("agent", "OCR command failed", map[string]any{ + "path": pdfPath, + "error": err.Error(), + }) + return fmt.Sprintf("[file:%s]", pdfPath) + } + + // Find the output markdown file + mdPath := findOCROutput(outputDir, pdfPath) + if mdPath == "" { + logger.WarnCF("agent", "OCR output not found", map[string]any{"output_dir": outputDir}) + return fmt.Sprintf("[file:%s]", pdfPath) + } + + // Read preview from first part of the markdown + mdData, err := os.ReadFile(mdPath) + if err != nil { + logger.WarnCF("agent", "Failed to read OCR output", map[string]any{"path": mdPath, "error": err.Error()}) + return fmt.Sprintf("[file:%s]", pdfPath) + } + + preview := extractPreview(string(mdData), maxPreviewRunes) + if totalPages == 0 { + totalPages = page // use detected page count as fallback + } + + // Store in cache + if al.mediaCache != nil { + _ = al.mediaCache.PutEntry(hash, mediacache.TypePDFOCR, mediacache.Entry{ + Result: preview, + FilePath: mdPath, + Pages: totalPages, + }) + } + + logger.InfoCF("agent", "PDF OCR completed", map[string]any{ + "path": pdfPath, + "pages": totalPages, + "md_path": mdPath, + }) + + return formatDocumentTag(preview, mdPath, totalPages) +} + +// formatDocumentTag creates the tag injected into message content. +func formatDocumentTag(preview, mdPath string, pages int) string { + pagesStr := mediacache.FormatPageCount(pages) + return fmt.Sprintf("[document: %s\n full: %s (%s pages)\n Use read_file to see the complete document.]", + preview, mdPath, pagesStr) +} + +// extractPreview returns the first maxRunes runes of text, appending "..." if truncated. +func extractPreview(text string, maxRunes int) string { + runes := []rune(strings.TrimSpace(text)) + if len(runes) <= maxRunes { + return string(runes) + } + return string(runes[:maxRunes]) + "..." +} + +// ocrOutputDir returns the directory for storing OCR markdown output files. +func (al *AgentLoop) ocrOutputDir() string { + registry := al.GetRegistry() + if agent := registry.GetDefaultAgent(); agent != nil { + return filepath.Join(agent.Workspace, ".ocr_cache") + } + return filepath.Join(os.TempDir(), "picoclaw-ocr") +} + +// findOCROutput locates the markdown file generated by yomitoku. +// yomitoku names output as .md or _combined.md in the output dir. +func findOCROutput(outputDir, pdfPath string) string { + base := strings.TrimSuffix(filepath.Base(pdfPath), filepath.Ext(pdfPath)) + + // Try common yomitoku output patterns + candidates := []string{ + filepath.Join(outputDir, base+".md"), + filepath.Join(outputDir, base+"_combined.md"), + } + + for _, c := range candidates { + if _, err := os.Stat(c); err == nil { + return c + } + } + + // Fallback: find any .md file in the output directory + entries, err := os.ReadDir(outputDir) + if err != nil { + return "" + } + for _, e := range entries { + if !e.IsDir() && strings.HasSuffix(e.Name(), ".md") { + return filepath.Join(outputDir, e.Name()) + } + } + + return "" +} + +// ocrEnvSlice converts a map to "KEY=VALUE" slice for exec.Cmd.Env. +func ocrEnvSlice(env map[string]string) []string { + result := make([]string, 0, len(env)) + for k, v := range env { + result = append(result, k+"="+v) + } + return result +} diff --git a/pkg/agent/loop_media_test.go b/pkg/agent/loop_media_test.go index 9afd647cb..e21e89b3b 100644 --- a/pkg/agent/loop_media_test.go +++ b/pkg/agent/loop_media_test.go @@ -1,6 +1,7 @@ package agent import ( + "strings" "testing" "github.com/sipeed/picoclaw/pkg/config" @@ -104,3 +105,75 @@ func TestResolveImageModel_FallsToPlanModel(t *testing.T) { t.Errorf("got %q, want %q", model, "openai/gpt-5.4-nano") } } + +func TestExtractPreview_Short(t *testing.T) { + text := "Hello world" + result := extractPreview(text, 500) + if result != "Hello world" { + t.Errorf("got %q", result) + } +} + +func TestExtractPreview_Truncated(t *testing.T) { + text := strings.Repeat("a", 600) + result := extractPreview(text, 500) + if len([]rune(result)) != 503 { // 500 + "..." + t.Errorf("len = %d, want 503", len([]rune(result))) + } + if !strings.HasSuffix(result, "...") { + t.Error("should end with ...") + } +} + +func TestFormatDocumentTag(t *testing.T) { + tag := formatDocumentTag("preview text", "/path/to/doc.md", 18) + if !strings.Contains(tag, "preview text") { + t.Error("should contain preview") + } + if !strings.Contains(tag, "/path/to/doc.md") { + t.Error("should contain file path") + } + if !strings.Contains(tag, "18 pages") { + t.Error("should contain page count") + } + if !strings.Contains(tag, "read_file") { + t.Error("should contain read_file hint") + } +} + +func TestFormatDocumentTag_UnknownPages(t *testing.T) { + tag := formatDocumentTag("preview", "/path.md", 0) + if !strings.Contains(tag, "? pages") { + t.Error("should show ? for unknown page count") + } +} + +func TestReplacePDFTags_NoPDF(t *testing.T) { + al := &AgentLoop{} + content := "Check this out [file:/path/to/audio.mp3]" + result := al.replacePDFTags(t.Context(), content, &config.OCRConfig{Command: "echo"}, "", "") + if result != content { + t.Errorf("non-PDF should be unchanged, got %q", result) + } +} + +func TestReplacePDFTags_NoTags(t *testing.T) { + al := &AgentLoop{} + content := "Hello world" + result := al.replacePDFTags(t.Context(), content, &config.OCRConfig{Command: "echo"}, "", "") + if result != content { + t.Errorf("no tags should be unchanged, got %q", result) + } +} + +func TestProcessPDFs_NilOCR(t *testing.T) { + // When OCR is not configured, messages should pass through unchanged + messages := []providers.Message{ + {Role: "user", Content: "Check [file:/tmp/doc.pdf]"}, + } + al := &AgentLoop{} + result := al.processPDFsInMessages(t.Context(), messages, nil, "", "") + if result[0].Content != messages[0].Content { + t.Errorf("content should be unchanged when OCR config is nil") + } +} diff --git a/pkg/agent/loop_run.go b/pkg/agent/loop_run.go index d0df75f43..66d4bed82 100644 --- a/pkg/agent/loop_run.go +++ b/pkg/agent/loop_run.go @@ -256,6 +256,11 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt messages = al.describeImagesInMessages(ctx, messages, agent, opts.Channel, opts.ChatID) } + // Process PDFs with OCR when configured + if ocrCfg := cfg.Agents.Defaults.OCR; ocrCfg != nil && ocrCfg.Command != "" { + messages = al.processPDFsInMessages(ctx, messages, ocrCfg, opts.Channel, opts.ChatID) + } + // 2b. Interview staleness nudge: if MEMORY.md hasn't been updated for // several consecutive turns, inject a reminder so the AI writes its findings. diff --git a/pkg/config/config.go b/pkg/config/config.go index 72a5aa656..f23b36f29 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -242,6 +242,23 @@ type AgentDefaults struct { TaskReminderInterval int `json:"task_reminder_interval" env:"PICOCLAW_AGENTS_DEFAULTS_TASK_REMINDER_INTERVAL"` Orchestration bool `json:"orchestration,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_ORCHESTRATION"` Routing *RoutingConfig `json:"routing,omitempty"` + OCR *OCRConfig `json:"ocr,omitempty"` +} + +// OCRConfig configures the external OCR command for PDF text extraction. +type OCRConfig struct { + Command string `json:"command"` // path to OCR binary (e.g. "/path/to/.venv/bin/yomitoku") + Args []string `json:"args,omitempty"` // static arguments (e.g. ["-f", "md", "--lite", ...]) + Env map[string]string `json:"env,omitempty"` // extra environment variables (e.g. {"HF_HOME": "/tmp/hf-home"}) + Timeout int `json:"timeout,omitempty"` // timeout in seconds (default: 600) +} + +// GetOCRTimeout returns the configured timeout or default (600s = 10min). +func (c *OCRConfig) GetOCRTimeout() int { + if c != nil && c.Timeout > 0 { + return c.Timeout + } + return 600 } const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB diff --git a/pkg/mediacache/pdf.go b/pkg/mediacache/pdf.go new file mode 100644 index 000000000..bb5e4acf3 --- /dev/null +++ b/pkg/mediacache/pdf.go @@ -0,0 +1,52 @@ +package mediacache + +import ( + "os" + "regexp" + "strconv" +) + +// pdfPageCountRe matches /Type /Pages ... /Count N in PDF cross-reference. +// This covers the vast majority of well-formed PDFs. +var pdfPageCountRe = regexp.MustCompile(`/Type\s*/Pages\b[^>]*/Count\s+(\d+)`) + +// PDFPageCount extracts the total page count from a PDF file by parsing +// the /Type /Pages dictionary. Returns 0 if the count cannot be determined +// (encrypted, malformed, or unusual structure). This is a best-effort +// extraction that avoids heavy PDF library dependencies. +func PDFPageCount(path string) int { + data, err := os.ReadFile(path) + if err != nil { + return 0 + } + // Search from the end of the file where the root Pages dict typically lives. + // Limit search to last 64KB for performance on large files. + searchStart := 0 + if len(data) > 64*1024 { + searchStart = len(data) - 64*1024 + } + matches := pdfPageCountRe.FindAllSubmatch(data[searchStart:], -1) + if len(matches) == 0 { + // Fallback: search entire file + matches = pdfPageCountRe.FindAllSubmatch(data, -1) + } + if len(matches) == 0 { + return 0 + } + // Use the largest /Count found (root Pages object has the total). + var maxCount int + for _, m := range matches { + if n, err := strconv.Atoi(string(m[1])); err == nil && n > maxCount { + maxCount = n + } + } + return maxCount +} + +// FormatPageCount returns pages as a string, or "?" if unknown. +func FormatPageCount(pages int) string { + if pages > 0 { + return strconv.Itoa(pages) + } + return "?" +} diff --git a/pkg/mediacache/pdf_test.go b/pkg/mediacache/pdf_test.go new file mode 100644 index 000000000..e0ab6b3be --- /dev/null +++ b/pkg/mediacache/pdf_test.go @@ -0,0 +1,81 @@ +package mediacache + +import ( + "os" + "path/filepath" + "testing" +) + +// minimalPDF is a valid PDF with 3 pages. +// This is the smallest possible multi-page PDF structure. +const minimalPDF = `%PDF-1.4 +1 0 obj <> endobj +2 0 obj <> endobj +3 0 obj <> endobj +4 0 obj <> endobj +5 0 obj <> endobj +xref +0 6 +trailer <> +startxref +0 +%%EOF` + +func TestPDFPageCount_ValidPDF(t *testing.T) { + path := writeTempFile(t, "test.pdf", minimalPDF) + count := PDFPageCount(path) + if count != 3 { + t.Errorf("PDFPageCount = %d, want 3", count) + } +} + +func TestPDFPageCount_NonExistent(t *testing.T) { + count := PDFPageCount("/nonexistent/file.pdf") + if count != 0 { + t.Errorf("PDFPageCount = %d, want 0 for missing file", count) + } +} + +func TestPDFPageCount_NotPDF(t *testing.T) { + path := writeTempFile(t, "test.txt", "hello world") + count := PDFPageCount(path) + if count != 0 { + t.Errorf("PDFPageCount = %d, want 0 for non-PDF", count) + } +} + +func TestPDFPageCount_SinglePage(t *testing.T) { + pdf := `%PDF-1.4 +1 0 obj <> endobj +2 0 obj <> endobj +3 0 obj <> endobj +xref +0 4 +trailer <> +startxref +0 +%%EOF` + path := writeTempFile(t, "single.pdf", pdf) + count := PDFPageCount(path) + if count != 1 { + t.Errorf("PDFPageCount = %d, want 1", count) + } +} + +func TestFormatPageCount(t *testing.T) { + if s := FormatPageCount(0); s != "?" { + t.Errorf("FormatPageCount(0) = %q, want %q", s, "?") + } + if s := FormatPageCount(18); s != "18" { + t.Errorf("FormatPageCount(18) = %q, want %q", s, "18") + } +} + +func writeTempFile(t *testing.T, name, content string) string { + t.Helper() + path := filepath.Join(t.TempDir(), name) + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write temp file: %v", err) + } + return path +}