quivr/backend/celery_task.py

20 lines
783 B
Python
Raw Normal View History

2023-09-15 16:52:06 +03:00
from celery import shared_task
from models.settings import get_supabase_db
from modules.brain.service.brain_vector_service import BrainVectorService
from packages.embeddings.vectors import Neurons
2023-09-15 16:52:06 +03:00
from repository.files.upload_file import DocumentSerializable
@shared_task
def create_embedding_for_document(brain_id, doc_with_metadata, file_sha1):
2023-09-15 16:52:06 +03:00
neurons = Neurons()
doc = DocumentSerializable.from_json(doc_with_metadata)
created_vector = neurons.create_vector(doc)
database = get_supabase_db()
database.set_file_sha_from_metadata(file_sha1)
2023-09-15 16:52:06 +03:00
created_vector_id = created_vector[0] # pyright: ignore reportPrivateUsage=none
brain_vector_service = BrainVectorService(brain_id)
brain_vector_service.create_brain_vector(created_vector_id, file_sha1)