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
 109         text, err := compactor.Summarize(t.Context(), task)
 110         require.NoError(t, err)
 111         require.Equal(t, "receipt, Ada, $12, red square", text)
 112         require.Contains(t, string(server.Call(0).raw), "input_image")
 113         require.Contains(t, string(server.Call(0).raw), "kind, visible text, amounts, names, and context")
 114         require.NotContains(t, string(server.Call(1).raw), "input_image")
 115
 116         _, err = compactor.Summarize(t.Context(), Task{
 117                 Level:   1,
 118                 Index:   0,
 119                 Context: "<chat>\n</chat>",
 120                 Input:   []string{text, "user: next"},
 121         })
 122         require.NoError(t, err)
 123         require.NotContains(t, string(server.Call(2).raw), "input_image")
 124
 125         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, compactor, discardLog())
 126         require.NoError(t, err)
 127         t.Cleanup(func() {
 128                 require.NoError(t, memory.Close())
 129         })
 130
 131         _, err = memory.Append("user", imageText([]string{name}, "A short caption."))
 132         require.NoError(t, err)
 133         require.NoError(t, memory.Settle(t.Context()))
 134         require.Equal(t, 4, server.Calls())
 135         require.Contains(t, string(server.Call(3).raw), "input_image")
 136         require.Contains(t, memory.View(), "receipt, Ada, $12")
 137 }
 138
 139 func TestZoomReturnsImagePartsButLogsOnlyText(t *testing.T) {
 140         server := newModelFake(t)
 141         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 142                 if call.number == 0 {
 143                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "zoom", `{"id":0,"n":1,"name":null}`)}, 100, 0, 5))
 144
 145                         return
 146                 }
 147
 148                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("A red square.")}, 100, 0, 5))
 149         }
 150
 151         test := newTestTurnLoop(t, server)
 152         name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
 153         require.NoError(t, err)
 154
 155         _, err = test.memory.Append("user", imageText([]string{name}, "Old image."))
 156         require.NoError(t, err)
 157
 158         _, err = test.loop.Submit(terminalSource, "Zoom the old image.")
 159         require.NoError(t, err)
 160         test.events.WaitFor(t, "idle", 1)
 161         require.NotContains(t, string(server.Call(0).raw), "input_image")
 162
 163         input := server.Call(1).body["input"].([]any)
 164         result := input[len(input)-1].(map[string]any)
 165         require.Equal(t, "function_call_output", result["type"])
 166
 167         parts := result["output"].([]any)
 168         require.Len(t, parts, 2)
 169         require.Equal(t, "input_text", parts[0].(map[string]any)["type"])
 170         require.Equal(t, "input_image", parts[1].(map[string]any)["type"])
 171         require.Contains(t, parts[0].(map[string]any)["text"], "Old image.")
 172         require.NotContains(t, strings.Join(test.Log(t), "\n"), "base64")
 173 }
 174
 175 func TestImageUserMessageKeepsTextPartsAndMarkerOrder(t *testing.T) {
 176         store := &ImageStore{dir: t.TempDir()}
 177         first, err := store.Save(imageTestPNG(t, 1, 2))
 178         require.NoError(t, err)
 179
 180         second, err := store.Save(imageTestPNG(t, 2, 1))
 181         require.NoError(t, err)
 182
 183         text := imageText([]string{first, second}, "Compare.")
 184         item, err := imageUserMessage(store, text)
 185         require.NoError(t, err)
 186
 187         data, err := json.Marshal(item)
 188         require.NoError(t, err)
 189
 190         var message struct {
 191                 Role    string           `json:"role"`
 192                 Content []ModelInputPart `json:"content"`
 193         }
 194
 195         require.NoError(t, json.Unmarshal(data, &message))
 196         require.Equal(t, "user", message.Role)
 197         require.Len(t, message.Content, 3)
 198         require.Equal(t, text, message.Content[0].Text)
 199         require.NotEqual(t, message.Content[1].ImageURL, message.Content[2].ImageURL)
 200 }