orchard

  1// Files a turn can carry.
  2//
  3// Nothing is kept. A file arrives with the turn, is read into text, and the
  4// bytes are dropped, so there is no upload directory to age out and incognito
  5// stays true to its name. What persists is the extracted text, which is already
  6// in the conversation, so a follow up question about a file works without the
  7// file being sent again.
  8package main
  9
 10import (
 11	"bytes"
 12	"context"
 13	"errors"
 14	"fmt"
 15	"io"
 16	"mime/multipart"
 17	"os"
 18	"os/exec"
 19	"path/filepath"
 20	"strings"
 21	"time"
 22	"unicode/utf8"
 23
 24	"github.com/ledongthuc/pdf"
 25)
 26
 27const (
 28	maxFiles      = 10
 29	maxFileBytes  = 20 << 20
 30	maxFileChars  = 40000
 31	maxTotalChars = 80000
 32)
 33
 34// Attachment is what the UI shows on the message and what the database keeps.
 35// The text itself is not here, it is folded into the message content.
 36type Attachment struct {
 37	Name  string `json:"name"`
 38	Size  int64  `json:"size"`
 39	Kind  string `json:"kind"`
 40	Chars int    `json:"chars,omitempty"`
 41	Err   string `json:"err,omitempty"`
 42}
 43
 44type filePart struct {
 45	Attachment
 46	Text string
 47}
 48
 49// readFiles turns the uploaded parts into text. A file that cannot be read is
 50// not an error for the turn: it comes back with Err set and the model is told
 51// which files it did not get, which is better than refusing the whole message
 52// because one of five was a screenshot.
 53func readFiles(headers []*multipart.FileHeader) []filePart {
 54	if len(headers) > maxFiles {
 55		headers = headers[:maxFiles]
 56	}
 57	out := make([]filePart, 0, len(headers))
 58	budget := maxTotalChars
 59	for _, h := range headers {
 60		p := filePart{Attachment: Attachment{Name: filepath.Base(h.Filename), Size: h.Size}}
 61		switch {
 62		case h.Size > maxFileBytes:
 63			p.Kind, p.Err = "skipped", fmt.Sprintf("larger than %s", humanSize(maxFileBytes))
 64		default:
 65			text, kind, err := extract(h)
 66			p.Kind = kind
 67			if err != nil {
 68				p.Err = err.Error()
 69			} else {
 70				if len(text) > maxFileChars {
 71					text = text[:maxFileChars] + fmt.Sprintf("\n\n[truncated, this is the first %d characters of %d]", maxFileChars, len(text))
 72				}
 73				if len(text) > budget {
 74					text = text[:budget] + "\n\n[truncated, the turn ran out of room for the rest]"
 75				}
 76				budget -= len(text)
 77				if budget < 0 {
 78					budget = 0
 79				}
 80				p.Text, p.Chars = text, len(text)
 81			}
 82		}
 83		out = append(out, p)
 84	}
 85	return out
 86}
 87
 88func extract(h *multipart.FileHeader) (string, string, error) {
 89	f, err := h.Open()
 90	if err != nil {
 91		return "", "unreadable", err
 92	}
 93	defer f.Close()
 94	raw, err := io.ReadAll(io.LimitReader(f, maxFileBytes))
 95	if err != nil {
 96		return "", "unreadable", err
 97	}
 98	if strings.EqualFold(filepath.Ext(h.Filename), ".pdf") || bytes.HasPrefix(raw, []byte("%PDF-")) {
 99		text, err := pdfText(raw)
100		return text, "pdf", err
101	}
102	if bytes.HasPrefix(raw, []byte("PK\x03\x04")) {
103		text, kind, err := officeText(raw)
104		if err == nil || kind != "" {
105			return text, kind, err
106		}
107		return "", "zip", fmt.Errorf("an archive cannot be read, send the files inside it instead")
108	}
109	if kind, ok := imageKind(raw); ok {
110		return "", kind, fmt.Errorf("this model reads text only, so an image cannot be looked at")
111	}
112	if !looksTextual(raw) {
113		return "", "binary", fmt.Errorf("not a text file, so there is nothing to read out of it")
114	}
115	return normalise(string(raw)), textKind(h.Filename), nil
116}
117
118// pdfText tries poppler first and the pure Go reader second.
119//
120// pdftotext is here because the Go reader fails on the pdfs that matter most. A
121// bank statement is normally encrypted with an empty owner password, which it
122// will not open at all, and even when it does open one it returns the words in
123// object order with no column alignment, so a statement comes out as a list of
124// numbers with nothing saying which row they belonged to. -layout keeps the
125// columns, which is the difference between a table and a pile of figures.
126func pdfText(raw []byte) (string, error) {
127	if out, err := popplerText(raw); err == nil {
128		return out, nil
129	} else if !errors.Is(err, errNoPoppler) {
130		// Poppler ran and could not read it, so the Go reader will not do
131		// better. Its own message is the useful one.
132		return "", err
133	}
134	return goPDFText(raw)
135}
136
137var errNoPoppler = errors.New("pdftotext is not installed")
138
139func popplerText(raw []byte) (string, error) {
140	bin, err := exec.LookPath("pdftotext")
141	if err != nil {
142		return "", errNoPoppler
143	}
144	dir, err := os.MkdirTemp("", "pdf")
145	if err != nil {
146		return "", errNoPoppler
147	}
148	defer os.RemoveAll(dir)
149	in := filepath.Join(dir, "in.pdf")
150	if err := os.WriteFile(in, raw, 0o600); err != nil {
151		return "", errNoPoppler
152	}
153
154	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
155	defer cancel()
156	// -layout keeps the columns, -nopgbrk drops the form feeds, and the empty
157	// -opw is what opens a statement encrypted with no password set.
158	cmd := exec.CommandContext(ctx, bin, "-layout", "-nopgbrk", "-enc", "UTF-8", "-opw", "", in, "-")
159	var out, errBuf bytes.Buffer
160	cmd.Stdout, cmd.Stderr = &out, &errBuf
161	if err := cmd.Run(); err != nil {
162		if ctx.Err() != nil {
163			return "", fmt.Errorf("this pdf took too long to read")
164		}
165		// poppler's exit codes are specific and its stderr wording is not, so
166		// the code decides and the text only ever adds detail.
167		var ee *exec.ExitError
168		if errors.As(err, &ee) {
169			switch ee.ExitCode() {
170			case 1:
171				return "", fmt.Errorf("this pdf could not be opened, it may be damaged")
172			case 3:
173				return "", fmt.Errorf("this pdf is locked against copying text")
174			}
175		}
176		if msg := strings.TrimSpace(errBuf.String()); msg != "" {
177			return "", fmt.Errorf("this pdf could not be read: %s", firstLine(msg))
178		}
179		return "", fmt.Errorf("this pdf could not be read")
180	}
181	text := normalise(out.String())
182	if strings.TrimSpace(text) == "" {
183		return "", fmt.Errorf("this pdf has no text in it, it is probably a scan and would need OCR")
184	}
185	return text, nil
186}
187
188// goPDFText is the fallback for a development run outside the container, where
189// poppler is not on the path. It recovers rather than propagating a panic,
190// because the reader panics on a malformed cross reference table instead of
191// returning an error.
192func goPDFText(raw []byte) (text string, err error) {
193	defer func() {
194		if r := recover(); r != nil {
195			text, err = "", fmt.Errorf("this pdf could not be parsed")
196		}
197	}()
198	r, err := pdf.NewReader(bytes.NewReader(raw), int64(len(raw)))
199	if err != nil {
200		return "", fmt.Errorf("this pdf could not be parsed")
201	}
202	rd, err := r.GetPlainText()
203	if err != nil {
204		return "", fmt.Errorf("this pdf has no extractable text, it is probably scanned")
205	}
206	var sb strings.Builder
207	if _, err := io.Copy(&sb, io.LimitReader(rd, 8<<20)); err != nil {
208		return "", fmt.Errorf("this pdf could not be read to the end")
209	}
210	out := normalise(sb.String())
211	if strings.TrimSpace(out) == "" {
212		return "", fmt.Errorf("this pdf has no extractable text, it is probably scanned")
213	}
214	return out, nil
215}
216
217// looksTextual decides by content and not by extension, so an unknown suffix on
218// a config file still reads and a .txt holding a binary blob still does not.
219func looksTextual(raw []byte) bool {
220	if len(raw) == 0 {
221		return false
222	}
223	head := raw
224	if len(head) > 8192 {
225		head = head[:8192]
226	}
227	if bytes.IndexByte(head, 0) >= 0 || !utf8.Valid(head) {
228		return false
229	}
230	odd := 0
231	for _, b := range head {
232		if b < 0x09 || (b > 0x0d && b < 0x20) {
233			odd++
234		}
235	}
236	return odd*20 < len(head)
237}
238
239func imageKind(raw []byte) (string, bool) {
240	switch {
241	case bytes.HasPrefix(raw, []byte("\x89PNG\r\n\x1a\n")):
242		return "png", true
243	case bytes.HasPrefix(raw, []byte{0xff, 0xd8, 0xff}):
244		return "jpeg", true
245	case bytes.HasPrefix(raw, []byte("GIF8")):
246		return "gif", true
247	case bytes.HasPrefix(raw, []byte("RIFF")) && len(raw) > 12 && bytes.Equal(raw[8:12], []byte("WEBP")):
248		return "webp", true
249	case bytes.HasPrefix(raw, []byte("BM")):
250		return "bmp", true
251	}
252	return "", false
253}
254
255func textKind(name string) string {
256	ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(name), "."))
257	if ext == "" {
258		return "text"
259	}
260	return ext
261}
262
263// normalise strips carriage returns and the byte order mark Windows editors
264// leave at the front, both of which reach the model as noise it has to spend
265// attention on.
266func normalise(s string) string {
267	s = strings.TrimPrefix(s, "\ufeff")
268	return strings.ReplaceAll(s, "\r\n", "\n")
269}
270
271// composeTurn builds what the model is handed. Files come first and the message
272// last, since the question is what the answer has to follow and the end of the
273// prompt is where a small model reads most carefully.
274func composeTurn(message string, parts []filePart) string {
275	if len(parts) == 0 {
276		return message
277	}
278	var sb strings.Builder
279	sb.WriteString("The user attached ")
280	if len(parts) == 1 {
281		sb.WriteString("a file. Its contents are below.\n\n")
282	} else {
283		fmt.Fprintf(&sb, "%d files. Their contents are below.\n\n", len(parts))
284	}
285	for _, p := range parts {
286		if p.Err != "" {
287			fmt.Fprintf(&sb, "--- %s (%s): not readable, %s ---\n\n", p.Name, humanSize(p.Size), p.Err)
288			continue
289		}
290		fmt.Fprintf(&sb, "--- %s (%s, %s) ---\n%s\n--- end of %s ---\n\n", p.Name, p.Kind, humanSize(p.Size), p.Text, p.Name)
291	}
292	sb.WriteString("The user's message about them:\n\n")
293	if strings.TrimSpace(message) == "" {
294		sb.WriteString("Read these and say what they are and what is in them.")
295	} else {
296		sb.WriteString(message)
297	}
298	return sb.String()
299}
300
301func attachments(parts []filePart) []Attachment {
302	out := make([]Attachment, 0, len(parts))
303	for _, p := range parts {
304		out = append(out, p.Attachment)
305	}
306	return out
307}
308
309func humanSize(n int64) string {
310	switch {
311	case n < 1024:
312		return fmt.Sprintf("%d B", n)
313	case n < 1024*1024:
314		return fmt.Sprintf("%.1f KB", float64(n)/1024)
315	}
316	return fmt.Sprintf("%.1f MB", float64(n)/(1024*1024))
317}
318
319// titleSeed keeps the text of a file out of the conversation title. A turn that
320// is only attachments is named after them instead.
321func titleSeed(message string, parts []filePart) string {
322	if strings.TrimSpace(message) != "" {
323		return message
324	}
325	if len(parts) == 0 {
326		return ""
327	}
328	names := make([]string, 0, len(parts))
329	for _, p := range parts {
330		names = append(names, p.Name)
331	}
332	return "A message attaching " + strings.Join(names, ", ")
333}
334
335func firstLine(s string) string {
336	if i := strings.IndexByte(s, '\n'); i >= 0 {
337		return s[:i]
338	}
339	return s
340}