Method persistEntities

api/service/rag/entity_extractor.go:222–290

(ctx context.Context, eid int64, chunk *model.DocumentChunk, entities []ExtractedEntity)

}

func (s *EntityExtractionService) persistEntities(ctx context.Context, eid int64, chunk *model.DocumentChunk, entities []ExtractedEntity) error {
	tx := s.db.Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
		}
	}()

	var library model.Library
	if err := tx.Select("id", "space_id").
		Where("eid = ? AND id = ?", eid, chunk.LibraryID).
		First(&library).Error; err != nil {
		tx.Rollback()
		return err
	}

	var relations []model.EntityChunkRelation
	const minAutoLLMConfidence = 0.85
	seenEntityIDs := make(map[int64]struct{})
	createdEntities := make(map[int64]*model.Entity)
	for _, e := range entities {
		if e.Confidence < minAutoLLMConfidence {
			continue
		}
		entityModel, created, err := model.GetOrCreateEntityWithDBAndCreated(tx, eid, e.Type, e.Name)
		if err != nil {
			tx.Rollback()
			return err
		}
		if created {
			createdEntities[entityModel.ID] = entityModel
		}
		if _, ok := seenEntityIDs[entityModel.ID]; ok {
			continue
		}
		seenEntityIDs[entityModel.ID] = struct{}{}

		relations = append(relations, model.EntityChunkRelation{
			Eid:        eid,
			EntityID:   entityModel.ID,
			SpaceID:    library.SpaceID,
			LibraryID:  chunk.LibraryID,
			FileID:     chunk.FileID,
			ChunkID:    chunk.ID,
			ChunkType:  chunk.ChunkType,
			Status:     model.EntityRelationStatusActive,
			Confidence: e.Confidence,
			Source:     model.EntityRelationSourceAutoLLM,
		})
	}

	if err := model.ReplaceEntityChunkRelationsBySourceWithDB(tx, eid, chunk.ID, model.EntityRelationSourceAutoLLM, relations); err != nil {
		tx.Rollback()
		return err
	}

	if err := tx.Commit().Error; err != nil {
		return err