BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

完善文档页文本向量与重排序调用支持

Browse Source
This commit is contained in:
wangbo
2026-05-31 21:18:41 +08:00
parent 8ee7a7969e
commit 644a6f9d17
24 changed files with 1945 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/runner/limits.go
+32 -12
View File
@@ -159,16 +159,14 @@ func hasRules(policy map[string]any) bool {
}
func estimateRequestTokens(body map[string]any) int {
text := ""
if prompt := stringFromMap(body, "prompt"); prompt != "" {
text += prompt
}
if input := stringFromMap(body, "input"); input != "" {
text += input
}
var text strings.Builder
appendTokenEstimateText(&text, body["prompt"])
appendTokenEstimateText(&text, body["input"])
appendTokenEstimateText(&text, body["query"])
appendTokenEstimateText(&text, body["documents"])
for _, item := range contentItems(body["content"]) {
if stringFromAny(item["type"]) == "text" {
text += stringFromAny(item["text"])
appendTokenEstimateText(&text, item["text"])
}
}
if messages, ok := body["messages"].([]any); ok {
@@ -176,19 +174,41 @@ func estimateRequestTokens(body map[string]any) int {
message, _ := raw.(map[string]any)
switch content := message["content"].(type) {
case string:
text += content
appendTokenEstimateText(&text, content)
case []any:
for _, rawPart := range content {
part, _ := rawPart.(map[string]any)
text += stringFromMap(part, "text")
appendTokenEstimateText(&text, part["text"])
}
}
}
}
if text == "" {
estimatedText := text.String()
if estimatedText == "" {
return 1
}
return len([]rune(text))/4 + 1
return len([]rune(estimatedText))/4 + 1
}
func appendTokenEstimateText(out *strings.Builder, value any) {
switch typed := value.(type) {
case string:
out.WriteString(typed)
case []string:
for _, item := range typed {
out.WriteString(item)
}
case []any:
for _, item := range typed {
appendTokenEstimateText(out, item)
}
case map[string]any:
for _, key := range []string{"text", "content", "query", "document"} {
if text := stringFromAny(typed[key]); text != "" {
out.WriteString(text)
}
}
}
}
func tokenUsageAmounts(usage clients.Usage) map[string]float64 {
apps/api/internal/runner/pricing.go
+17 -8
View File
@@ -40,8 +40,11 @@ func (s *Service) Estimate(ctx context.Context, kind string, model string, body
}
func (s *Service) estimatedBillings(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate) []any {
usage := clients.Usage{InputTokens: estimateRequestTokens(body), OutputTokens: int(floatFromAny(body["max_tokens"]))}
if usage.OutputTokens == 0 {
usage := clients.Usage{InputTokens: estimateRequestTokens(body)}
if isTextGenerationKind(kind) {
usage.OutputTokens = int(floatFromAny(body["max_tokens"]))
}
if isTextGenerationKind(kind) && usage.OutputTokens == 0 {
usage.OutputTokens = 64
}
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
@@ -56,19 +59,25 @@ func (s *Service) estimatedBillings(ctx context.Context, user *auth.User, kind s
func (s *Service) billings(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate, response clients.Response, simulated bool) []any {
config := s.effectiveBillingConfig(ctx, candidate)
discount := effectiveDiscount(ctx, s.store, user, candidate)
if isTextGenerationKind(kind) {
if isTextBillingKind(kind) {
inputTokens := response.Usage.InputTokens
outputTokens := response.Usage.OutputTokens
if isTextInputOnlyKind(kind) && inputTokens == 0 && response.Usage.TotalTokens > 0 {
inputTokens = response.Usage.TotalTokens
}
if inputTokens == 0 && outputTokens == 0 {
inputTokens = estimateRequestTokens(body)
outputTokens = 1
if isTextGenerationKind(kind) {
outputTokens = 1
}
}
inputAmount := roundPrice(float64(inputTokens) / 1000 * resourcePrice(config, "text", "textInputPer1k", "inputTokenPrice", "basePrice") * discount)
outputAmount := roundPrice(float64(outputTokens) / 1000 * resourcePrice(config, "text", "textOutputPer1k", "outputTokenPrice", "basePrice") * discount)
return []any{
billingLine(candidate, "text_input", "1k_tokens", inputTokens, inputAmount, discount, simulated),
billingLine(candidate, "text_output", "1k_tokens", outputTokens, outputAmount, discount, simulated),
lines := []any{billingLine(candidate, "text_input", "1k_tokens", inputTokens, inputAmount, discount, simulated)}
if isTextGenerationKind(kind) {
outputAmount := roundPrice(float64(outputTokens) / 1000 * resourcePrice(config, "text", "textOutputPer1k", "outputTokenPrice", "basePrice") * discount)
lines = append(lines, billingLine(candidate, "text_output", "1k_tokens", outputTokens, outputAmount, discount, simulated))
}
return lines
}
count := requestOutputCount(body)
resource := "image"
apps/api/internal/runner/service.go
+49 -3
View File
@@ -698,7 +698,7 @@ func (s *Service) recordTaskParameterPreprocessing(ctx context.Context, taskID s
func skipTaskParameterPreprocessingLog(modelType string) bool {
switch strings.TrimSpace(modelType) {
case "text_generate", "chat", "responses", "text":
case "text_generate", "text_embedding", "text_rerank", "chat", "responses", "text":
return true
default:
return false
@@ -923,6 +923,10 @@ func modelTypeFromKind(kind string, body map[string]any) string {
switch kind {
case "chat.completions", "responses":
return "text_generate"
case "embeddings":
return "text_embedding"
case "reranks":
return "text_rerank"
case "images.generations", "images.edits":
if kind == "images.edits" {
return "image_edit"
@@ -943,7 +947,7 @@ func modelTypeFromKind(kind string, body map[string]any) string {
func requestedModelTypeFromBody(body map[string]any) string {
for _, key := range []string{"modelType", "model_type", "capability", "capabilityType"} {
value := strings.TrimSpace(stringFromMap(body, key))
value := canonicalModelType(strings.TrimSpace(stringFromMap(body, key)))
if isKnownModelType(value) {
return value
}
@@ -951,9 +955,21 @@ func requestedModelTypeFromBody(body map[string]any) string {
return ""
}
func canonicalModelType(value string) string {
normalized := strings.ReplaceAll(strings.ToLower(strings.TrimSpace(value)), "-", "_")
switch normalized {
case "embedding":
return "text_embedding"
case "rerank", "reranks":
return "text_rerank"
default:
return normalized
}
}
func isKnownModelType(value string) bool {
switch value {
case "text_generate", "image_generate", "image_edit", "video_generate", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni":
case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "video_generate", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni":
return true
default:
return false
@@ -1001,6 +1017,14 @@ func isTextGenerationKind(kind string) bool {
return kind == "chat.completions" || kind == "responses"
}
func isTextInputOnlyKind(kind string) bool {
return kind == "embeddings" || kind == "reranks"
}
func isTextBillingKind(kind string) bool {
return isTextGenerationKind(kind) || isTextInputOnlyKind(kind)
}
func isSimulation(task store.GatewayTask, candidate store.RuntimeModelCandidate) bool {
if task.RunMode == "simulation" {
return true
@@ -1115,6 +1139,17 @@ func validateRequest(kind string, body map[string]any) error {
if body["input"] == nil && body["messages"] == nil {
return errors.New("input or messages is required")
}
case "embeddings":
if body["input"] == nil {
return errors.New("input is required")
}
case "reranks":
if body["query"] == nil {
return errors.New("query is required")
}
if !hasRerankDocuments(body["documents"]) {
return errors.New("documents is required")
}
case "images.generations", "images.edits":
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
return errors.New("prompt is required")
@@ -1123,6 +1158,17 @@ func validateRequest(kind string, body map[string]any) error {
return nil
}
func hasRerankDocuments(value any) bool {
switch typed := value.(type) {
case []any:
return len(typed) > 0
case []string:
return len(typed) > 0
default:
return false
}
}
func parameterPreprocessClientError(err error) *clients.ClientError {
if err == nil {
return nil