diff --git a/internal/application/service/agent_service.go b/internal/application/service/agent_service.go index e8200ede3..1f3ad8895 100644 --- a/internal/application/service/agent_service.go +++ b/internal/application/service/agent_service.go @@ -687,7 +687,9 @@ func (s *agentService) getKnowledgeBaseInfos(ctx context.Context, kbIDs []string pageResult, err := s.knowledgeService.ListPagedKnowledgeByKnowledgeBaseID(ctx, kbID, &types.Pagination{ Page: 1, PageSize: 10, - }, types.KnowledgeListFilter{}) + }, types.KnowledgeListFilter{ + ParseStatus: types.ParseStatusCompleted, + }) if err == nil && pageResult != nil { docCount = int(pageResult.Total) diff --git a/internal/application/service/agent_service_test.go b/internal/application/service/agent_service_test.go new file mode 100644 index 000000000..d74600ec4 --- /dev/null +++ b/internal/application/service/agent_service_test.go @@ -0,0 +1,95 @@ +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type fakeAgentKnowledgeBaseService struct { + interfaces.KnowledgeBaseService + kb *types.KnowledgeBase +} + +func (s *fakeAgentKnowledgeBaseService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) { + if s.kb == nil { + return nil, errors.New("knowledge base not found") + } + return s.kb, nil +} + +type fakeAgentKnowledgeService struct { + interfaces.KnowledgeService + knowledges []*types.Knowledge + lastFilter types.KnowledgeListFilter +} + +func (s *fakeAgentKnowledgeService) ListPagedKnowledgeByKnowledgeBaseID( + _ context.Context, + _ string, + page *types.Pagination, + filter types.KnowledgeListFilter, +) (*types.PageResult, error) { + s.lastFilter = filter + + filtered := make([]*types.Knowledge, 0, len(s.knowledges)) + for _, knowledge := range s.knowledges { + if filter.ParseStatus != "" && knowledge.ParseStatus != filter.ParseStatus { + continue + } + filtered = append(filtered, knowledge) + } + return types.NewPageResult(int64(len(filtered)), page, filtered), nil +} + +func TestGetKnowledgeBaseInfos_ExcludesUnprocessedDocuments(t *testing.T) { + now := time.Now() + knowledgeService := &fakeAgentKnowledgeService{ + knowledges: []*types.Knowledge{ + { + ID: "doc-processing", + KnowledgeBaseID: "kb-1", + Title: "still parsing", + FileName: "processing.pdf", + FileType: "pdf", + ParseStatus: types.ParseStatusProcessing, + CreatedAt: now, + }, + { + ID: "doc-completed", + KnowledgeBaseID: "kb-1", + Title: "ready document", + FileName: "ready.pdf", + FileType: "pdf", + ParseStatus: types.ParseStatusCompleted, + CreatedAt: now.Add(-time.Minute), + }, + }, + } + service := &agentService{ + knowledgeBaseService: &fakeAgentKnowledgeBaseService{ + kb: &types.KnowledgeBase{ + ID: "kb-1", + Name: "KB", + Type: types.KnowledgeBaseTypeDocument, + TenantID: 1, + }, + }, + knowledgeService: knowledgeService, + } + + infos, err := service.getKnowledgeBaseInfos(context.Background(), []string{"kb-1"}) + + require.NoError(t, err) + require.Len(t, infos, 1) + assert.Equal(t, types.ParseStatusCompleted, knowledgeService.lastFilter.ParseStatus) + assert.Equal(t, 1, infos[0].DocCount) + require.Len(t, infos[0].RecentDocs, 1) + assert.Equal(t, "doc-completed", infos[0].RecentDocs[0].KnowledgeID) +}