xinmi/internal/application/service/custom_agent_draft_test.go

package service

import (
	"context"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

type agentDraftModels struct {
	interfaces.ModelService
	model   *types.Model
	llm     *agentDraftChat
	lookups int
}

func (m *agentDraftModels) GetModelByID(context.Context, string) (*types.Model, error) {
	m.lookups++
	return m.model, nil
}

func (m *agentDraftModels) GetChatModel(context.Context, string) (chat.Chat, error) {
	return m.llm, nil
}

type agentDraftChat struct {
	response *types.ChatResponse
	calls    int
	ctx      context.Context
	messages []chat.Message
}

func (m *agentDraftChat) ChatStream(context.Context, []chat.Message, *chat.ChatOptions) (<-chan types.StreamResponse, error) {
	panic("draft generation must not use streaming")
}

func (m *agentDraftChat) GetModelName() string { return "selected" }
func (m *agentDraftChat) GetModelID() string   { return "selected" }

func (m *agentDraftChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
	m.calls++
	m.ctx, m.messages = ctx, messages
	return m.response, nil
}

func TestAgentDraftUsesSelectedModelWithoutPersistingBindings(t *testing.T) {
	llm := &agentDraftChat{response: &types.ChatResponse{Content: `{"name":"Course Planner","description":"Creates connected daily frameworks","role":"training_map","content_tags":["planning"],"id":"injected","system_prompt":"injected","knowledge_bases":["private"],"config":{"allowed_tools":["shell"],"model_id":"other"}}`}}
	models := &agentDraftModels{model: &types.Model{ID: "selected", Type: types.ModelTypeKnowledgeQA, Status: types.ModelStatusActive}, llm: llm}
	service := &customAgentService{models: models}
	draft, err := service.DraftAgent(instructionContext(1, "alice", types.TenantRoleContributor), "selected", "Plan a course")
	require.NoError(t, err)
	require.Equal(t, 1, llm.calls)
	purpose, correlation := types.LLMCallMetadataFromContext(llm.ctx)
	require.Equal(t, "agent_draft_generation", purpose)
	require.NotEmpty(t, correlation)
	_, deadline := llm.ctx.Deadline()
	require.True(t, deadline)
	require.Equal(t, "Plan a course", llm.messages[1].Content)
	require.Empty(t, draft.ID)
	require.Equal(t, "private", draft.Visibility)
	require.Equal(t, "selected", draft.Config.ModelID)
	require.Equal(t, types.AgentRoleTrainingMap, draft.Config.Role)
	require.Empty(t, draft.Config.SystemPrompt)
	require.Empty(t, draft.Config.KnowledgeBases)
	require.Empty(t, draft.Config.AllowedTools)
	require.Equal(t, "none", draft.Config.KBSelectionMode)
}

func TestAgentDraftRejectsUnauthorizedOrInvalidRequestsBeforeCallingModel(t *testing.T) {
	ctx := instructionContext(1, "alice", types.TenantRoleContributor)
	for _, test := range []struct {
		name        string
		ctx         context.Context
		requirement string
		model       *types.Model
	}{
		{"denied", accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true}), "plan", &types.Model{Type: types.ModelTypeKnowledgeQA, Status: types.ModelStatusActive}},
		{"empty", ctx, " ", nil},
		{"too-long", ctx, strings.Repeat("a", 4001), nil},
		{"missing-model", ctx, "plan", nil},
		{"embedding-model", ctx, "plan", &types.Model{Type: types.ModelTypeEmbedding, Status: types.ModelStatusActive}},
	} {
		t.Run(test.name, func(t *testing.T) {
			llm := &agentDraftChat{}
			models := &agentDraftModels{model: test.model, llm: llm}
			_, err := (&customAgentService{models: models}).DraftAgent(test.ctx, "selected", test.requirement)
			require.Error(t, err)
			require.Zero(t, llm.calls)
		})
	}
}

func TestAgentDraftRejectsMalformedAndIncompleteResponses(t *testing.T) {
	for _, content := range []string{"not-json", `{"name":"A","role":"admin"}`, `{"name":"A","content_tags":["same","same"]}`, `{"name":""}`, strings.Repeat("x", 16385)} {
		_, err := parseAgentDraft(content, "selected")
		require.Error(t, err)
	}
	draft, err := parseAgentDraft("```json\n{\"name\":\"Valid\"}\n```", "selected")
	require.NoError(t, err)
	require.Equal(t, types.AgentRoleGeneral, draft.Config.Role)
	llm := &agentDraftChat{response: &types.ChatResponse{Content: `{"name":"A"}`, FinishReason: "length"}}
	models := &agentDraftModels{model: &types.Model{Type: types.ModelTypeKnowledgeQA, Status: types.ModelStatusActive}, llm: llm}
	_, err = (&customAgentService{models: models}).DraftAgent(instructionContext(1, "alice", types.TenantRoleContributor), "selected", "plan")
	require.ErrorContains(t, err, "完整")
}