feat: agent callbacks for adk

Author: N3kox

adk/agent_middleware.go

@@ -0,0 +1,164 @@
+/*
+ * Copyright 2025 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package adk
+
+import (
+	"context"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/compose"
+)
+
+// AgentMiddleware provides hooks to customize agent behavior at various stages of execution.
+type AgentMiddleware struct {
+	// Name of the middleware, default empty. This will be used for middleware deduplication.
+	Name string
+
+	// AdditionalInstruction adds supplementary text to the agent's system instruction.
+	// This instruction is concatenated with the base instruction before each chat model call.
+	AdditionalInstruction string
+
+	// AdditionalTools adds supplementary tools to the agent's available toolset.
+	// These tools are combined with the tools configured for the agent.
+	AdditionalTools []tool.BaseTool
+
+	// BeforeChatModel is called before each ChatModel invocation, allowing modification of the agent state.
+	BeforeChatModel func(context.Context, *ChatModelAgentState) error
+
+	// AfterChatModel is called after each ChatModel invocation, allowing modification of the agent state.
+	AfterChatModel func(context.Context, *ChatModelAgentState) error
+
+	// WrapToolCall wraps tool calls with custom middleware logic.
+	// Each middleware contains Invokable and/or Streamable functions for tool calls.
+	WrapToolCall compose.ToolMiddleware
+
+	// BeforeAgent is called before the agent starts executing. It allows modifying the context
+	// or performing any setup actions before the agent begins processing.
+	// When an error is returned:
+	// 1. The framework will immediately return an AsyncIterator containing only this error
+	// 2. Subsequent BeforeAgent steps in other middlewares will be interrupted
+	// 3. The OnEvents handlers in previously executed middlewares will be invoked
+	BeforeAgent func(ctx context.Context, ac *AgentContext) (nextContext context.Context, err error)
+
+	// OnEvents is called to handle events generated by the agent during execution.
+	// - iter: The iterator contains the original output from the agent or the processed output from the previous middlewares.
+	// - gen: The generator is used to send the processed events to the next middleware or directly as output.
+	// This allows for filtering, transforming, or adding events in the middleware chain.
+	OnEvents func(ctx context.Context, ac *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent])
+}
+
+// AgentMiddlewareChecker is an interface that agents can implement to indicate
+// whether they support and enable middleware functionality.
+// Agents implementing this interface will execute middlewares internally;
+// otherwise, middlewares will be executed outside the agent by Runner.
+type AgentMiddlewareChecker interface {
+	IsAgentMiddlewareEnabled() bool
+}
+
+// ChatModelAgentState represents the state of a chat model agent during conversation.
+type ChatModelAgentState struct {
+	// Messages contains all messages in the current conversation session.
+	Messages []Message
+}
+
+type EntranceType string
+
+const (
+	// EntranceTypeRun indicates the agent is starting a new execution from scratch.
+	EntranceTypeRun EntranceType = "Run"
+	// EntranceTypeResume indicates the agent is resuming a previously interrupted execution.
+	EntranceTypeResume EntranceType = "Resume"
+)
+
+// AgentContext contains the context information for an agent's execution.
+// It provides access to input data, resume information, and execution options.
+type AgentContext struct {
+	// AgentInput contains the input data for the agent's execution.
+	AgentInput *AgentInput
+	// ResumeInfo contains information needed to resume a previously interrupted execution.
+	ResumeInfo *ResumeInfo
+	// AgentRunOptions contains options for configuring the agent's execution.
+	AgentRunOptions []AgentRunOption
+
+	// internal properties, read only
+	agentName string
+	entrance  EntranceType
+}
+
+func (a *AgentContext) AgentName() string {
+	return a.agentName
+}
+
+func (a *AgentContext) EntranceType() EntranceType {
+	return a.entrance
+}
+
+// GetGlobalAgentMiddlewares try to get agent middlewares set by RunnerMiddleware
+func GetGlobalAgentMiddlewares(ctx context.Context) []AgentMiddleware {
+	if v, ok := ctx.Value(globalAgentMiddlewareCtxKey{}).([]AgentMiddleware); ok && v != nil {
+		return v
+	}
+	return nil
+}
+
+type globalAgentMiddlewareCtxKey struct{}
+
+func isAgentMiddlewareEnabled(a Agent) bool {
+	if c, ok := a.(AgentMiddlewareChecker); ok && c.IsAgentMiddlewareEnabled() {
+		return true
+	}
+	return false
+}
+
+type agentMWHelper struct {
+	beforeAgentFns []func(ctx context.Context, ac *AgentContext) (nextContext context.Context, err error)
+	onEventsFns    []func(ctx context.Context, ac *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent])
+}
+
+func (a *agentMWHelper) execBeforeAgents(ctx context.Context, ac *AgentContext) (context.Context, *AsyncIterator[*AgentEvent]) {
+	var err error
+	for i, beforeAgent := range a.beforeAgentFns {
+		if beforeAgent == nil {
+			continue
+		}
+		ctx, err = beforeAgent(ctx, ac)
+		if err != nil {
+			iter, gen := NewAsyncIteratorPair[*AgentEvent]()
+			gen.Send(&AgentEvent{Err: err})
+			gen.Close()
+			return ctx, a.execOnEventsFromIndex(ctx, ac, i-1, iter)
+		}
+	}
+	return ctx, nil
+}
+
+func (a *agentMWHelper) execOnEvents(ctx context.Context, ac *AgentContext, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {
+	return a.execOnEventsFromIndex(ctx, ac, len(a.onEventsFns)-1, iter)
+}
+
+func (a *agentMWHelper) execOnEventsFromIndex(ctx context.Context, ac *AgentContext, fromIdx int, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {
+	for idx := fromIdx; idx >= 0; idx-- {
+		onEvents := a.onEventsFns[idx]
+		if onEvents == nil {
+			continue
+		}
+		i, g := NewAsyncIteratorPair[*AgentEvent]()
+		onEvents(ctx, ac, iter, g)
+		iter = i
+	}
+	return iter
+}