feat: sync deployed v1.0.0-rc.38-custom.67

pkg/model_benchmark/response.go

@@ -0,0 +1,144 @@
+package model_benchmark
+
+import (
+	"bufio"
+	"errors"
+	"io"
+	"strings"
+
+	"github.com/tidwall/gjson"
+)
+
+var ErrInvalidSample = errors.New("invalid probe answer")
+var ErrIncomplete = errors.New("probe response incomplete")
+
+type ProbeResponse struct {
+	Text             string `json:"-"`
+	DeclaredModel    string `json:"declared_model,omitempty"`
+	InputTokens      int64  `json:"input_tokens"`
+	OutputTokens     int64  `json:"output_tokens"`
+	CachedTokens     int64  `json:"cached_tokens"`
+	CacheWriteTokens int64  `json:"cache_write_tokens"`
+	HasUsage         bool   `json:"has_usage"`
+}
+
+func completedResponse(response gjson.Result) (ProbeResponse, error) {
+	r := ProbeResponse{}
+	r.DeclaredModel = response.Get("model").String()
+	if len(r.DeclaredModel) > 128 {
+		r.DeclaredModel = ""
+	}
+	usage := response.Get("usage")
+	r.HasUsage = usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
+	r.InputTokens = usage.Get("input_tokens").Int()
+	r.OutputTokens = usage.Get("output_tokens").Int()
+	r.CachedTokens = usage.Get("input_tokens_details.cached_tokens").Int()
+	r.CacheWriteTokens = max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int())
+	if response.Get("status").String() != "completed" {
+		return r, ErrIncomplete
+	}
+	for _, item := range response.Get("output").Array() {
+		switch item.Get("type").String() {
+		case "reasoning":
+		case "message":
+			for _, content := range item.Get("content").Array() {
+				if content.Get("type").String() != "output_text" {
+					return r, ErrInvalidSample
+				}
+				r.Text += content.Get("text").String()
+				if len(r.Text) > 16384 {
+					return ProbeResponse{}, ErrInvalidSample
+				}
+			}
+		default:
+			return r, ErrInvalidSample
+		}
+	}
+	return r, nil
+}
+
+// Only a complete response is a sample. Partial deltas, tool calls and a
+// truncated stream cannot become a model verdict. Total bytes are bounded.
+func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
+	limited := &io.LimitedReader{R: reader, N: 8 << 20}
+	buffered := bufio.NewReader(limited)
+	peek, _ := buffered.Peek(1)
+	if !strings.Contains(contentType, "text/event-stream") && len(peek) > 0 && peek[0] == '{' {
+		data, err := io.ReadAll(buffered)
+		if err != nil {
+			return ProbeResponse{}, err
+		}
+		if limited.N == 0 || !gjson.ValidBytes(data) {
+			return ProbeResponse{}, ErrInvalidSample
+		}
+		return completedResponse(gjson.ParseBytes(data))
+	}
+	scanner := bufio.NewScanner(buffered)
+	scanner.Buffer(make([]byte, 4096), 2<<20)
+	var data strings.Builder
+	var deltas strings.Builder
+	refused := false
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "data:") {
+			if data.Len() > 0 {
+				data.WriteByte('\n')
+			}
+			data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
+			if data.Len() > 2<<20 {
+				return ProbeResponse{}, ErrInvalidSample
+			}
+			continue
+		}
+		if line != "" || data.Len() == 0 {
+			continue
+		}
+		payload := data.String()
+		data.Reset()
+		if payload == "[DONE]" {
+			break
+		}
+		if !gjson.Valid(payload) {
+			return ProbeResponse{}, ErrInvalidSample
+		}
+		event := gjson.Parse(payload)
+		typ := event.Get("type").String()
+		if typ == "response.output_text.delta" {
+			deltas.WriteString(event.Get("delta").String())
+			if deltas.Len() > 65536 {
+				return ProbeResponse{}, ErrInvalidSample
+			}
+		}
+		if typ == "response.refusal.delta" || typ == "response.refusal.done" {
+			refused = true
+		}
+		if typ == "response.output_item.added" || typ == "response.output_item.done" {
+			itemType := event.Get("item.type").String()
+			if itemType != "message" && itemType != "reasoning" {
+				return ProbeResponse{}, ErrInvalidSample
+			}
+		}
+		if typ == "response.completed" {
+			response := event.Get("response")
+			answer, err := completedResponse(response)
+			if refused {
+				return answer, ErrInvalidSample
+			}
+			// Meow 4.5.4 accepts deltas only when completion omits output entirely.
+			if err == nil && !response.Get("output").Exists() {
+				answer.Text = deltas.String()
+			}
+			return answer, err
+		}
+		if typ == "response.failed" || typ == "response.incomplete" {
+			return completedResponse(event.Get("response"))
+		}
+		if typ == "error" {
+			return ProbeResponse{}, ErrIncomplete
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return ProbeResponse{}, err
+	}
+	return ProbeResponse{}, ErrIncomplete
+}