feat(artifact): use retrievable to decide if chunk can be return (#54)

Because

1. we wanna monitoring the latency of retrieval test in easy way
2. we wanna enable retrievable button of chunks

This commit

1. add temporary log
2. add the logic to skip the unretirevable chunk

pkg/handler/retrieval.go

@@ -3,6 +3,7 @@ package handler
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/gofrs/uuid"
 	"github.com/instill-ai/artifact-backend/pkg/customerror"
@@ -31,17 +32,22 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 		log.Error("failed to parse user id", zap.Error(err))
 		return nil, fmt.Errorf("failed to parse user id: %v. err: %w", err, customerror.ErrUnauthenticated)
 	}
+	t := time.Now()
 	ns, err := ph.service.GetNamespaceByNsID(ctx, req.GetNamespaceId())
 	if err != nil {
 		log.Error("failed to get namespace by ns id", zap.Error(err))
 		return nil, fmt.Errorf("failed to get namespace by ns id. err: %w", err)
 	}
+	log.Info("get namespace by ns id", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	ownerUID := ns.NsUID
 	kb, err := ph.service.Repository.GetKnowledgeBaseByOwnerAndKbID(ctx, ownerUID.String(), req.KbId)
 	if err != nil {
 		log.Error("failed to get knowledge base by owner and kb id", zap.Error(err))
 		return nil, fmt.Errorf("failed to get knowledge base by owner and kb id. err: %w", err)
 	}
+	log.Info("get knowledge base by owner and kb id", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	// ACL : check user has access to the knowledge base
 	granted, err := ph.service.ACLClient.CheckPermission(ctx, "knowledgebase", kb.UID, "reader")
 	if err != nil {
@@ -52,7 +58,8 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 		log.Error("permission denied", zap.String("user_id", uid), zap.String("kb_id", kb.UID.String()))
 		return nil, fmt.Errorf("SimilarityChunksSearch permission denied. err: %w", service.ErrNoPermission)
 	}
-
+	log.Info("check permission", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	// retrieve the chunks based on the similarity
 	simChunksScroes, err := ph.service.SimilarityChunksSearch(ctx, uidUUID, ownerUID.String(), req)
 	if err != nil {
@@ -63,6 +70,8 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 	for _, simChunk := range simChunksScroes {
 		chunkUIDs = append(chunkUIDs, simChunk.ChunkUID)
 	}
+	log.Info("get similarity chunks", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	// fetch the chunks metadata
 	chunks, err := ph.service.Repository.GetChunksByUIDs(ctx, chunkUIDs)
 	if err != nil {
@@ -74,12 +83,15 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 	for _, chunk := range chunks {
 		chunkFilePaths = append(chunkFilePaths, chunk.ContentDest)
 	}
+	log.Info("get chunks by uids", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	// fetch the chunks content from minio
 	chunkContents, err := ph.service.MinIO.GetFilesByPaths(ctx, chunkFilePaths)
 	if err != nil {
 		log.Error("failed to get chunks content", zap.Error(err))
 		return nil, fmt.Errorf("failed to get chunks content. err: %w", err)
 	}
+	log.Info("get chunks content from minIO", zap.Duration("duration", time.Since(t)))
 
 	// fetch the file names
 	fileUIDMapName := make(map[uuid.UUID]string)
@@ -90,6 +102,7 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 	for fileUID := range fileUIDMapName {
 		fileUids = append(fileUids, fileUID)
 	}
+	t = time.Now()
 	files, err := ph.service.Repository.GetKnowledgeBaseFilesByFileUIDs(
 		ctx, fileUids, repository.KnowledgeBaseFileColumn.UID, repository.KnowledgeBaseFileColumn.Name)
 	if err != nil {
@@ -99,10 +112,14 @@ func (ph *PublicHandler) SimilarityChunksSearch(
 	for _, file := range files {
 		fileUIDMapName[file.UID] = file.Name
 	}
+	log.Info("get knowledge base files by file uids", zap.Duration("duration", time.Since(t)))
 
 	// prepare the response
 	simChunks := make([]*artifactv1alpha.SimilarityChunk, 0, len(chunks))
 	for i, chunk := range chunks {
+		if !chunk.Retrievable {
+			continue
+		}
 		simChunks = append(simChunks, &artifactv1alpha.SimilarityChunk{
 			ChunkUid:        chunk.UID.String(),
 			SimilarityScore: float32(simChunksScroes[i].Score),

pkg/milvus/milvus.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/instill-ai/artifact-backend/pkg/logger"
 	"github.com/milvus-io/milvus-sdk-go/v2/client"
@@ -315,6 +316,7 @@ func (m *MilvusClient) SearchSimilarEmbeddings(ctx context.Context, collectionNa
 		log.Error("failed to get logger", zap.Error(err))
 		return nil, fmt.Errorf("failed to get logger: %w", err)
 	}
+	t := time.Now()
 	// Check if the collection exists
 	has, err := m.c.HasCollection(ctx, collectionName)
 	if err != nil {
@@ -325,6 +327,8 @@ func (m *MilvusClient) SearchSimilarEmbeddings(ctx context.Context, collectionNa
 		log.Error("collection does not exist", zap.String("collection_name", collectionName))
 		return nil, fmt.Errorf("collection %s does not exist", collectionName)
 	}
+	log.Info("check collection existence", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 
 	// Load the collection if it's not already loaded
 	err = m.c.LoadCollection(ctx, collectionName, false)
@@ -338,6 +342,8 @@ func (m *MilvusClient) SearchSimilarEmbeddings(ctx context.Context, collectionNa
 	for i, v := range vectors {
 		milvusVectors[i] = entity.FloatVector(v)
 	}
+	log.Info("load collection", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 
 	// Perform the search
 	sp, err := entity.NewIndexSCANNSearchParam(Nprobe, ReorderK)
@@ -356,7 +362,7 @@ func (m *MilvusClient) SearchSimilarEmbeddings(ctx context.Context, collectionNa
 		log.Error("failed to search embeddings", zap.Error(err))
 		return nil, fmt.Errorf("failed to search embeddings: %w", err)
 	}
-
+	log.Info("search embeddings", zap.Duration("duration", time.Since(t)))
 	// Extract the embeddings from the search results
 	var embeddings [][]SimilarEmbedding
 	for _, result := range results {

pkg/service/retrieval.go

@@ -3,6 +3,7 @@ package service
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/gofrs/uuid"
 	"github.com/instill-ai/artifact-backend/pkg/logger"
@@ -17,28 +18,30 @@ type SimChunk struct {
 
 func (s *Service) SimilarityChunksSearch(ctx context.Context, caller uuid.UUID, ownerUID string, req *artifactv1alpha.SimilarityChunksSearchRequest) ([]SimChunk, error) {
 	log, _ := logger.GetZapLogger(ctx)
-	log.Info("SimilarityChunksSearch")
+	t := time.Now()
 	textVector, err := s.VectorizeText(ctx, caller, []string{req.TextPrompt})
-	fmt.Println("textVector", textVector[0][:10])
 	if err != nil {
 		log.Error("failed to vectorize text", zap.Error(err))
 		return nil, fmt.Errorf("failed to vectorize text. err: %w", err)
 	}
-
+	log.Info("vectorize text", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 	// get kb by kb_id and owner uid
 	kb, err := s.Repository.GetKnowledgeBaseByOwnerAndKbID(ctx, ownerUID, req.KbId)
 	if err != nil {
 		log.Error("failed to get knowledge base by owner and id", zap.Error(err))
 		return nil, fmt.Errorf("failed to get knowledge base by owner and id. err: %w", err)
-
 	}
+	log.Info("get knowledge base by owner and id", zap.Duration("duration", time.Since(t)))
+	t = time.Now()
 
 	// search similar embeddings in kb
 	simEmbeddings, err := s.MilvusClient.SearchSimilarEmbeddingsInKB(ctx, kb.UID.String(), textVector, int(req.Topk))
 	if err != nil {
 		log.Error("failed to search similar embeddings in kb", zap.Error(err))
 		return nil, fmt.Errorf("failed to search similar embeddings in kb. err: %w", err)
 	}
+	log.Info("search similar embeddings in kb", zap.Duration("duration", time.Since(t)))
 
 	// fetch chunks by their UIDs
 	res := make([]SimChunk, 0, len(simEmbeddings))