diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index 216c8a75ff..28f47c8445 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -443,40 +443,11 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map hasRuntimeUserPrompt = !shouldFallbackToSysQuery(v) } - // v3.6.1: derive the driver and bare model name from the - // composite llm_id when the Agent DSL didn't set `driver`. The - // Python side does the same in split_model_name at - // api/db/joint_services/tenant_model_service.py:163-178. We - // also use this opportunity to look up the tenant's LLM - // credentials from `tenant_llm` when the DSL omitted `api_key` - // — mirrors Python's get_model_config_from_provider_instance, - // which is how the Python canvas finds the tenant's - // provider-specific API key + base URL without storing them - // in the canvas DSL. - // Save the original composite llm_id before the split drops the - // instance-name segment. We need it for the tenant_model_instance - // fallback path below. - originalModelID := p.ModelID - - if p.Driver == "" && p.ModelID != "" { - if m, prov, ok := agentProviderLastSegmentSplit(p.ModelID); ok { - p.Driver = prov - p.ModelID = m - } + var err error + p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL) + if err != nil { + return nil, err } - if p.Driver == "" && p.ModelID != "" { - modelID, driver, apiKey, baseURL, ok, err := resolveTenantChatModelByID(ctx, p.ModelID, p.APIKey, p.BaseURL) - if err != nil { - return nil, err - } - if ok { - p.ModelID = modelID - p.Driver = driver - p.APIKey = apiKey - p.BaseURL = baseURL - } - } - p.APIKey, p.BaseURL = resolveTenantLLMConfig(ctx, p.Driver, p.ModelID, p.APIKey, p.BaseURL, originalModelID) var state *runtime.CanvasState if s, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && s != nil { diff --git a/internal/agent/component/categorize.go b/internal/agent/component/categorize.go index c81dad985e..9c8168c9e9 100644 --- a/internal/agent/component/categorize.go +++ b/internal/agent/component/categorize.go @@ -15,6 +15,8 @@ import ( "strings" "github.com/cloudwego/eino/schema" + + "ragflow/internal/agent/runtime" ) // CategorizeComponent is an LLM classifier. @@ -24,15 +26,19 @@ type CategorizeComponent struct { // CategorizeParam captures the (resolved) DSL parameters for a Categorize node. type CategorizeParam struct { - ModelID string - Items []string - Categories []string - CategoryRoutes map[string]string - SysPrompt string - DefaultCategory string - Driver string - APIKey string - BaseURL string + ModelID string + Query string + Items []string + Categories []string + CategoryRoutes map[string]string + CategoryDescriptions map[string]string + CategoryExamples map[string][]string + SysPrompt string + DefaultCategory string + MessageHistoryWindowSize int + Driver string + APIKey string + BaseURL string } // CategorizeOutput mirrors the outputs map (per plan §2.11.3 row 6): @@ -60,20 +66,20 @@ func (c *CategorizeComponent) Name() string { return "Categorize" } // something outside the configured set). func (c *CategorizeComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { p := mergeCategorizeParam(c.param, inputs) - originalModelID := p.ModelID - if p.Driver == "" && p.ModelID != "" { - if modelID, driver, ok := splitCompositeLLMID(p.ModelID); ok { - p.Driver = driver - p.ModelID = modelID - } + var err error + p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL) + if err != nil { + return nil, err } - p.APIKey, p.BaseURL = resolveTenantLLMConfig(ctx, p.Driver, p.ModelID, p.APIKey, p.BaseURL, originalModelID) if p.ModelID == "" { return nil, &ParamError{Field: "model_id", Reason: "required"} } if len(p.Categories) == 0 { return nil, &ParamError{Field: "categories", Reason: "at least one category is required"} } + if p.MessageHistoryWindowSize < 0 { + return nil, &ParamError{Field: "message_history_window_size", Reason: "cannot be negative"} + } if p.DefaultCategory == "" { // Fall back to the first category so the run never fails purely // because the user omitted the default. @@ -83,9 +89,10 @@ func (c *CategorizeComponent) Invoke(ctx context.Context, inputs map[string]any) inv := getDefaultChatInvoker() sysPrompt := p.SysPrompt if sysPrompt == "" { - sysPrompt = "You are a strict classifier." + sysPrompt = buildCategorizeSystemPrompt(p) } - userPrompt := buildCategorizePrompt(p) + query := resolveCategorizeQuery(ctx, p, inputs) + userPrompt := buildCategorizePrompt(categorizeHistory(ctx, p.MessageHistoryWindowSize, query)) msgs := []schema.Message{ {Role: schema.System, Content: sysPrompt}, {Role: schema.User, Content: userPrompt}, @@ -132,13 +139,24 @@ func (c *CategorizeComponent) Stream(ctx context.Context, inputs map[string]any) // Inputs returns parameter metadata for tooling. func (c *CategorizeComponent) Inputs() map[string]string { return map[string]string{ - "model_id": "Provider-side model identifier", - "items": "Optional list of items to classify (added to the prompt as context)", - "categories": "List of allowed category names (response must match one)", - "sys_prompt": "Optional system prompt; defaults to a strict classifier instruction", - "default_category": "Category returned if the model's answer is not in `categories` (defaults to categories[0])", - "driver": "Provider driver name", - "api_key": "Override API key", + "model_id": "Provider-side model identifier", + "query": "Variable reference or literal text to classify. Defaults to sys.query.", + "items": "Optional list of items to classify (added to the prompt as context)", + "categories": "List of allowed category names (response must match one)", + "sys_prompt": "Optional system prompt; defaults to a strict classifier instruction", + "default_category": "Category returned if the model's answer is not in `categories` (defaults to categories[0])", + "message_history_window_size": "How many prior conversation turns to include in the Real Data prompt.", + "driver": "Provider driver name", + "api_key": "Override API key", + } +} + +func (c *CategorizeComponent) GetInputForm() map[string]any { + return map[string]any{ + "query": map[string]any{ + "type": "line", + "name": "Query", + }, } } @@ -152,31 +170,158 @@ func (c *CategorizeComponent) Outputs() map[string]string { } } -// buildCategorizePrompt assembles a prompt that asks the model to pick a -// category. The categories are listed deterministically (sorted) so the -// prompt is stable across runs. -func buildCategorizePrompt(p CategorizeParam) string { +func buildCategorizeSystemPrompt(p CategorizeParam) string { cats := append([]string(nil), p.Categories...) sort.Strings(cats) var b strings.Builder - b.WriteString("Classify the following item into exactly one of these categories:\n") + b.WriteString("You are an advanced classification system that categorizes user questions into specific types. Analyze the input question and classify it into ONE of the following categories:\n") for _, c := range cats { b.WriteString("- ") b.WriteString(c) b.WriteString("\n") } - if len(p.Items) > 0 { - b.WriteString("\nItems:\n") - for _, it := range p.Items { - b.WriteString("- ") - b.WriteString(it) + + if len(p.CategoryDescriptions) > 0 { + b.WriteString("\nHere's description of each category:\n") + for _, c := range cats { + desc := strings.TrimSpace(p.CategoryDescriptions[c]) + if desc == "" { + continue + } + b.WriteString("\n------\nCategory: ") + b.WriteString(c) + b.WriteString("\nDescription: ") + b.WriteString(desc) + b.WriteString("\n") + } + } + + b.WriteString("\n---- Instructions ----\n") + b.WriteString("- Consider both explicit mentions and implied context\n") + b.WriteString("- Prioritize the most specific applicable category\n") + b.WriteString("- Return only the category name without explanations\n") + if containsCategory(cats, "Other") { + b.WriteString("- Use \"Other\" only when no other category fits\n") + } + + examples := categorizeExamples(p, cats) + if len(examples) > 0 { + b.WriteString("\n---- Examples ----\n") + for _, line := range examples { + b.WriteString(line) b.WriteString("\n") } } - b.WriteString("\nRespond with ONLY the category name, no other text.") return b.String() } +func containsCategory(categories []string, target string) bool { + for _, c := range categories { + if c == target { + return true + } + } + return false +} + +func categorizeExamples(p CategorizeParam, cats []string) []string { + var lines []string + for _, c := range cats { + for _, example := range p.CategoryExamples[c] { + example = strings.TrimSpace(strings.ReplaceAll(example, "\n", " ")) + if example == "" { + continue + } + lines = append(lines, fmt.Sprintf("USER: %q -> %s", example, c)) + } + } + for _, it := range p.Items { + it = strings.TrimSpace(strings.ReplaceAll(it, "\n", " ")) + if it == "" { + continue + } + lines = append(lines, fmt.Sprintf("USER: %q", it)) + } + return lines +} + +type categorizeHistoryMessage struct { + Role string + Content string +} + +func categorizeHistory(ctx context.Context, window int, query string) []categorizeHistoryMessage { + msgs := []categorizeHistoryMessage{} + if window > 0 { + if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { + start := 0 + if len(state.History) > window { + start = len(state.History) - window + } + for i := start; i < len(state.History); i++ { + entry := state.History[i] + role, _ := entry["role"].(string) + content, _ := entry["content"].(string) + if role == "" || content == "" { + continue + } + msgs = append(msgs, categorizeHistoryMessage{Role: role, Content: content}) + } + } + } + if len(msgs) == 0 { + msgs = append(msgs, categorizeHistoryMessage{Role: "user"}) + } + msgs[len(msgs)-1].Content = query + return msgs +} + +func buildCategorizePrompt(history []categorizeHistoryMessage) string { + parts := make([]string, 0, len(history)) + for _, msg := range history { + role := strings.ToUpper(strings.TrimSpace(msg.Role)) + if role == "" { + role = "USER" + } + content := strings.NewReplacer("\r\n", "", "\n", "", "\r", "").Replace(msg.Content) + parts = append(parts, fmt.Sprintf("%s: %q", role, content)) + } + return fmt.Sprintf("\n---- Real Data ----\n%s ->\n", strings.Join(parts, " | ")) +} + +func resolveCategorizeQuery(ctx context.Context, p CategorizeParam, inputs map[string]any) string { + if v, ok := stringValueFromAny(inputs["query"]); ok { + return v + } + queryRef := strings.TrimSpace(p.Query) + if queryRef == "" { + queryRef = "sys.query" + } + if v, ok := stringValueFromAny(inputs[queryRef]); ok && strings.TrimSpace(v) != "" { + return v + } + if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { + if v, err := state.GetVar(queryRef); err == nil { + if s, ok := stringValueFromAny(v); ok { + return s + } + } + } + return "" +} + +func stringValueFromAny(v any) (string, bool) { + if v == nil { + return "", false + } + switch t := v.(type) { + case string: + return t, true + default: + return fmt.Sprint(t), true + } +} + // pickCategory extracts a category from the model's response. Strategy: // 1. exact match (case-sensitive) // 2. case-insensitive match @@ -229,23 +374,24 @@ func mergeCategorizeParam(base CategorizeParam, inputs map[string]any) Categoriz } else if v, ok := stringFrom(inputs, "llm_id"); ok { p.ModelID = v } + if v, ok := stringFrom(inputs, "query"); ok { + p.Query = v + } if v, ok := sliceFrom(inputs, "items"); ok { p.Items = v } if v, ok := sliceFrom(inputs, "categories"); ok { p.Categories = v - } else if m, ok := stringMapFrom(inputs, "category_description"); ok && len(m) > 0 { - // v1 stores the categories as a map of {name: description}. - // We only need the keys to drive the picker. - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) + } + if meta, ok := categoryMetadataFrom(inputs, "category_description"); ok && len(meta.Names) > 0 { + if len(p.Categories) == 0 { + p.Categories = meta.Names } - sort.Strings(keys) - p.Categories = keys + p.CategoryDescriptions = mergeStringMap(p.CategoryDescriptions, meta.Descriptions) + p.CategoryExamples = mergeStringSliceMap(p.CategoryExamples, meta.Examples) } if routes, ok := categoryRoutesFrom(inputs, "category_description"); ok { - p.CategoryRoutes = routes + p.CategoryRoutes = mergeStringMap(p.CategoryRoutes, routes) } if v, ok := stringFrom(inputs, "sys_prompt"); ok { p.SysPrompt = v @@ -255,6 +401,9 @@ func mergeCategorizeParam(base CategorizeParam, inputs map[string]any) Categoriz if v, ok := stringFrom(inputs, "default_category"); ok { p.DefaultCategory = v } + if v, ok := intFrom(inputs, "message_history_window_size"); ok { + p.MessageHistoryWindowSize = v + } if v, ok := stringFrom(inputs, "driver"); ok { p.Driver = v } @@ -267,35 +416,91 @@ func mergeCategorizeParam(base CategorizeParam, inputs map[string]any) Categoriz return p } -// stringMapFrom extracts map[string]string from inputs[name]. The v1 -// "category_description" field is shaped this way (name → human -// description); we only consume the keys. -func stringMapFrom(inputs map[string]any, name string) (map[string]string, bool) { +type categorizeMetadata struct { + Names []string + Descriptions map[string]string + Examples map[string][]string +} + +func categoryMetadataFrom(inputs map[string]any, name string) (categorizeMetadata, bool) { v, ok := inputs[name] if !ok { - return nil, false + return categorizeMetadata{}, false } raw, ok := v.(map[string]any) if !ok { - return nil, false + return categorizeMetadata{}, false + } + out := categorizeMetadata{ + Names: make([]string, 0, len(raw)), + Descriptions: make(map[string]string, len(raw)), + Examples: make(map[string][]string, len(raw)), } - out := make(map[string]string, len(raw)) for k, child := range raw { + out.Names = append(out.Names, k) if s, ok := child.(string); ok { - out[k] = s + out.Descriptions[k] = s continue } - // Some encoders nest the description under a "description" - // key; handle that fallback defensively. if nested, ok := child.(map[string]any); ok { if s, ok := nested["description"].(string); ok { - out[k] = s - continue + out.Descriptions[k] = s + } + out.Examples[k] = examplesFromAny(nested["examples"]) + continue + } + } + sort.Strings(out.Names) + return out, true +} + +func examplesFromAny(v any) []string { + raw, ok := v.([]any) + if !ok { + return nil + } + out := make([]string, 0, len(raw)) + for _, item := range raw { + switch t := item.(type) { + case string: + if t != "" { + out = append(out, t) + } + case map[string]any: + if s, _ := t["value"].(string); s != "" { + out = append(out, s) } } - out[k] = "" } - return out, true + return out +} + +func mergeStringMap(base, override map[string]string) map[string]string { + if len(override) == 0 { + return base + } + out := make(map[string]string, len(base)+len(override)) + for k, v := range base { + out[k] = v + } + for k, v := range override { + out[k] = v + } + return out +} + +func mergeStringSliceMap(base, override map[string][]string) map[string][]string { + if len(override) == 0 { + return base + } + out := make(map[string][]string, len(base)+len(override)) + for k, v := range base { + out[k] = append([]string(nil), v...) + } + for k, v := range override { + out[k] = append([]string(nil), v...) + } + return out } func categoryRoutesFrom(inputs map[string]any, name string) (map[string]string, bool) { @@ -327,12 +532,15 @@ func categoryRoutesFrom(inputs map[string]any, name string) (map[string]string, // init registers CategorizeComponent with the orchestrator-owned registry. func init() { Register("Categorize", func(params map[string]any) (Component, error) { - var p CategorizeParam + p := CategorizeParam{MessageHistoryWindowSize: 1} if v, ok := stringFrom(params, "model_id"); ok { p.ModelID = v } else if v, ok := stringFrom(params, "llm_id"); ok { p.ModelID = v } + if v, ok := stringFrom(params, "query"); ok { + p.Query = v + } // Check the object-style []any of maps first. sliceFrom would // otherwise match the same []any input and return (empty, true) // for non-string elements, making the object branch unreachable. @@ -349,6 +557,12 @@ func init() { continue } names = append(names, name) + if desc, _ := m["description"].(string); desc != "" { + if p.CategoryDescriptions == nil { + p.CategoryDescriptions = map[string]string{} + } + p.CategoryDescriptions[name] = desc + } if route, ok := firstRouteTarget(m["to"]); ok { routes[name] = route } else if uuid, _ := m["uuid"].(string); uuid != "" { @@ -356,12 +570,18 @@ func init() { } if examples, ok := m["examples"].([]any); ok { for _, example := range examples { - em, ok := example.(map[string]any) - if !ok { - continue + if p.CategoryExamples == nil { + p.CategoryExamples = map[string][]string{} } - if v, _ := em["value"].(string); v != "" { - p.Items = append(p.Items, v) + switch em := example.(type) { + case map[string]any: + if v, _ := em["value"].(string); v != "" { + p.CategoryExamples[name] = append(p.CategoryExamples[name], v) + } + case string: + if em != "" { + p.CategoryExamples[name] = append(p.CategoryExamples[name], em) + } } } } @@ -377,27 +597,15 @@ func init() { } if v, ok := sliceFrom(params, "categories"); ok { p.Categories = v - } else if m, ok := params["category_description"].(map[string]any); ok && len(m) > 0 { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) + } + if meta, ok := categoryMetadataFrom(params, "category_description"); ok && len(meta.Names) > 0 { + if len(p.Categories) == 0 { + p.Categories = meta.Names } - sort.Strings(keys) - p.Categories = keys - routes := make(map[string]string, len(m)) - for k, child := range m { - nested, ok := child.(map[string]any) - if !ok { - continue - } - if route, ok := firstRouteTarget(nested["to"]); ok { - routes[k] = route - } else if uuid, _ := nested["uuid"].(string); uuid != "" { - routes[k] = uuid - } - } - if len(routes) > 0 { - p.CategoryRoutes = routes + p.CategoryDescriptions = mergeStringMap(p.CategoryDescriptions, meta.Descriptions) + p.CategoryExamples = mergeStringSliceMap(p.CategoryExamples, meta.Examples) + if routes, ok := categoryRoutesFrom(params, "category_description"); ok { + p.CategoryRoutes = mergeStringMap(p.CategoryRoutes, routes) } } if v, ok := stringFrom(params, "sys_prompt"); ok { @@ -408,6 +616,9 @@ func init() { if v, ok := stringFrom(params, "default_category"); ok { p.DefaultCategory = v } + if v, ok := intFrom(params, "message_history_window_size"); ok { + p.MessageHistoryWindowSize = v + } if v, ok := stringFrom(params, "driver"); ok { p.Driver = v } diff --git a/internal/agent/component/categorize_test.go b/internal/agent/component/categorize_test.go index fcb75bc1eb..943e76d73b 100644 --- a/internal/agent/component/categorize_test.go +++ b/internal/agent/component/categorize_test.go @@ -125,20 +125,190 @@ func TestCategorize_PromptListsCategories(t *testing.T) { if stub.captured == nil { t.Fatal("invoker not called") } + var systemContent string + for _, m := range stub.captured.Messages { + if m.Role == "system" { + systemContent = m.Content + } + } + if systemContent == "" { + t.Fatal("no system message in captured invoker request") + } + for _, want := range []string{"x", "y", "z", "foo", "bar"} { + if !strings.Contains(systemContent, want) { + t.Errorf("prompt missing %q; got: %s", want, systemContent) + } + } +} + +func TestCategorize_PromptIncludesRuntimeQuery(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "English", Model: "stub"}} + withStubInvoker(t, stub) + + state := canvas.NewCanvasState("run-1", "task-1") + state.Sys["query"] = "he who desires but acts not" + ctx := canvas.WithState(context.Background(), state) + + c := NewCategorizeComponent(CategorizeParam{ + ModelID: "stub", + Query: "sys.query", + Categories: []string{"Number", "chinese", "English"}, + CategoryDescriptions: map[string]string{ + "Number": "This query has only a number", + "chinese": "this query only has chinese", + "English": "this query has english letter", + }, + }) + _, err := c.Invoke(ctx, map[string]any{}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if stub.captured == nil { + t.Fatal("invoker not called") + } var userContent string for _, m := range stub.captured.Messages { if m.Role == "user" { userContent = m.Content } } - if userContent == "" { - t.Fatal("no user message in captured invoker request") + if !strings.Contains(userContent, "he who desires but acts not") { + t.Fatalf("user prompt = %q, want runtime query", userContent) } - for _, want := range []string{"x", "y", "z", "foo", "bar"} { - if !strings.Contains(userContent, want) { - t.Errorf("prompt missing %q; got: %s", want, userContent) + if strings.Contains(userContent, "Number") || strings.Contains(userContent, "chinese") { + t.Fatalf("user prompt should carry real data only, got %q", userContent) + } +} + +func TestCategorize_PromptUsesInputQueryValue(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "chinese", Model: "stub"}} + withStubInvoker(t, stub) + + c := NewCategorizeComponent(CategorizeParam{ + ModelID: "stub", + Query: "sys.query", + Categories: []string{"Number", "chinese", "English"}, + }) + _, err := c.Invoke(context.Background(), map[string]any{"query": "测试"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if stub.captured == nil { + t.Fatal("invoker not called") + } + var userContent string + for _, m := range stub.captured.Messages { + if m.Role == "user" { + userContent = m.Content } } + if !strings.Contains(userContent, "测试") { + t.Fatalf("user prompt = %q, want input query value", userContent) + } +} + +func TestCategorize_HistoryWindowRealData(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "English", Model: "stub"}} + withStubInvoker(t, stub) + + state := canvas.NewCanvasState("run-1", "task-1") + state.History = []map[string]any{ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "prior answer"}, + {"role": "user", "content": "stale latest"}, + } + ctx := canvas.WithState(context.Background(), state) + + c := NewCategorizeComponent(CategorizeParam{ + ModelID: "stub", + Query: "sys.query", + Categories: []string{"Number", "chinese", "English"}, + MessageHistoryWindowSize: 2, + }) + _, err := c.Invoke(ctx, map[string]any{"query": "current question"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + var userContent string + for _, m := range stub.captured.Messages { + if m.Role == "user" { + userContent = m.Content + } + } + if !strings.Contains(userContent, `ASSISTANT: "prior answer" | USER: "current question"`) { + t.Fatalf("user prompt = %q, want bounded history plus current query", userContent) + } + if strings.Contains(userContent, "old user") || strings.Contains(userContent, "stale latest") { + t.Fatalf("user prompt should truncate old history and overwrite latest content, got %q", userContent) + } +} + +func TestCategorizeRegistered_DefaultHistoryWindow(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "English", Model: "stub"}} + withStubInvoker(t, stub) + + state := canvas.NewCanvasState("run-1", "task-1") + state.History = []map[string]any{ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "stale latest"}, + } + ctx := canvas.WithState(context.Background(), state) + + c, err := New("Categorize", map[string]any{ + "model_id": "stub", + "query": "sys.query", + "categories": []any{"Number", "chinese", "English"}, + }) + if err != nil { + t.Fatalf("New(Categorize): %v", err) + } + _, err = c.Invoke(ctx, map[string]any{"query": "current question"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + var userContent string + for _, m := range stub.captured.Messages { + if m.Role == "user" { + userContent = m.Content + } + } + if !strings.Contains(userContent, `ASSISTANT: "current question"`) { + t.Fatalf("user prompt = %q, want default one-message history with current query", userContent) + } + if strings.Contains(userContent, "old user") || strings.Contains(userContent, "stale latest") { + t.Fatalf("user prompt should only include default one-message window with overwritten content, got %q", userContent) + } +} + +func TestCategorize_HistoryWindowRejectsNegative(t *testing.T) { + c := NewCategorizeComponent(CategorizeParam{ + ModelID: "stub", + Categories: []string{"Number", "chinese", "English"}, + MessageHistoryWindowSize: -1, + }) + _, err := c.Invoke(context.Background(), map[string]any{"query": "current question"}) + if err == nil { + t.Fatal("expected negative message_history_window_size error") + } + if !strings.Contains(err.Error(), "message_history_window_size") { + t.Fatalf("error = %v, want message_history_window_size", err) + } +} + +func TestCategorize_OtherInstructionOnlyWhenAllowed(t *testing.T) { + withoutOther := buildCategorizeSystemPrompt(CategorizeParam{ + Categories: []string{"Number", "Chinese", "English"}, + }) + if strings.Contains(withoutOther, `Use "Other"`) { + t.Fatalf("prompt should not mention Other when category is absent: %s", withoutOther) + } + + withOther := buildCategorizeSystemPrompt(CategorizeParam{ + Categories: []string{"Number", "Other"}, + }) + if !strings.Contains(withOther, `Use "Other" only when no other category fits`) { + t.Fatalf("prompt should mention Other when category exists: %s", withOther) + } } func TestCategorize_Registered(t *testing.T) { @@ -243,6 +413,66 @@ func TestCategorize_ResolvesTenantModelInstanceCredentials(t *testing.T) { } } +func TestCategorize_ResolvesTenantModelID(t *testing.T) { + db := setupComponentTestDB(t) + pushComponentDB(t, db) + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-1", + TenantID: "tenant-1", + ProviderName: "SILICONFLOW", + }).Error; err != nil { + t.Fatalf("create provider: %v", err) + } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-1", + ProviderID: "provider-1", + InstanceName: "prod-east", + APIKey: "instance-key", + Status: "active", + Extra: `{"base_url":"https://instance.example"}`, + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "3d2d824e7e5d11f1a845455b140cef90", + ProviderID: "provider-1", + InstanceID: "instance-1", + ModelName: "Qwen/Qwen3-8B", + ModelType: int(entity.ModelTypeChat), + Status: "active", + }).Error; err != nil { + t.Fatalf("create model: %v", err) + } + + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "support", Model: "stub"}} + withStubInvoker(t, stub) + + c := NewCategorizeComponent(CategorizeParam{ + ModelID: "3d2d824e7e5d11f1a845455b140cef90", + Categories: []string{"sales", "support"}, + DefaultCategory: "support", + }) + _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if stub.captured == nil { + t.Fatal("invoker not called") + } + if stub.captured.Driver != "SILICONFLOW" { + t.Fatalf("Driver=%q, want SILICONFLOW", stub.captured.Driver) + } + if stub.captured.ModelName != "Qwen/Qwen3-8B" { + t.Fatalf("ModelName=%q, want Qwen/Qwen3-8B", stub.captured.ModelName) + } + if stub.captured.APIKey != "instance-key" { + t.Fatalf("APIKey=%q, want instance-key", stub.captured.APIKey) + } + if stub.captured.BaseURL != "https://instance.example" { + t.Fatalf("BaseURL=%q, want https://instance.example", stub.captured.BaseURL) + } +} + func TestCategorize_ResolvesSoleActiveInstanceWhenDefaultMissing(t *testing.T) { db := setupComponentTestDB(t) pushComponentDB(t, db) @@ -334,6 +564,83 @@ func TestCategorize_RoutesFromCategoryDescriptionToList(t *testing.T) { } } +func TestCategorize_ExplicitCategoriesKeepCategoryDescriptionMetadata(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "English", Model: "stub"}} + withStubInvoker(t, stub) + + c := NewCategorizeComponent(CategorizeParam{ModelID: "stub"}) + out, err := c.Invoke(context.Background(), map[string]any{ + "categories": []any{"Number", "chinese", "English"}, + "category_description": map[string]any{ + "Number": map[string]any{"description": "This query has only a number", "examples": []any{"4321"}, "to": []any{"Message:Number"}}, + "chinese": map[string]any{"description": "this query only has chinese", "examples": []any{"测试"}, "to": []any{"Message:Chinese"}}, + "English": map[string]any{"description": "this query has english letter", "examples": []any{"hello"}, "to": []any{"Message:English"}}, + }, + "query": "hello", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + next, ok := out["_next"].([]string) + if !ok { + t.Fatalf("_next missing or wrong type: %T", out["_next"]) + } + if len(next) != 1 || next[0] != "Message:English" { + t.Fatalf("_next=%v, want [\"Message:English\"]", next) + } + var systemContent string + for _, m := range stub.captured.Messages { + if m.Role == "system" { + systemContent = m.Content + } + } + for _, want := range []string{"this query has english letter", `USER: "hello" -> English`} { + if !strings.Contains(systemContent, want) { + t.Fatalf("system prompt missing %q; got %s", want, systemContent) + } + } +} + +func TestCategorizeRegistered_ExplicitCategoriesKeepCategoryDescriptionMetadata(t *testing.T) { + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "English", Model: "stub"}} + withStubInvoker(t, stub) + + c, err := New("Categorize", map[string]any{ + "model_id": "stub", + "categories": []any{"Number", "chinese", "English"}, + "category_description": map[string]any{ + "Number": map[string]any{"description": "This query has only a number", "examples": []any{"4321"}, "to": []any{"Message:Number"}}, + "chinese": map[string]any{"description": "this query only has chinese", "examples": []any{"测试"}, "to": []any{"Message:Chinese"}}, + "English": map[string]any{"description": "this query has english letter", "examples": []any{"hello"}, "to": []any{"Message:English"}}, + }, + }) + if err != nil { + t.Fatalf("New(Categorize): %v", err) + } + out, err := c.Invoke(context.Background(), map[string]any{"query": "hello"}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + next, ok := out["_next"].([]string) + if !ok { + t.Fatalf("_next missing or wrong type: %T", out["_next"]) + } + if len(next) != 1 || next[0] != "Message:English" { + t.Fatalf("_next=%v, want [\"Message:English\"]", next) + } + var systemContent string + for _, m := range stub.captured.Messages { + if m.Role == "system" { + systemContent = m.Content + } + } + for _, want := range []string{"this query has english letter", `USER: "hello" -> English`} { + if !strings.Contains(systemContent, want) { + t.Fatalf("system prompt missing %q; got %s", want, systemContent) + } + } +} + func stateWithTenant(tenantID string) context.Context { state := canvas.NewCanvasState("run-1", "task-1") state.Sys["tenant_id"] = tenantID diff --git a/internal/agent/component/llm_credentials.go b/internal/agent/component/llm_credentials.go index f1391af3bd..71f26b7a46 100644 --- a/internal/agent/component/llm_credentials.go +++ b/internal/agent/component/llm_credentials.go @@ -47,6 +47,30 @@ func resolveTenantLLMConfig(ctx context.Context, driver, modelID, apiKey, baseUR return apiKey, baseURL } +func resolveChatModelRef(ctx context.Context, modelID, driver, apiKey, baseURL string) (string, string, string, string, error) { + originalModelID := modelID + if driver == "" && modelID != "" { + if m, prov, ok := agentProviderLastSegmentSplit(modelID); ok { + driver = prov + modelID = m + } + } + if driver == "" && modelID != "" { + resolvedModelID, resolvedDriver, resolvedAPIKey, resolvedBaseURL, ok, err := resolveTenantChatModelByID(ctx, modelID, apiKey, baseURL) + if err != nil { + return "", "", "", "", err + } + if ok { + modelID = resolvedModelID + driver = resolvedDriver + apiKey = resolvedAPIKey + baseURL = resolvedBaseURL + } + } + apiKey, baseURL = resolveTenantLLMConfig(ctx, driver, modelID, apiKey, baseURL, originalModelID) + return modelID, driver, apiKey, baseURL, nil +} + // resolveTenantLLMCredentials looks up the old tenant_llm table for the given // tenant / factory / model. Returns true when credentials were found. func resolveTenantLLMCredentials(tid, driver, modelID, baseURL string) (string, string, bool) {