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

   1 package main
   2
   3 import (
   4         "encoding/base64"
   5         "encoding/json"
   6         "net/http"
   7         "path/filepath"
   8         "testing"
   9
  10         "github.com/stretchr/testify/require"
  11 )
  12
  13 func TestImageShowUsesStoredImageForTelegramAndTerminal(t *testing.T) {
  14         picture := imageTestPNG(t, 5, 4)
  15         store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
  16         name, err := store.Save(picture)
  17         require.NoError(t, err)
  18
  19         telegram := newTelegramFake(t)
  20         bot, logs := telegramStartTestBot(t, telegram)
  21         box := &Toolbox{imageShow: &ImageShow{
  22                 store: store,
  23                 bot:   bot,
  24         }}
  25
  26         args, err := json.Marshal(ToolImageShowArguments{
  27                 Image:   name,
  28                 Caption: "A red point.",
  29         })
  30         require.NoError(t, err)
  31
  32         output := box.run(t.Context(), "image_show", string(args))
  33         require.Equal(t, imageText([]string{name}, "A red point.")+"\n"+filepath.Join(store.dir, name), output)
  34
  35         calls := telegram.Calls()
  36         require.Len(t, calls, 1)
  37         require.Equal(t, "sendPhoto", calls[0].method)
  38         require.Equal(t, "4242", calls[0].params["chat_id"])
  39         require.Equal(t, "A red point.", calls[0].params["caption"])
  40         require.Equal(t, name, calls[0].params["photo_name"])
  41         require.Equal(t, picture, calls[0].params["photo"])
  42         require.Equal(t, "<6>image sent\n", logs.String())
  43 }
  44
  45 func TestImageShowRefusesUnknownNamesBeforeSending(t *testing.T) {
  46         telegram := newTelegramFake(t)
  47         bot, _ := telegramStartTestBot(t, telegram)
  48         show := &ImageShow{
  49                 store: &ImageStore{dir: t.TempDir()},
  50                 bot:   bot,
  51         }
  52
  53         for _, name := range []string{"../image.png", "unknown", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.png"} {
  54                 _, err := show.Show(t.Context(), name, "Do not send this.")
  55                 require.Error(t, err)
  56         }
  57
  58         require.Empty(t, telegram.Calls())
  59
  60         box := &Toolbox{}
  61         require.Equal(t, `error: unknown tool "image_show"`, box.run(t.Context(), "image_show", `{"image":"unknown","caption":""}`))
  62 }
  63
  64 func TestImageShowFallsBackOnlyWhenTelegramRefusesPhoto(t *testing.T) {
  65         testCases := []struct {
  66                 name    string
  67                 status  int
  68                 methods []string
  69                 failure bool
  70         }{
  71                 {
  72                         name:    "refused photo sends document",
  73                         status:  http.StatusBadRequest,
  74                         methods: []string{"sendPhoto", "sendDocument"},
  75                 },
  76                 {
  77                         name:    "uncertain server error does not send twice",
  78                         status:  http.StatusBadGateway,
  79                         methods: []string{"sendPhoto"},
  80                         failure: true,
  81                 },
  82         }
  83
  84         for _, tc := range testCases {
  85                 t.Run(tc.name, func(t *testing.T) {
  86                         telegram := newTelegramFake(t)
  87                         telegram.fail = func(call TelegramFakeCall) (int, string) {
  88                                 if call.method == "sendPhoto" {
  89                                         return tc.status, "photo refused"
  90                                 }
  91
  92                                 return 0, ""
  93                         }
  94
  95                         store := &ImageStore{dir: t.TempDir()}
  96                         picture := imageTestPNG(t, 3, 3)
  97                         name, err := store.Save(picture)
  98                         require.NoError(t, err)
  99
 100                         bot, _ := telegramStartTestBot(t, telegram)
 101                         show := &ImageShow{
 102                                 store: store,
 103                                 bot:   bot,
 104                         }
 105
 106                         _, err = show.Show(t.Context(), name, "Caption")
 107                         if tc.failure {
 108                                 require.Error(t, err)
 109                         } else {
 110                                 require.NoError(t, err)
 111                         }
 112
 113                         var methods []string
 114
 115                         for _, call := range telegram.Calls() {
 116                                 methods = append(methods, call.method)
 117                                 if call.method == "sendDocument" {
 118                                         require.Equal(t, picture, call.params["document"])
 119                                         require.Equal(t, name, call.params["document_name"])
 120                                         require.Equal(t, "Caption", call.params["caption"])
 121                                 }
 122                         }
 123
 124                         require.Equal(t, tc.methods, methods)
 125                 })
 126         }
 127 }
 128
 129 func TestAgentShowsImageInTerminalWithoutTelegram(t *testing.T) {
 130         server := newModelFake(t)
 131         test := newTestTurnLoop(t, server)
 132         name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
 133         require.NoError(t, err)
 134
 135         output := test.loop.tools.run(t.Context(), "image_show", `{"image":"`+name+`","caption":"Here it is."}`)
 136         require.Equal(t, imageText([]string{name}, "Here it is.")+"\n"+filepath.Join(test.agent.Images().dir, name), output)
 137         require.Contains(t, terminalEventText(Event{
 138                 Kind: "echo",
 139                 Text: output,
 140         }), filepath.Join(test.agent.Images().dir, name))
 141 }
 142
 143 func TestTellImageMarkerGivesTheChildPixels(t *testing.T) {
 144         server := newModelFake(t)
 145         test := newDaemonMessageTest(t, t.TempDir(), server)
 146         picture := imageTestPNG(t, 4, 2)
 147         name, err := test.agents[0].agent.Images().Save(picture)
 148         require.NoError(t, err)
 149         test.Start(t, 0, 1)
 150
 151         arguments, err := json.Marshal(ToolTellArguments{
 152                 Name:    "Garden",
 153                 Message: imageText([]string{name}, "Review this plan."),
 154         })
 155         require.NoError(t, err)
 156         require.Contains(t, test.loops[0].tools.run(t.Context(), "tell", string(arguments)), "Sent.")
 157         test.events[1].WaitFor(t, "idle", 1)
 158
 159         var childCall *ModelFakeCall
 160
 161         for i := range server.Calls() {
 162                 call := server.Call(i)
 163                 if call.body["model"] == "gpt-test-garden" {
 164                         childCall = &call
 165                         break
 166                 }
 167         }
 168
 169         require.NotNil(t, childCall)
 170         require.Contains(t, string(childCall.raw), "input_image")
 171         require.Contains(t, string(childCall.raw), base64.StdEncoding.EncodeToString(picture))
 172         require.Contains(t, test.events[1].Last("user"), "[image "+name+"]")
 173         require.NotContains(t, promptToolNames(test.agents[1].agent.Tools()), "image_show")
 174 }
 175
 176 func TestAgentShowsTelegramReceivedImageToItsOwner(t *testing.T) {
 177         picture := imageTestPNG(t, 6, 4)
 178         store := &ImageStore{dir: t.TempDir()}
 179         name, err := store.Save(picture)
 180         require.NoError(t, err)
 181
 182         telegram := newTelegramFake(t)
 183         telegram.Image("owner-photo", picture)
 184
 185         bot, logs := telegramStartTestBot(t, telegram)
 186         server := newModelFake(t)
 187         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 188                 if call.number == 0 {
 189                         arguments, err := json.Marshal(ToolImageShowArguments{
 190                                 Image:   name,
 191                                 Caption: "Here is your photo.",
 192                         })
 193                         require.NoError(t, err)
 194                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "image_show", string(arguments))}, 100, 0, 5))
 195
 196                         return
 197                 }
 198
 199                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("Shown.")}, 100, 0, 5))
 200         }
 201
 202         socket := daemonStartTestWithBot(t, server, bot)
 203         terminal := terminalStartTestClient(t, socket)
 204         terminal.WaitFor(t, "[connected to OneChat]")
 205
 206         message := telegramOwnerMessage("")
 207         message.Caption = "Show this photo."
 208         message.Photo = []TelegramPhoto{{
 209                 FileID: "owner-photo",
 210                 Width:  6,
 211                 Height: 4,
 212         }}
 213
 214         telegram.Add(TelegramUpdate{
 215                 UpdateID: 1,
 216                 Message:  message,
 217         })
 218
 219         terminal.WaitFor(t, "Shown.\n[done]")
 220         require.Contains(t, terminal.Output(), "[result] [image "+name+"]")
 221         require.Contains(t, terminal.Output(), filepath.Join(imageDirName, name))
 222
 223         var photos []TelegramFakeCall
 224
 225         for _, call := range telegram.Calls() {
 226                 if call.method == "sendPhoto" {
 227                         photos = append(photos, call)
 228                 }
 229         }
 230
 231         require.Len(t, photos, 1)
 232         require.Equal(t, picture, photos[0].params["photo"])
 233         require.Equal(t, "4242", photos[0].params["chat_id"])
 234         require.Equal(t, "Here is your photo.", photos[0].params["caption"])
 235         require.NotContains(t, logs.String(), "Here is your photo.")
 236         require.NotContains(t, logs.String(), "Show this photo.")
 237         require.NotContains(t, logs.String(), telegramTestToken)
 238 }