hhs
ad700743ef
- 移除 adapter.go 中的 ctxKeySessionID 定义 - 移除 callback.go 中的 ctxKeyRequestID 定义 - 统一使用 trace.WithSessionID/WithRequestID - adapter.go 使用 trace.FromContext 替换 logger.Log - callback.go 使用 trace.FromContext - 移除双重日志，SetActiveRequest 失败直接返回错误 - 更新测试文件导入 trace 包
132 lines
3.8 KiB
Go
132 lines
3.8 KiB
Go
package eino
|
||
|
||
import (
|
||
"context"
|
||
"io"
|
||
|
||
"github.com/cloudwego/eino/callbacks"
|
||
"github.com/cloudwego/eino/components/model"
|
||
"github.com/cloudwego/eino/schema"
|
||
callbacksHelper "github.com/cloudwego/eino/utils/callbacks"
|
||
|
||
"github.com/hhs/camtalk/internal/models"
|
||
"github.com/hhs/camtalk/internal/orchestrator"
|
||
"github.com/hhs/camtalk/internal/trace"
|
||
)
|
||
|
||
// context key 类型，避免与其他包冲突。
|
||
type ctxKeySender struct{}
|
||
type ctxKeyState struct{}
|
||
|
||
// WithSender 将 Sender 注入 context。
|
||
func WithSender(ctx context.Context, sender orchestrator.Sender) context.Context {
|
||
return context.WithValue(ctx, ctxKeySender{}, sender)
|
||
}
|
||
|
||
// WithRequestID 将 requestID 注入 context（使用 trace 包）。
|
||
func WithRequestID(ctx context.Context, requestID string) context.Context {
|
||
return trace.WithRequestID(ctx, requestID)
|
||
}
|
||
|
||
// WithPipelineState 将 PipelineState 注入 context。
|
||
func WithPipelineState(ctx context.Context, state *PipelineState) context.Context {
|
||
return context.WithValue(ctx, ctxKeyState{}, state)
|
||
}
|
||
|
||
// senderFromCtx 从 context 获取 Sender。
|
||
func senderFromCtx(ctx context.Context) orchestrator.Sender {
|
||
s, _ := ctx.Value(ctxKeySender{}).(orchestrator.Sender)
|
||
return s
|
||
}
|
||
|
||
// requestIDFromCtx 从 context 获取 requestID（使用 trace 包）。
|
||
func requestIDFromCtx(ctx context.Context) string {
|
||
return trace.GetRequestID(ctx)
|
||
}
|
||
|
||
// stateFromCtx 从 context 获取 PipelineState。
|
||
func stateFromCtx(ctx context.Context) *PipelineState {
|
||
s, _ := ctx.Value(ctxKeyState{}).(*PipelineState)
|
||
return s
|
||
}
|
||
|
||
// BuildCallbackHandler 构建 Eino Callback Handler。
|
||
//
|
||
// 核心职责：ChatModel 节点通过 OnEndWithStreamOutput 逐 token 推送 llm_chunk 到客户端，
|
||
// 同时累积完整文本到 PipelineState。
|
||
//
|
||
// 其他节点的消息推送（stt_result、tts_audio、llm_done）由各 Lambda 内部直接调用 Sender。
|
||
func BuildCallbackHandler() callbacks.Handler {
|
||
return callbacksHelper.NewHandlerHelper().
|
||
ChatModel(&callbacksHelper.ModelCallbackHandler{
|
||
OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {
|
||
log := trace.FromContext(ctx)
|
||
sender := senderFromCtx(ctx)
|
||
requestID := requestIDFromCtx(ctx)
|
||
state := stateFromCtx(ctx)
|
||
|
||
if sender == nil || requestID == "" {
|
||
log.Warnw("ModelCallback: missing sender or request_id in context",
|
||
"node", info.Name)
|
||
return ctx
|
||
}
|
||
|
||
// 异步消费流，避免阻塞框架的下游处理。
|
||
// 框架对流做了内部拷贝，此 goroutine 读取独立副本。
|
||
go func() {
|
||
defer output.Close()
|
||
|
||
for {
|
||
chunk, err := output.Recv()
|
||
if err != nil {
|
||
if err == io.EOF {
|
||
return
|
||
}
|
||
log.Errorw("ModelCallback: stream recv error",
|
||
"node", info.Name, "error", err)
|
||
return
|
||
}
|
||
|
||
if chunk == nil || chunk.Message == nil {
|
||
continue
|
||
}
|
||
|
||
delta := chunk.Message.Content
|
||
if delta == "" {
|
||
continue
|
||
}
|
||
|
||
// 推送 llm_chunk 到客户端
|
||
if err := sender.SendLLMChunk(models.WsLLMChunk{
|
||
Type: "llm_chunk",
|
||
RequestID: requestID,
|
||
Delta: delta,
|
||
Role: "assistant",
|
||
}); err != nil {
|
||
log.Errorw("ModelCallback: send llm_chunk failed", "error", err)
|
||
}
|
||
|
||
// 累积完整文本到 State
|
||
if state != nil {
|
||
state.AppendText(delta)
|
||
}
|
||
|
||
// 记录 token 用量（流的最后一帧携带）
|
||
if chunk.TokenUsage != nil && state != nil {
|
||
state.mu.Lock()
|
||
state.TokenUsage = &TokenUsage{
|
||
Prompt: chunk.TokenUsage.PromptTokens,
|
||
Completion: chunk.TokenUsage.CompletionTokens,
|
||
Total: chunk.TokenUsage.TotalTokens,
|
||
}
|
||
state.mu.Unlock()
|
||
}
|
||
}
|
||
}()
|
||
|
||
return ctx
|
||
},
|
||
}).
|
||
Handler()
|
||
}
|