X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/6d1d9d9a7cf90b25a2e10af1c5bc56faee7e494c..8d9547f6e7ffa1e95ad25d9f8f2e7cfd23d992d0:/model_test.go?ds=inline diff --git a/model_test.go b/model_test.go index 2197eca..11a7d00 100644 --- a/model_test.go +++ b/model_test.go @@ -1,24 +1,16 @@ package main import ( - "bufio" - "encoding/json" - "fmt" "io" "net/http" - "net/http/httptest" "os" "path/filepath" - "sync" "testing" "time" "github.com/stretchr/testify/require" ) -// modelTestKey is a fake API key. -const modelTestKey = "sk-test-not-a-real-key" - func TestClientSendsTheFixedRequestFields(t *testing.T) { server := newModelFake(t) client := newTestModelClient(t, server) @@ -93,6 +85,19 @@ func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) { call := server.Call(0) require.Equal(t, tc.want, call.header.Get("session_id")) require.Equal(t, "agent-key", call.body["prompt_cache_key"]) + + // A subagent has its own key. + request := modelTestRequest() + request.Agent = "GreenOtter" + _, err = client.Create(t.Context(), request) + require.NoError(t, err) + + call = server.Call(1) + if tc.want != "" { + require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id")) + } + + require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"]) }) } } @@ -152,11 +157,13 @@ func TestClientRecordsEachCall(t *testing.T) { client := newTestModelClient(t, server) request := modelTestRequest() - request.Node = &nodeKey{2, 5} + request.Node = &NodeKey{2, 5} + request.Agent = "GreenOtter" _, err := client.Create(t.Context(), request) require.NoError(t, err) require.NotContains(t, server.Call(0).body, "node") require.NotContains(t, string(server.Call(0).raw), "[2,5]") + require.NotContains(t, server.Call(0).body, "agent") _, err = client.Create(t.Context(), modelTestRequest()) require.ErrorContains(t, err, "HTTP 429") @@ -176,12 +183,14 @@ func TestClientRecordsEachCall(t *testing.T) { require.Equal(t, 7, records[0].OutputTokens) require.Equal(t, quota, records[0].Quota) require.Equal(t, &[2]int{2, 5}, records[0].Node) + require.Equal(t, "GreenOtter", records[0].Agent) require.Equal(t, len("ok"), records[0].TextBytes) require.Greater(t, records[0].Seconds, 0.0) require.Empty(t, records[0].Error) require.NotEmpty(t, records[0].Date) require.Nil(t, records[1].Node) + require.Empty(t, records[1].Agent) require.Zero(t, records[1].TextBytes) require.Equal(t, http.StatusTooManyRequests, records[1].Status) @@ -261,177 +270,8 @@ func TestClientFailsWhenTheResponseFails(t *testing.T) { } } -// ModelFakeCall is one call to a fake endpoint, numbered from 0. -type ModelFakeCall struct { - number int - header http.Header - raw []byte - body map[string]any -} - -// ModelFake is an httptest Responses endpoint. It keeps each call, and reply -// writes the response. The default reply is one short message. -type ModelFake struct { - *httptest.Server - mutex sync.Mutex - calls []ModelFakeCall - reply func(http.ResponseWriter, ModelFakeCall) -} - -// Call returns call number n. -func (f *ModelFake) Call(n int) ModelFakeCall { - f.mutex.Lock() - defer f.mutex.Unlock() - - return f.calls[n] -} - -// Calls returns the number of calls. -func (f *ModelFake) Calls() int { - f.mutex.Lock() - defer f.mutex.Unlock() - - return len(f.calls) -} - -func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) { - raw, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - - call := ModelFakeCall{ - header: r.Header.Clone(), - raw: raw, - } - - err = json.Unmarshal(raw, &call.body) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - - f.mutex.Lock() - call.number = len(f.calls) - f.calls = append(f.calls, call) - reply := f.reply - f.mutex.Unlock() - - reply(w, call) -} - -func newModelFake(t *testing.T) *ModelFake { - t.Helper() - - f := &ModelFake{ - reply: func(w http.ResponseWriter, call ModelFakeCall) { - modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5)) - }, - } - - f.Server = httptest.NewServer(http.HandlerFunc(f.serve)) - t.Cleanup(f.Close) - - return f -} - func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient { t.Helper() return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl")) } - -func modelTestRequest() ModelRequest { - return ModelRequest{ - Model: "gpt-test", - Effort: "low", - Input: []any{modelUserMessage("hello")}, - } -} - -// modelWriteEvents writes events as a server-sent event stream. -func modelWriteEvents(w http.ResponseWriter, events ...any) { - w.Header().Set("Content-Type", "text/event-stream") - for _, event := range events { - data, err := json.Marshal(event) - if err != nil { - panic(err) - } - - _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data) - if err != nil { - panic(err) - } - } -} - -func modelMessageItem(text string) map[string]any { - return map[string]any{ - "type": "message", - "role": "assistant", - "content": []any{ - map[string]any{"type": "output_text", "text": text}, - }, - } -} - -func modelCompletedEvent(output []any, input, cached, out int) map[string]any { - if output == nil { - output = []any{} - } - - return map[string]any{ - "type": "response.completed", - "response": map[string]any{ - "output": output, - "usage": map[string]any{ - "input_tokens": input, - "input_tokens_details": map[string]any{"cached_tokens": cached}, - "output_tokens": out, - }, - }, - } -} - -func modelReadUsage(t *testing.T, path string) []ModelUsageRecord { - t.Helper() - - file, err := os.Open(path) - require.NoError(t, err) - - defer func() { - require.NoError(t, file.Close()) - }() - - var records []ModelUsageRecord - scanner := bufio.NewScanner(file) - for scanner.Scan() { - var record ModelUsageRecord - require.NoError(t, json.Unmarshal(scanner.Bytes(), &record)) - records = append(records, record) - } - - require.NoError(t, scanner.Err()) - - return records -} - -func modelReasoningItem(id, encrypted string) map[string]any { - return map[string]any{ - "type": "reasoning", - "id": id, - "encrypted_content": encrypted, - "summary": []any{}, - } -} - -func modelFunctionCallItem(id, callID, name, arguments string) map[string]any { - return map[string]any{ - "type": "function_call", - "id": id, - "call_id": callID, - "name": name, - "arguments": arguments, - } -}