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

   1 package main
   2
   3 import (
   4         "net/http"
   5         "path/filepath"
   6         "testing"
   7
   8         "github.com/stretchr/testify/require"
   9 )
  10
  11 func TestChainReturnsTheRepliesOfTheFinalAnswer(t *testing.T) {
  12         server := newModelFake(t)
  13         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  14                 if call.number == 0 {
  15                         modelWriteEvents(w, modelCompletedEvent([]any{
  16                                 modelMessageItem("Looking."),
  17                                 modelFunctionCallItem("fc_1", "call_1", "shell", `{"command":"echo hi"}`),
  18                         }, 100, 0, 5))
  19
  20                         return
  21                 }
  22
  23                 modelWriteEvents(w, modelCompletedEvent([]any{
  24                         modelMessageItem("It said hi."),
  25                         modelMessageItem("Done."),
  26                 }, 120, 100, 5))
  27         }
  28
  29         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
  30         require.NoError(t, err)
  31
  32         t.Cleanup(func() {
  33                 require.NoError(t, memory.Close())
  34         })
  35
  36         chain := &Chain{
  37                 client: newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl")),
  38                 tools: &Toolbox{
  39                         log:       memory,
  40                         chat:      memory,
  41                         subagents: newTestSubagents(t),
  42                         workdir:   t.TempDir(),
  43                 },
  44                 request: ModelRequest{
  45                         Model:  "gpt-test",
  46                         Effort: "low",
  47                 },
  48                 take: func() []Message {
  49                         return nil
  50                 },
  51         }
  52
  53         text, err := chain.Run(t.Context(), []any{modelUserMessage("Run echo hi.")})
  54         require.NoError(t, err)
  55         require.Equal(t, "It said hi.\n\nDone.", text)
  56         require.Equal(t, 2, server.Calls())
  57
  58         var kinds []string
  59         for id := range memory.Len() {
  60                 message, err := memory.Message(id)
  61                 require.NoError(t, err)
  62
  63                 kinds = append(kinds, message.Kind)
  64         }
  65
  66         require.Equal(t, []string{"reply", "tool", "echo", "reply", "reply"}, kinds)
  67 }