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                         workdir: t.TempDir(),
  42                 },
  43                 request: ModelRequest{
  44                         Model:  "gpt-test",
  45                         Effort: "low",
  46                 },
  47                 take: func() []Message {
  48                         return nil
  49                 },
  50         }
  51
  52         text, err := chain.Run(t.Context(), []any{modelUserMessage("Run echo hi.")})
  53         require.NoError(t, err)
  54         require.Equal(t, "It said hi.\n\nDone.", text)
  55         require.Equal(t, 2, server.Calls())
  56
  57         var kinds []string
  58         for id := range memory.Len() {
  59                 message, err := memory.Message(id)
  60                 require.NoError(t, err)
  61
  62                 kinds = append(kinds, message.Kind)
  63         }
  64
  65         require.Equal(t, []string{"reply", "tool", "echo", "reply", "reply"}, kinds)
  66 }