Eino ADK: ChatModelAgentMiddleware

type ChatModelAgentMiddleware interface {
    // BeforeAgent 在每次 agent 运行前调用，允许修改 instruction 和 tools 配置
    BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)

    // BeforeModelRewriteState 在每次模型调用前调用
    // 返回的 state 会被持久化到 agent 内部状态并传递给模型
    // 返回的 context 会传播到模型调用和后续 handler
    BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)

    // AfterModelRewriteState 在每次模型调用后调用
    // 输入的 state 包含模型响应作为最后一条消息
    AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)

    // WrapInvokableToolCall 用自定义行为包装工具的同步执行
    // 如果不需要包装，返回原始 endpoint 和 nil error
    // 仅对实现了 InvokableTool 的工具调用此方法
    WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error)

    // WrapStreamableToolCall 用自定义行为包装工具的流式执行
    // 如果不需要包装，返回原始 endpoint 和 nil error
    // 仅对实现了 StreamableTool 的工具调用此方法
    WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error)

    // WrapEnhancedInvokableToolCall 用自定义行为包装增强型工具的同步执行
    WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error)

    // WrapEnhancedStreamableToolCall 用自定义行为包装增强型工具的流式执行
    WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error)

    // WrapModel 用自定义行为包装聊天模型
    // 如果不需要包装，返回原始 model 和 nil error
    // 在请求时调用，每次模型调用前都会执行
    WrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error)
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    return ctx, state, nil
}

// InvokableToolCallEndpoint 是同步工具调用的函数签名
type InvokableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error)

// StreamableToolCallEndpoint 是流式工具调用的函数签名
type StreamableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error)

// EnhancedInvokableToolCallEndpoint 是增强型同步工具调用的函数签名
type EnhancedInvokableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)

// EnhancedStreamableToolCallEndpoint 是增强型流式工具调用的函数签名
type EnhancedStreamableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)

type ChatModelAgentContext struct {
    // Instruction 是当前 Agent 执行的指令
    // 包括 agent 配置的指令、框架和 AgentMiddleware 追加的额外指令，
    // 以及之前 BeforeAgent handler 应用的修改
    Instruction string

    // Tools 是当前为 Agent 执行配置的原始工具（无任何 wrapper 或 tool middleware）
    // 包括 AgentConfig 中传入的工具、框架隐式添加的工具（如 transfer/exit 工具），
    // 以及 middleware 已添加的其他工具
    Tools []tool.BaseTool

    // ReturnDirectly 是当前配置为使 Agent 直接返回的工具名称集合
    ReturnDirectly map[string]bool
}

type ChatModelAgentState struct {
    // Messages 包含当前对话会话中的所有消息
    Messages []Message
}

type ToolContext struct {
    // Name 是工具名称
    Name string

    // CallID 是此特定工具调用的唯一标识符
    CallID string
}

func (h *MyHandler) WrapInvokableToolCall(ctx context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {
    return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
        log.Printf("Tool %s (call %s) starting with args: %s", tCtx.Name, tCtx.CallID, argumentsInJSON)
        
        result, err := endpoint(ctx, argumentsInJSON, opts...)
        
        if err != nil {
            log.Printf("Tool %s failed: %v", tCtx.Name, err)
            return "", err
        }
        
        log.Printf("Tool %s completed with result: %s", tCtx.Name, result)
        return result, nil
    }, nil
}

type ModelContext struct {
    // Tools 是当前配置给 agent 的工具列表
    // 在请求时填充，包含将发送给模型的工具
    Tools []*schema.ToolInfo

    // ModelRetryConfig 包含模型的重试配置
    // 在请求时从 agent 的 ModelRetryConfig 填充
    // 用于 EventSenderModelWrapper 适当地包装流错误
    ModelRetryConfig *ModelRetryConfig
}

func (h *MyHandler) WrapModel(ctx context.Context, m model.BaseChatModel, mc *adk.ModelContext) (model.BaseChatModel, error) {
    return &myModelWrapper{
        inner: m,
        tools: mc.Tools,
    }, nil
}

type myModelWrapper struct {
    inner model.BaseChatModel
    tools []*schema.ToolInfo
}

func (w *myModelWrapper) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
    log.Printf("Model called with %d tools", len(w.tools))
    return w.inner.Generate(ctx, msgs, opts...)
}

