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

   1 package main
   2
   3 import (
   4         "bufio"
   5         "encoding/json"
   6         "fmt"
   7         "io"
   8         "net/http"
   9         "net/http/httptest"
  10         "os"
  11         "sync"
  12         "testing"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 // modelTestKey is a fake API key.
  18 const modelTestKey = "sk-test-not-a-real-key"
  19
  20 // ModelFakeCall is one call to a fake endpoint, numbered from 0.
  21 type ModelFakeCall struct {
  22         number int
  23         header http.Header
  24         raw    []byte
  25         body   map[string]any
  26 }
  27
  28 // ModelFake is an httptest Responses endpoint. It keeps each call, and reply
  29 // writes the response. The default reply is one short message.
  30 type ModelFake struct {
  31         *httptest.Server
  32         mutex sync.Mutex
  33         calls []ModelFakeCall
  34         reply func(http.ResponseWriter, ModelFakeCall)
  35 }
  36
  37 func newModelFake(t *testing.T) *ModelFake {
  38         t.Helper()
  39
  40         f := &ModelFake{
  41                 reply: func(w http.ResponseWriter, _ ModelFakeCall) {
  42                         modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
  43                 },
  44         }
  45
  46         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
  47         t.Cleanup(f.Close)
  48
  49         return f
  50 }
  51
  52 // Call returns call number n.
  53 func (f *ModelFake) Call(n int) ModelFakeCall {
  54         f.mutex.Lock()
  55         defer f.mutex.Unlock()
  56
  57         return f.calls[n]
  58 }
  59
  60 // Calls returns the number of calls.
  61 func (f *ModelFake) Calls() int {
  62         f.mutex.Lock()
  63         defer f.mutex.Unlock()
  64
  65         return len(f.calls)
  66 }
  67
  68 func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
  69         raw, err := io.ReadAll(r.Body)
  70         if err != nil {
  71                 http.Error(w, err.Error(), http.StatusBadRequest)
  72
  73                 return
  74         }
  75
  76         call := ModelFakeCall{
  77                 header: r.Header.Clone(),
  78                 raw:    raw,
  79         }
  80
  81         err = json.Unmarshal(raw, &call.body)
  82         if err != nil {
  83                 http.Error(w, err.Error(), http.StatusBadRequest)
  84
  85                 return
  86         }
  87
  88         f.mutex.Lock()
  89
  90         call.number = len(f.calls)
  91         f.calls = append(f.calls, call)
  92         reply := f.reply
  93         f.mutex.Unlock()
  94
  95         reply(w, call)
  96 }
  97
  98 func modelTestRequest() ModelRequest {
  99         return ModelRequest{
 100                 Model:  "gpt-test",
 101                 Effort: "low",
 102                 Input:  []any{modelUserMessage("hello")},
 103         }
 104 }
 105
 106 // modelWriteEvents writes events as a server-sent event stream.
 107 func modelWriteEvents(w http.ResponseWriter, events ...any) {
 108         w.Header().Set("Content-Type", "text/event-stream")
 109
 110         for _, event := range events {
 111                 data, err := json.Marshal(event)
 112                 if err != nil {
 113                         panic(err)
 114                 }
 115
 116                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 117                 if err != nil {
 118                         panic(err)
 119                 }
 120         }
 121 }
 122
 123 func modelMessageItem(text string) map[string]any {
 124         return map[string]any{
 125                 "type": "message",
 126                 "role": "assistant",
 127                 "content": []any{
 128                         map[string]any{"type": "output_text", "text": text},
 129                 },
 130         }
 131 }
 132
 133 func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
 134         if output == nil {
 135                 output = []any{}
 136         }
 137
 138         return map[string]any{
 139                 "type": "response.completed",
 140                 "response": map[string]any{
 141                         "output": output,
 142                         "usage": map[string]any{
 143                                 "input_tokens":         input,
 144                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 145                                 "output_tokens":        out,
 146                         },
 147                 },
 148         }
 149 }
 150
 151 func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
 152         t.Helper()
 153
 154         file, err := os.Open(path)
 155         require.NoError(t, err)
 156
 157         defer func() {
 158                 require.NoError(t, file.Close())
 159         }()
 160
 161         var records []ModelUsageRecord
 162
 163         scanner := bufio.NewScanner(file)
 164         for scanner.Scan() {
 165                 var record ModelUsageRecord
 166
 167                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 168
 169                 records = append(records, record)
 170         }
 171
 172         require.NoError(t, scanner.Err())
 173
 174         return records
 175 }
 176
 177 func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
 178         return map[string]any{
 179                 "type":      "function_call",
 180                 "id":        id,
 181                 "call_id":   callID,
 182                 "name":      name,
 183                 "arguments": arguments,
 184         }
 185 }