CamTalk/backend/internal/eino/nodes_done.go

package eino

import (
	"context"
	"time"

	"github.com/cloudwego/eino/compose"

	"github.com/hhs/camtalk/internal/logger"
	"github.com/hhs/camtalk/internal/models"
)

// ctxKeyStartTime 请求开始时间的 context key。
type ctxKeyStartTime struct{}

// WithStartTime 将请求开始时间注入 context。
func WithStartTime(ctx context.Context, t time.Time) context.Context {
	return context.WithValue(ctx, ctxKeyStartTime{}, t)
}

// latencyFromCtx 从 context 获取开始时间并计算延迟（毫秒）。
func latencyFromCtx(ctx context.Context) int64 {
	if startTime, ok := ctx.Value(ctxKeyStartTime{}).(time.Time); ok {
		return time.Since(startTime).Milliseconds()
	}
	return 0
}

// NewDoneLambda 创建 Done Lambda 节点。
// 输入: struct{}（TTS 完成信号）→ 输出: *PipelineOutput
//
// 从 PipelineState 读取完整回复和 token 用量，发送 llm_done 到客户端。
// 历史消息追加由适配器负责（避免重复写入）。
func NewDoneLambda(defaultModel string) *compose.Lambda {
	return compose.InvokableLambda(func(ctx context.Context, _ struct{}) (PipelineOutput, error) {
		log := logger.Log
		sender := senderFromCtx(ctx)
		state := stateFromCtx(ctx)

		if state == nil {
			return PipelineOutput{}, nil
		}

		state.mu.Lock()
		fullResponse := state.FullResponse.String()
		transcribedText := state.TranscribedText
		tokenUsage := state.TokenUsage
		requestID := state.RequestID
		modelName := defaultModel
		state.mu.Unlock()

		// 发送 llm_done
		if sender != nil && requestID != "" {
			done := models.WsLLMDone{
				Type:      "llm_done",
				RequestID: requestID,
				FullText:  fullResponse,
				Model:     modelName,
				LatencyMs: latencyFromCtx(ctx),
			}
			if tokenUsage != nil {
				done.TokensUsed = struct {
					Prompt     int `json:"prompt"`
					Completion int `json:"completion"`
					Total      int `json:"total"`
				}{
					Prompt:     tokenUsage.Prompt,
					Completion: tokenUsage.Completion,
					Total:      tokenUsage.Total,
				}
			}
			if err := sender.SendLLMDone(done); err != nil {
				log.Errorw("发送 llm_done 失败", "error", err)
			}
		}

		log.Infow("查询处理完成",
			"request_id", requestID,
			"response_length", len(fullResponse))

		return PipelineOutput{
			TranscribedText: transcribedText,
			FullResponse:    fullResponse,
			Model:           modelName,
			TokenUsage:      tokenUsage,
		}, nil
	})
}