mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-15 05:04:27 +08:00
feat(agent): Go port — canvas engine, 22 components, DSL v2, 13 endpoints (#15952)
Ports the agent canvas subsystem from Python to Go.
## What's included
### Canvas Engine (Phase 0/1)
- State engine, scheduler, variable resolver, Redis checkpoint store,
cancel protocol
- **209 tests** across canvas / component / io packages
### 22 Components (P0–P4)
| Tier | Components |
|---|---|
| P0 T1+T2+T3 | LLM, Agent, ExitLoop, Switch, Categorize, Begin,
Message, Invoke |
| P1 T3 | VariableAggregator, VariableAssigner, StringTransform,
ListOperations, DataOperations |
| P2 T3 | Iteration, IterationItem, Loop, LoopItem |
| P3 T3 | UserFillUp, Fillup |
| P4 T5 | Browser, ExcelProcessor, DocsGenerator |
### DSL v2 Schema (Phase 2.5)
- Typed v2 in-memory model with v1-to-v2 auto-detect converter
- v1 legacy field stripping per plan §2.11.7
### HTTP Endpoints & Bug Fixes (Plans PR1–PR3)
- **DELETE SQL bug fix**: gorm v2 `Where("id = ?", id).Delete(...)`
pattern
- **CreateAgent validation**: title/DSL required, duplicate check, 103
envelope
- **13 new endpoints**: templates, prompts, tags, sessions CRUD,
chat/completions (SSE + non-stream stubs), rerun, test_db_connection,
logs, webhook/logs
- **756 Go unit tests** (745 → 756, +18)
- **17 → 0 Python integration test failures** (test_agents.py +
test_session_management/)
### Tools
21 eino tools: HTTPHelper, search tools, financial/data tools, mandatory
stubs
### Infrastructure
OTel observability, NATS message queue, DeepDoc gRPC client, SSRF
guards, IDOR mitigation
This commit is contained in:
398
internal/agent/component/agent_test.go
Normal file
398
internal/agent/component/agent_test.go
Normal file
@@ -0,0 +1,398 @@
|
||||
// Package component — Agent unit tests (Phase 2 P0, plan §2.11.3 row 8).
|
||||
//
|
||||
// Tests inject a canned agentRunner to verify the component contract
|
||||
// without requiring a real model or eino react agent runtime:
|
||||
//
|
||||
// 1. NoToolsReAct: the runner returns a plain answer → component
|
||||
// surfaces content with empty tool_calls.
|
||||
// 2. ToolCallRound: the runner returns a message with ToolCalls →
|
||||
// component extracts them into the tool_calls output.
|
||||
// 3. ExhaustRoundsError: the runner returns an error → component
|
||||
// propagates it.
|
||||
// 4. MissingModelID: the component rejects before calling the runner.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/flow/agent/react"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
|
||||
agenttool "ragflow/internal/agent/tool"
|
||||
)
|
||||
|
||||
// withAgentRunner replaces the package-level agentRunner for the duration
|
||||
// of t.
|
||||
func withAgentRunner(t *testing.T, fn func(context.Context, AgentParam) (*schema.Message, error)) {
|
||||
t.Helper()
|
||||
prev := agentRunner
|
||||
agentRunner = fn
|
||||
t.Cleanup(func() { agentRunner = prev })
|
||||
}
|
||||
|
||||
func TestAgent_NoToolsReAct(t *testing.T) {
|
||||
var calls int
|
||||
withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) {
|
||||
calls++
|
||||
return &schema.Message{Role: schema.Assistant, Content: "the answer is 42"}, nil
|
||||
})
|
||||
|
||||
c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 3})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"user_prompt": "what is 6*7?",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["content"], "the answer is 42"; got != want {
|
||||
t.Errorf("content=%v, want %v", got, want)
|
||||
}
|
||||
toolCalls, ok := out["tool_calls"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("tool_calls missing or wrong type: %T", out["tool_calls"])
|
||||
}
|
||||
if len(toolCalls) != 0 {
|
||||
t.Errorf("tool_calls=%d, want 0", len(toolCalls))
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Errorf("runner called %d times, want 1", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_ToolCallRound(t *testing.T) {
|
||||
var calls int
|
||||
withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) {
|
||||
calls++
|
||||
return &schema.Message{
|
||||
Role: schema.Assistant,
|
||||
Content: "final answer based on tool",
|
||||
ToolCalls: []schema.ToolCall{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "search",
|
||||
Arguments: `{"q": "ragflow"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
|
||||
c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 3})
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"user_prompt": "find out about ragflow",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got, want := out["content"], "final answer based on tool"; got != want {
|
||||
t.Errorf("content=%v, want %v", got, want)
|
||||
}
|
||||
toolCalls, ok := out["tool_calls"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("tool_calls missing or wrong type: %T", out["tool_calls"])
|
||||
}
|
||||
if len(toolCalls) != 1 {
|
||||
t.Fatalf("tool_calls=%d, want 1", len(toolCalls))
|
||||
}
|
||||
if toolCalls[0]["name"] != "search" {
|
||||
t.Errorf("tool name=%v, want search", toolCalls[0]["name"])
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Errorf("runner called %d times, want 1", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_ExhaustRoundsError(t *testing.T) {
|
||||
withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) {
|
||||
return nil, errors.New("agent: exhausted rounds without final answer")
|
||||
})
|
||||
|
||||
c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 2})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"user_prompt": "x",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error when loop exhausts without a final answer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_MissingModelID(t *testing.T) {
|
||||
c := NewAgentComponent(AgentParam{MaxRounds: 1})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{"user_prompt": "x"})
|
||||
if err == nil {
|
||||
t.Fatal("expected ParamError for missing model_id")
|
||||
}
|
||||
var pe *ParamError
|
||||
if !errors.As(err, &pe) {
|
||||
t.Errorf("err type=%T, want *ParamError", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_UnknownToolName(t *testing.T) {
|
||||
c := NewAgentComponent(AgentParam{
|
||||
ModelID: "stub",
|
||||
MaxRounds: 1,
|
||||
Tools: []string{"does_not_exist"},
|
||||
})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"user_prompt": "x",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown tool")
|
||||
}
|
||||
if !strings.Contains(err.Error(), `build tools: agent tool: unsupported tool "does_not_exist"`) {
|
||||
t.Fatalf("err = %q, want unsupported tool message", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_AllRegisteredToolsConfigPassesToRunner(t *testing.T) {
|
||||
var captured AgentParam
|
||||
withAgentRunner(t, func(_ context.Context, p AgentParam) (*schema.Message, error) {
|
||||
captured = p
|
||||
return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil
|
||||
})
|
||||
|
||||
c := NewAgentComponent(AgentParam{ModelID: "stub", MaxRounds: 1})
|
||||
_, err := c.Invoke(context.Background(), map[string]any{
|
||||
"user_prompt": "x",
|
||||
"tools": []any{
|
||||
"akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo",
|
||||
"email", "github", "google", "google_scholar", "jin10", "pubmed",
|
||||
"qweather", "retrieval", "searxng", "tavily", "tushare", "wencai",
|
||||
"wikipedia", "yahoo_finance", "execute_sql",
|
||||
},
|
||||
"tool_params": map[string]any{
|
||||
"execute_sql": map[string]any{
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if len(captured.Tools) != 21 {
|
||||
t.Fatalf("captured.Tools len = %d, want 21", len(captured.Tools))
|
||||
}
|
||||
if captured.ToolParams == nil || captured.ToolParams["execute_sql"] == nil {
|
||||
t.Fatalf("captured.ToolParams missing execute_sql: %#v", captured.ToolParams)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeToolCallingChatModel struct {
|
||||
tools []*schema.ToolInfo
|
||||
}
|
||||
|
||||
func (m *fakeToolCallingChatModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
||||
return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func (m *fakeToolCallingChatModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
sr, sw := schema.Pipe[*schema.Message](1)
|
||||
go func() {
|
||||
defer sw.Close()
|
||||
_ = sw.Send(&schema.Message{Role: schema.Assistant, Content: "ok"}, io.EOF)
|
||||
}()
|
||||
return sr, nil
|
||||
}
|
||||
|
||||
func (m *fakeToolCallingChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
||||
cp := *m
|
||||
cp.tools = append([]*schema.ToolInfo(nil), tools...)
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func TestAgent_CanCreateReactAgentWithAllRegisteredTools(t *testing.T) {
|
||||
p := AgentParam{
|
||||
Tools: []string{
|
||||
"akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo",
|
||||
"email", "github", "google", "google_scholar", "jin10", "pubmed",
|
||||
"qweather", "retrieval", "searxng", "tavily", "tushare", "wencai",
|
||||
"wikipedia", "yahoo_finance", "execute_sql",
|
||||
},
|
||||
ToolParams: map[string]map[string]any{
|
||||
"execute_sql": {
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
},
|
||||
},
|
||||
MaxRounds: 1,
|
||||
}
|
||||
tools, err := buildAgentTools(p)
|
||||
if err != nil {
|
||||
t.Fatalf("buildAgentTools: %v", err)
|
||||
}
|
||||
if len(tools) != len(p.Tools) {
|
||||
t.Fatalf("len(tools) = %d, want %d", len(tools), len(p.Tools))
|
||||
}
|
||||
_, err = react.NewAgent(context.Background(), &react.AgentConfig{
|
||||
ToolCallingModel: &fakeToolCallingChatModel{},
|
||||
ToolsConfig: compose.ToolsNodeConfig{
|
||||
Tools: tools,
|
||||
},
|
||||
MaxStep: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("react.NewAgent(all tools): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_Registered(t *testing.T) {
|
||||
c, err := New("Agent", map[string]any{"model_id": "stub", "user_prompt": "x"})
|
||||
if err != nil {
|
||||
t.Fatalf("New(Agent): %v", err)
|
||||
}
|
||||
if c.Name() != "Agent" {
|
||||
t.Errorf("Name()=%q, want Agent", c.Name())
|
||||
}
|
||||
}
|
||||
|
||||
// exhaustStepsModel is a scripted ToolCallingChatModel that emits a
|
||||
// tool_call on every Generate and never returns final content. It
|
||||
// is the input driver for TestAgent_ReActExhaustsSteps, which needs
|
||||
// the eino ReAct loop to hit its MaxStep ceiling.
|
||||
type exhaustStepsModel struct {
|
||||
turn int
|
||||
rounds [][]*schema.Message
|
||||
boundTools []*schema.ToolInfo
|
||||
toolName string
|
||||
toolArgs string
|
||||
}
|
||||
|
||||
func (m *exhaustStepsModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
||||
m.boundTools = tools
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m *exhaustStepsModel) Generate(_ context.Context, in []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
||||
cp := make([]*schema.Message, len(in))
|
||||
copy(cp, in)
|
||||
m.rounds = append(m.rounds, cp)
|
||||
m.turn++
|
||||
return &schema.Message{
|
||||
Role: schema.Assistant,
|
||||
ToolCalls: []schema.ToolCall{{
|
||||
ID: fmt.Sprintf("call_%d", m.turn),
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: m.toolName,
|
||||
Arguments: m.toolArgs,
|
||||
},
|
||||
}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *exhaustStepsModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
sr, sw := schema.Pipe[*schema.Message](1)
|
||||
sw.Close()
|
||||
return sr, nil
|
||||
}
|
||||
|
||||
// TestAgent_ReActExhaustsSteps drives a real react.NewAgent whose
|
||||
// scripted model always returns a tool_call and never returns final
|
||||
// content. With MaxStep: 2 the loop must terminate with an error
|
||||
// from eino's MaxStep guard, while the real ExeSQLTool is invoked
|
||||
// at least once on the way. This is the eino error-path counterpart
|
||||
// to TestExeSQL_RealReactAgent_ExecutesTool: the latter proves the
|
||||
// happy path (model returns tool_call, framework runs tool, model
|
||||
// returns final); this one proves the loop guard.
|
||||
func TestAgent_ReActExhaustsSteps(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Real ExeSQLTool with sqlmock. The query is identical across
|
||||
// turns; sqlmock's QueryMatcherEqual will accept each call.
|
||||
// eino's MaxStep=2 with a tool_call-only model invokes the tool
|
||||
// exactly once before the loop guard fires (per eino's react
|
||||
// internals — the second iteration is the MaxStep check itself,
|
||||
// not a new tool call), so stage one ping + one query.
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
|
||||
if err != nil {
|
||||
t.Fatalf("sqlmock.New: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"x"}).AddRow(1))
|
||||
|
||||
// Default sql.Open would try to connect to a real MySQL; the
|
||||
// dialer stub makes the tool talk to sqlmock instead.
|
||||
dialer := func(_, _ string) (*sql.DB, error) { return db, nil }
|
||||
// BuildByName goes through the public registry — the same path
|
||||
// AgentComponent.buildAgentTools takes. This proves the agent's
|
||||
// own wiring (ToolsConfig -> real BaseTool) works under the
|
||||
// MaxStep guard, not a backdoor constructor.
|
||||
built, err := agenttool.BuildByName("execute_sql", map[string]any{
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("agenttool.BuildByName(execute_sql): %v", err)
|
||||
}
|
||||
exeSQLTool, ok := built.(*agenttool.ExeSQLTool)
|
||||
if !ok {
|
||||
t.Fatalf("BuildByName(execute_sql) returned %T, want *ExeSQLTool", built)
|
||||
}
|
||||
realTool := exeSQLTool.WithExeSQLDialer(dialer)
|
||||
|
||||
mdl := &exhaustStepsModel{
|
||||
toolName: "execute_sql",
|
||||
toolArgs: `{"sql": "SELECT 1"}`,
|
||||
}
|
||||
|
||||
agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
|
||||
ToolCallingModel: mdl,
|
||||
ToolsConfig: compose.ToolsNodeConfig{
|
||||
Tools: []einotool.BaseTool{realTool},
|
||||
},
|
||||
MaxStep: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("react.NewAgent: %v", err)
|
||||
}
|
||||
|
||||
out, err := agent.Generate(context.Background(), []*schema.Message{
|
||||
schema.UserMessage("loop forever"),
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatalf("agent.Generate returned no error; out=%+v — expected MaxStep exhaustion", out)
|
||||
}
|
||||
if mdl.turn < 1 {
|
||||
t.Errorf("model.Generate called %d times, want >= 1 (the loop should have invoked it before giving up)", mdl.turn)
|
||||
}
|
||||
if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "execute_sql" {
|
||||
names := make([]string, 0, len(mdl.boundTools))
|
||||
for _, ti := range mdl.boundTools {
|
||||
names = append(names, ti.Name)
|
||||
}
|
||||
t.Errorf("tools bound to model = %v, want [execute_sql]", names)
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user