71 lines
2.8 KiB
Go
71 lines
2.8 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestAgentDutyEnforcedOnSave(t *testing.T) {
|
||
|
|
s := sqliteWorkbench(t)
|
||
|
|
ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
|
||
|
|
for _, role := range []string{"", types.AgentRoleGeneral, types.AgentRoleTrainingMap, types.AgentRoleProduction, types.AgentRoleQC} {
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: role}}
|
||
|
|
m := &Resource{Name: "map", Kind: "training-maps", Data: asJSON(MapData{Topic: "topic", AgentID: "agent"})}
|
||
|
|
f := &Resource{Name: "flow", Kind: "flows", Data: asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})}
|
||
|
|
for _, item := range []struct {
|
||
|
|
resource *Resource
|
||
|
|
duty string
|
||
|
|
}{{m, types.AgentRoleTrainingMap}, {f, types.AgentRoleProduction}} {
|
||
|
|
err := s.validate(ctx, item.resource)
|
||
|
|
want := (types.CustomAgentConfig{Role: role}).SupportsDuty(item.duty)
|
||
|
|
if (err == nil) != want {
|
||
|
|
t.Fatalf("role %q duty %q: %v", role, item.duty, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAgentContentTagsReachTaskDecomposition(t *testing.T) {
|
||
|
|
s, ctx, job, model := generationHarness(t)
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ContentTags: []string{"课程标签"}}}
|
||
|
|
f := seedResource(t, s, 1, "writer", "flows")
|
||
|
|
f.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})
|
||
|
|
if err := s.db.Save(f).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
job.Action = "decompose"
|
||
|
|
model.output = `{"segments":[{"id":"s","title":"segment","node_id":"n","agent_id":"agent"}]}`
|
||
|
|
r := &Resource{Kind: "content-tasks", Data: asJSON(TaskData{FlowID: f.ID, AgentID: "agent", SourceText: "framework"})}
|
||
|
|
if err := s.executeTask(ctx, job, r, ActionRequest{}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(model.messages) < 2 || !strings.Contains(model.messages[1].Content, "课程标签") {
|
||
|
|
t.Fatal("decomposition omitted native content tags")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAgentDutyChangeBlocksSavedMapAndFlowExecution(t *testing.T) {
|
||
|
|
s, ctx, job, model := generationHarness(t)
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: types.AgentRoleQC}}
|
||
|
|
job.Action = "generate"
|
||
|
|
m := &Resource{Kind: "training-maps", Data: asJSON(MapData{Topic: "topic", AgentID: "agent"})}
|
||
|
|
if err := s.executeMap(ctx, job, m, ActionRequest{}); !errors.Is(err, ErrInvalid) {
|
||
|
|
t.Fatalf("map: %v", err)
|
||
|
|
}
|
||
|
|
if len(model.messages) != 0 {
|
||
|
|
t.Fatal("incompatible agent invoked model")
|
||
|
|
}
|
||
|
|
f := seedResource(t, s, 1, "writer", "flows")
|
||
|
|
f.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})
|
||
|
|
if err := s.db.Save(f).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
snapshot := FlowData{Nodes: []FlowNode{{ID: "frozen", AgentID: "agent"}}}
|
||
|
|
if _, err := s.taskFlow(ctx, TaskData{FlowID: f.ID, FlowSnapshot: &snapshot}); !errors.Is(err, ErrInvalid) {
|
||
|
|
t.Fatalf("frozen flow: %v", err)
|
||
|
|
}
|
||
|
|
}
|