1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

package main

// photoid.go handles the "photo button" flow: when the user snaps a photo, we
// mint a short id, keep the pristine full-res JPEG for archival, and send the
// model a DOWNSCALED copy. The id is delivered to the model as text (in the
// clientContent turn), so the model references it via a tool arg (image_id) to
// attach the stored full-res photo to an entity.

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"image"
	"image/jpeg"
	"sync"

	xdraw "golang.org/x/image/draw"
)

// idAlphabet excludes look-alike characters (0/O, 1/I/L, etc.) so the id is
// unambiguous when it appears in text and tool arguments.
const idAlphabet = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"

// photoIDLen is the number of characters in a minted photo id.
const photoIDLen = 4

// downscaleWidth is the width of the copy sent to the model.
const downscaleWidth = 640

// pendingPhotos is a per-session buffer of full-res photos awaiting attachment,
// keyed by their minted id. Bounded; oldest evicted. The tool runtime pops
// from here when the model references an id.
type pendingPhotos struct {
	mu    sync.Mutex
	byID  map[string][]byte
	order []string
	cap   int
}

func newPendingPhotos(capacity int) *pendingPhotos {
	return &pendingPhotos{byID: map[string][]byte{}, cap: capacity}
}

// add stores full-res bytes under a freshly minted, currently-unused id and
// returns the id, evicting the oldest entry if over capacity.
func (p *pendingPhotos) add(fullres []byte) string {
	p.mu.Lock()
	defer p.mu.Unlock()
	var id string
	for {
		id = mintPhotoID()
		if _, exists := p.byID[id]; !exists {
			break
		}
	}
	p.byID[id] = fullres
	p.order = append(p.order, id)
	for len(p.order) > p.cap {
		old := p.order[0]
		p.order = p.order[1:]
		delete(p.byID, old)
	}
	return id
}

// pop returns and removes the full-res bytes for id (case-insensitive-ish: the
// caller should upper-case). Returns false if unknown/evicted.
func (p *pendingPhotos) pop(id string) ([]byte, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	b, ok := p.byID[id]
	if !ok {
		return nil, false
	}
	delete(p.byID, id)
	for i, o := range p.order {
		if o == id {
			p.order = append(p.order[:i], p.order[i+1:]...)
			break
		}
	}
	return b, true
}

// peek returns the full-res bytes for id WITHOUT removing it, so the same photo
// can be attached to more than one entity and survives model retries (e.g. after
// a transcription glitch). Returns false if unknown/evicted.
func (p *pendingPhotos) peek(id string) ([]byte, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	b, ok := p.byID[id]
	return b, ok
}

// ids returns the currently-buffered ids, for debug logging.
func (p *pendingPhotos) ids() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]string, len(p.order))
	copy(out, p.order)
	return out
}

// mintPhotoID returns a short random id from the unambiguous alphabet, e.g. "A3F9".
func mintPhotoID() string {
	b := make([]byte, photoIDLen)
	if _, err := rand.Read(b); err != nil {
		// extremely unlikely; fall back to a fixed-but-unique-ish value
		return "AAAA"
	}
	out := make([]byte, photoIDLen)
	for i := range out {
		out[i] = idAlphabet[int(b[i])%len(idAlphabet)]
	}
	return string(out)
}

// downscaleJPEG decodes a JPEG and scales it to targetWidth (keeping aspect
// ratio). If the source is already narrower, it is scaled to its original size.
// Returns the scaled RGBA image.
func downscaleJPEG(src []byte, targetWidth int) (*image.RGBA, error) {
	img, err := jpeg.Decode(bytes.NewReader(src))
	if err != nil {
		return nil, fmt.Errorf("decode jpeg: %w", err)
	}
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	if w <= 0 || h <= 0 {
		return nil, fmt.Errorf("empty image")
	}
	tw := targetWidth
	if tw > w {
		tw = w
	}
	th := h * tw / w
	if th < 1 {
		th = 1
	}
	dst := image.NewRGBA(image.Rect(0, 0, tw, th))
	xdraw.CatmullRom.Scale(dst, dst.Bounds(), img, b, xdraw.Over, nil)
	return dst, nil
}

// encodeJPEG encodes an image to JPEG bytes at the given quality.
func encodeJPEG(img image.Image, quality int) ([]byte, error) {
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: quality}); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// makeDownscaledFrame produces a downscaled JPEG suitable for sending to the
// model. The original full-res bytes are left untouched (stored elsewhere for
// archival). The photo id is delivered to the model as text, not drawn here.
func makeDownscaledFrame(fullres []byte) ([]byte, error) {
	small, err := downscaleJPEG(fullres, downscaleWidth)
	if err != nil {
		return nil, err
	}
	return encodeJPEG(small, 80)
}