integrated celery

README.md

@@ -26,6 +26,12 @@ Additionally, it handles various word forms and multi-term descriptions to maint
 
 ## Running
 
+```
+# Start celery worker
+celery -A tasks worker --loglevel=info
+python run.py
+```
+
 ```
 docker build -t brightly-ai .
 docker run -p 7860:7860 brightly-ai

algo.py

@@ -13,8 +13,10 @@ from db.db_utils import store_mapping_to_db, cached_get_mapping_from_db, get_dic
 from ask_gpt import query_gpt
 from multi_food_item_detector import extract_items, has_delimiters
 from mapping_template import empty_template
+from tasks import insert_result_async
 
-logging.basicConfig(level=logging.WARNING, format='%(asctime)s - %(levelname)s - %(message)s')
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 similarity_threshold = 0.75
 
 
@@ -385,7 +387,9 @@ class Algo:
                 mapping.update({
                     'run_row': input_row_num
                 })
-                results.append(mapping)
-                store_result_to_db(self.db_cursor, self.db_conn, self.run_key, mapping)
+                insert_result_async.delay(self.run_key, mapping)
+                # TODO: push this to celery queue
+                # results.append(mapping)
+                # store_result_to_db(self.db_cursor, self.db_conn, self.run_key, mapping)
 
         return results

requirements.txt

@@ -1,5 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
 autocorrect==2.6.1
+celery==5.4.0
 gradio==4.36.1
 mistralai==0.4.0
 numpy==1.26.4
@@ -10,6 +11,7 @@ psutil==6.0.0
 psycopg2-binary==2.9.9
 python-dotenv==1.0.1
 python-Levenshtein==0.25.1
+redis==5.0.7
 requests==2.32.3
 sentence_transformers==3.0.1
 spacy==3.7.5

tasks.py

@@ -0,0 +1,24 @@
+# tasks.py
+import logging
+from celery import Celery
+from db.db_utils import get_connection, store_result_to_db
+
+
+app = Celery('tasks', broker='redis://localhost:6379/0', backend='redis://localhost:6379/0')
+
+app.conf.update(
+    result_expires=3600,
+    task_serializer='json',
+    result_serializer='json',
+    accept_content=['json'],
+    timezone='UTC',
+    enable_utc=True,
+)
+
+@app.task
+def insert_result_async(run_key, mapping):
+    logging.info(f"Inserting result for run_key: {run_key}, mapping: {mapping['input_word']}")
+    db_conn = get_connection()
+    db_cursor = db_conn.cursor()
+    store_result_to_db(db_cursor, db_conn, run_key, mapping)
+    db_conn.close()