refactor(model): 重构模型网关结构和流式响应解析

2026-07-06 13:20:19 +08:00
parent 1ce1913437
commit d3a8043337
6 changed files with 199 additions and 107 deletions
@@ -1,40 +1,128 @@
 package util

 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"model-gateway/service/gateway"
 	"sort"
 	"strings"

 	"github.com/gogf/gf/v2/encoding/gjson"
 )

-// ================================================================
-
-// ParseStreamResponse 流式响应解析（通用入口）
-func ParseStreamResponse(rawBytes []byte, streamConfig map[string]any) (map[string]any, error) {
+// ParseStreamResponse 流式响应解析
+func ParseStreamResponse(ctx context.Context, rawBytes []byte, streamConfig map[string]any) (map[string]any, error) {
 	enabled, _ := streamConfig["enabled"].(bool)
 	if !enabled {
 		return gjson.New(string(rawBytes)).Map(), nil
 	}

-	parser, _ := streamConfig["parser"].(string)
-	if parser == "base64_concat" {
-		return parseBase64Stream(rawBytes)
+	outputType, _ := streamConfig["output_type"].(string)
+	streamToClient, _ := streamConfig["stream_to_client"].(bool)
+	events, _ := streamConfig["events"].([]any)
+
+	if len(events) == 0 {
+		return gjson.New(string(rawBytes)).Map(), nil
 	}

-	return parseSSEStream(rawBytes, streamConfig)
+	// 如果业务需要流式返回，直接透传原始数据
+	if streamToClient {
+		return map[string]any{"stream_data": rawBytes}, nil
+	}
+
+	result := make(map[string]any)
+	lines := strings.Split(string(rawBytes), "\n")
+
+	for _, evt := range events {
+		e, _ := evt.(map[string]any)
+		evtType, _ := e["type"].(string)
+
+		switch evtType {
+		case "concat":
+			processConcat(lines, e, result)
+		case "base64_concat":
+			processBase64Concat(lines, e, result)
+		case "collect":
+			processCollect(lines, e, result)
+		case "final":
+			processFinal(lines, e, result)
+		}
+	}
+
+	switch outputType {
+	case "audio":
+		if audioBytes, ok := result["audio"].([]byte); ok {
+			oss, err := gateway.UploadByTask(ctx, audioBytes, "mp3")
+			if err != nil {
+				return nil, err
+			}
+			result["content"] = oss.FileAddressPrefix + oss.FileURL
+		}
+	case "text":
+		if v, ok := result["content"]; ok {
+			return map[string]any{"content": v, "usage": result["usage"]}, nil
+		}
+	case "image":
+		if v, ok := result["urls"]; ok {
+			return map[string]any{"content": v, "usage": result["usage"]}, nil
+		}
+	}
+
+	return result, nil
 }

-// parseBase64Stream 拼接流式 base64 并解码为二进制（TTS 等音频模型）
-func parseBase64Stream(rawBytes []byte) (map[string]any, error) {
-	lines := strings.Split(string(rawBytes), "\n")
-	var audioBase64 strings.Builder
+// processConcat 文本拼接
+func processConcat(lines []string, event map[string]any, result map[string]any) {
+	match, _ := event["match"].(string)
+	aggregateTo, _ := event["aggregate_to"].(string)
+	fields, _ := event["fields"].(map[string]any)

+	var parts []string
 	for _, line := range lines {
 		line = strings.TrimSpace(line)
-		if line == "" {
+		if line == "" || line == "[DONE]" {
+			continue
+		}
+		if strings.HasPrefix(line, "event:") {
+			continue
+		}
+		if strings.HasPrefix(line, "data:") {
+			line = strings.TrimPrefix(line, "data:")
+			line = strings.TrimSpace(line)
+		}
+
+		var chunk map[string]any
+		if err := json.Unmarshal([]byte(line), &chunk); err != nil {
+			continue
+		}
+
+		chunkType, _ := chunk["type"].(string)
+		if match != "" && !strings.Contains(chunkType, match) {
+			continue
+		}
+
+		for _, chunkPath := range fields {
+			val := gjson.New(chunk).Get(chunkPath.(string)).String()
+			if val != "" {
+				parts = append(parts, val)
+			}
+		}
+	}
+
+	result[aggregateTo] = strings.Join(parts, "")
+}
+
+// processBase64Concat base64 拼接
+func processBase64Concat(lines []string, event map[string]any, result map[string]any) {
+	aggregateTo, _ := event["aggregate_to"].(string)
+	fields, _ := event["fields"].(map[string]any)
+
+	var builder strings.Builder
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if line == "" || line == "[DONE]" {
 			continue
 		}

@@ -43,8 +131,10 @@ func parseBase64Stream(rawBytes []byte) (map[string]any, error) {
 			continue
 		}

-		if data, ok := chunk["data"].(string); ok && data != "" {
-			audioBase64.WriteString(data)
+		for _, chunkPath := range fields {
+			if data := gjson.New(chunk).Get(chunkPath.(string)).String(); data != "" {
+				builder.WriteString(data)
+			}
 		}
 	}

@@ -53,29 +143,73 @@ func parseBase64Stream(rawBytes []byte) (map[string]any, error) {
 			return -1
 		}
 		return r
-	}, audioBase64.String())
+	}, builder.String())

 	audioBytes, err := base64.StdEncoding.DecodeString(cleanBase64)
 	if err != nil {
-		audioBytes, err = base64.RawStdEncoding.DecodeString(cleanBase64)
-		if err != nil {
-			return nil, fmt.Errorf("base64 解码失败: %w", err)
-		}
+		audioBytes, _ = base64.RawStdEncoding.DecodeString(cleanBase64)
 	}
-
-	return map[string]any{"audio": audioBytes}, nil
+	result[aggregateTo] = audioBytes
 }

