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

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "io"
   6         "net/http"
   7         "path/filepath"
   8         "strconv"
   9         "strings"
  10         "testing"
  11         "unicode/utf8"
  12
  13         "github.com/stretchr/testify/require"
  14 )
  15
  16 func TestCompactorSendsThePromptToolsAndTask(t *testing.T) {
  17         server := newFakeResponses(t)
  18         compactor := newTestCompactor(t, server)
  19
  20         line, err := compactor.Summarize(t.Context(), Task{
  21                 Level:   0,
  22                 Index:   7,
  23                 Context: "<chat>\n0+4|a\n4+1|b\n5+1|c\n6+1|d\n7+1|e\n</chat>",
  24                 Input:   []string{"user: " + strings.Repeat("x", 600)},
  25         })
  26         require.NoError(t, err)
  27         require.Equal(t, "a line", line)
  28         require.Equal(t, 1, server.Calls())
  29
  30         body := server.Call(0).body
  31         require.Equal(t, prompt("Ada"), body["instructions"])
  32         require.Equal(t, "gpt-test", body["model"])
  33         require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
  34         require.Equal(t, "agent-key", body["prompt_cache_key"])
  35
  36         tools, err := json.Marshal(agentTools())
  37         require.NoError(t, err)
  38
  39         sent, err := json.Marshal(body["tools"])
  40         require.NoError(t, err)
  41         require.JSONEq(t, string(tools), string(sent))
  42
  43         parts := inputParts(t, server.Call(0))
  44         require.Equal(t, []string{
  45                 "<chat>\n",
  46                 "0+4|a\n4+1|b\n5+1|c\n6+1|d\n",
  47                 "7+1|e\n</chat>\n" +
  48                         "Compaction: compress message 7 into one line of at most 512 bytes\n" +
  49                         "(about 70 words), the length of this ruler:\n" +
  50                         strings.Repeat("-", 512) + "\n" +
  51                         "<input>\nuser: " + strings.Repeat("x", 600) + "\n</input>",
  52         }, parts)
  53 }
  54
  55 func TestCompactionPartsEndWithTheRestAndTheTask(t *testing.T) {
  56         testCases := []struct {
  57                 name    string
  58                 context string
  59                 want    []string
  60         }{
  61                 {
  62                         name:    "empty view has the opening tag and the last part",
  63                         context: "<chat>\n</chat>",
  64                         want:    []string{"<chat>\n", "</chat>\nTASK"},
  65                 },
  66                 {
  67                         name:    "lines short of a block go into the last part",
  68                         context: "<chat>\n0+1|a\n1+1|b\n2+1|c\n</chat>",
  69                         want:    []string{"<chat>\n", "0+1|a\n1+1|b\n2+1|c\n</chat>\nTASK"},
  70                 },
  71                 {
  72                         name:    "one whole block",
  73                         context: "<chat>\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n</chat>",
  74                         want:    []string{"<chat>\n", "0+1|a\n1+1|b\n2+1|c\n3+1|d\n", "</chat>\nTASK"},
  75                 },
  76                 {
  77                         name:    "two whole blocks and a rest",
  78                         context: "<chat>\n0+2|a\n2+1|b\n3+1|c\n4+1|d\n5+1|e\n6+1|f\n7+1|g\n8+1|h\n9+1|i\n</chat>",
  79                         want: []string{
  80                                 "<chat>\n",
  81                                 "0+2|a\n2+1|b\n3+1|c\n4+1|d\n",
  82                                 "5+1|e\n6+1|f\n7+1|g\n8+1|h\n",
  83                                 "9+1|i\n</chat>\nTASK",
  84                         },
  85                 },
  86         }
  87
  88         for _, tc := range testCases {
  89                 t.Run(tc.name, func(t *testing.T) {
  90                         parts := compactionParts(tc.context, "TASK")
  91                         require.Equal(t, tc.want, parts)
  92                         require.Equal(t, tc.context+"\nTASK", strings.Join(parts, ""))
  93                 })
  94         }
  95 }
  96
  97 func TestCompactionPartsOfALongerViewBeginWithTheBlocksOfAShorterOne(t *testing.T) {
  98         shorter := compactionParts("<chat>\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n4+1|e\n</chat>", "task 5")
  99         longer := compactionParts("<chat>\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n4+1|e\n5+1|f\n6+1|g\n7+1|h\n</chat>", "task 8")
 100         require.Len(t, shorter, 3)
 101         require.Equal(t, shorter[:2], longer[:2])
 102         require.NotEqual(t, shorter[2], longer[2])
 103 }
 104
 105 func TestCompactorWritesTheMergeTask(t *testing.T) {
 106         text, err := compactionTask(Task{
 107                 Level: 2,
 108                 Index: 3,
 109                 Input: []string{"left line", "right line"},
 110         })
 111         require.NoError(t, err)
 112         require.Equal(t, "Compaction: merge lines 12+2 and 14+2, adjacent, into one line of at most\n"+
 113                 "512 bytes (about 70 words), the length of this ruler:\n"+
 114                 strings.Repeat("-", 512)+"\n"+
 115                 "<chat> may hold their messages, 12 to 15, in more detail: take details\n"+
 116                 "of them from there too.\n"+
 117                 "<input>\nleft line\nright line\n</input>", text)
 118 }
 119
 120 func TestCompactorRetriesALongLineWithTheNewestNote(t *testing.T) {
 121         long := []string{strings.Repeat("a", 600), strings.Repeat("b", 550), strings.Repeat("c", 400)}
 122         server := newFakeResponses(t)
 123         server.reply = func(w http.ResponseWriter, call fakeCall) {
 124                 writeEvents(w, completedEvent([]any{messageItem(long[call.number])}, 100, 0, 5))
 125         }
 126
 127         compactor := newTestCompactor(t, server)
 128         line, err := compactor.Summarize(t.Context(), leafTask())
 129         require.NoError(t, err)
 130         require.Equal(t, long[2], line)
 131         require.Equal(t, 3, server.Calls())
 132
 133         records := readUsage(t, compactor.client.usage.path)
 134         require.Len(t, records, 3)
 135         for n, record := range records {
 136                 require.Equal(t, &[2]int{0, 0}, record.Node)
 137                 require.Equal(t, len(long[n]), record.TextBytes)
 138         }
 139
 140         first := inputParts(t, server.Call(0))
 141         require.Len(t, first, 2)
 142         for n, want := range []string{long[0], long[1]} {
 143                 call := server.Call(n + 1)
 144                 parts := inputParts(t, call)
 145                 require.Equal(t, first, parts[:2])
 146                 require.Equal(t, []string{tooLong(want)}, parts[2:])
 147                 require.Contains(t, parts[2], "was "+strconv.Itoa(len(want))+" bytes,")
 148                 require.Contains(t, parts[2], "\n"+want[:512]+"| <- LIMIT")
 149
 150                 note, err := json.Marshal(map[string]string{
 151                         "type": "input_text",
 152                         "text": parts[2],
 153                 })
 154                 require.NoError(t, err)
 155                 require.LessOrEqual(t, len(call.raw), len(server.Call(0).raw)+len(",")+len(note))
 156         }
 157 }
 158
 159 func TestCompactorKeepsTheShortestOfFiveLines(t *testing.T) {
 160         sizes := []int{700, 530, 600, 520, 560, 513}
 161         server := newFakeResponses(t)
 162         server.reply = func(w http.ResponseWriter, call fakeCall) {
 163                 writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("z", sizes[call.number]))}, 100, 0, 5))
 164         }
 165
 166         line, err := newTestCompactor(t, server).Summarize(t.Context(), leafTask())
 167         require.NoError(t, err)
 168         require.Equal(t, compactionTries, server.Calls())
 169         require.Len(t, line, 520)
 170 }
 171
 172 func TestCompactorFailsWithoutALine(t *testing.T) {
 173         testCases := []struct {
 174                 name  string
 175                 reply func(http.ResponseWriter, fakeCall)
 176                 want  string
 177         }{
 178                 {
 179                         name: "failed call",
 180                         reply: func(w http.ResponseWriter, call fakeCall) {
 181                                 w.WriteHeader(http.StatusInternalServerError)
 182                                 _, _ = io.WriteString(w, "down")
 183                         },
 184                         want: "HTTP 500",
 185                 },
 186                 {
 187                         name: "tool call and no text",
 188                         reply: func(w http.ResponseWriter, call fakeCall) {
 189                                 writeEvents(w, completedEvent([]any{map[string]any{"type": "function_call", "name": "zoom", "arguments": "{}"}}, 100, 0, 5))
 190                         },
 191                         want: "no line",
 192                 },
 193         }
 194
 195         for _, tc := range testCases {
 196                 t.Run(tc.name, func(t *testing.T) {
 197                         server := newFakeResponses(t)
 198                         server.reply = tc.reply
 199
 200                         _, err := newTestCompactor(t, server).Summarize(t.Context(), leafTask())
 201                         require.ErrorContains(t, err, tc.want)
 202                         require.Equal(t, 1, server.Calls())
 203                 })
 204         }
 205 }
 206
 207 func TestCutLineKeepsWholeCharacters(t *testing.T) {
 208         line := strings.Repeat("a", 511) + "é" + "tail"
 209         cut := cutLine(line)
 210         require.Equal(t, strings.Repeat("a", 511), cut)
 211         require.True(t, utf8.ValidString(cut))
 212         require.Equal(t, "short", cutLine("short"))
 213 }
 214
 215 func TestCompactorBuildsNodesInTheCore(t *testing.T) {
 216         server := newFakeResponses(t)
 217         server.reply = func(w http.ResponseWriter, call fakeCall) {
 218                 if call.number == 0 {
 219                         w.WriteHeader(http.StatusServiceUnavailable)
 220                         return
 221                 }
 222
 223                 writeEvents(w, completedEvent([]any{messageItem("summary of the long output")}, 100, 0, 5))
 224         }
 225
 226         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newTestCompactor(t, server))
 227         require.NoError(t, err)
 228
 229         defer func() {
 230                 require.NoError(t, memory.Close())
 231         }()
 232
 233         _, err = memory.Append("echo", strings.Repeat("output ", 200))
 234         require.NoError(t, err)
 235         require.NoError(t, memory.Wait(t.Context()))
 236         require.False(t, isBuilt(t, memory.tree, 0, 0))
 237
 238         _, err = memory.Append("user", "next")
 239         require.NoError(t, err)
 240         require.NoError(t, memory.Wait(t.Context()))
 241
 242         node := readNode(t, memory.tree, 0, 0)
 243         require.Equal(t, "summary of the long output", node.Text)
 244         require.Equal(t, "<chat>\n0+1|summary of the long output\n1+1|user: next\n</chat>", memory.View())
 245         require.Equal(t, 2, server.Calls())
 246 }
 247
 248 func newTestCompactor(t *testing.T, server *fakeResponses) *Compactor {
 249         t.Helper()
 250
 251         client := newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 252
 253         return newCompactor(client, "Ada", ModelConfig{
 254                 Model:  "gpt-test",
 255                 Effort: "high",
 256         })
 257 }
 258
 259 func leafTask() Task {
 260         return Task{
 261                 Level:   0,
 262                 Index:   0,
 263                 Context: "<chat>\n</chat>",
 264                 Input:   []string{"user: " + strings.Repeat("y", 600)},
 265         }
 266 }
 267
 268 // inputParts returns the texts of the one user message of call.
 269 func inputParts(t *testing.T, call fakeCall) []string {
 270         t.Helper()
 271
 272         var body struct {
 273                 Input []struct {
 274                         Role    string `json:"role"`
 275                         Content []struct {
 276                                 Type string `json:"type"`
 277                                 Text string `json:"text"`
 278                         } `json:"content"`
 279                 } `json:"input"`
 280         }
 281
 282         require.NoError(t, json.Unmarshal(call.raw, &body))
 283         require.Len(t, body.Input, 1)
 284         require.Equal(t, "user", body.Input[0].Role)
 285
 286         texts := make([]string, len(body.Input[0].Content))
 287         for n, part := range body.Input[0].Content {
 288                 require.Equal(t, "input_text", part.Type)
 289                 texts[n] = part.Text
 290         }
 291
 292         return texts
 293 }