diff --git a/cmd/server/main.go b/cmd/server/main.go index f8a84c6..507c699 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -64,17 +64,6 @@ func buildRepositories(cfg config.Config) (repository.PaperRepository, repositor return nil, nil, nil, err } return database.Papers, database.Chunks, database.Users, nil - - case "json": - repo, err := repository.NewJSONRepository(cfg.DBPath) - if err != nil { - return nil, nil, nil, err - } - return repo, repository.NewJSONChunkRepository(repo), nil, nil - - case "memory": - return repository.NewMemoryPaperRepository(), repository.NewMemoryChunkRepository(), nil, nil - default: return nil, nil, nil, fmt.Errorf("unsupported CERULEAN_DB_DRIVER=%q; supported: mysql, json, memory", cfg.DBDriver) } diff --git a/internal/dao/chunk.go b/internal/dao/chunk.go index a830a99..ffbf635 100644 --- a/internal/dao/chunk.go +++ b/internal/dao/chunk.go @@ -3,6 +3,7 @@ package dao import ( "context" "fmt" + "strings" "github.com/CeruleanFlow/cerulean/internal/domain" "github.com/CeruleanFlow/cerulean/internal/entity" @@ -135,3 +136,56 @@ func fromJSONMap(values datatypes.JSONMap) map[string]string { return out } + +func (d *ChunkDAO) ReplaceByPaperID(ctx context.Context, paperID string, chunks []domain.Chunk) error { + if strings.TrimSpace(paperID) == "" { + return fmt.Errorf("invalid paper id") + } + + return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := tx.Where("paper_id = ?", paperID).Delete(&entity.Chunk{}).Error; err != nil { + return fmt.Errorf("delete old chunks: %w", err) + } + if len(chunks) == 0 { + return nil + } + + entities := make([]entity.Chunk, 0, len(chunks)) + + for _, chunk := range chunks { + if chunk.PaperID != paperID { + return fmt.Errorf("invalid paper id") + } + if chunk.PaperID == "" { + chunk.PaperID = paperID + } + + entities = append(entities, domainChunkToEntity(chunk)) + } + + if err := tx.Create(&entities).Error; err != nil { + return fmt.Errorf("insert new chunks: %w", err) + } + return nil + }) +} + +func domainChunkToEntity(chunk domain.Chunk) entity.Chunk { + metadata := map[string]any{} + for k, v := range chunk.Metadata { + metadata[k] = v + } + + return entity.Chunk{ + ID: chunk.ID, + PaperID: chunk.PaperID, + PageNo: chunk.PageNo, + ChunkIndex: chunk.Index, + Text: chunk.Text, + ObjectKey: chunk.ObjectKey, + VectorID: chunk.VectorID, + Metadata: metadata, + CreatedAt: chunk.CreatedAt, + UpdatedAt: chunk.UpdatedAt, + } +} diff --git a/internal/ingest/service.go b/internal/ingest/service.go index 9e86bfb..4aaf71b 100644 --- a/internal/ingest/service.go +++ b/internal/ingest/service.go @@ -140,22 +140,15 @@ func (s *Service) runPDFTextIngest(ctx context.Context, job task.Task, paper dom return fmt.Errorf("pdf text parser produced no chunks") } - // delete old chunks in mysql - if err := s.chunks.DeleteByPaperID(ctx, paper.ID); err != nil { - return fmt.Errorf("delete old chunks from mysql: %w", err) + if err := s.chunks.ReplaceByPaperID(ctx, paper.ID, chunks); err != nil { + return fmt.Errorf("replace chunks in mysql: %w", err) } - // delete old chunk index in es + if s.search != nil { if err := s.search.DeleteByPaperID(ctx, paper.ID); err != nil { return fmt.Errorf("delete old chunks from elasticsearch: %w", err) } - } - // save new chunks in mysql - if err := s.chunks.UpsertMany(ctx, chunks); err != nil { - return fmt.Errorf("save chunks to mysql: %w", err) - } - // save new chunks index in es - if s.search != nil { + if err := s.search.IndexChunks(ctx, chunks); err != nil { return fmt.Errorf("index chunks to elasticsearch: %w", err) } diff --git a/internal/repository/chunk_repository.go b/internal/repository/chunk_repository.go index 2f8e2dd..490a873 100644 --- a/internal/repository/chunk_repository.go +++ b/internal/repository/chunk_repository.go @@ -8,13 +8,6 @@ import ( "github.com/CeruleanFlow/cerulean/internal/domain" ) -type ChunkRepository interface { - UpsertMany(ctx context.Context, chunks []domain.Chunk) error - List(ctx context.Context, filters map[string]string) ([]domain.Chunk, error) - ListByPaperID(ctx context.Context, paperID string) ([]domain.Chunk, error) - DeleteByPaperID(ctx context.Context, paperID string) error -} - type MemoryChunkRepository struct { mu sync.RWMutex chunks map[string]domain.Chunk diff --git a/internal/repository/repository.go b/internal/repository/repository.go new file mode 100644 index 0000000..9c4e807 --- /dev/null +++ b/internal/repository/repository.go @@ -0,0 +1,15 @@ +package repository + +import ( + "context" + + "github.com/CeruleanFlow/cerulean/internal/domain" +) + +type ChunkRepository interface { + UpsertMany(ctx context.Context, chunks []domain.Chunk) error + List(ctx context.Context, filters map[string]string) ([]domain.Chunk, error) + ListByPaperID(ctx context.Context, paperID string) ([]domain.Chunk, error) + DeleteByPaperID(ctx context.Context, paperID string) error + ReplaceByPaperID(ctx context.Context, paperID string, chunks []domain.Chunk) error +}