192 lines
8.6 KiB
Go
192 lines
8.6 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"reflect"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
)
|
||
|
|
|
||
|
|
type flowKnowledgeBases struct {
|
||
|
|
interfaces.KnowledgeBaseService
|
||
|
|
ids []string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s flowKnowledgeBases) GetKnowledgeBaseByID(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||
|
|
for _, allowed := range s.ids {
|
||
|
|
if id == allowed {
|
||
|
|
return &types.KnowledgeBase{ID: id, TenantID: 1}, nil
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil, ErrForbidden
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s flowKnowledgeBases) ListKnowledgeBases(context.Context) ([]*types.KnowledgeBase, error) {
|
||
|
|
var bases []*types.KnowledgeBase
|
||
|
|
for _, id := range s.ids {
|
||
|
|
bases = append(bases, &types.KnowledgeBase{ID: id, TenantID: 1})
|
||
|
|
}
|
||
|
|
return bases, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type flowKnowledgeFiles struct{ interfaces.KnowledgeService }
|
||
|
|
|
||
|
|
func (flowKnowledgeFiles) GetKnowledgeBatchWithSharedAccess(_ context.Context, _ uint64, ids []string) ([]*types.Knowledge, error) {
|
||
|
|
if len(ids) > 0 {
|
||
|
|
return nil, ErrForbidden
|
||
|
|
}
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestFlowPersistenceDropsLegacyKnowledgeAndPreservesMetadata(t *testing.T) {
|
||
|
|
for _, operation := range []string{"create", "save_supplied_scope", "save_omitted_scope", "rename"} {
|
||
|
|
t.Run(operation, func(t *testing.T) {
|
||
|
|
s := sqliteWorkbench(t)
|
||
|
|
s.agents = generationAgentStub{}
|
||
|
|
ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
|
||
|
|
data := asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "writer"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"folder"}}, KnowledgeTags: map[string][]string{"deleted": {"tag"}}})
|
||
|
|
data["opaque_metadata"] = JSON{"nested": []any{"preserved", 42}, "enabled": true}
|
||
|
|
data["layout_metadata"] = JSON{"zoom": 0.75}
|
||
|
|
r := &Resource{ID: "flow", TenantID: 1, OwnerID: "alice", Kind: "flows", Name: "Original", Status: "draft", Version: 1, Data: data}
|
||
|
|
original := compact(data)
|
||
|
|
want := asJSON(data)
|
||
|
|
if operation == "create" {
|
||
|
|
if err := s.validate(ctx, r); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := s.create(ctx, r); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
if err := s.db.Create(r).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
proposed := asJSON(data)
|
||
|
|
if operation != "rename" {
|
||
|
|
proposed["nodes"] = []FlowNode{{ID: "n", Label: "Updated", AgentID: "writer"}}
|
||
|
|
}
|
||
|
|
if operation == "save_omitted_scope" {
|
||
|
|
for _, key := range []string{"knowledge_base_ids", "knowledge_ids", "knowledge_folders", "knowledge_tags"} {
|
||
|
|
delete(proposed, key)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
want = asJSON(proposed)
|
||
|
|
if response := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: "Updated", Version: 1, Data: proposed}); response.Code != 200 {
|
||
|
|
t.Fatal(response.Code, response.Body.String())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, key := range []string{"knowledge_base_ids", "knowledge_ids", "knowledge_folders", "knowledge_tags"} {
|
||
|
|
delete(want, key)
|
||
|
|
}
|
||
|
|
stored, err := s.get(ctx, "flows", r.ID)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if compact(stored.Data) != compact(want) {
|
||
|
|
t.Fatalf("persisted canvas retained legacy scope or changed other metadata: got %s, want %s", compact(stored.Data), compact(want))
|
||
|
|
}
|
||
|
|
if operation != "create" {
|
||
|
|
var before, after Revision
|
||
|
|
if err := s.db.First(&before, "resource_id = ? AND version = 1", r.ID).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := s.db.First(&after, "resource_id = ? AND version = 2", r.ID).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if stored.Version != 2 || compact(before.Data) != original || compact(after.Data) != compact(want) {
|
||
|
|
t.Fatal("canvas save lost its original revision or recorded unsanitized data")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestFlowKnowledgeModesIgnoreStaleCanvasAndTaskScopes(t *testing.T) {
|
||
|
|
for _, mode := range []string{"selected", "all", "none"} {
|
||
|
|
t.Run(mode, func(t *testing.T) {
|
||
|
|
s := NewService(nil)
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{KBSelectionMode: mode, KnowledgeBases: []string{"selected"}}}
|
||
|
|
s.knowledgeBases = flowKnowledgeBases{ids: []string{"selected", "other"}}
|
||
|
|
s.knowledge = flowKnowledgeFiles{}
|
||
|
|
flow := FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"empty-folder"}}, KnowledgeTags: map[string][]string{"deleted": {"deleted-tag"}}}
|
||
|
|
d := TaskData{KBIDs: flow.KBIDs, KnowledgeIDs: flow.KnowledgeIDs, KnowledgeFolders: flow.KnowledgeFolders, KnowledgeTags: flow.KnowledgeTags}
|
||
|
|
if err := s.inheritFlow(workbenchContext(1, "writer", types.TenantRoleContributor), &d, flow); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
var want []string
|
||
|
|
if mode == "selected" {
|
||
|
|
want = []string{"selected"}
|
||
|
|
} else if mode == "all" {
|
||
|
|
want = []string{"selected", "other"}
|
||
|
|
}
|
||
|
|
if !reflect.DeepEqual(d.KBIDs, want) || len(d.KnowledgeIDs) != 0 || len(d.KnowledgeFolders) != 0 || len(d.KnowledgeTags) != 0 {
|
||
|
|
t.Fatalf("unexpected agent scope: %+v", d)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTaskAgentCategoriesAndGenerationIgnoreLegacyFlowKnowledge(t *testing.T) {
|
||
|
|
s, ctx, job, _ := generationHarness(t)
|
||
|
|
s.agents = &fusionAgents{rows: map[string]*types.CustomAgent{
|
||
|
|
"writer": {ID: "writer", TenantID: 1, Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"writer-live"}}},
|
||
|
|
"editor": {ID: "editor", TenantID: 1, Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"editor-kb"}}},
|
||
|
|
}}
|
||
|
|
s.knowledgeBases = flowKnowledgeBases{ids: []string{"writer-live", "writer-frozen", "editor-kb"}}
|
||
|
|
s.knowledge = flowKnowledgeFiles{}
|
||
|
|
if err := s.db.AutoMigrate(&types.KnowledgeTag{}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
for _, baseID := range []string{"writer-live", "writer-frozen", "editor-kb"} {
|
||
|
|
if err := s.db.Create(&types.KnowledgeTag{ID: baseID + "-category", TenantID: 1, KnowledgeBaseID: baseID, Name: baseID}).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
flow := FlowData{Nodes: []FlowNode{{ID: "n1", AgentID: "writer"}, {ID: "n2", AgentID: "editor"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"empty-folder"}}, KnowledgeTags: map[string][]string{"deleted": {"deleted-category"}}}
|
||
|
|
r := &Resource{ID: "flow", TenantID: 1, OwnerID: "writer", Kind: "flows", Name: "Flow", Version: 1, Data: asJSON(flow)}
|
||
|
|
if err := s.validate(ctx, r); err != nil {
|
||
|
|
t.Fatalf("legacy scope prevented saving a valid flow: %v", err)
|
||
|
|
}
|
||
|
|
if err := s.db.Create(r).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
d := TaskData{FlowID: r.ID, FlowSnapshot: &flow, AgentSnapshots: map[string]AgentSnapshot{"writer": {ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"writer-frozen"}}}}
|
||
|
|
resolved, err := s.taskFlow(ctx, d)
|
||
|
|
if err != nil || len(resolved.KBIDs) != 0 || len(resolved.KnowledgeIDs) != 0 || len(resolved.KnowledgeFolders) != 0 || len(resolved.KnowledgeTags) != 0 {
|
||
|
|
t.Fatalf("legacy scope obstructed the flow snapshot: %+v, %v", resolved, err)
|
||
|
|
}
|
||
|
|
if err := s.inheritFlow(ctx, &d, resolved); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if !reflect.DeepEqual(d.KBIDs, []string{"writer-frozen", "editor-kb"}) {
|
||
|
|
t.Fatalf("planning scope did not retain each selected Agent's snapshot: %v", d.KBIDs)
|
||
|
|
}
|
||
|
|
options, err := s.agentCategoryOptions(ctx, d, resolved)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(options["writer"]) != 1 || options["writer"][0].ID != "writer-frozen-category" || len(options["editor"]) != 1 || options["editor"][0].ID != "editor-kb-category" {
|
||
|
|
t.Fatalf("categories leaked between Agents or ignored the snapshot: %+v", options)
|
||
|
|
}
|
||
|
|
if err := s.db.Exec("CREATE TABLE knowledges (id TEXT PRIMARY KEY, knowledge_base_id TEXT, folder_path TEXT)").Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := s.db.Exec("INSERT INTO knowledges VALUES ('folder-only-file','editor-kb','')").Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
search := &selectionSearchStub{}
|
||
|
|
s.sessions = search
|
||
|
|
d.Segments = []Segment{{ID: "s1", Title: "Draft", AgentID: "writer", CategoryID: "editor-kb-category", Status: "pending"}}
|
||
|
|
resource := &Resource{Kind: "content-tasks", Name: "Task", Data: asJSON(d)}
|
||
|
|
job.Action = "generate"
|
||
|
|
if err := s.executeTask(ctx, job, resource, ActionRequest{SegmentID: "s1", AgentID: "editor"}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
wantTags := []types.TagScope{{KnowledgeBaseID: "editor-kb", TagIDs: []string{"editor-kb-category"}}}
|
||
|
|
if len(search.kbs) != 0 || len(search.files) != 0 || !reflect.DeepEqual(search.tags, wantTags) {
|
||
|
|
t.Fatalf("reassigned segment lost its Agent/category scope: kbs=%v files=%v tags=%v", search.kbs, search.files, search.tags)
|
||
|
|
}
|
||
|
|
}
|