ai-agent/internal/services/knowledge_base_service_test.go

package services

import (
	"encoding/json"
	"strings"
	"testing"

	"code.tczkiot.com/wlw/ai-agent/internal/ai/workflow/dsl"
	workflowregistry "code.tczkiot.com/wlw/ai-agent/internal/ai/workflow/registry"
	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/dto/request"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
	"code.tczkiot.com/wlw/ai-agent/internal/repositories"

	"github.com/glebarez/sqlite"
	"github.com/mlogclub/simple/sqls"
	"gorm.io/gorm"
)

func TestBuildKnowledgeBaseModelUsesLowerDefaultScoreThreshold(t *testing.T) {
	item, err := KnowledgeBaseService.buildKnowledgeBaseModel(request.CreateKnowledgeBaseRequest{})
	if err != nil {
		t.Fatalf("build knowledge base model failed: %v", err)
	}
	if item.DefaultScoreThreshold != 0.2 {
		t.Fatalf("expected default score threshold 0.2, got %v", item.DefaultScoreThreshold)
	}
}

func TestDeleteKnowledgeBaseRejectsWorkflowDraftReference(t *testing.T) {
	setupKnowledgeBaseServiceTestDB(t)
	kb := createKnowledgeBaseServiceTestBase(t, "Referenced KB")
	otherKB := createKnowledgeBaseServiceTestBase(t, "Other KB")
	createKnowledgeBaseServiceTestWorkflow(t, "Support Workflow", knowledgeBaseServiceTestWorkflowDefinition([]int64{12, otherKB.ID}))
	createKnowledgeBaseServiceTestWorkflow(t, "Knowledge Workflow", knowledgeBaseServiceTestWorkflowDefinition([]int64{12, kb.ID, otherKB.ID}))

	err := KnowledgeBaseService.DeleteKnowledgeBase(kb.ID)
	if err == nil {
		t.Fatal("DeleteKnowledgeBase() error is nil, want referenced workflow error")
	}
	if got := err.Error(); !strings.Contains(got, "Knowledge Workflow") {
		t.Fatalf("DeleteKnowledgeBase() error = %q, want workflow name", got)
	}
	if repositories.KnowledgeBaseRepository.Get(sqls.DB(), kb.ID) == nil {
		t.Fatal("knowledge base was deleted despite workflow reference")
	}
}

func TestDeleteKnowledgeBaseRejectsWorkflowVersionReference(t *testing.T) {
	setupKnowledgeBaseServiceTestDB(t)
	kb := createKnowledgeBaseServiceTestBase(t, "Version KB")
	workflow := createKnowledgeBaseServiceTestWorkflow(t, "Published Workflow", knowledgeBaseServiceTestWorkflowDefinition([]int64{999}))
	raw, err := json.Marshal(knowledgeBaseServiceTestWorkflowDefinition([]int64{kb.ID}))
	if err != nil {
		t.Fatalf("marshal workflow version definition: %v", err)
	}
	if err := repositories.AIWorkflowVersionRepository.Create(sqls.DB(), &models.AIWorkflowVersion{
		WorkflowID: workflow.ID,
		Version:    1,
		Status:     enums.StatusOk,
		Definition: string(raw),
	}); err != nil {
		t.Fatalf("create workflow version: %v", err)
	}

	err = KnowledgeBaseService.DeleteKnowledgeBase(kb.ID)
	if err == nil {
		t.Fatal("DeleteKnowledgeBase() error is nil, want referenced workflow version error")
	}
	if got := err.Error(); !strings.Contains(got, "Published Workflow") {
		t.Fatalf("DeleteKnowledgeBase() error = %q, want workflow name", got)
	}
}

