X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/4b9b22619008f0829b7f5a8685e7b1d2a9aeb6bb..e1e48ba388c28105fff3b70c8ca02304f5d57f53:/compactor_test.go diff --git a/compactor_test.go b/compactor_test.go index 20862c7..a128bc7 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,82 +28,30 @@ 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)) + parts := compactorInputParts(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) } -func TestCompactionPartsEndWithTheRestAndTheTask(t *testing.T) { - testCases := []struct { - name string - context string - want []string - }{ - { - name: "empty view has the opening tag and the last part", - context: "\n", - want: []string{"\n", "\nTASK"}, - }, - { - name: "lines short of a block go into the last part", - context: "\n0+1|a\n1+1|b\n2+1|c\n", - want: []string{"\n", "0+1|a\n1+1|b\n2+1|c\n\nTASK"}, - }, - { - name: "one whole block", - 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", "\nTASK"}, - }, - { - name: "two whole blocks and a rest", - 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", - "5+1|e\n6+1|f\n7+1|g\n8+1|h\n", - "9+1|i\n\nTASK", - }, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - parts := compactionParts(tc.context, "TASK") - require.Equal(t, tc.want, parts) - require.Equal(t, tc.context+"\nTASK", strings.Join(parts, "")) - }) - } -} - -func TestCompactionPartsOfALongerViewBeginWithTheBlocksOfAShorterOne(t *testing.T) { - shorter := compactionParts("\n0+1|a\n1+1|b\n2+1|c\n3+1|d\n4+1|e\n", "task 5") - longer := compactionParts("\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", "task 8") - require.Len(t, shorter, 3) - require.Equal(t, shorter[:2], longer[:2]) - require.NotEqual(t, shorter[2], longer[2]) -} - func TestCompactorWritesTheMergeTask(t *testing.T) { - text, err := compactionTask(Task{ + text, err := compactorTask(Task{ Level: 2, Index: 3, Input: []string{"left line", "right line"}, @@ -119,31 +67,31 @@ 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)) } compactor := newTestCompactor(t, server) - line, err := compactor.Summarize(t.Context(), leafTask()) + line, err := compactor.Summarize(t.Context(), compactorLeafTask()) 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)) + first := compactorInputParts(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) + parts := compactorInputParts(t, call) require.Equal(t, first, parts[:2]) - require.Equal(t, []string{tooLong(want)}, parts[2:]) + require.Equal(t, []string{compactorTooLong(want)}, parts[2:]) require.Contains(t, parts[2], "was "+strconv.Itoa(len(want))+" bytes,") require.Contains(t, parts[2], "\n"+want[:512]+"| <- LIMIT") @@ -158,26 +106,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") }, @@ -185,8 +133,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", }, @@ -194,10 +142,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()) }) @@ -206,21 +154,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)) @@ -233,40 +181,31 @@ 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 { +// compactorInputParts returns the texts of the one user message of call. +func compactorInputParts(t *testing.T, call ModelFakeCall) []string { t.Helper() var body struct { @@ -291,3 +230,17 @@ func inputParts(t *testing.T, call fakeCall) []string { 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 +}