RAG/service/dataset_service.py

from datetime import datetime

from core.database import DBHelper
from data_models.contents import Contents
from data_models.dataset import Dataset


def get_datasets():
    db_helper = DBHelper()
    entities = db_helper.get_all(Dataset, status=1)
    data = []
    for entity in entities:
        count = db_helper.count(Contents, dataset_id=entity.id, doc_status = 1)
        data.append({'dataset_id': entity.id, 'name': entity.name, 'count': count,
                     'created_at': entity.created_at.strftime('%Y-%m-%d')})
    return data


def add_datasets(dataset_param):
    db_helper = DBHelper()
    dataset = Dataset(name=dataset_param.name, created_at=datetime.now(), updated_at=datetime.now(), status=1)
    db_helper.add(dataset)
    return True