git.giacomolli.uk - onechat/blob - image_model_test.go

   1 package main
   2
   3 import (
   4         "encoding/base64"
   5         "encoding/json"
   6         "net/http"
   7         "path/filepath"
   8         "strings"
   9         "testing"
  10
  11         "github.com/stretchr/testify/require"
  12 )
  13
  14 func TestNewTurnKeepsImagesInEveryChainRequestButNotLaterTurns(t *testing.T) {
  15         server := newModelFake(t)
  16         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  17                 if call.number == 0 {
  18                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "date", `{"id":0}`)}, 100, 0, 5))
  19
  20                         return
  21                 }
  22
  23                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("A red square.")}, 100, 0, 5))
  24         }
  25
  26         test := newTestTurnLoop(t, server)
  27         data := imageTestPNG(t, 3, 3)
  28         name, err := test.agent.Images().Save(data)
  29         require.NoError(t, err)
  30
  31         marker := imageText([]string{name}, "What is this?")
  32         _, err = test.loop.Submit(terminalSource, marker)
  33         require.NoError(t, err)
  34         test.events.WaitFor(t, "idle", 1)
  35         require.Equal(t, 2, server.Calls())
  36
  37         first := server.Call(0).body["input"].([]any)
  38         second := server.Call(1).body["input"].([]any)
  39         require.Equal(t, first, second[:len(first)])
  40
  41         expected := "data:image/png;base64," + base64.StdEncoding.EncodeToString(data)
  42         require.Contains(t, string(server.Call(0).raw), expected)
  43         require.Contains(t, string(server.Call(1).raw), expected)
  44
  45         _, err = test.loop.Submit(terminalSource, "Continue with text.")
  46         require.NoError(t, err)
  47         test.events.WaitFor(t, "idle", 2)
  48         require.NotContains(t, string(server.Call(2).raw), "input_image")
  49         require.NotContains(t, strings.Join(test.Log(t), "\n"), "base64")
  50 }
  51
  52 func TestImageArrivingDuringAChainJoinsItsNextRequest(t *testing.T) {
  53         server := newModelFake(t)
  54         arrived := make(chan struct{})
  55         release := make(chan struct{})
  56         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  57                 if call.number == 0 {
  58                         close(arrived)
  59                         <-release
  60                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "date", `{"id":0}`)}, 100, 0, 5))
  61
  62                         return
  63                 }
  64
  65                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("done")}, 100, 0, 5))
  66         }
  67
  68         test := newTestTurnLoop(t, server)
  69         _, err := test.loop.Submit(terminalSource, "Start.")
  70         require.NoError(t, err)
  71         <-arrived
  72
  73         name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
  74         require.NoError(t, err)
  75
  76         _, err = test.loop.Submit(terminalSource, imageText([]string{name}, "Read this."))
  77         require.NoError(t, err)
  78         close(release)
  79         test.events.WaitFor(t, "idle", 1)
  80         require.NotContains(t, string(server.Call(0).raw), "input_image")
  81         require.Contains(t, string(server.Call(1).raw), "input_image")
  82 }
  83
  84 func TestLeafCompactionDescribesImagesOnceAndMergesOnlyText(t *testing.T) {
  85         server := newModelFake(t)
  86         store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
  87         name, err := store.Save(imageTestPNG(t, 2, 2))
  88         require.NoError(t, err)
  89
  90         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  91                 text := "receipt, Ada, $12, red square"
  92                 if call.number == 0 {
  93                         text = strings.Repeat("detail ", 100)
  94                 }
  95
  96                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(text)}, 100, 0, 5))
  97         }
  98
  99         compactor := newTestCompactor(t, server)
 100         compactor.images = store
 101         task := Task{
 102                 Level:   0,
 103                 Index:   0,
 104                 Context: "<chat>\n</chat>",
 105                 Input:   []string{"user: " + imageText([]string{name}, "Read the receipt.")},
 106                 Images:  []string{name},
 107         }
 108         text, err := compactor.Summarize(t.Context(), task)
 109         require.NoError(t, err)
 110         require.Equal(t, "receipt, Ada, $12, red square", text)
 111         require.Contains(t, string(server.Call(0).raw), "input_image")
 112         require.Contains(t, string(server.Call(0).raw), "kind, visible text, amounts, names, and context")
 113         require.NotContains(t, string(server.Call(1).raw), "input_image")
 114
 115         _, err = compactor.Summarize(t.Context(), Task{
 116                 Level:   1,
 117                 Index:   0,
 118                 Context: "<chat>\n</chat>",
 119                 Input:   []string{text, "user: next"},
 120         })
 121         require.NoError(t, err)
 122         require.NotContains(t, string(server.Call(2).raw), "input_image")
 123
 124         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, compactor, discardLog())
 125         require.NoError(t, err)
 126         t.Cleanup(func() {
 127                 require.NoError(t, memory.Close())
 128         })
 129
 130         _, err = memory.Append("user", imageText([]string{name}, "A short caption."))
 131         require.NoError(t, err)
 132         require.NoError(t, memory.Settle(t.Context()))
 133         require.Equal(t, 4, server.Calls())
 134         require.Contains(t, string(server.Call(3).raw), "input_image")
 135         require.Contains(t, memory.View(), "receipt, Ada, $12")
 136 }
 137
 138 func TestZoomReturnsImagePartsButLogsOnlyText(t *testing.T) {
 139         server := newModelFake(t)
 140         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 141                 if call.number == 0 {
 142                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "zoom", `{"id":0,"n":1,"name":null}`)}, 100, 0, 5))
 143
 144                         return
 145                 }
 146
 147                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("A red square.")}, 100, 0, 5))
 148         }
 149
 150         test := newTestTurnLoop(t, server)
 151         name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
 152         require.NoError(t, err)
 153
 154         _, err = test.memory.Append("user", imageText([]string{name}, "Old image."))
 155         require.NoError(t, err)
 156
 157         _, err = test.loop.Submit(terminalSource, "Zoom the old image.")
 158         require.NoError(t, err)
 159         test.events.WaitFor(t, "idle", 1)
 160         require.NotContains(t, string(server.Call(0).raw), "input_image")
 161
 162         input := server.Call(1).body["input"].([]any)
 163         result := input[len(input)-1].(map[string]any)
 164         require.Equal(t, "function_call_output", result["type"])
 165
 166         parts := result["output"].([]any)
 167         require.Len(t, parts, 2)
 168         require.Equal(t, "input_text", parts[0].(map[string]any)["type"])
 169         require.Equal(t, "input_image", parts[1].(map[string]any)["type"])
 170         require.Contains(t, parts[0].(map[string]any)["text"], "Old image.")
 171         require.NotContains(t, strings.Join(test.Log(t), "\n"), "base64")
 172 }
 173
 174 func TestImageUserMessageKeepsTextPartsAndMarkerOrder(t *testing.T) {
 175         store := &ImageStore{dir: t.TempDir()}
 176         first, err := store.Save(imageTestPNG(t, 1, 2))
 177         require.NoError(t, err)
 178
 179         second, err := store.Save(imageTestPNG(t, 2, 1))
 180         require.NoError(t, err)
 181
 182         text := imageText([]string{first, second}, "Compare.")
 183         item, err := imageUserMessage(store, text)
 184         require.NoError(t, err)
 185
 186         data, err := json.Marshal(item)
 187         require.NoError(t, err)
 188
 189         var message struct {
 190                 Role    string           `json:"role"`
 191                 Content []ModelInputPart `json:"content"`
 192         }
 193
 194         require.NoError(t, json.Unmarshal(data, &message))
 195         require.Equal(t, "user", message.Role)
 196         require.Len(t, message.Content, 3)
 197         require.Equal(t, text, message.Content[0].Text)
 198         require.NotEqual(t, message.Content[1].ImageURL, message.Content[2].ImageURL)
 199 }