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