diff --git a/apiproxy/proxy.go b/apiproxy/proxy.go index 933e1e4..76a61a2 100644 --- a/apiproxy/proxy.go +++ b/apiproxy/proxy.go @@ -2,6 +2,7 @@ package apiproxy import ( "bytes" + "encoding/json" "fmt" "io" "log" @@ -150,6 +151,7 @@ func (h *baseHandle) HandleAzure(w http.ResponseWriter, r *http.Request, backend r.URL.Path = strings.TrimPrefix(r.URL.Path, "/api") // Preserve incoming query parameters unchanged. r.URL.RawQuery = r.URL.RawQuery + ensureStreamUsageForChatCompletions(r) backendProxy[r.Host] = proxy @@ -186,6 +188,62 @@ func (h *baseHandle) HandleAzure(w http.ResponseWriter, r *http.Request, backend } +// ensureStreamUsageForChatCompletions enables usage emission for streaming chat +// completions by adding stream_options.include_usage=true when it is missing. +func ensureStreamUsageForChatCompletions(r *http.Request) { + if r == nil || r.URL == nil || r.Body == nil { + return + } + path := strings.ToLower(r.URL.Path) + if !(strings.HasSuffix(path, "/chat/completions") || strings.Contains(path, "/chat/completions/")) { + return + } + + bodyBytes, err := io.ReadAll(r.Body) + if err != nil { + return + } + defer func() { + r.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + }() + if len(bytes.TrimSpace(bodyBytes)) == 0 { + return + } + + var payload map[string]interface{} + if err := json.Unmarshal(bodyBytes, &payload); err != nil { + return + } + + streamEnabled, _ := payload["stream"].(bool) + if !streamEnabled { + return + } + + changed := false + streamOptions, hasOptions := payload["stream_options"].(map[string]interface{}) + if !hasOptions { + streamOptions = map[string]interface{}{} + payload["stream_options"] = streamOptions + changed = true + } + if _, exists := streamOptions["include_usage"]; !exists { + streamOptions["include_usage"] = true + changed = true + } + if !changed { + return + } + + updatedBody, err := json.Marshal(payload) + if err != nil { + return + } + bodyBytes = updatedBody + r.ContentLength = int64(len(bodyBytes)) + r.Header.Set("Content-Length", fmt.Sprintf("%d", len(bodyBytes))) +} + // singleJoiningSlash joins two URL paths with a single slash between them. func singleJoiningSlash(a, b string) string { aslash := strings.HasSuffix(a, "/") diff --git a/apiproxy/proxy_test.go b/apiproxy/proxy_test.go index 08b02fd..9951f86 100644 --- a/apiproxy/proxy_test.go +++ b/apiproxy/proxy_test.go @@ -1,6 +1,7 @@ package apiproxy import ( + "encoding/json" "io" "net/http" "net/http/httptest" @@ -92,3 +93,55 @@ func TestForwarding_V1ChatCompletions(t *testing.T) { t.Fatalf("unexpected body forwarded: %s", gotBody) } } + +func TestEnsureStreamUsageForChatCompletions_AddsIncludeUsageWhenMissing(t *testing.T) { + body := `{"model":"gpt-5-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}` + req := httptest.NewRequest("POST", "http://localhost/api/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + + ensureStreamUsageForChatCompletions(req) + + gotBytes, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]interface{} + if err := json.Unmarshal(gotBytes, &payload); err != nil { + t.Fatalf("failed to decode patched body: %v", err) + } + + streamOptions, ok := payload["stream_options"].(map[string]interface{}) + if !ok { + t.Fatalf("expected stream_options object in patched payload, got: %v", payload["stream_options"]) + } + includeUsage, ok := streamOptions["include_usage"].(bool) + if !ok || !includeUsage { + t.Fatalf("expected stream_options.include_usage=true, got: %v", streamOptions["include_usage"]) + } +} + +func TestEnsureStreamUsageForChatCompletions_KeepsExistingIncludeUsage(t *testing.T) { + body := `{"model":"gpt-5-mini","stream":true,"stream_options":{"include_usage":false},"messages":[{"role":"user","content":"hi"}]}` + req := httptest.NewRequest("POST", "http://localhost/api/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + + ensureStreamUsageForChatCompletions(req) + + gotBytes, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var payload map[string]interface{} + if err := json.Unmarshal(gotBytes, &payload); err != nil { + t.Fatalf("failed to decode patched body: %v", err) + } + + streamOptions, ok := payload["stream_options"].(map[string]interface{}) + if !ok { + t.Fatalf("expected stream_options object, got: %v", payload["stream_options"]) + } + includeUsage, ok := streamOptions["include_usage"].(bool) + if !ok || includeUsage { + t.Fatalf("expected existing include_usage=false to remain unchanged, got: %v", streamOptions["include_usage"]) + } +} diff --git a/apiproxy/response.go b/apiproxy/response.go index 7aa2262..966fb19 100644 --- a/apiproxy/response.go +++ b/apiproxy/response.go @@ -308,17 +308,15 @@ func dedupPromptTokens(total, cached int) int { return total - cached } -var snapshotSuffixPattern = regexp.MustCompile(`\d{4}-\d{2}-\d{2}$`) +var snapshotSuffixPattern = regexp.MustCompile(`^(.*)-(\d{4}-\d{2}-\d{2})$`) func splitModelSnapshot(model string) (string, string) { if model == "" { return "", "" } - if idx := strings.LastIndex(model, "-"); idx > 0 { - suffix := model[idx+1:] - if snapshotSuffixPattern.MatchString(suffix) { - return model[:idx], suffix - } + matches := snapshotSuffixPattern.FindStringSubmatch(model) + if len(matches) == 3 && matches[1] != "" { + return matches[1], matches[2] } return model, "" } diff --git a/apiproxy/response_model_snapshot_test.go b/apiproxy/response_model_snapshot_test.go new file mode 100644 index 0000000..203b166 --- /dev/null +++ b/apiproxy/response_model_snapshot_test.go @@ -0,0 +1,40 @@ +package apiproxy + +import "testing" + +func TestSplitModelSnapshot(t *testing.T) { + tests := []struct { + name string + model string + wantModel string + wantSnapshot string + }{ + { + name: "model with snapshot suffix", + model: "gpt-5-mini-2025-08-07", + wantModel: "gpt-5-mini", + wantSnapshot: "2025-08-07", + }, + { + name: "model without snapshot suffix", + model: "gpt-5-mini", + wantModel: "gpt-5-mini", + wantSnapshot: "", + }, + { + name: "empty model", + model: "", + wantModel: "", + wantSnapshot: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotModel, gotSnapshot := splitModelSnapshot(tt.model) + if gotModel != tt.wantModel || gotSnapshot != tt.wantSnapshot { + t.Fatalf("splitModelSnapshot(%q) = (%q, %q), want (%q, %q)", tt.model, gotModel, gotSnapshot, tt.wantModel, tt.wantSnapshot) + } + }) + } +} diff --git a/apiproxy/response_sse_nested_usage_test.go b/apiproxy/response_sse_nested_usage_test.go index 638a753..372d18b 100644 --- a/apiproxy/response_sse_nested_usage_test.go +++ b/apiproxy/response_sse_nested_usage_test.go @@ -192,6 +192,62 @@ func TestNewResponse_SSEOverCapEventApproximatesCompletionTokens(t *testing.T) { } } +func TestNewResponse_ChatCompletionStoresInputTokens(t *testing.T) { + jsonBody := `{ + "id":"chatcmpl-DHWFPHKozPO2ImSWSs4Nc7aGEoyMg", + "object":"chat.completion", + "model":"gpt-5-mini-2025-08-07", + "usage":{ + "completion_tokens":200, + "completion_tokens_details":{ + "accepted_prediction_tokens":0, + "audio_tokens":0, + "reasoning_tokens":200, + "rejected_prediction_tokens":0 + }, + "prompt_tokens":23, + "prompt_tokens_details":{ + "audio_tokens":0, + "cached_tokens":0 + }, + "total_tokens":223 + } + }` + + req, _ := http.NewRequest("POST", "https://example.local/openai/v1/chat/completions", nil) + req.Header.Set(authHeader, "Bearer TESTTOKEN") + + resp := &http.Response{ + Request: req, + Header: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}}, + Body: io.NopCloser(strings.NewReader(jsonBody)), + } + + fb := &fakeDBForTest{} + hash, _ := bcrypt.GenerateFromPassword([]byte("TESTTOKEN"), 5) + fb.apiKeys = []db.ApiKey{{UUID: "uid-1", ApiKey: string(hash), Owner: "owner1"}} + + rc := &ResponseConf{db: fb} + if err := rc.NewResponse(resp); err != nil { + t.Fatalf("NewResponse error: %v", err) + } + + if len(fb.writes) != 1 { + t.Fatalf("expected 1 DB write, got %d", len(fb.writes)) + } + + w := fb.writes[0] + if w.InputTokenCount != 23 { + t.Fatalf("expected input tokens 23, got %d", w.InputTokenCount) + } + if w.CachedInputTokenCount != 0 { + t.Fatalf("expected cached input tokens 0, got %d", w.CachedInputTokenCount) + } + if w.OutputTokenCount != 200 { + t.Fatalf("expected output tokens 200, got %d", w.OutputTokenCount) + } +} + // fake DB implementation for test type fakeDBForTest struct { apiKeys []db.ApiKey diff --git a/local-dev/example_requests/requests.http b/local-dev/example_requests/requests.http index 5f32813..5530898 100644 --- a/local-dev/example_requests/requests.http +++ b/local-dev/example_requests/requests.http @@ -1,25 +1,23 @@ POST http://localhost:8082/api/v1/chat/completions Content-Type: application/json -Authorization: Bearer +Authorization: Bearer { - "model": "gpt-4o", + "model": "gpt-5-mini", "messages": [ {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]}, {"role": "user", "content": [{"type": "text", "text": "Please summarize the cost collection fix."}]} ], - "temperature": 0.7, - "max_tokens": 200, - "top_p": 0.95 + "max_completion_tokens": 200 } ### POST http://localhost:8082/api/v1/responses Content-Type: application/json -Authorization: Bearer +Authorization: Bearer { - "model": "gpt-4o", + "model": "gpt-5-mini", "input": [ { "role": "system", @@ -39,8 +37,5 @@ Authorization: Bearer } ] } - ], - "temperature": 0.7, - "max_output_tokens": 200, - "top_p": 0.95 + ] } \ No newline at end of file