70 lines
2.4 KiB
Go
70 lines
2.4 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
)
|
||
|
|
|
||
|
|
type instructionStub struct {
|
||
|
|
text string
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
type disabledInstructionAgentStub struct{ interfaces.CustomAgentService }
|
||
|
|
|
||
|
|
func (disabledInstructionAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
|
||
|
|
tenant, _ := types.TenantIDFromContext(ctx)
|
||
|
|
return &types.CustomAgent{ID: id, TenantID: tenant, Visibility: "global", Enabled: false, Config: types.CustomAgentConfig{ModelID: "model"}}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestDisabledAgentCannotGenerateWithoutInstructionResolver(t *testing.T) {
|
||
|
|
s, ctx, _, model := generationHarness(t)
|
||
|
|
s.agents = disabledInstructionAgentStub{}
|
||
|
|
if _, err := s.complete(ctx, "disabled", "system", "prompt"); !errors.Is(err, ErrInvalid) {
|
||
|
|
t.Fatalf("expected disabled Agent rejection, got %v", err)
|
||
|
|
}
|
||
|
|
if len(model.messages) != 0 {
|
||
|
|
t.Fatal("disabled Agent reached provider")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s instructionStub) ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error) {
|
||
|
|
return s.text, s.err
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMountedInstructionsParticipateInGenerationAndBudget(t *testing.T) {
|
||
|
|
s, ctx, _, model := generationHarness(t)
|
||
|
|
s.SetInstructionResolver(instructionStub{text: "技能规则：先列判断依据，再给案例。"})
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if !strings.Contains(model.messages[0].Content, "技能规则：先列判断依据") {
|
||
|
|
t.Fatal("mounted instruction missing from provider request")
|
||
|
|
}
|
||
|
|
s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096, MaxOutputTokens: 512}}
|
||
|
|
s.SetInstructionResolver(instructionStub{text: strings.Repeat("很长的技能规则", 3000)})
|
||
|
|
model.messages = nil
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrInvalid) {
|
||
|
|
t.Fatalf("expected budget rejection, got %v", err)
|
||
|
|
}
|
||
|
|
if len(model.messages) != 0 {
|
||
|
|
t.Fatal("oversized skills reached model")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestUnavailableInstructionStopsGeneration(t *testing.T) {
|
||
|
|
s, ctx, _, model := generationHarness(t)
|
||
|
|
s.SetInstructionResolver(instructionStub{err: ErrForbidden})
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrForbidden) {
|
||
|
|
t.Fatalf("got %v", err)
|
||
|
|
}
|
||
|
|
if len(model.messages) != 0 {
|
||
|
|
t.Fatal("generation silently ignored unavailable instruction")
|
||
|
|
}
|
||
|
|
}
|