feat: add PDF OCR processing with yomitoku CLI integration
PDF files sent via chat are now processed with an external OCR command (configurable via config.json `ocr` section). The OCR output (markdown) is cached in media_cache.db with a first-page preview stored inline. The LLM receives a [document: preview (full: path, N pages)] tag and can use read_file to access the complete OCR result on demand. - OCRConfig: command, args, env, timeout in AgentDefaults - PDFPageCount: lightweight /Count N parser (fallback to "?" on failure) - processPDFsInMessages: replaces [file:*.pdf] tags with OCR results - progressIndicator: upgraded to dynamic labels via atomic.Value for real-time page progress from stderr (TextDetector cycle counting) - ocrPDF: exec.CommandContext with stderr parsing, cache integration - OCR output stored in workspace/.ocr_cache/ Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dj-oyu
parent
aaa9876acf
commit
bc45e31227
6 changed files with 522 additions and 12 deletions
|
|
@ -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{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
|
||||
|
||||
// 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()) {
|
||||
if al.bus == nil || channel == "" || chatID == "" {
|
||||
return func() {}
|
||||
// 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 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 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 <basename>.md or <basename>_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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
52
pkg/mediacache/pdf.go
Normal file
52
pkg/mediacache/pdf.go
Normal file
|
|
@ -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 "?"
|
||||
}
|
||||
81
pkg/mediacache/pdf_test.go
Normal file
81
pkg/mediacache/pdf_test.go
Normal file
|
|
@ -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 <</Type /Catalog /Pages 2 0 R>> endobj
|
||||
2 0 obj <</Type /Pages /Kids [3 0 R 4 0 R 5 0 R] /Count 3>> endobj
|
||||
3 0 obj <</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>> endobj
|
||||
4 0 obj <</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>> endobj
|
||||
5 0 obj <</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>> endobj
|
||||
xref
|
||||
0 6
|
||||
trailer <</Size 6 /Root 1 0 R>>
|
||||
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 <</Type /Catalog /Pages 2 0 R>> endobj
|
||||
2 0 obj <</Type /Pages /Kids [3 0 R] /Count 1>> endobj
|
||||
3 0 obj <</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>> endobj
|
||||
xref
|
||||
0 4
|
||||
trailer <</Size 4 /Root 1 0 R>>
|
||||
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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue