Source code for nlpmed_portal.nlp.tasks

# SPDX-FileCopyrightText: Copyright (C) 2026 Omid Jafari <omidjafari.com>
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


from django.db import transaction
from django.db.models import F

from nlpmed_portal.annotations.models import Patient
from nlpmed_portal.nlp.models import NlpProcessJob
from nlpmed_portal.nlp.models import NLPSetting
from nlpmed_portal.nlp.utils import process_patient_api_call
from nlpmed_portal.nlp.utils import process_patient_write_db

CHUNK_SIZE = 10




def chunked_nlp_master_task(job_id: int) -> None:
    try:
        job = NlpProcessJob.objects.get(pk=job_id)
    except NlpProcessJob.DoesNotExist:
        return

    job.mark_in_progress()

    patient_list = job.patient_pks
    job.total_count = len(patient_list)
    job.save(update_fields=["total_count", "updated_at"])

    project = job.project

    try:
        nlp_setting = NLPSetting.objects.get(project=project)
    except NLPSetting.DoesNotExist:
        job.job_status = "failed"
        job.error_message = f"NLPSetting does not exist for project: {project.id}"
        job.save()
        return

    chunked_lists = [
        patient_list[i : i + CHUNK_SIZE] for i in range(0, len(patient_list), CHUNK_SIZE)
    ]

    for patient_chunk in chunked_lists:
        process_patient_chunk(
            job_id,
            patient_chunk,
            nlp_setting.id,
        )

    finalize_nlp_job(job_id)





def process_patient_chunk(
    job_id: int,
    patient_chunk: list[int],
    nlp_setting_id: int,
) -> None:
    try:  # ruff: ignore[too-many-statements-in-try-clause]
        job = NlpProcessJob.objects.get(pk=job_id)
        if job.job_status == "failed":
            return

        nlp_setting = NLPSetting.objects.get(pk=nlp_setting_id)

        for pat_pk in patient_chunk:
            if job.job_status == "failed":
                break

            try:
                patient = Patient.objects.get(pk=pat_pk)
            except Patient.DoesNotExist:
                continue

            notes = patient.notes.all()

            if not notes.exists():
                continue

            response_json = process_patient_api_call(nlp_setting, pat_pk, notes)

            with transaction.atomic():
                process_patient_write_db(nlp_setting, response_json, notes)

        job.completed_count = F("completed_count") + len(patient_chunk)
        job.save(update_fields=["completed_count", "updated_at"])

    except Exception as exc:
        NlpProcessJob.objects.filter(pk=job_id).update(
            job_status="failed",
            error_message=str(exc),
        )
        raise





def finalize_nlp_job(job_id: int) -> None:
    try:
        job = NlpProcessJob.objects.get(pk=job_id)
        if job.job_status != "failed":
            job.mark_completed()

    except NlpProcessJob.DoesNotExist:
        pass