X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/56129a5fcbcbe51a7ac7d6b009d4551ed9f71109..459656ed8c15847d325f07758e4a775935411c3d:/compactor_test.go?ds=inline diff --git a/compactor_test.go b/compactor_test.go index cc99f33..0d591ee 100644 --- a/compactor_test.go +++ b/compactor_test.go @@ -14,13 +14,13 @@ import ( ) func TestCompactorSendsThePromptToolsAndTask(t *testing.T) { - server := newFakeResponses(t) + server := newModelFake(t) compactor := newTestCompactor(t, server) line, err := compactor.Summarize(t.Context(), Task{ Level: 0, Index: 7, - Context: "\n0+4|a\n4+1|b\n5+1|c\n6+1|d\n7+1|e\n", + Context: "\n0+4|user: earlier\n", Input: []string{"user: " + strings.Repeat("x", 600)}, }) require.NoError(t, err) @@ -28,81 +28,114 @@ func TestCompactorSendsThePromptToolsAndTask(t *testing.T) { require.Equal(t, 1, server.Calls()) body := server.Call(0).body - require.Equal(t, prompt("Ada"), body["instructions"]) + require.Equal(t, prompt("Ada", ""), body["instructions"]) require.Equal(t, "gpt-test", body["model"]) require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"]) require.Equal(t, "agent-key", body["prompt_cache_key"]) - tools, err := json.Marshal(agentTools()) + tools, err := json.Marshal(promptTools()) require.NoError(t, err) sent, err := json.Marshal(body["tools"]) require.NoError(t, err) require.JSONEq(t, string(tools), string(sent)) - parts := inputParts(t, server.Call(0)) + texts := compactorInputTexts(t, server.Call(0)) require.Equal(t, []string{ - "\n", - "0+4|a\n4+1|b\n5+1|c\n6+1|d\n", - "7+1|e\n\n", - "Compaction: compress message 7 into one line of at most 512 bytes\n" + + "\n0+4|user: earlier\n\n" + + "Compaction: compress message 7 into one line of at most 512 bytes\n" + "(about 70 words), the length of this ruler:\n" + strings.Repeat("-", 512) + "\n" + "\nuser: " + strings.Repeat("x", 600) + "\n", - }, parts) + }, texts) } -func TestViewPartsSendWholeBlocksOfLines(t *testing.T) { +func TestCompactorSendsEachWholeBlockOfTheViewAsAMessage(t *testing.T) { testCases := []struct { name string context string want []string }{ { - name: "empty view has the tags only", + name: "empty view goes with the task", context: "\n", - want: []string{"\n", "\n"}, + want: []string{"\n\nTASK"}, }, { - name: "lines short of a block go with the closing tag", + name: "lines short of a block go with the task", context: "\n0+1|a\n1+1|b\n2+1|c\n", - want: []string{"\n", "0+1|a\n1+1|b\n2+1|c\n\n"}, + want: []string{"\n0+1|a\n1+1|b\n2+1|c\n\nTASK"}, }, { - name: "one whole block", + name: "one whole block and the task", context: "\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n", - want: []string{"\n", "0+1|a\n1+1|b\n2+1|c\n3+1|d\n", "\n"}, + want: []string{"\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n", "\nTASK"}, }, { - name: "two whole blocks and a rest", + name: "two whole blocks and the rest with the task", context: "\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", want: []string{ - "\n", - "0+2|a\n2+1|b\n3+1|c\n4+1|d\n", + "\n0+2|a\n2+1|b\n3+1|c\n4+1|d\n", "5+1|e\n6+1|f\n7+1|g\n8+1|h\n", - "9+1|i\n\n", + "9+1|i\n\nTASK", }, }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - parts := viewParts(tc.context) - require.Equal(t, tc.want, parts) - require.Equal(t, tc.context+"\n", strings.Join(parts, "")) + server := newModelFake(t) + task := compactorLeafTask() + task.Context = tc.context + + _, err := newTestCompactor(t, server).Summarize(t.Context(), task) + require.NoError(t, err) + + text, err := compactorTask(task) + require.NoError(t, err) + + want := make([]string, len(tc.want)) + for n, message := range tc.want { + want[n] = strings.Replace(message, "TASK", text, 1) + } + + texts := compactorInputTexts(t, server.Call(0)) + require.Equal(t, want, texts) + require.Equal(t, tc.context+"\n"+text, strings.Join(texts, "")) }) } } -func TestViewPartsKeepTheBlocksOfAShorterView(t *testing.T) { - shorter := viewParts("\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n4+1|e\n") - longer := viewParts("\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") - require.Equal(t, shorter[:2], longer[:2]) - require.NotEqual(t, shorter[2], longer[2]) +func TestCompactorSendsTheSameBlocksForViewsWithTheSameFirstLines(t *testing.T) { + shared := "\n0+4|a\n4+1|b\n5+1|c\n6+1|d\n7+1|e\n8+1|f\n9+1|g\n10+1|h\n" + server := newModelFake(t) + compactor := newTestCompactor(t, server) + _, err := compactor.Summarize(t.Context(), Task{ + Level: 0, + Index: 11, + Context: shared + "11+1|i\n", + Input: []string{"user: eleven"}, + }) + require.NoError(t, err) + + _, err = compactor.Summarize(t.Context(), Task{ + Level: 1, + Index: 5, + Context: shared + "11+1|j\n12+1|k\n", + Input: []string{"left line", "right line"}, + }) + require.NoError(t, err) + + first := compactorInputTexts(t, server.Call(0)) + second := compactorInputTexts(t, server.Call(1)) + require.Len(t, first, 3) + require.Len(t, second, 3) + require.Equal(t, first[:2], second[:2]) + require.NotEqual(t, first[2], second[2]) } func TestCompactorWritesTheMergeTask(t *testing.T) { - text, err := compactionTask(Task{ + text, err := compactorTask(Task{ Level: 2, Index: 3, Input: []string{"left line", "right line"}, @@ -118,38 +151,38 @@ func TestCompactorWritesTheMergeTask(t *testing.T) { func TestCompactorRetriesALongLineWithTheNewestNote(t *testing.T) { long := []string{strings.Repeat("a", 600), strings.Repeat("b", 550), strings.Repeat("c", 400)} - server := newFakeResponses(t) - server.reply = func(w http.ResponseWriter, call fakeCall) { - writeEvents(w, completedEvent([]any{messageItem(long[call.number])}, 100, 0, 5)) + server := newModelFake(t) + server.reply = func(w http.ResponseWriter, call ModelFakeCall) { + modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(long[call.number])}, 100, 0, 5)) } + task := compactorLeafTask() + task.Context = "\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n4+1|e\n" + compactor := newTestCompactor(t, server) - line, err := compactor.Summarize(t.Context(), leafTask()) + line, err := compactor.Summarize(t.Context(), task) require.NoError(t, err) require.Equal(t, long[2], line) require.Equal(t, 3, server.Calls()) - records := readUsage(t, compactor.client.usage.path) + records := modelReadUsage(t, compactor.client.usage.path) require.Len(t, records, 3) for n, record := range records { require.Equal(t, &[2]int{0, 0}, record.Node) require.Equal(t, len(long[n]), record.TextBytes) } - first := inputParts(t, server.Call(0)) - require.Len(t, first, 3) + first := compactorInputTexts(t, server.Call(0)) + require.Len(t, first, 2) for n, want := range []string{long[0], long[1]} { call := server.Call(n + 1) - parts := inputParts(t, call) - require.Equal(t, first, parts[:3]) - require.Equal(t, []string{tooLong(want)}, parts[3:]) - require.Contains(t, parts[3], "was "+strconv.Itoa(len(want))+" bytes,") - require.Contains(t, parts[3], "\n"+want[:512]+"| <- LIMIT") - - note, err := json.Marshal(map[string]string{ - "type": "input_text", - "text": parts[3], - }) + texts := compactorInputTexts(t, call) + require.Equal(t, first, texts[:2]) + require.Equal(t, []string{compactorTooLong(want)}, texts[2:]) + require.Contains(t, texts[2], "was "+strconv.Itoa(len(want))+" bytes,") + require.Contains(t, texts[2], "\n"+want[:512]+"| <- LIMIT") + + note, err := json.Marshal(modelUserMessage(texts[2])) require.NoError(t, err) require.LessOrEqual(t, len(call.raw), len(server.Call(0).raw)+len(",")+len(note)) } @@ -157,26 +190,26 @@ func TestCompactorRetriesALongLineWithTheNewestNote(t *testing.T) { func TestCompactorKeepsTheShortestOfFiveLines(t *testing.T) { sizes := []int{700, 530, 600, 520, 560, 513} - server := newFakeResponses(t) - server.reply = func(w http.ResponseWriter, call fakeCall) { - writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("z", sizes[call.number]))}, 100, 0, 5)) + server := newModelFake(t) + server.reply = func(w http.ResponseWriter, call ModelFakeCall) { + modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("z", sizes[call.number]))}, 100, 0, 5)) } - line, err := newTestCompactor(t, server).Summarize(t.Context(), leafTask()) + line, err := newTestCompactor(t, server).Summarize(t.Context(), compactorLeafTask()) require.NoError(t, err) - require.Equal(t, compactionTries, server.Calls()) + require.Equal(t, compactorTries, server.Calls()) require.Len(t, line, 520) } func TestCompactorFailsWithoutALine(t *testing.T) { testCases := []struct { name string - reply func(http.ResponseWriter, fakeCall) + reply func(http.ResponseWriter, ModelFakeCall) want string }{ { name: "failed call", - reply: func(w http.ResponseWriter, call fakeCall) { + reply: func(w http.ResponseWriter, call ModelFakeCall) { w.WriteHeader(http.StatusInternalServerError) _, _ = io.WriteString(w, "down") }, @@ -184,8 +217,8 @@ func TestCompactorFailsWithoutALine(t *testing.T) { }, { name: "tool call and no text", - reply: func(w http.ResponseWriter, call fakeCall) { - writeEvents(w, completedEvent([]any{map[string]any{"type": "function_call", "name": "zoom", "arguments": "{}"}}, 100, 0, 5)) + reply: func(w http.ResponseWriter, call ModelFakeCall) { + modelWriteEvents(w, modelCompletedEvent([]any{map[string]any{"type": "function_call", "name": "zoom", "arguments": "{}"}}, 100, 0, 5)) }, want: "no line", }, @@ -193,10 +226,10 @@ func TestCompactorFailsWithoutALine(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - server := newFakeResponses(t) + server := newModelFake(t) server.reply = tc.reply - _, err := newTestCompactor(t, server).Summarize(t.Context(), leafTask()) + _, err := newTestCompactor(t, server).Summarize(t.Context(), compactorLeafTask()) require.ErrorContains(t, err, tc.want) require.Equal(t, 1, server.Calls()) }) @@ -205,21 +238,21 @@ func TestCompactorFailsWithoutALine(t *testing.T) { func TestCutLineKeepsWholeCharacters(t *testing.T) { line := strings.Repeat("a", 511) + "é" + "tail" - cut := cutLine(line) + cut := compactorCutLine(line) require.Equal(t, strings.Repeat("a", 511), cut) require.True(t, utf8.ValidString(cut)) - require.Equal(t, "short", cutLine("short")) + require.Equal(t, "short", compactorCutLine("short")) } func TestCompactorBuildsNodesInTheCore(t *testing.T) { - server := newFakeResponses(t) - server.reply = func(w http.ResponseWriter, call fakeCall) { + server := newModelFake(t) + server.reply = func(w http.ResponseWriter, call ModelFakeCall) { if call.number == 0 { w.WriteHeader(http.StatusServiceUnavailable) return } - writeEvents(w, completedEvent([]any{messageItem("summary of the long output")}, 100, 0, 5)) + modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("summary of the long output")}, 100, 0, 5)) } memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newTestCompactor(t, server)) @@ -232,40 +265,32 @@ func TestCompactorBuildsNodesInTheCore(t *testing.T) { _, err = memory.Append("echo", strings.Repeat("output ", 200)) require.NoError(t, err) require.NoError(t, memory.Wait(t.Context())) - require.False(t, isBuilt(t, memory.tree, 0, 0)) + require.False(t, treeIsBuilt(t, memory.tree, 0, 0)) _, err = memory.Append("user", "next") require.NoError(t, err) require.NoError(t, memory.Wait(t.Context())) - node := readNode(t, memory.tree, 0, 0) + node := treeReadNode(t, memory.tree, 0, 0) require.Equal(t, "summary of the long output", node.Text) require.Equal(t, "\n0+1|summary of the long output\n1+1|user: next\n", memory.View()) require.Equal(t, 2, server.Calls()) } -func newTestCompactor(t *testing.T, server *fakeResponses) *Compactor { +func newTestCompactor(t *testing.T, server *ModelFake) *Compactor { t.Helper() - client := newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl")) + client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl")) - return newCompactor(client, "Ada", ModelConfig{ + return newCompactor(client, compactorAdaPrompt, ModelConfig{ Model: "gpt-test", Effort: "high", }) } -func leafTask() Task { - return Task{ - Level: 0, - Index: 0, - Context: "\n", - Input: []string{"user: " + strings.Repeat("y", 600)}, - } -} - -// inputParts returns the texts of the one user message of call. -func inputParts(t *testing.T, call fakeCall) []string { +// compactorInputTexts returns the text of each user message of call. Each +// message has one text part. +func compactorInputTexts(t *testing.T, call ModelFakeCall) []string { t.Helper() var body struct { @@ -279,14 +304,28 @@ func inputParts(t *testing.T, call fakeCall) []string { } require.NoError(t, json.Unmarshal(call.raw, &body)) - require.Len(t, body.Input, 1) - require.Equal(t, "user", body.Input[0].Role) - texts := make([]string, len(body.Input[0].Content)) - for n, part := range body.Input[0].Content { - require.Equal(t, "input_text", part.Type) - texts[n] = part.Text + texts := make([]string, len(body.Input)) + for n, message := range body.Input { + require.Equal(t, "user", message.Role) + require.Len(t, message.Content, 1) + require.Equal(t, "input_text", message.Content[0].Type) + texts[n] = message.Content[0].Text } return texts } + +func compactorLeafTask() Task { + return Task{ + Level: 0, + Index: 0, + Context: "\n", + Input: []string{"user: " + strings.Repeat("y", 600)}, + } +} + +// compactorAdaPrompt returns the prompt of the agent Ada, without instructions. +func compactorAdaPrompt() (string, error) { + return prompt("Ada", ""), nil +}