func (w *myModelWrapper) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
    return w.inner.Stream(ctx, msgs, opts...)
}

// SetRunLocalValue 设置一个在当前 agent Run() 调用期间持久化的键值对
// 值的作用域限于此特定执行，不会在不同的 Run() 调用或 agent 实例之间共享
//
// 存储在这里的值与中断/恢复周期兼容 - 它们会被序列化并在 agent 恢复时还原
// 对于自定义类型，必须在 init() 函数中使用 schema.RegisterName[T]() 注册以确保正确序列化
//
// 此函数只能在 agent 执行期间从 ChatModelAgentMiddleware 内部调用
// 如果在 agent 执行上下文之外调用，返回错误
func SetRunLocalValue(ctx context.Context, key string, value any) error

// GetRunLocalValue 获取在当前 agent Run() 调用期间设置的值
// 值的作用域限于此特定执行，不会在不同的 Run() 调用或 agent 实例之间共享
//
// 通过 SetRunLocalValue 存储的值与中断/恢复周期兼容 - 它们会被序列化并在 agent 恢复时还原
// 对于自定义类型，必须在 init() 函数中使用 schema.RegisterName[T]() 注册以确保正确序列化
//
// 此函数只能在 agent 执行期间从 ChatModelAgentMiddleware 内部调用
// 如果找到值返回 (value, true, nil)，如果未找到返回 (nil, false, nil)，
// 如果在 agent 执行上下文之外调用返回错误
func GetRunLocalValue(ctx context.Context, key string) (any, bool, error)

// DeleteRunLocalValue 删除在当前 agent Run() 调用期间设置的值
//
// 此函数只能在 agent 执行期间从 ChatModelAgentMiddleware 内部调用
// 如果在 agent 执行上下文之外调用，返回错误
func DeleteRunLocalValue(ctx context.Context, key string) error

func init() {
    schema.RegisterName[*MyCustomData]("my_package.MyCustomData")
}

type MyCustomData struct {
    Count int
    Name  string
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) WrapInvokableToolCall(ctx context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {
    return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
        result, err := endpoint(ctx, argumentsInJSON, opts...)
        
        data := &MyCustomData{Count: 1, Name: tCtx.Name}
        if err := adk.SetRunLocalValue(ctx, "my_handler.last_tool", data); err != nil {
            log.Printf("Failed to set run local value: %v", err)
        }
        
        return result, err
    }, nil
}

func (h *MyHandler) AfterModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    if val, found, err := adk.GetRunLocalValue(ctx, "my_handler.last_tool"); err == nil && found {
        if data, ok := val.(*MyCustomData); ok {
            log.Printf("Last tool was: %s (count: %d)", data.Name, data.Count)
        }
    }
    return ctx, state, nil
}

// SendEvent 在 agent 执行期间向事件流发送自定义 AgentEvent
// 允许 ChatModelAgentMiddleware 实现发出自定义事件，
// 这些事件将被遍历 agent 事件流的调用者接收
//
// 此函数只能在 agent 执行期间从 ChatModelAgentMiddleware 内部调用
// 如果在 agent 执行上下文之外调用，返回错误
func SendEvent(ctx context.Context, event *AgentEvent) error

type State struct {
    Messages []Message
    extra    map[string]any  // 未导出，通过 SetRunLocalValue/GetRunLocalValue 访问

    // 以下为内部字段 - 请勿直接访问
    // 为与现有 checkpoint 向后兼容而保持导出
    ReturnDirectlyToolCallID string
    ToolGenActions           map[string]*AgentAction
    AgentName                string
    RemainingIterations      int

    internals map[string]any
}

middleware := adk.AgentMiddleware{
    BeforeChatModel: func(ctx context.Context, state *adk.ChatModelAgentState) error {
        return nil
    },
}

type MyHandler struct {
    *adk.BaseChatModelAgentMiddleware
}

func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) {
    newCtx := context.WithValue(ctx, myKey, myValue)
    return newCtx, state, nil
}

compose.ProcessState(ctx, func(_ context.Context, st *adk.State) error {
    st.Extra["myKey"] = myValue
    return nil
})

if err := adk.SetRunLocalValue(ctx, "myKey", myValue); err != nil {
    return ctx, state, err
}

if val, found, err := adk.GetRunLocalValue(ctx, "myKey"); err == nil && found {
}