feat: agent callbacks for adk

adk/agent_middleware.go

@@ -0,0 +1,234 @@
+/*
+ * 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"
+	"runtime/debug"
+	"sync/atomic"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/compose"
+	"github.com/cloudwego/eino/internal/safe"
+)
+
+// AgentMiddleware provides hooks to customize agent behavior at various stages of execution.
+type AgentMiddleware struct {
+	// 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, arc *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, arc *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
+	isRootAgent bool
+	entrance    EntranceType
+}
+
+func (a *AgentContext) AgentName() string {
+	return a.agentName
+}
+
+func (a *AgentContext) IsRootAgent() bool {
+	return a.isRootAgent
+}
+
+func (a *AgentContext) EntranceType() EntranceType {
+	return a.entrance
+}
+
+type (
+	runnerPassedMiddlewaresCtxKey struct{}
+	runnerPassedMiddlewaresInfo   struct {
+		middlewares []AgentMiddleware
+		isRootAgent int32
+	}
+)
+
+func isRootAgent(ctx context.Context) bool {
+	if v, ok := ctx.Value(runnerPassedMiddlewaresCtxKey{}).(*runnerPassedMiddlewaresInfo); ok && v != nil {
+		val := atomic.SwapInt32(&v.isRootAgent, 1)
+		return val == 0
+	}
+	return false
+}
+
+func getRunnerPassedAgentMWs(ctx context.Context) []AgentMiddleware {
+	if v, ok := ctx.Value(runnerPassedMiddlewaresCtxKey{}).(*runnerPassedMiddlewaresInfo); ok && v != nil {
+		return v.middlewares
+	}
+	return nil
+}
+
+func isAgentMiddlewareEnabled(a Agent) bool {
+	if c, ok := a.(AgentMiddlewareChecker); ok && c.IsAgentMiddlewareEnabled() {
+		return true
+	}
+	return false
+}
+
+type agentMWRunner struct {
+	beforeAgentFns []func(ctx context.Context, arc *AgentContext) (nextContext context.Context, err error)
+	onEventsFns    []func(ctx context.Context, arc *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent])
+}
+
+func (a *agentMWRunner) 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 *agentMWRunner) execOnEvents(ctx context.Context, ac *AgentContext, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] {
+	return a.execOnEventsFromIndex(ctx, ac, len(a.onEventsFns)-1, iter)
+}
+
+func (a *agentMWRunner) 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
+}
+
+// NewAsyncOnSingleEventHandler creates an OnEvents middleware function that handles each event asynchronously.
+// It wraps the synchronous single event handler in a goroutine.
+func NewAsyncOnSingleEventHandler(onEvent func(ctx context.Context, arc *AgentContext, fromEvent *AgentEvent) (toEvent *AgentEvent)) (
+	onEventsFn func(ctx context.Context, arc *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent])) {
+	return func(ctx context.Context, arc *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent]) {
+		go NewSyncOnSingleEventHandler(onEvent)(ctx, arc, iter, gen)
+	}
+}
+
+// NewSyncOnSingleEventHandler creates an OnEvents middleware function that handles each event synchronously.
+// It applies the given onEvent function to each event in the iterator.
+func NewSyncOnSingleEventHandler(onEvent func(ctx context.Context, arc *AgentContext, fromEvent *AgentEvent) (toEvent *AgentEvent)) (
+	onEventsFn func(ctx context.Context, arc *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent])) {
+	return func(ctx context.Context, arc *AgentContext, iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent]) {
+		defer func() {
+			panicErr := recover()
+			if panicErr != nil {
+				e := safe.NewPanicErr(panicErr, debug.Stack())
+				gen.Send(&AgentEvent{Err: e})
+			}
+			gen.Close()
+		}()
+
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			toEvent := onEvent(ctx, arc, event)
+			if toEvent == nil {
+				continue
+			}
+			gen.Send(toEvent)
+		}
+	}
+}
+
+// BypassIterator creates a goroutine that simply passes events from the input iterator to the output generator.
+// This is useful when you need to do something without modifying events.
+func BypassIterator(iter *AsyncIterator[*AgentEvent], gen *AsyncGenerator[*AgentEvent]) {
+	go func() {
+		defer gen.Close()
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			gen.Send(event)
+		}
+	}()
+}