ragflow/internal/agent/harness/route_test.go

package harness

import (
	"context"
	"testing"

	"gorm.io/gorm"

	"ragflow/internal/agent/component"
)

// fakeChatInvoker returns a fixed content for the chat call, so tests can drive
// the route/planner LLM output deterministically.
type fakeChatInvoker struct{ content string }

func (f *fakeChatInvoker) Invoke(_ context.Context, _ *gorm.DB, _ component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) {
	return &component.ChatInvokeResponse{Content: f.content}, nil
}

func installChat(t *testing.T, content string) {
	t.Helper()
	component.SetDefaultChatInvoker(&fakeChatInvoker{content: content})
	t.Cleanup(func() { component.SetDefaultChatInvoker(nil) })
}

// TestRouteNode_Classifies asserts route classification drives the execution
// strategy from the mode and the LLM's question_type.
func TestRouteNode_Classifies(t *testing.T) {
	installChat(t, `{"question_type":"comparative","requires_decomposition":true,"reasoning":"cmp"}`)
	r := RouteNode(context.Background(), nil, "Compare A and B", "medium")
	if r.QuestionType != "comparative" {
		t.Errorf("question_type = %q, want comparative", r.QuestionType)
	}
	// medium mode requires decomposition AND LLM says true -> true.
	if !r.RequiresDecomposition {
		t.Errorf("requires_decomposition = false, want true")
	}
	if r.ExecutionStrategy != "decompose_and_search" {
		t.Errorf("execution_strategy = %q, want decompose_and_search", r.ExecutionStrategy)
	}
}

// TestRouteNode_LowModeDisablesDecomposition asserts low mode never decomposes
// even if the LLM requests it.
func TestRouteNode_LowModeDisablesDecomposition(t *testing.T) {
	installChat(t, `{"question_type":"analytical","requires_decomposition":true}`)
	r := RouteNode(context.Background(), nil, "Analyze X", "low")
	if r.RequiresDecomposition {
		t.Errorf("low mode must disable decomposition, got true")
	}
	if r.ExecutionStrategy != "direct_search" {
		t.Errorf("low mode execution_strategy = %q, want direct_search", r.ExecutionStrategy)
	}
}

// TestRouteNode_FencedJSON asserts think-tag/fence stripping works.
func TestRouteNode_FencedJSON(t *testing.T) {
	installChat(t, "Sure!\n```json\n{\"question_type\":\"factual\",\"requires_decomposition\":false}\n```")
	r := RouteNode(context.Background(), nil, "What is X?", "medium")
	if r.QuestionType != "factual" {
		t.Errorf("question_type = %q, want factual", r.QuestionType)
	}
	if r.RequiresDecomposition {
		t.Errorf("requires_decomposition = true, want false")
	}
}

// TestRouteNode_EmptyQuestionFallsBack asserts an empty question yields a
// direct factual decision without calling the LLM.
func TestRouteNode_EmptyQuestionFallsBack(t *testing.T) {
	r := RouteNode(context.Background(), nil, "", "medium")
	if r.QuestionType != "factual" || r.RequiresDecomposition {
		t.Errorf("empty question must fall back to direct factual, got %+v", r)
	}
}

// TestPlannerNode_DirectMode asserts a non-decomposed route yields one coarse
// claim without calling the LLM.
func TestPlannerNode_DirectMode(t *testing.T) {
	plan := PlannerNode(context.Background(), nil, RouteDecision{
		Question: "What is X?", RequiresDecomposition: false,
	}, nil)
	if plan.PlanType != "direct" || len(plan.Claims) != 1 {
		t.Fatalf("direct plan = %+v, want 1 direct claim", plan)
	}
}

// TestPlannerNode_Decomposes asserts the planner builds claims from the LLM
// output and applies the mode's max iterations.
func TestPlannerNode_Decomposes(t *testing.T) {
	installChat(t, `{"claims":[
		{"claim_id":"c0","description":"fact one","priority":0},
		{"claim_id":"c1","description":"fact two","priority":1}
	]}`)
	plan := PlannerNode(context.Background(), nil, RouteDecision{
		Question: "Compare A and B", QuestionType: "comparative", RequiresDecomposition: true, ThinkingMode: "medium",
	}, nil)
	if plan.PlanType != "comparative_decomposition" {
		t.Errorf("plan_type = %q, want comparative_decomposition", plan.PlanType)
	}
	if len(plan.Claims) != 2 {
		t.Fatalf("claims = %d, want 2", len(plan.Claims))
	}
	if plan.Claims[1].Priority != 1 {
		t.Errorf("claim priority = %d, want 1", plan.Claims[1].Priority)
	}
	// medium maxOrchestratorCycles = 3
	if plan.MaxIterations != 3 {
		t.Errorf("max_iterations = %d, want 3", plan.MaxIterations)
	}
}

// TestPlannerNode_UnknownModeFallsBack asserts an unknown (non-empty) mode label
// falls back to medium so the planner is not driven by a zero-valued mode
// (which would produce a degenerate plan with max_claims=0).
func TestPlannerNode_UnknownModeFallsBack(t *testing.T) {
	installChat(t, `{"claims":[{"claim_id":"c0","description":"fact one","priority":0}]}`)
	plan := PlannerNode(context.Background(), nil, RouteDecision{
		Question: "Q", RequiresDecomposition: true, ThinkingMode: "turbo-unknown",
	}, nil)
	// medium maxOrchestratorCycles = 3, and claims must still be built.
	if plan.MaxIterations != 3 {
		t.Errorf("max_iterations = %d, want 3 (medium fallback)", plan.MaxIterations)
	}
	if len(plan.Claims) != 1 {
		t.Errorf("claims = %d, want 1", len(plan.Claims))
	}
}

// TestPlannerNode_BadJSONFallsBack asserts unparseable planner output falls back
// to the direct plan.
func TestPlannerNode_BadJSONFallsBack(t *testing.T) {
	installChat(t, "not json at all")
	plan := PlannerNode(context.Background(), nil, RouteDecision{
		Question: "Q", RequiresDecomposition: true, ThinkingMode: "medium",
	}, nil)
	if plan.PlanType != "direct" || len(plan.Claims) != 1 {
		t.Fatalf("fallback plan = %+v, want direct", plan)
	}
}