func TestDeleteKnowledgeBaseCascadesContentWhenNotReferenced(t *testing.T) {
	setupKnowledgeBaseServiceTestDB(t)
	kb := createKnowledgeBaseServiceTestBase(t, "Delete KB")
	document := &models.KnowledgeDocument{
		KnowledgeBaseID: kb.ID,
		Title:           "Doc",
		ContentType:     enums.KnowledgeDocumentContentTypeMarkdown,
		Content:         "content",
		Status:          enums.StatusOk,
		IndexStatus:     enums.KnowledgeDocumentIndexStatusIndexed,
	}
	if err := repositories.KnowledgeDocumentRepository.Create(sqls.DB(), document); err != nil {
		t.Fatalf("create document: %v", err)
	}
	faq := &models.KnowledgeFAQ{
		KnowledgeBaseID: kb.ID,
		Question:        "Question",
		Answer:          "Answer",
		Status:          enums.StatusOk,
		IndexStatus:     enums.KnowledgeDocumentIndexStatusIndexed,
	}
	if err := repositories.KnowledgeFAQRepository.Create(sqls.DB(), faq); err != nil {
		t.Fatalf("create faq: %v", err)
	}
	if err := repositories.KnowledgeChunkRepository.BatchCreate(sqls.DB(), []models.KnowledgeChunk{
		{KnowledgeBaseID: kb.ID, DocumentID: document.ID, ChunkNo: 1, Status: enums.StatusOk},
		{KnowledgeBaseID: kb.ID, FaqID: faq.ID, ChunkNo: 1, Status: enums.StatusOk},
	}); err != nil {
		t.Fatalf("create chunks: %v", err)
	}

	if err := KnowledgeBaseService.DeleteKnowledgeBase(kb.ID); err != nil {
		t.Fatalf("DeleteKnowledgeBase() error = %v", err)
	}

	assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeBase{}, "id = ?", kb.ID, 0)
	assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeDocument{}, "knowledge_base_id = ?", kb.ID, 0)
	assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeFAQ{}, "knowledge_base_id = ?", kb.ID, 0)
	assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeChunk{}, "knowledge_base_id = ?", kb.ID, 0)
}

func setupKnowledgeBaseServiceTestDB(t *testing.T) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	if err != nil {
		t.Fatalf("open sqlite db: %v", err)
	}
	if err := db.AutoMigrate(&models.KnowledgeBase{}, &models.KnowledgeDocument{}, &models.KnowledgeFAQ{}, &models.KnowledgeChunk{}, &models.AIAgent{}, &models.AIWorkflow{}, &models.AIWorkflowVersion{}); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)
}

func createKnowledgeBaseServiceTestWorkflow(t *testing.T, name string, definition dsl.Definition) *models.AIWorkflow {
	t.Helper()
	raw, err := json.Marshal(definition)
	if err != nil {
		t.Fatalf("marshal workflow definition: %v", err)
	}
	item := &models.AIWorkflow{
		Name:            name,
		Status:          enums.StatusOk,
		DraftDefinition: string(raw),
	}
	if err := repositories.AIWorkflowRepository.Create(sqls.DB(), item); err != nil {
		t.Fatalf("create workflow: %v", err)
	}
	return item
}

func knowledgeBaseServiceTestWorkflowDefinition(knowledgeBaseIDs []int64) dsl.Definition {
	return dsl.Definition{
		SchemaVersion: dsl.SchemaVersion,
		Nodes: []dsl.Node{
			{
				ID:   "start_1",
				Type: workflowregistry.NodeTypeStart,
			},
			{
				ID:   "retrieve_1",
				Type: workflowregistry.NodeTypeKnowledgeRetrieve,
				Data: dsl.NodeData{
					Config: mustKnowledgeBaseServiceTestJSON(map[string]any{"knowledgeBaseIds": knowledgeBaseIDs}),
					InputsValues: map[string]dsl.Value{
						"query": dsl.RefValue("start_1", "userMessage"),
					},
				},
			},
			{
				ID:   "end_1",
				Type: workflowregistry.NodeTypeEnd,
			},
		},
		Edges: []dsl.Edge{
			{SourceNodeID: "start_1", TargetNodeID: "retrieve_1"},
			{SourceNodeID: "retrieve_1", TargetNodeID: "end_1"},
		},
	}
}

func mustKnowledgeBaseServiceTestJSON(value any) json.RawMessage {
	raw, err := json.Marshal(value)
	if err != nil {
		panic(err)
	}
	return raw
}

func createKnowledgeBaseServiceTestBase(t *testing.T, name string) *models.KnowledgeBase {
	t.Helper()
	item := &models.KnowledgeBase{
		Name:          name,
		KnowledgeType: string(enums.KnowledgeBaseTypeDocument),
		Status:        enums.StatusOk,
	}
	if err := repositories.KnowledgeBaseRepository.Create(sqls.DB(), item); err != nil {
		t.Fatalf("create knowledge base: %v", err)
	}
	return item
}

func assertKnowledgeBaseServiceTestCount(t *testing.T, model any, query string, arg any, want int64) {
	t.Helper()
	var count int64
	if err := sqls.DB().Model(model).Where(query, arg).Count(&count).Error; err != nil {
		t.Fatalf("count %T: %v", model, err)
	}
	if count != want {
		t.Fatalf("count %T = %d, want %d", model, count, want)
	}
}