diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index fd5edaa8c1..aef5318290 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -13,6 +13,7 @@ package component import ( "context" + "encoding/json" "fmt" "strings" @@ -160,7 +161,13 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro } input := []*schema.Message{schema.UserMessage(p.UserPrompt)} - return agent.Generate(ctx, input) + opt, future := react.WithMessageFuture() + ctx = setArtifactCollector(ctx, future) + msg, err := agent.Generate(ctx, input, opt) + if err != nil { + return nil, err + } + return msg, nil } // addToolCallMemory summarizes the tool calls observed in msg via @@ -530,7 +537,7 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map } } } - artifacts := collectArtifactsFromToolCalls(msg) + artifacts := collectArtifactsFromToolCalls(ctx, msg) artifactMD := formatArtifactMarkdown(artifacts, content) out := map[string]any{ "content": content + artifactMD, @@ -653,24 +660,132 @@ type artifactEntry struct { URL string `json:"url"` } -// collectArtifactsFromToolCalls extracts artifact entries from a -// ReAct final message. Today, eino's react.Agent.Generate returns only -// the final message — the tool-response state lives inside the agent -// runtime and is not directly accessible. So this v1 returns an empty -// slice; the wiring lives in a follow-up that hoists the eino state -// into a place AgentComponent can read. -// -// When the tool response state becomes accessible (a future phase), -// the entry point to wire it is here: scan the eino conversation -// messages for entries whose `Extra["_ARTIFACTS"]` carries the -// per-tool artifact metadata, decode the JSON, and append to the -// returned slice. The shape expected from each tool is: -// -// { "name": "report.pdf", "url": "https://..." } -func collectArtifactsFromToolCalls(_ *schema.Message) []artifactEntry { +// artifactCollectorKey is the context key used to stash the +// MessageFuture from react.WithMessageFuture() so the AgentComponent +// can collect artifacts after the ReAct loop finishes. The collector +// is created per-invocation in runEinoReActAgent. +type artifactCollectorKey struct{} + +// setArtifactCollector registers the MessageFuture for this agent run +// in the context. It is called from runEinoReActAgent after +// react.WithMessageFuture() returns a future. +func setArtifactCollector(ctx context.Context, future react.MessageFuture) context.Context { + return context.WithValue(ctx, artifactCollectorKey{}, future) +} + +// getArtifactCollector retrieves the MessageFuture registered for the +// current agent run. Returns nil when no collector was registered +// (e.g., tests that stub agentRunner). +func getArtifactCollector(ctx context.Context) react.MessageFuture { + v := ctx.Value(artifactCollectorKey{}) + if v == nil { + return nil + } + if f, ok := v.(react.MessageFuture); ok { + return f + } return nil } +// collectArtifactsFromToolCalls drains the MessageFuture stored in +// ctx (if any) and extracts artifact entries from every tool response +// message that carries a `_ARTIFACTS` payload in its Extra field. +// The final message is ignored because it is an assistant message and +// does not contain tool results. Returns a de-duplicated list ordered +// by first appearance. +// +// The expected payload shape in each tool response is: +// +// { "_ARTIFACTS": [{ "name": "report.pdf", "url": "https://..." }, ...] } +func collectArtifactsFromToolCalls(ctx context.Context, _ *schema.Message) []artifactEntry { + future := getArtifactCollector(ctx) + if future == nil { + return nil + } + + seen := make(map[string]struct{}) + var out []artifactEntry + + iter := future.GetMessages() + for { + msg, ok, err := iter.Next() + if err != nil { + common.Debug("agent: artifact collection iterator error", zap.Error(err)) + break + } + if !ok { + break + } + if msg == nil || msg.Role != schema.Tool { + continue + } + rawArtifacts := extractArtifactsFromToolMessage(msg) + for _, a := range rawArtifacts { + if a.URL == "" || a.Name == "" { + continue + } + if _, exists := seen[a.URL]; exists { + continue + } + seen[a.URL] = struct{}{} + out = append(out, a) + } + } + return out +} + +// extractArtifactsFromToolMessage parses the JSON payload of a tool +// response message and returns the `_ARTIFACTS` list. The payload is +// read from msg.Content when it is non-empty; otherwise the first text +// element of msg.UserInputMultiContent is used. This matches the eino +// tool contract where tool results are delivered as a string. +func extractArtifactsFromToolMessage(msg *schema.Message) []artifactEntry { + payload := msg.Content + if payload == "" && len(msg.UserInputMultiContent) > 0 { + payload = toolMessageTextContent(msg) + } + if payload == "" { + return nil + } + + var envelope map[string]any + if err := json.Unmarshal([]byte(payload), &envelope); err != nil { + return nil + } + + raw, ok := envelope["_ARTIFACTS"].([]any) + if !ok { + return nil + } + + out := make([]artifactEntry, 0, len(raw)) + for _, item := range raw { + m, ok := item.(map[string]any) + if !ok { + continue + } + name, _ := m["name"].(string) + url, _ := m["url"].(string) + if name == "" || url == "" { + continue + } + out = append(out, artifactEntry{Name: name, URL: url}) + } + return out +} + +// toolMessageTextContent returns the first text content part of a tool +// message, or an empty string if no text part is found. +func toolMessageTextContent(msg *schema.Message) string { + for i := range msg.UserInputMultiContent { + part := &msg.UserInputMultiContent[i] + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + return part.Text + } + } + return "" +} + // formatArtifactMarkdown renders a slice of artifacts as markdown // links, omitting URLs already present in the existing text (Python's // `_collect_tool_artifact_markdown` does the same de-duplication). diff --git a/internal/agent/component/agent_artifact_test.go b/internal/agent/component/agent_artifact_test.go new file mode 100644 index 0000000000..f84cbf7be0 --- /dev/null +++ b/internal/agent/component/agent_artifact_test.go @@ -0,0 +1,249 @@ +package component + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent/react" + "github.com/cloudwego/eino/schema" +) + +// artifactTool is an invokable tool that returns a JSON envelope with _ARTIFACTS. +type artifactTool struct { + result string +} + +func (t *artifactTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "artifact_tool", + Desc: "returns artifacts", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "unused": {Type: schema.String}, + }), + }, nil +} + +func (t *artifactTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return t.result, nil +} + +// artifactModel is a scripted ToolCallingChatModel that emits one tool call +// and then a final answer. +type artifactModel struct { + turn int + callID string + toolName string + toolArgs string + final string +} + +func (m *artifactModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.turn++ + if m.turn == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: m.callID, + Type: "function", + Function: schema.FunctionCall{ + Name: m.toolName, + Arguments: m.toolArgs, + }, + }}, + }, nil + } + return &schema.Message{Role: schema.Assistant, Content: m.final}, nil +} + +func (m *artifactModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + sw.Close() + return sr, nil +} + +func (m *artifactModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func TestAgent_ReActAgent_CollectsArtifactsFromCodeExecTool(t *testing.T) { + toolResult, err := json.Marshal(map[string]any{ + "tool_called": true, + "message": "CodeExec executed", + "_ARTIFACTS": []map[string]any{ + { + "name": "agent_artifact_bug_demo.png", + "url": "/api/v1/documents/artifact/1ae8d553478544628bb8be267d502371.png", + }, + }, + }) + if err != nil { + t.Fatalf("marshal tool result: %v", err) + } + + opt, future := react.WithMessageFuture() + agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + ToolCallingModel: &artifactModel{ + callID: "call_1", + toolName: "artifact_tool", + toolArgs: `{"unused":"x"}`, + final: "The image has been generated.", + }, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&artifactTool{result: string(toolResult)}}, + }, + MaxStep: 3, + }) + if err != nil { + t.Fatalf("react.NewAgent: %v", err) + } + + _, err = agent.Generate(context.Background(), []*schema.Message{ + schema.UserMessage("generate a test image"), + }, opt) + if err != nil { + t.Fatalf("agent.Generate: %v", err) + } + + // Drain the future so collectArtifactsFromToolCalls can iterate synchronously. + msgs := drainFutureMessages(t, future) + if len(msgs) == 0 { + t.Fatal("MessageFuture produced no messages") + } + + // Re-create the same sequence in a context and call the collector. + fakeFuture := newSliceFuture(msgs) + ctx := setArtifactCollector(context.Background(), fakeFuture) + got := collectArtifactsFromToolCalls(ctx, nil) + + if len(got) != 1 { + t.Fatalf("collected %d artifacts, want 1", len(got)) + } + if got[0].Name != "agent_artifact_bug_demo.png" { + t.Errorf("name=%q, want agent_artifact_bug_demo.png", got[0].Name) + } + if got[0].URL == "" { + t.Error("artifact URL is empty") + } + + md := formatArtifactMarkdown(got, "done") + want := "![agent_artifact_bug_demo.png](/api/v1/documents/artifact/1ae8d553478544628bb8be267d502371.png)" + if !strings.Contains(md, want) { + t.Errorf("markdown=%q, want substring %q", md, want) + } +} + +func drainFutureMessages(t *testing.T, future react.MessageFuture) []*schema.Message { + t.Helper() + var out []*schema.Message + iter := future.GetMessages() + for { + msg, ok, err := iter.Next() + if err != nil { + t.Fatalf("MessageFuture.Next: %v", err) + } + if !ok { + break + } + out = append(out, msg) + } + return out +} + +// sliceFuture is a react.MessageFuture backed by a slice. +type sliceFuture struct { + iter *react.Iterator[*schema.Message] +} + +func newSliceFuture(messages []*schema.Message) *sliceFuture { + // Re-run a real react agent whose model returns the supplied messages as + // tool-call / tool-response / final-answer sequence. This gives us a real + // react.Iterator populated by eino's own callback plumbing. + opt, future := react.WithMessageFuture() + + // Extract any tool name referenced by the replayed messages so the agent's + // tool node can dispatch it. + toolName := "passthrough" + for _, m := range messages { + for _, tc := range m.ToolCalls { + if tc.Function.Name != "" { + toolName = tc.Function.Name + } + } + } + + model := &replayModel{messages: messages} + agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + ToolCallingModel: model, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&passthroughTool{name: toolName}}, + }, + MaxStep: len(messages) + 1, + }) + if err != nil { + panic(err) + } + _, err = agent.Generate(context.Background(), []*schema.Message{ + schema.UserMessage("replay"), + }, opt) + if err != nil { + panic(err) + } + return &sliceFuture{iter: future.GetMessages()} +} + +func (f *sliceFuture) GetMessages() *react.Iterator[*schema.Message] { + return f.iter +} + +func (f *sliceFuture) GetMessageStreams() *react.Iterator[*schema.StreamReader[*schema.Message]] { + return nil +} + +// replayModel replays a scripted sequence of messages on successive Generate calls. +type replayModel struct { + messages []*schema.Message + pos int +} + +func (m *replayModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if m.pos >= len(m.messages) { + return &schema.Message{Role: schema.Assistant, Content: "done"}, nil + } + msg := m.messages[m.pos] + m.pos++ + return msg, nil +} + +func (m *replayModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + sw.Close() + return sr, nil +} + +func (m *replayModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// passthroughTool echoes its input as a tool result. +type passthroughTool struct { + name string +} + +func (t *passthroughTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "echoes arguments", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "unused": {Type: schema.String}, + }), + }, nil +} + +func (t *passthroughTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return argumentsInJSON, nil +}