sift/internal/api/extract.go

package api

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"
	"unicode/utf8"

	"git.wntrmute.dev/kyle/sift/internal/clean"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
	"git.wntrmute.dev/kyle/sift/internal/digest"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/pdfdoc"
	"git.wntrmute.dev/kyle/sift/internal/pdftext"
)

// v1extract answers a PDF the caller POSTs in the request body, exactly as /v1/fetch
// answers the same PDF by URL. Nothing is fetched: the bytes are hashed, the hash is the
// cache key, and identical content is extracted once (content-addressed, so no TTL and no
// revalidation).
func (s *Server) v1extract(w http.ResponseWriter, r *http.Request) {
	s.reqN.Add(1)
	c := s.caller(w, r)
	if c == nil {
		return
	}

	select {
	case s.uploads <- struct{}{}:
		defer func() { <-s.uploads }()
	case <-time.After(uploadWait):
		w.Header().Set("Retry-After", "10")
		writeErr(w, 429, "busy", "")
		return
	case <-r.Context().Done():
		return
	}

	body, rerr := io.ReadAll(r.Body)
	if rerr != nil {
		var me *http.MaxBytesError
		if errors.As(rerr, &me) {
			writeErr(w, 413, fmt.Sprintf("too large (>%d MB)", s.pdfMaxUpload()>>20), "")
			return
		}
		writeErr(w, 400, "empty body", "")
		return
	}
	if len(body) == 0 {
		writeErr(w, 400, "empty body", "")
		return
	}
	if clean.Kind(r.Header.Get("Content-Type"), body) != "pdf" {
		writeErr(w, 415, "extract accepts PDFs only", "")
		return
	}

	q := r.URL.Query()
	pages := q.Get("pages")
	grep := q.Get("grep")
	question := q.Get("question")
	mode := q.Get("mode")
	if mode == "" {
		if question != "" {
			mode = "digest"
		} else {
			mode = "clean"
		}
	}
	if mode != "clean" && mode != "digest" {
		writeErr(w, 400, "invalid mode", "")
		return
	}

	hash := "sha256:" + HashBody(body)
	key := dcache.PDFKey("upload:"+hash, s.callerClass(c))

	t0 := time.Now()
	var res *fetch.Resource
	var doc *pdftext.Doc
	cached := false
	if cache := s.cache(); cache != nil {
		if e, ok := cache.GetPDF(key); ok {
			res, doc = s.pdfFromEntry("upload", e)
			cached = true
			s.Prom.PDFCacheHit()
		}
	}
	if res == nil {
		res = &fetch.Resource{URL: "upload", FinalURL: "upload", Status: 200, ContentType: "application/pdf", Body: body, Hash: hash, FetchedAt: t0}
		var eerr error
		doc, eerr = s.extractPDF(r.Context(), c, res)
		if eerr != nil {
			s.writeFetchErr(w, c, "", eerr)
			return
		}
		s.storePDF(key, res, doc)
	}
	fetchMS := time.Since(t0).Milliseconds()

	rp, err := s.renderPDF(doc, pdfdoc.Options{Pages: pages, Grep: grep, OutCap: s.OutCap})
	if err != nil {
		if errors.Is(err, pdfdoc.ErrBadPages) || errors.Is(err, pdfdoc.ErrBadGrep) {
			s.writeFetchErr(w, c, "", &fetch.Error{Status: 400, Msg: err.Error()})
			return
		}
		s.writeFetchErr(w, c, "", err)
		return
	}
	full := pdfdoc.Full(doc)
	docLen := utf8.RuneCountInString(full)
	content := rp.Content
	if mode == "digest" {
		content = full
	}

	warn := []string{}
	var digMeta map[string]any
	var digestMS int64
	if mode == "digest" {
		n := len(digest.Split(full, digest.ChunkTokens))
		limit := s.PDFMaxQuestionChunks
		if limit == 0 {
			limit = DefaultPDFMaxQuestionChunks
		}
		if n > limit {
			s.writeFetchErr(w, c, "", &fetch.Error{Status: 422, Msg: fmt.Sprintf("document too long for question (%d chunks); use grep or pages", n)})
			return
		}
		if s.Digest == nil {
			warn = append(warn, "digest fell back to clean: unset")
			mode = "clean"
			content, warn = capDigestFallback(content, warn)
			digMeta = map[string]any{"fallback": "unset"}
		} else {
			dr := s.Digest.Run(r.Context(), content, question, res.FinalURL, s.callerClass(c))
			digestMS = dr.DigestMS
			if !dr.Cached {
				s.Prom.ObserveDigest(time.Duration(dr.DigestMS) * time.Millisecond)
			}
			if dr.Fallback != "" {
				warn = append(warn, "digest fell back to clean: "+dr.Fallback)
				mode = "clean"
				content, warn = capDigestFallback(content, warn)
				digMeta = map[string]any{"fallback": dr.Fallback, "model": dr.Model}
			} else {
				content = dr.Text
				digMeta = map[string]any{
					"model":        dr.Model,
					"chunks_used":  dr.ChunksUsed,
					"chunks_total": dr.ChunksTotal,
					"retrieval":    dr.Retrieval,
					"embedded":     dr.Embedded,
					"fallback":     false,
					"cached":       dr.Cached,
				}
			}
		}
	}

	pdfMap := map[string]any{
		"pages":     doc.Pages,
		"title":     doc.Title,
		"outline":   doc.Outline,
		"extracted": len(doc.Text),
		"first":     rp.First,
		"last":      rp.Last,
		"map":       rp.Map,
		"grep_hits": rp.GrepHits,
		"note":      rp.Note,
	}

	s.Log.Info("extract", "caller", c.Name, "bytes", len(body), "pdf_pages", doc.Pages, "mode", mode, "cached", cached)

	src := map[string]any{
		"kind":       "upload",
		"hash":       hash,
		"bytes":      len(body),
		"cached":     cached,
		"fetched_at": res.FetchedAt.Format(time.RFC3339),
	}
	out := map[string]any{
		"content":   content,
		"mode":      mode,
		"untrusted": true,
		"source":    src,
		"warnings":  warn,
		"sizes": map[string]int{
			"raw_tokens":      docLen / 4,
			"clean_tokens":    digest.EstTokens(full),
			"returned_tokens": digest.EstTokens(content),
		},
		"timings": map[string]int64{"fetch_ms": fetchMS, "digest_ms": digestMS},
	}
	if pdfMap != nil {
		out["pdf"] = pdfMap
	}
	if digMeta != nil {
		out["digest"] = digMeta
	}
	writeJSON(w, 200, out)
}

// maxUploads is how many uploads may be in memory at once, and uploadWait how long a further
// one waits for a slot before a 429: at 100 MiB each, a few are all a 1 GiB service can hold
// beside the extraction children.
const (
	maxUploads = 2
	uploadWait = 5 * time.Second
)

// pdfMaxUpload is the /v1/extract upload cap: the configured value, or the default when unset.
func (s *Server) pdfMaxUpload() int64 {
	if s.PDFMaxUpload > 0 {
		return s.PDFMaxUpload
	}
	return DefaultPDFMaxUpload
}