Underground-Digital/Workflow-Engine

1import datetime2import time3 4import click5from sqlalchemy import text6from werkzeug.exceptions import NotFound7 8import app9from configs import dify_config10from extensions.ext_database import db11from models.dataset import Embedding12 13 14@app.celery.task(queue="dataset")15def clean_embedding_cache_task():16    click.echo(click.style("Start clean embedding cache.", fg="green"))17    clean_days = int(dify_config.PLAN_SANDBOX_CLEAN_DAY_SETTING)18    start_at = time.perf_counter()19    thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days)20    while True:21        try:22            embedding_ids = (23                db.session.query(Embedding.id)24                .filter(Embedding.created_at < thirty_days_ago)25                .order_by(Embedding.created_at.desc())26                .limit(100)27                .all()28            )29            embedding_ids = [embedding_id[0] for embedding_id in embedding_ids]30        except NotFound:31            break32        if embedding_ids:33            for embedding_id in embedding_ids:34                db.session.execute(35                    text("DELETE FROM embeddings WHERE id = :embedding_id"), {"embedding_id": embedding_id}36                )37 38            db.session.commit()39        else:40            break41    end_at = time.perf_counter()42    click.echo(click.style("Cleaned embedding cache from db success latency: {}".format(end_at - start_at), fg="green"))43