tjdify/api/schedule/clean_embedding_cache_task.py

import datetime
import time

import click
from sqlalchemy import text
from werkzeug.exceptions import NotFound

import app
from configs import dify_config
from extensions.ext_database import db
from models.dataset import Embedding


@app.celery.task(queue='dataset')
def clean_embedding_cache_task():
    click.echo(click.style('Start clean embedding cache.', fg='green'))
    clean_days = int(dify_config.CLEAN_DAY_SETTING)
    start_at = time.perf_counter()
    thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days)
    while True:
        try:
            embedding_ids = db.session.query(Embedding.id).filter(Embedding.created_at < thirty_days_ago) \
                .order_by(Embedding.created_at.desc()).limit(100).all()
            embedding_ids = [embedding_id[0] for embedding_id in embedding_ids]
        except NotFound:
            break
        if embedding_ids:
            for embedding_id in embedding_ids:
                db.session.execute(text(
                    "DELETE FROM embeddings WHERE id = :embedding_id"
                ), {'embedding_id': embedding_id})

            db.session.commit()
        else:
            break
    end_at = time.perf_counter()
    click.echo(click.style('Cleaned embedding cache from db success latency: {}'.format(end_at - start_at), fg='green'))