mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-14 00:48:26 +08:00
706e0d2d06
### What problem does this PR solve? - Tools management - Pregel engine wrapper for better usage - UT race - Coding style ### Type of change - [x] Refactoring
548 lines
15 KiB
Go
548 lines
15 KiB
Go
package profile
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"testing"
|
|
|
|
"ragflow/internal/harness/core"
|
|
"ragflow/internal/harness/core/middlewares/subagent"
|
|
"ragflow/internal/harness/core/schema"
|
|
)
|
|
|
|
// ---- Mock model for testing ----
|
|
|
|
type mockProfileModel struct {
|
|
response string
|
|
}
|
|
|
|
func (m *mockProfileModel) Generate(ctx context.Context, msgs []*schema.Message, opts ...core.ModelOption) (*schema.Message, error) {
|
|
return &schema.Message{Role: schema.RoleAssistant, Content: m.response}, nil
|
|
}
|
|
func (m *mockProfileModel) Stream(ctx context.Context, msgs []*schema.Message, opts ...core.ModelOption) (*schema.StreamReader[*schema.Message], error) {
|
|
return schema.StreamReaderFromArray([]*schema.Message{{Role: schema.RoleAssistant, Content: m.response}}), nil
|
|
}
|
|
func (m *mockProfileModel) BindTools(tools []*schema.ToolInfo) error { return nil }
|
|
|
|
// ---- Tests ----
|
|
|
|
// TestRegisterAndParse verifies provider registration and model spec parsing.
|
|
func TestRegisterAndParse(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test-provider",
|
|
DefaultModel: "test-model",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "ok"}, nil
|
|
},
|
|
})
|
|
|
|
p := LookupProvider("test-provider")
|
|
if p == nil {
|
|
t.Fatal("expected provider to be found")
|
|
}
|
|
if p.DefaultModel != "test-model" {
|
|
t.Errorf("expected DefaultModel=test-model, got %s", p.DefaultModel)
|
|
}
|
|
|
|
// ParseModelSpec
|
|
prov, model, err := ParseModelSpec("test-provider:my-model")
|
|
if err != nil {
|
|
t.Fatalf("ParseModelSpec error: %v", err)
|
|
}
|
|
if prov != "test-provider" || model != "my-model" {
|
|
t.Errorf("expected (test-provider, my-model), got (%s, %s)", prov, model)
|
|
}
|
|
|
|
// Invalid spec
|
|
_, _, err = ParseModelSpec("no-colon")
|
|
if err == nil {
|
|
t.Error("expected error for invalid spec")
|
|
}
|
|
|
|
// Duplicate registration should panic
|
|
defer func() {
|
|
if r := recover(); r == nil {
|
|
t.Error("expected panic on duplicate registration")
|
|
}
|
|
}()
|
|
RegisterProvider(&ProviderProfile{Name: "test-provider"})
|
|
}
|
|
|
|
// TestHarnessRegistration verifies harness profile registration.
|
|
func TestHarnessRegistration(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "coding",
|
|
BaseSystemPrompt: StrPtr("You are a coder."),
|
|
MaxIterations: 20,
|
|
RecursionDepth: 5,
|
|
})
|
|
|
|
h := LookupHarness("coding")
|
|
if h == nil {
|
|
t.Fatal("expected harness to be found")
|
|
}
|
|
if h.BaseSystemPrompt == nil || *h.BaseSystemPrompt != "You are a coder." {
|
|
t.Errorf("unexpected BaseSystemPrompt")
|
|
}
|
|
if h.MaxIterations != 20 {
|
|
t.Errorf("expected MaxIterations=20, got %d", h.MaxIterations)
|
|
}
|
|
if h.RecursionDepth != 5 {
|
|
t.Errorf("expected RecursionDepth=5, got %d", h.RecursionDepth)
|
|
}
|
|
}
|
|
|
|
// TestNewAgent_Basic verifies the full NewAgent flow: provider + model + agent.
|
|
func TestNewAgent_Basic(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "hello"}, nil
|
|
},
|
|
})
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
if agent == nil {
|
|
t.Fatal("expected non-nil agent")
|
|
}
|
|
|
|
// Run and check output.
|
|
runner := core.NewTypedRunner(core.RunnerConfig[*schema.Message]{Agent: agent})
|
|
iter := runner.Run(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
var final string
|
|
for {
|
|
ev, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if ev.Err != nil {
|
|
t.Fatalf("run error: %v", ev.Err)
|
|
}
|
|
if ev.Output != nil && ev.Output.MessageOutput != nil &&
|
|
!ev.Output.MessageOutput.IsStreaming &&
|
|
ev.Output.MessageOutput.Message != nil {
|
|
final = ev.Output.MessageOutput.Message.Content
|
|
}
|
|
}
|
|
if final != "hello" {
|
|
t.Errorf("expected 'hello', got %q", final)
|
|
}
|
|
}
|
|
|
|
// TestNewAgent_WithHarness verifies harness profile overrides the system prompt.
|
|
func TestNewAgent_WithHarness(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "ok"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "research",
|
|
BaseSystemPrompt: StrPtr("Research mode."),
|
|
MaxIterations: 15,
|
|
})
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
HarnessProfileName: "research",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
// Verify the system prompt was set correctly by running a query.
|
|
runner := core.NewTypedRunner(core.RunnerConfig[*schema.Message]{Agent: agent})
|
|
iter := runner.Run(context.Background(), []*schema.Message{schema.UserMessage("test")})
|
|
for {
|
|
ev, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if ev.Err != nil {
|
|
t.Logf("event error (expected with mock model): %v", ev.Err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestNewAgent_UserOverrides verifies user explicit config takes precedence.
|
|
func TestNewAgent_UserOverrides(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "ok"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "chat",
|
|
BaseSystemPrompt: StrPtr("Harness prompt."),
|
|
MaxIterations: 5,
|
|
})
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
HarnessProfileName: "chat",
|
|
Instruction: StrPtr("User instruction."),
|
|
MaxIterations: 20,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
// Just verify no error - user instruction takes precedence.
|
|
runner := core.NewTypedRunner(core.RunnerConfig[*schema.Message]{Agent: agent})
|
|
iter := runner.Run(context.Background(), []*schema.Message{schema.UserMessage("t")})
|
|
for {
|
|
ev, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
_ = ev
|
|
}
|
|
}
|
|
|
|
// TestNewAgent_WithSubAgents verifies SubAgentSpecs integration.
|
|
func TestNewAgent_WithSubAgents(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "ok"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "agentic",
|
|
RecursionDepth: 3,
|
|
})
|
|
|
|
// Create a sub-agent spec (uses mock model for simplicity).
|
|
subAgentModel := &mockProfileModel{response: "sub result"}
|
|
subSpec := subagent.SubAgentSpec{
|
|
Name: "helper",
|
|
Description: "Helper agent",
|
|
AgentConfig: &subagent.AgentConfig{
|
|
Model: subAgentModel,
|
|
},
|
|
}
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
HarnessProfileName: "agentic",
|
|
SubAgentSpecs: []subagent.SubAgentSpec{subSpec},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
if agent == nil {
|
|
t.Fatal("expected non-nil agent")
|
|
}
|
|
t.Log("NewAgent with sub-agents: OK")
|
|
}
|
|
|
|
// TestNewAgent_ProviderNotFound verifies error for unknown provider.
|
|
func TestNewAgent_ProviderNotFound(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
_, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "nonexistent:model",
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected error for unknown provider")
|
|
}
|
|
t.Logf("expected error: %v", err)
|
|
}
|
|
|
|
// TestNewAgent_InvalidSpec verifies error for malformed model spec.
|
|
func TestNewAgent_InvalidSpec(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
_, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "bad",
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected error for invalid model spec")
|
|
}
|
|
}
|
|
|
|
// TestRegisterPanics verifies nil/empty registration panics.
|
|
func TestRegisterPanics(t *testing.T) {
|
|
defer func() {
|
|
if r := recover(); r == nil {
|
|
t.Error("expected panic for nil provider")
|
|
}
|
|
}()
|
|
RegisterProvider(nil)
|
|
}
|
|
|
|
// TestValidate verifies the Validate helper.
|
|
func TestValidate(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
// Valid config: no errors.
|
|
errs := Validate(&AgentConfig{ModelSpec: "a:b"})
|
|
if len(errs) != 0 {
|
|
t.Errorf("expected no errors, got %v", errs)
|
|
}
|
|
|
|
// Missing ModelSpec.
|
|
errs = Validate(&AgentConfig{})
|
|
if len(errs) == 0 {
|
|
t.Error("expected error for missing ModelSpec")
|
|
}
|
|
|
|
// Invalid ModelSpec.
|
|
errs = Validate(&AgentConfig{ModelSpec: "no-colon"})
|
|
if len(errs) == 0 {
|
|
t.Error("expected error for invalid ModelSpec")
|
|
}
|
|
|
|
// Unknown harness profile.
|
|
RegisterHarness(&HarnessProfile{Name: "real-profile"})
|
|
errs = Validate(&AgentConfig{ModelSpec: "a:b", HarnessProfileName: "fake-profile"})
|
|
if len(errs) == 0 {
|
|
t.Error("expected error for unknown harness profile")
|
|
}
|
|
}
|
|
|
|
// TestFilterMiddlewareByTypeName verifies middleware exclusion.
|
|
func TestFilterMiddlewareByTypeName(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
type mockMW struct {
|
|
core.BaseMiddleware[*schema.Message]
|
|
}
|
|
type excludedMW struct {
|
|
core.BaseMiddleware[*schema.Message]
|
|
}
|
|
|
|
mws := []core.ReActMiddleware{
|
|
&mockMW{},
|
|
&excludedMW{},
|
|
&mockMW{},
|
|
}
|
|
|
|
// Exclude excludedMW by type name.
|
|
typeName := fmt.Sprintf("%T", &excludedMW{})
|
|
filtered := filterMiddlewareByTypeName(mws, []string{typeName})
|
|
|
|
if len(filtered) != 2 {
|
|
t.Errorf("expected 2 middlewares after exclusion, got %d", len(filtered))
|
|
}
|
|
}
|
|
|
|
// TestDualTrack verifies the full ProviderProfile + HarnessProfile flow.
|
|
func TestDualTrack(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "dual",
|
|
DefaultModel: "test-model",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "dual"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "coding",
|
|
MaxIterations: 25,
|
|
ExcludedMiddlewareNames: []string{
|
|
"*some.ExcludedMiddleware",
|
|
},
|
|
})
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "dual:test-model",
|
|
HarnessProfileName: "coding",
|
|
Tools: []core.Tool{&simpleTool{name: "t1"}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
if agent == nil {
|
|
t.Fatal("expected non-nil agent")
|
|
}
|
|
t.Log("dual track: OK")
|
|
}
|
|
|
|
// TestDescriptionOverride verifies tool description overrides are applied.
|
|
func TestDescriptionOverride(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "x"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "overrider",
|
|
ToolDescriptionOverrides: map[string]string{
|
|
"search": "Search the web (custom description)",
|
|
},
|
|
})
|
|
|
|
searchTool := &simpleTool{name: "search", desc: "Original description"}
|
|
_, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
HarnessProfileName: "overrider",
|
|
Tools: []core.Tool{searchTool},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
// The tool should be wrapped with the new description.
|
|
// We verify indirectly by checking the harness was applied without error.
|
|
t.Log("description override: OK")
|
|
}
|
|
|
|
// TestProviderOpts verifies custom provider options are passed through.
|
|
func TestProviderOpts(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
var capturedOpts map[string]any
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "opts-test",
|
|
DefaultOpts: map[string]any{
|
|
"temperature": 0.7,
|
|
},
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
capturedOpts = opts
|
|
return &mockProfileModel{response: "ok"}, nil
|
|
},
|
|
})
|
|
|
|
_, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "opts-test:m1",
|
|
ProviderOpts: map[string]any{
|
|
"temperature": 0.2,
|
|
"max_tokens": 2048,
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
|
|
if capturedOpts == nil {
|
|
t.Fatal("expected opts to be captured")
|
|
}
|
|
// User opts should override defaults.
|
|
if temp, ok := toFloat64(capturedOpts["temperature"]); !ok || temp != 0.2 {
|
|
t.Errorf("expected temperature=0.2, got %v", capturedOpts["temperature"])
|
|
}
|
|
if tokens, ok := toFloat64(capturedOpts["max_tokens"]); !ok || tokens != 2048 {
|
|
t.Errorf("expected max_tokens=2048, got %v", capturedOpts["max_tokens"])
|
|
}
|
|
t.Logf("captured opts: %v", capturedOpts)
|
|
}
|
|
|
|
// TestInitModelError verifies InitModel errors are propagated.
|
|
func TestInitModelError(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
expectedErr := errors.New("api failure")
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "faulty",
|
|
DefaultModel: "x",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return nil, expectedErr
|
|
},
|
|
})
|
|
|
|
_, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "faulty:x",
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected error from InitModel")
|
|
}
|
|
if !errors.Is(err, expectedErr) && !strings.Contains(err.Error(), expectedErr.Error()) {
|
|
t.Errorf("expected error containing %q, got %q", expectedErr.Error(), err.Error())
|
|
}
|
|
t.Logf("InitModel error propagated: %v", err)
|
|
}
|
|
|
|
// TestExcludedToolNames verifies tools are removed from the agent's tool list.
|
|
func TestExcludedToolNames(t *testing.T) {
|
|
defer resetRegistries()
|
|
|
|
RegisterProvider(&ProviderProfile{
|
|
Name: "test",
|
|
DefaultModel: "m1",
|
|
InitModel: func(ctx context.Context, modelName string, opts map[string]any) (core.Model[*schema.Message], error) {
|
|
return &mockProfileModel{response: "x"}, nil
|
|
},
|
|
})
|
|
RegisterHarness(&HarnessProfile{
|
|
Name: "limited",
|
|
ExcludedToolNames: []string{"dangerous_tool"},
|
|
})
|
|
|
|
agent, err := NewAgent(context.Background(), &AgentConfig{
|
|
ModelSpec: "test:m1",
|
|
HarnessProfileName: "limited",
|
|
Tools: []core.Tool{
|
|
&simpleTool{name: "safe_tool"},
|
|
&simpleTool{name: "dangerous_tool"},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewAgent error: %v", err)
|
|
}
|
|
_ = agent
|
|
t.Log("excluded tool names: OK")
|
|
}
|
|
|
|
// ---- Helpers ----
|
|
|
|
type simpleTool struct {
|
|
name string
|
|
desc string
|
|
}
|
|
|
|
func (t *simpleTool) Name() string { return t.name }
|
|
func (t *simpleTool) Description() string { return t.desc }
|
|
func (t *simpleTool) Invoke(ctx context.Context, args string, opts ...core.ToolOption) (string, error) {
|
|
return "result", nil
|
|
}
|
|
func (t *simpleTool) Stream(ctx context.Context, args string, opts ...core.ToolOption) (*schema.StreamReader[string], error) {
|
|
return schema.StreamReaderFromArray([]string{"result"}), nil
|
|
}
|
|
|
|
func toFloat64(v any) (float64, bool) {
|
|
switch n := v.(type) {
|
|
case float64:
|
|
return n, true
|
|
case int:
|
|
return float64(n), true
|
|
case int64:
|
|
return float64(n), true
|
|
default:
|
|
return 0, false
|
|
}
|
|
}
|
|
|
|
// resetRegistries clears the global registries for test isolation.
|
|
func resetRegistries() {
|
|
ClearProviders()
|
|
ClearHarnesses()
|
|
}
|