-// parseSSEStream SSE 流式解析（图片模型等）
-func parseSSEStream(rawBytes []byte, streamConfig map[string]any) (map[string]any, error) {
-	events, _ := streamConfig["events"].([]any)
-	if len(events) == 0 {
-		return gjson.New(string(rawBytes)).Map(), nil
+// processCollect 数组收集
+func processCollect(lines []string, event map[string]any, result map[string]any) {
+	match, _ := event["match"].(string)
+	aggregateTo, _ := event["aggregate_to"].(string)
+	orderBy, _ := event["order_by"].(string)
+	fields, _ := event["fields"].(map[string]any)
+
+	var items []map[string]any
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if line == "" || line == "[DONE]" {
+			continue
+		}
+
+		var chunk map[string]any
+		if err := json.Unmarshal([]byte(line), &chunk); err != nil {
+			continue
+		}
+
+		chunkType, _ := chunk["type"].(string)
+		if match != "" && !strings.Contains(chunkType, match) {
+			continue
+		}
+
+		item := make(map[string]any)
+		for localKey, chunkPath := range fields {
+			item[localKey] = gjson.New(chunk).Get(chunkPath.(string)).Val()
+		}
+		items = append(items, item)
 	}

-	lines := strings.Split(string(rawBytes), "\n")
-	result := make(map[string]any)
-	var partials []map[string]any
+	if orderBy != "" {
+		sort.Slice(items, func(i, j int) bool {
+			return fmt.Sprint(items[i][orderBy]) < fmt.Sprint(items[j][orderBy])
+		})
+	}
+
+	// 如果只有一个字段，直接存值数组
+	if len(fields) == 1 {
+		var vals []any
+		for _, item := range items {
+			for _, v := range item {
+				vals = append(vals, v)
+			}
+		}
+		result[aggregateTo] = vals
+	} else {
+		result[aggregateTo] = items
+	}
+}
+
+// processFinal 取最后一条匹配的数据
+func processFinal(lines []string, event map[string]any, result map[string]any) {
+	match, _ := event["match"].(string)
+	aggregateTo, _ := event["aggregate_to"].(string)
+	fields, _ := event["fields"].(map[string]any)
+
+	var lastMatch map[string]any

 	for _, line := range lines {
 		line = strings.TrimSpace(line)
@@ -96,55 +230,26 @@ func parseSSEStream(rawBytes []byte, streamConfig map[string]any) (map[string]an
 		}

 		chunkType, _ := chunk["type"].(string)
+		if match != "" && !strings.Contains(chunkType, match) {
+			continue
+		}

-		for _, evt := range events {
-			e, _ := evt.(map[string]any)
-			match, _ := e["match"].(string)
-			if !strings.Contains(chunkType, match) {
-				continue
-			}
+		lastMatch = chunk
+	}

-			fields, _ := e["fields"].(map[string]any)
-			aggregateTo, _ := e["aggregate_to"].(string)
-			evtType, _ := e["type"].(string)
+	if lastMatch == nil {
+		return
+	}

-			switch evtType {
-			case "partial":
-				item := make(map[string]any)
-				for localKey, chunkKey := range fields {
-					item[localKey] = chunk[chunkKey.(string)]
-				}
-				partials = append(partials, item)
-
-			case "final":
-				for localKey, chunkKey := range fields {
-					val := gjson.New(chunk).Get(chunkKey.(string))
-					if !val.IsNil() {
-						if _, exists := result[aggregateTo]; !exists {
-							result[aggregateTo] = make(map[string]any)
-						}
-						result[aggregateTo].(map[string]any)[localKey] = val.Val()
-					}
-				}
-			}
+	data := make(map[string]any)
+	for localKey, chunkPath := range fields {
+		val := gjson.New(lastMatch).Get(chunkPath.(string)).Val()
+		if val != nil {
+			data[localKey] = val
 		}
 	}

-	if len(partials) > 0 {
-		for _, evt := range events {
-			e, _ := evt.(map[string]any)
-			if e["type"] == "partial" {
-				if orderBy, ok := e["order_by"].(string); ok {
-					sort.Slice(partials, func(i, j int) bool {
-						return fmt.Sprint(partials[i][orderBy]) < fmt.Sprint(partials[j][orderBy])
-					})
-				}
-				result[e["aggregate_to"].(string)] = partials
-				break
-			}
-		}
+	if len(data) > 0 {
+		result[aggregateTo] = data
 	}
-
-	mergedBytes, _ := json.Marshal(result)
-	return gjson.New(mergedBytes).Map(), nil
 }