106 lines
3.3 KiB
Go
106 lines
3.3 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func processVisionDetection(ctx context.Context, run model.DetectionRun) {
|
|
if run.UserID > 0 {
|
|
var price service.DetectionPrice
|
|
if common.UnmarshalJsonStr(run.Price, &price) != nil || gjson.GetBytes(price.Request.Body, "max_output_tokens").Int() != service.VisionOutputTokens || gjson.GetBytes(price.Request.Body, "reasoning.effort").String() != "medium" {
|
|
finishVisionDetection(run, "", "pricing_changed")
|
|
return
|
|
}
|
|
}
|
|
channel, err := model.GetChannelById(run.ChannelID, true)
|
|
if err != nil || channel.Status != common.ChannelStatusEnabled {
|
|
finishVisionDetection(run, "", "channel_unavailable")
|
|
return
|
|
}
|
|
var previous []model.DetectionCall
|
|
if model.DB.Where("id = ?", run.ID).Find(&previous).Error != nil {
|
|
return
|
|
}
|
|
var call model.DetectionCall
|
|
if len(previous) > 0 {
|
|
call = previous[0]
|
|
if call.State == "running" {
|
|
return
|
|
}
|
|
} else {
|
|
call = model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(run.ID), State: "running", StartedAt: time.Now().Unix()}
|
|
if model.DB.Create(&call).Error != nil {
|
|
return
|
|
}
|
|
probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
|
|
detectionCalls.Store(call.ID, cancel)
|
|
answer, _, probeErr := runModelDetectionProbe(probeCtx, channel, 0, run.Kind)
|
|
cancel()
|
|
detectionCalls.Delete(call.ID)
|
|
success := probeErr == nil
|
|
call.Available = &success
|
|
call.FirstTokenMS = answer.FirstTokenMS
|
|
if answer.HasUsage {
|
|
usage, _ := common.Marshal(answer)
|
|
call.Usage = string(usage)
|
|
if run.UserID > 0 {
|
|
quota, quotaErr := service.DetectionUsageQuota(run.Price, answer.InputTokens, answer.OutputTokens, answer.CachedTokens, answer.CacheWriteTokens)
|
|
if quotaErr != nil || quota > run.Reserved {
|
|
call.Error = "settlement_unavailable"
|
|
} else {
|
|
call.Quota = quota
|
|
}
|
|
}
|
|
}
|
|
if run.UserID > 0 && !answer.HasUsage {
|
|
call.Error = "usage_unavailable"
|
|
}
|
|
if probeErr != nil {
|
|
call.Error = "vision_request_failed"
|
|
}
|
|
sample, _ := common.Marshal(fingerprint.Sample{Text: answer.Text, Count: 1})
|
|
call.Sample = string(sample)
|
|
if model.FinishDetectionCall(&call) != nil {
|
|
return
|
|
}
|
|
}
|
|
var sample fingerprint.Sample
|
|
if common.UnmarshalJsonStr(call.Sample, &sample) != nil {
|
|
call.Error = "vision_response_invalid"
|
|
}
|
|
finishVisionDetection(run, sample.Text, call.Error)
|
|
}
|
|
|
|
func finishVisionDetection(run model.DetectionRun, answer, reason string) {
|
|
quality := service.VisionQuality(answer)
|
|
if run.Kind == "ciallo" {
|
|
quality = service.CialloQuality(answer)
|
|
}
|
|
if reason == "" && strings.TrimSpace(answer) == "" {
|
|
reason = "vision_response_invalid"
|
|
}
|
|
valid := 1
|
|
if reason != "" {
|
|
quality = "unknown"
|
|
valid = 0
|
|
}
|
|
result, _ := common.Marshal(fingerprint.Result{Engine: run.Kind + "-v1", Status: quality, Closest: answer, ValidSamples: valid})
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "result": string(result)}).Error != nil {
|
|
return
|
|
}
|
|
if reason != "" {
|
|
finishDetection(run, "failed", reason)
|
|
} else {
|
|
finishDetection(run, "succeeded", "")
|
|
}
|
|
service.InvalidateGroupQualityCache()
|
|
}
|