diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index aef531829..d4b117c2c 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -15,6 +15,7 @@ import ( "context" "encoding/json" "fmt" + "sort" "strings" einotool "github.com/cloudwego/eino/components/tool" @@ -282,17 +283,16 @@ func chunksFromState(ctx context.Context) []prompts.CitationSource { return out } -// GetInputForm aggregates the Agent's own meta-schema with each -// sub-tool's input form. Mirrors Python's `Agent.get_input_form`. +// GetInputForm aggregates the Agent's own user-callable parameters with each +// sub-tool's input form. Mirrors Python's `Agent.get_input_form`, which +// returns a flat field-definition map keyed by input name. // // Today the sub-tool input forms are aggregated via an optional // `InputForm() map[string]any` method on the eino tool (when tools // implement it); tools that don't expose a structured input form // are skipped silently. func (c *AgentComponent) GetInputForm() map[string]any { - out := map[string]any{ - "self": c.param.Meta, - } + out := extractAgentPromptInputForm(c.param.SystemPrompt, c.param.UserPrompt) tools, err := buildAgentTools(c.param) if err != nil { return out @@ -314,6 +314,41 @@ func (c *AgentComponent) GetInputForm() map[string]any { return out } +func extractAgentPromptInputForm(systemPrompt, userPrompt string) map[string]any { + out := map[string]any{} + seen := map[string]struct{}{} + matches := append(runtime.VarRefPattern.FindAllStringSubmatch(systemPrompt, -1), runtime.VarRefPattern.FindAllStringSubmatch(userPrompt, -1)...) + for _, match := range matches { + if len(match) < 2 { + continue + } + key := strings.TrimSpace(match[1]) + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + out[key] = map[string]any{ + "type": "line", + "name": key, + "optional": false, + } + } + return out +} + +func sortedAgentPromptInputKeys(systemPrompt, userPrompt string) []string { + form := extractAgentPromptInputForm(systemPrompt, userPrompt) + keys := make([]string, 0, len(form)) + for key := range form { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} + // Reset calls Reset on every sub-tool that implements the Resetter // interface. Mirrors Python's per-tool reset() — useful for clearing // per-invocation state (caches, scratch buffers) between calls. diff --git a/internal/agent/component/agent_meta_test.go b/internal/agent/component/agent_meta_test.go index ddf087f82..94b6e5365 100644 --- a/internal/agent/component/agent_meta_test.go +++ b/internal/agent/component/agent_meta_test.go @@ -22,33 +22,43 @@ import ( "github.com/cloudwego/eino/schema" ) -// TestAgent_GetInputForm_IncludesSelfMeta: GetInputForm returns the -// Agent's own meta-schema under the "self" key, even when no tools -// are configured. -func TestAgent_GetInputForm_IncludesSelfMeta(t *testing.T) { +// TestAgent_GetInputForm_UsesPromptReferences verifies the Agent input form +// follows Python's behavior: it is derived from prompt template references. +func TestAgent_GetInputForm_UsesPromptReferences(t *testing.T) { c := NewAgentComponent(AgentParam{ - ModelID: "echo", - Meta: AgentMeta{ - Name: "research_agent", - Description: "Performs multi-step research", - Parameters: map[string]AgentMetaParam{ - "user_prompt": {Type: "string", Description: "The question", Required: true}, - }, - }, + ModelID: "echo", + SystemPrompt: "Use {sys.query} and {{sys.user_id}} to answer.", + UserPrompt: "Question: {sys.query}", }) form := c.GetInputForm() if form == nil { t.Fatal("GetInputForm returned nil") } - self, ok := form["self"].(AgentMeta) + if _, ok := form["self"]; ok { + t.Fatalf("form unexpectedly contains synthetic self key: %#v", form["self"]) + } + queryField, ok := form["sys.query"].(map[string]any) if !ok { - t.Fatalf("form['self'] type=%T, want AgentMeta", form["self"]) + t.Fatalf("form['sys.query'] type=%T, want map[string]any", form["sys.query"]) } - if self.Name != "research_agent" { - t.Errorf("Name=%q, want research_agent", self.Name) + if queryField["type"] != "line" { + t.Errorf("sys.query.type=%v, want line", queryField["type"]) } - if self.Parameters["user_prompt"].Type != "string" { - t.Errorf("user_prompt type=%q, want string", self.Parameters["user_prompt"].Type) + if queryField["name"] != "sys.query" { + t.Errorf("sys.query.name=%v, want 'sys.query'", queryField["name"]) + } + if queryField["optional"] != false { + t.Errorf("sys.query.optional=%v, want false", queryField["optional"]) + } + userIDField, ok := form["sys.user_id"].(map[string]any) + if !ok { + t.Fatalf("form['sys.user_id'] type=%T, want map[string]any", form["sys.user_id"]) + } + if userIDField["name"] != "sys.user_id" { + t.Errorf("sys.user_id.name=%v, want sys.user_id", userIDField["name"]) + } + if userIDField["optional"] != false { + t.Errorf("sys.user_id.optional=%v, want false", userIDField["optional"]) } } @@ -59,17 +69,27 @@ func TestAgent_Reset_NoTools(t *testing.T) { c.Reset() // should not panic } -// TestAgent_Meta_DefaultsToEmpty: zero-value AgentParam.Meta is the -// empty AgentMeta struct (not a nil map dereference). +// TestAgent_GetInputForm_DeduplicatesPromptReferences ensures repeated refs +// across system and user prompts collapse to one input field. +func TestAgent_GetInputForm_DeduplicatesPromptReferences(t *testing.T) { + c := NewAgentComponent(AgentParam{ + ModelID: "echo", + SystemPrompt: "A {sys.query}", + UserPrompt: "B {{sys.query}}", + }) + keys := sortedAgentPromptInputKeys(c.param.SystemPrompt, c.param.UserPrompt) + if len(keys) != 1 || keys[0] != "sys.query" { + t.Fatalf("keys=%v, want [sys.query]", keys) + } +} + +// TestAgent_Meta_DefaultsToEmpty: prompts without variable references yield +// an empty input-form map. func TestAgent_Meta_DefaultsToEmpty(t *testing.T) { c := NewAgentComponent(AgentParam{ModelID: "echo"}) form := c.GetInputForm() - self, ok := form["self"].(AgentMeta) - if !ok { - t.Fatalf("form['self'] type=%T, want AgentMeta", form["self"]) - } - if self.Name != "" || self.Description != "" { - t.Errorf("expected empty meta, got %+v", self) + if len(form) != 0 { + t.Fatalf("expected empty input form, got %+v", form) } } diff --git a/internal/handler/agent_component.go b/internal/handler/agent_component.go index 4e67278ac..3698a48bd 100644 --- a/internal/handler/agent_component.go +++ b/internal/handler/agent_component.go @@ -32,6 +32,7 @@ import ( "context" "errors" "fmt" + "strings" "github.com/gin-gonic/gin" @@ -211,6 +212,11 @@ func (h *AgentHandler) DebugComponent(c *gin.Context) { // decorator for the debug endpoint. debugState := canvas.NewCanvasState("debug-"+componentID, "debug-task") debugState.Sys["tenant_id"] = user.ID + for key, value := range inputs { + if strings.HasPrefix(key, "sys.") && key != "sys.tenant_id" { + debugState.Sys[strings.TrimPrefix(key, "sys.")] = value + } + } invokeCtx := runtime.WithState(c.Request.Context(), debugState) outputs, err := comp.Invoke(invokeCtx, inputs) diff --git a/internal/handler/agent_component_test.go b/internal/handler/agent_component_test.go index c92a9bf3f..725ac1a9d 100644 --- a/internal/handler/agent_component_test.go +++ b/internal/handler/agent_component_test.go @@ -17,6 +17,7 @@ package handler import ( + "context" "encoding/json" "errors" "net/http/httptest" @@ -25,6 +26,7 @@ import ( "github.com/gin-gonic/gin" + agentruntime "ragflow/internal/agent/runtime" _ "ragflow/internal/agent/component" // registers the production factory "ragflow/internal/dao" "ragflow/internal/entity" @@ -210,6 +212,115 @@ func TestDebugComponent_HappyPath_Begin(t *testing.T) { } } +type sysEchoComponent struct{} + +func (s *sysEchoComponent) Invoke(ctx context.Context, _ map[string]any) (map[string]any, error) { + state, _, err := agentruntime.GetStateFromContext[*agentruntime.CanvasState](ctx) + if err != nil { + return nil, err + } + return map[string]any{ + "query": state.Sys["query"], + "tenant_id": state.Sys["tenant_id"], + }, nil +} + +func TestDebugComponent_SeedsSysInputsIntoCanvasState(t *testing.T) { + origFactory := agentruntime.DefaultFactory() + agentruntime.SetDefaultFactory(func(name string, _ map[string]any) (agentruntime.Component, error) { + if name != "Probe" { + return nil, errors.New("unexpected component name: " + name) + } + return &sysEchoComponent{}, nil + }) + t.Cleanup(func() { agentruntime.SetDefaultFactory(origFactory) }) + + cv := &entity.UserCanvas{ + ID: "c1", + DSL: map[string]any{ + "components": map[string]any{ + "probe": map[string]any{ + "obj": map[string]any{ + "component_name": "Probe", + "params": map[string]any{}, + }, + }, + }, + }, + } + h := &AgentHandler{loader: &fakeCanvasLoader{canvas: cv}} + + body := `{"params":{"sys.query":{"value":"hello sys"}}}` + c, w := componentCtx(t, "POST", "/api/v1/agents/c1/components/probe/debug", body) + c.Params = gin.Params{ + {Key: "canvas_id", Value: "c1"}, + {Key: "component_id", Value: "probe"}, + } + h.DebugComponent(c) + + if w.Code != 200 { + t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String()) + } + var env struct { + Code int `json:"code"` + Data map[string]interface{} `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &env); err != nil { + t.Fatalf("decode: %v (body=%s)", err, w.Body.String()) + } + if env.Data["query"] != "hello sys" { + t.Fatalf("data.query = %v, want 'hello sys'", env.Data["query"]) + } +} + +func TestDebugComponent_RejectsSysTenantIDOverride(t *testing.T) { + origFactory := agentruntime.DefaultFactory() + agentruntime.SetDefaultFactory(func(name string, _ map[string]any) (agentruntime.Component, error) { + if name != "Probe" { + return nil, errors.New("unexpected component name: " + name) + } + return &sysEchoComponent{}, nil + }) + t.Cleanup(func() { agentruntime.SetDefaultFactory(origFactory) }) + + cv := &entity.UserCanvas{ + ID: "c1", + DSL: map[string]any{ + "components": map[string]any{ + "probe": map[string]any{ + "obj": map[string]any{ + "component_name": "Probe", + "params": map[string]any{}, + }, + }, + }, + }, + } + h := &AgentHandler{loader: &fakeCanvasLoader{canvas: cv}} + + body := `{"params":{"sys.query":{"value":"hello sys"},"sys.tenant_id":{"value":"attacker-tenant"}}}` + c, w := componentCtx(t, "POST", "/api/v1/agents/c1/components/probe/debug", body) + c.Params = gin.Params{ + {Key: "canvas_id", Value: "c1"}, + {Key: "component_id", Value: "probe"}, + } + h.DebugComponent(c) + + if w.Code != 200 { + t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String()) + } + var env struct { + Code int `json:"code"` + Data map[string]interface{} `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &env); err != nil { + t.Fatalf("decode: %v (body=%s)", err, w.Body.String()) + } + if env.Data["tenant_id"] != "u-1" { + t.Fatalf("data.tenant_id = %v, want u-1", env.Data["tenant_id"]) + } +} + func TestDebugComponent_InvalidParams_MissingField(t *testing.T) { cv := beginCanvas() h := &AgentHandler{loader: &fakeCanvasLoader{canvas: cv}}