ai-agent/gateway/model_stream.go

package gateway

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"strings"

	commonHttp "gitea.redpowerfuture.com/red-future/common/http"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

// ModelCallStreamReq 调用模型网关流式（client 侧请求体）
type ModelCallStreamReq struct {
	ModelId        int64          `json:"modelId" dc:"模型ID"`
	BizName        string         `json:"bizName" dc:"业务名称"`
	SessionId      string         `json:"sessionId" dc:"会话ID"`
	RequestParams  map[string]any `json:"requestParams" dc:"请求参数（模板字段）"`
	BusinessParams map[string]any `json:"businessParams" dc:"业务参数"`
}

// ModelCallStream 流式调用模型网关：POST /modelCallStream，逐 chunk 调 onChunk（文本增量转发），
// 流结束返回 ModelCallRes：Content 为累加全量文本（{"respBody": ...}），Tools 为流末工具列表。
func ModelCallStream(ctx context.Context, modelId int64, sessionId string, requestParams, businessParams map[string]any, onChunk func(chunk map[string]any) error) (*ModelCallRes, error) {
	req := ModelCallStreamReq{
		ModelId:        modelId,
		BizName:        g.Cfg().MustGet(ctx, "server.name").String(),
		SessionId:      sessionId,
		RequestParams:  requestParams,
		BusinessParams: businessParams,
	}
	body, err := commonHttp.PostStream(ctx, "model-gateway/model/call/modelCallStream", requestHeaders(ctx), &req)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	res := new(ModelCallRes)
	var contentBuf strings.Builder
	scanErr := parseSSEStream(ctx, body, func(chunk map[string]any) error {
		tools, done := processStreamChunk(chunk)
		if done {
			res.Tools = tools
			// 流末 done 事件携带该次调用累计 token 与费用：纯工具调用步骤无文本 chunk，靠它取数（>0 才覆盖，兼容旧网关无 token/cost 的 done 事件）
			if total := gconv.Int64(chunk["totalTokens"]); total > 0 {
				res.TotalTokens = total
				res.PromptTokens = gconv.Int64(chunk["promptTokens"])
				res.CompletionTokens = gconv.Int64(chunk["completionTokens"])
			}
			if cost := gconv.Float64(chunk["cost"]); cost > 0 {
				res.Cost = cost
			}
			return nil
		}
		if onChunk != nil {
			if err := onChunk(chunk); err != nil {
				return err
			}
		}
		contentBuf.WriteString(streamTextDelta(chunk))
		// 网关端每个 chunk 携带累计 token 值，取最后一个文本 chunk 即为该步总消耗
		res.TotalTokens = gconv.Int64(chunk["totalTokens"])
		res.PromptTokens = gconv.Int64(chunk["promptTokens"])
		res.CompletionTokens = gconv.Int64(chunk["completionTokens"])
		res.Cost = gconv.Float64(chunk["cost"])
		return nil
	})
	if scanErr != nil {
		// 流被中断（如 ctx 取消）时仍返回已累计的 res，调用方可读取已产生的 token 正常落库
		return res, scanErr
	}
	if contentBuf.Len() > 0 {
		res.Content = map[string]any{"respBody": contentBuf.String()}
	}
	return res, nil
}

// processStreamChunk 处理单个流式 chunk。done 事件（{"type":"done","tools":[...]}）返回解析出的
// tools（可为空）且 done=true；其余为文本增量 chunk，返回 done=false。
func processStreamChunk(chunk map[string]any) (tools []ModelTool, done bool) {
	if typ, _ := chunk["type"].(string); typ == "done" {
		raw, _ := chunk["tools"].([]any)
		for _, t := range raw {
			m := gconv.Map(t)
			if m == nil {
				continue
			}
			var tool ModelTool
			tool.Id = gconv.String(m["id"])
			tool.Type = gconv.String(m["type"])
			if fn := gconv.Map(m["function"]); fn != nil {
				tool.Function.Name = gconv.String(fn["name"])
				tool.Function.Arguments = gconv.String(fn["arguments"])
			}
			tools = append(tools, tool)
		}
		return tools, true
	}
	return nil, false
}

// streamTextDelta 提取文本增量 chunk 的文本（读 content 子对象，字段名由网关端结构体统一管理）
func streamTextDelta(chunk map[string]any) string {
	var sb strings.Builder
	if content, ok := chunk["content"].(map[string]any); ok {
		for _, v := range content {
			if s := gconv.String(v); s != "" {
				sb.WriteString(s)
			}
		}
	}
	return sb.String()
}

// parseSSEStream 标准 SSE 解析：逐 data 行 JSON 回调，支持 [DONE] 与上下文取消。
// 对齐 model-gateway ParseSSEStream 语义：跳过空行、[DONE]、非 data 行。
func parseSSEStream(ctx context.Context, reader io.Reader, onChunk func(chunk map[string]any) error) error {
	scanner := bufio.NewScanner(reader)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	var sb strings.Builder
	for scanner.Scan() {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			if sb.Len() == 0 {
				continue
			}
			data := sb.String()
			sb.Reset()
			if data == "[DONE]" {
				continue
			}
			var chunk map[string]any
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				continue
			}
			if onChunk != nil {
				if err := onChunk(chunk); err != nil {
					return err
				}
			}
			continue
		}
		if strings.HasPrefix(line, "data:") {
			sb.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
		}
	}
	return scanner.Err()
}