117 lines
4.5 KiB
Go
117 lines
4.5 KiB
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 {
|
|
RoutingPolicyFingerprint string `json:"routing_policy_fingerprint,omitempty"`
|
|
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{RoutingPolicyFingerprint: c.GetString("pool_affinity_policy_fingerprint"), 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)
|
|
}
|