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 // Call returns call number n.
  38 func (f *ModelFake) Call(n int) ModelFakeCall {
  39         f.mutex.Lock()
  40         defer f.mutex.Unlock()
  41
  42         return f.calls[n]
  43 }
  44
  45 // Calls returns the number of calls.
  46 func (f *ModelFake) Calls() int {
  47         f.mutex.Lock()
  48         defer f.mutex.Unlock()
  49
  50         return len(f.calls)
  51 }
  52
  53 func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
  54         raw, err := io.ReadAll(r.Body)
  55         if err != nil {
  56                 http.Error(w, err.Error(), http.StatusBadRequest)
  57                 return
  58         }
  59
  60         call := ModelFakeCall{
  61                 header: r.Header.Clone(),
  62                 raw:    raw,
  63         }
  64
  65         err = json.Unmarshal(raw, &call.body)
  66         if err != nil {
  67                 http.Error(w, err.Error(), http.StatusBadRequest)
  68                 return
  69         }
  70
  71         f.mutex.Lock()
  72         call.number = len(f.calls)
  73         f.calls = append(f.calls, call)
  74         reply := f.reply
  75         f.mutex.Unlock()
  76
  77         reply(w, call)
  78 }
  79
  80 func newModelFake(t *testing.T) *ModelFake {
  81         t.Helper()
  82
  83         f := &ModelFake{
  84                 reply: func(w http.ResponseWriter, call ModelFakeCall) {
  85                         modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
  86                 },
  87         }
  88
  89         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
  90         t.Cleanup(f.Close)
  91
  92         return f
  93 }
  94
  95 func modelTestRequest() ModelRequest {
  96         return ModelRequest{
  97                 Model:  "gpt-test",
  98                 Effort: "low",
  99                 Input:  []any{modelUserMessage("hello")},
 100         }
 101 }
 102
 103 // modelWriteEvents writes events as a server-sent event stream.
 104 func modelWriteEvents(w http.ResponseWriter, events ...any) {
 105         w.Header().Set("Content-Type", "text/event-stream")
 106         for _, event := range events {
 107                 data, err := json.Marshal(event)
 108                 if err != nil {
 109                         panic(err)
 110                 }
 111
 112                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 113                 if err != nil {
 114                         panic(err)
 115                 }
 116         }
 117 }
 118
 119 func modelMessageItem(text string) map[string]any {
 120         return map[string]any{
 121                 "type": "message",
 122                 "role": "assistant",
 123                 "content": []any{
 124                         map[string]any{"type": "output_text", "text": text},
 125                 },
 126         }
 127 }
 128
 129 func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
 130         if output == nil {
 131                 output = []any{}
 132         }
 133
 134         return map[string]any{
 135                 "type": "response.completed",
 136                 "response": map[string]any{
 137                         "output": output,
 138                         "usage": map[string]any{
 139                                 "input_tokens":         input,
 140                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 141                                 "output_tokens":        out,
 142                         },
 143                 },
 144         }
 145 }
 146
 147 func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
 148         t.Helper()
 149
 150         file, err := os.Open(path)
 151         require.NoError(t, err)
 152
 153         defer func() {
 154                 require.NoError(t, file.Close())
 155         }()
 156
 157         var records []ModelUsageRecord
 158         scanner := bufio.NewScanner(file)
 159         for scanner.Scan() {
 160                 var record ModelUsageRecord
 161                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 162                 records = append(records, record)
 163         }
 164
 165         require.NoError(t, scanner.Err())
 166
 167         return records
 168 }
 169
 170 func modelReasoningItem(id, encrypted string) map[string]any {
 171         return map[string]any{
 172                 "type":              "reasoning",
 173                 "id":                id,
 174                 "encrypted_content": encrypted,
 175                 "summary":           []any{},
 176         }
 177 }
 178
 179 func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
 180         return map[string]any{
 181                 "type":      "function_call",
 182                 "id":        id,
 183                 "call_id":   callID,
 184                 "name":      name,
 185                 "arguments": arguments,
 186         }
 187 }