new-api-custom/service/request_details.go

package service

import (
	"crypto/sha256"
	"fmt"
	"strings"

	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/gin-gonic/gin"
)

type RequestDetails struct {
	SessionID       string                          `json:"session_id,omitempty"`
	SessionSource   string                          `json:"session_source"`
	CorrelationKey  string                          `json:"correlation_key,omitempty"`
	RequestedModel  string                          `json:"requested_model,omitempty"`
	UpstreamModel   string                          `json:"upstream_model,omitempty"`
	ReturnedModel   string                          `json:"returned_model,omitempty"`
	ReasoningEffort string                          `json:"reasoning_effort,omitempty"`
	Status          int                             `json:"status"`
	ModelGuard      *relaycommon.ResponseModelGuard `json:"model_guard,omitempty"`
}

func appendRequestDetails(c *gin.Context, info *relaycommon.RelayInfo, other *model.LogOther) {
	d := RequestDetails{SessionSource: "none", RequestedModel: c.GetString("original_model"), Status: c.Writer.Status()}
	if status, ok := other.Snapshot()["status_code"].(int); ok {
		d.Status = status
	}
	if c.Request != nil {
		for _, h := range []struct{ name, source string }{{"session_id", "session"}, {"Session-Id", "session"}} {
			if value := strings.TrimSpace(c.GetHeader(h.name)); value != "" {
				d.SessionID, d.SessionSource = value, h.source
				break
			}
		}
	}
	if d.SessionID == "" {
		if value := strings.TrimSpace(c.GetString("request_session_id")); value != "" {
			d.SessionID = value
			d.SessionSource = "session_body"
		}
	}
	if d.SessionID == "" && c.Request != nil {
		for _, key := range []string{"thread_id", "Thread-Id"} {
			if value := strings.TrimSpace(c.GetHeader(key)); value != "" {
				d.SessionID = value
				d.SessionSource = "thread"
				break
			}
		}
	}
	if d.SessionID == "" {
		if value := strings.TrimSpace(c.GetString("request_thread_id")); value != "" {
			d.SessionID = value
			d.SessionSource = "thread_body"
		}
	}
	if d.SessionID == "" {
		if meta, ok := getChannelAffinityMeta(c); ok && meta.KeyFingerprint != "" {
			d.SessionID = meta.KeyFingerprint
			d.SessionSource = "affinity"
		}
	}
	if d.SessionID != "" {
		d.CorrelationKey = fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprintf("%d\x00%s\x00%s", c.GetInt("id"), d.SessionSource, d.SessionID))))
		// Hash the complete identifier before limiting its display size.
		if len(d.SessionID) > 256 {
			d.SessionID = string([]rune(d.SessionID)[:min(256, len([]rune(d.SessionID)))]) + "…"
		}
	}
	c.Set("log_correlation_key", d.CorrelationKey)
	if info != nil {
		d.ModelGuard = info.ResponseGuard
		d.RequestedModel = info.OriginModelName
		d.UpstreamModel = info.GetUpstreamModelName()
		d.ReasoningEffort = info.ReasoningEffort
		if info.ResponseModel != nil {
			d.ReturnedModel = info.ResponseModel.ReturnedModel
		}
	}
	other.SetAdmin("request_details", d)
}

func recordSelection(c *gin.Context, group string, trace *model.SelectionTrace, reason string) {
	state := RequestPolicy(c)
	remaining := max(0, 100-state.candidateCount)
	items := trace.Candidates[:min(remaining, len(trace.Candidates))]
	state.candidateCount += len(items)
	channelID := trace.SelectedID
	for _, item := range trace.Candidates {
		if item.Reason == "selected" || item.Reason == "affinity_hit" || item.Reason == "pinned_channel" {
			channelID = item.ChannelID
			break
		}
	}
	state.AddEvent(PolicyEvent{Attempt: state.Attempts + 1, ChannelID: channelID, Group: group, Candidates: items, CandidateTotal: trace.Total, Decision: PolicyDecision{Action: "selection", Reason: reason, Source: "routing"}})
}

func selectWithTrace(param *RetryParam, group string, retry int) (*model.Channel, error) {
	trace := &model.SelectionTrace{}
	channel, err := model.GetRandomSatisfiedChannel(group, param.ModelName, retry, GetChannelConstraints(param.Ctx).Filters, trace)
	reason := "priority_weight_selection"
	if channel == nil {
		reason = "no_available_channel"
	}
	recordSelection(param.Ctx, group, trace, reason)
	return channel, err
}

func recordDirectSelection(c *gin.Context, ch *model.Channel, id int, group, modelName, reason string) {
	trace := &model.SelectionTrace{}
	trace.Add(ch, id, group, modelName, reason)
	recordSelection(c, group, trace, reason)
}