new-api-custom/pkg/model_fingerprint/response.go

package model_fingerprint

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
	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_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" {
			return completedResponse(event.Get("response"))
		}
		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
}