Eino 流式处理与 Callback 切面系统实战

// Invoke：等待完整结果
result, err := chatModel.Generate(ctx, messages)

// Stream：返回 StreamReader，逐片段读取
streamReader, err := chatModel.Stream(ctx, messages)
for {
    chunk, err := streamReader.Recv()
    if err == io.EOF {
        break
    }
    fmt.Print(chunk.Content)
}

// 创建流对
reader, writer := schema.Pipe[*schema.Message](bufferSize)

// 写入端（通常在组件内部）
go func() {
    defer writer.Close()
    for _, chunk := range chunks {
        if closed := writer.Send(chunk, nil); closed {
            return
        }
    }
}()

// 读取端（由下游消费）
for {
    msg, err := reader.Recv()
    if err == io.EOF {
        break
    }
    process(msg)
}

// 引擎内部等效逻辑
fullMessage := concatenate(streamReader) // 收集所有 chunk
nextNode.Invoke(ctx, fullMessage)         // 传递完整值

// 一个 StreamReader 被复制为多份
readers := stream.Split(originalReader, consumerCount)
// 每个 consumer 独立消费完整的流数据

// 多个上游流合并为一个
mergedReader := stream.Merge(reader1, reader2, reader3)

handler := callbacks.NewHandlerBuilder().
    OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
        log.Printf("[%s] 开始执行, 组件: %s, 输入: %v",
            info.Name, info.Type, input)
        return ctx
    }).
    OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
        log.Printf("[%s] 执行完成, 组件: %s, 输出: %v",
            info.Name, info.Type, output)
        return ctx
    }).
    OnEndWithStreamFn(func(ctx context.Context, info *callbacks.RunInfo,
        sr *schema.StreamReader[callbacks.CallbackOutput]) context.Context {
        log.Printf("[%s] 流式输出开始, 组件: %s", info.Name, info.Type)
        return ctx
    }).
    Build()

// 1. 全局：对 Graph 中所有节点生效
result, err := compiledGraph.Invoke(ctx, input,
    compose.WithCallbacks(handler))

// 2. 按组件类型：仅对特定类型组件生效
result, err := compiledGraph.Invoke(ctx, input,
    compose.WithCallbacks(handler).DesignateType(components.ComponentOfChatModel))

// 3. 按节点名称：仅对指定节点生效
result, err := compiledGraph.Invoke(ctx, input,
    compose.WithCallbacks(handler).DesignateNode("llm_node"))

import (
    "go.opentelemetry.io/otel"
    "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
    sdktrace "go.opentelemetry.io/otel/sdk/trace"
    einootel "github.com/cloudwego/eino/callbacks/otel"
)

func initTracer() func() {
    exporter, _ := otlptracehttp.New(ctx,
        otlptracehttp.WithEndpoint("otel-collector:4318"),
    )
    tp := sdktrace.NewTracerProvider(
        sdktrace.WithBatcher(exporter),
        sdktrace.WithResource(resource.NewWithAttributes(
            semconv.SchemaURL,
            semconv.ServiceName("my-agent-service"),
        )),
    )
    otel.SetTracerProvider(tp)
    return func() { tp.Shutdown(ctx) }
}

// 创建 OTel Callback Handler
otelHandler := einootel.NewHandler()

// 注入到 Graph 执行
result, err := graph.Invoke(ctx, input,
    compose.WithCallbacks(otelHandler))

package main

import (
    "context"
    "fmt"
    "io"
    "log"
    "time"

    "github.com/cloudwego/eino/callbacks"
    "github.com/cloudwego/eino/compose"
    "github.com/cloudwego/eino/components/model"
    "github.com/cloudwego/eino/schema"
)

// 构建带耗时追踪的 Callback Handler
func buildTracingHandler() callbacks.Handler {
    return callbacks.NewHandlerBuilder().
        OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
            start := time.Now()
            ctx = context.WithValue(ctx, "start_"+info.Name, start)
            log.Printf("[TRACE] ▶ Node=%s Type=%s started", info.Name, info.Type)
            return ctx
        }).
        OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
            if start, ok := ctx.Value("start_" + info.Name).(time.Time); ok {
                duration := time.Since(start)
                log.Printf("[TRACE] ◀ Node=%s Type=%s duration=%v",
                    info.Name, info.Type, duration)
            }
            return ctx
        }).
        OnEndWithStreamFn(func(ctx context.Context, info *callbacks.RunInfo,
            sr *schema.StreamReader[callbacks.CallbackOutput]) context.Context {
            if start, ok := ctx.Value("start_" + info.Name).(time.Time); ok {
                ttft := time.Since(start)
                log.Printf("[TRACE] ⇥ Node=%s first_chunk_latency=%v", info.Name, ttft)
            }
            return ctx
        }).
        Build()
}

func main() {
    ctx := context.Background()

    // 构建 Graph（参考上一篇文章的编排方式）
    g := compose.NewGraph[string, string]()

    // ... 添加节点和边 ...

    compiled, err := g.Compile(ctx)
    if err != nil {
        log.Fatal(err)
    }

    // 注入 tracing handler
    handler := buildTracingHandler()

    // 流式执行
    sr, err := compiled.Stream(ctx, "你好，请帮我分析这段代码的性能问题",
        compose.WithCallbacks(handler))
    if err != nil {
        log.Fatal(err)
    }

    // 消费流式输出
    for {
        chunk, err := sr.Recv()
        if err == io.EOF {
            break
        }
        if err != nil {
            log.Fatal(err)
        }
        fmt.Print(chunk)
    }
    fmt.Println()
}

[TRACE] ▶ Node=prompt_template Type=PromptTemplate started
[TRACE] ◀ Node=prompt_template Type=PromptTemplate duration=1.2ms
[TRACE] ▶ Node=llm Type=ChatModel started
[TRACE] ⇥ Node=llm first_chunk_latency=156ms
[TRACE] ▶ Node=output_parser Type=Lambda started
[TRACE] ◀ Node=output_parser Type=Lambda duration=0.3ms