heuristic.py

from fast_api.models import CreateHeuristicBody, UpdateHeuristicBody
from fast_api.routes.client_response import pack_json_result, get_silent_success
from fastapi import APIRouter, Body, Depends, Request
from controller.information_source import manager
from submodules.model.business_objects import weak_supervision
from controller.auth import manager as auth_manager
from controller.payload import manager as payload_manager
from submodules.model.business_objects import information_source
from submodules.model.business_objects.payload import get_payload_with_heuristic_type
from submodules.model.util import sql_alchemy_to_dict
from util import notification
from controller.task_master import manager as task_master_manager
from submodules.model import enums
from controller.auth import kratos

router = APIRouter()


@router.get(
    "/{project_id}/information-sources-overview-data",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def get_information_sources_overview_data(
    request: Request,
    project_id: str,
):
    data = manager.get_overview_data(project_id)
    return pack_json_result(data)


@router.get(
    "/{project_id}/weak-supervision-run",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def get_weak_supervision_run(
    request: Request,
    project_id: str,
):
    result = weak_supervision.get_current_weak_supervision_run(project_id)
    if result is None:
        return pack_json_result(None)

    user_id = auth_manager.get_user_id_by_info(request.state.info)
    names, mail = kratos.resolve_user_name_and_email_by_id(user_id)
    first_name = names.get("first", "")
    last_name = names.get("last", "")

    data = {
        "id": str(result.id),
        "state": result.state,
        "createdAt": result.created_at,
        "finishedAt": result.finished_at,
        "selectedInformationSources": result.selected_information_sources,
        "selectedLabelingTasks": result.selected_labeling_tasks,
        "distinctRecords": result.distinct_records,
        "result_count": result.result_count,
        "user": {
            "id": str(user_id),
            "firstName": first_name,
            "lastName": last_name,
            "email": mail,
        },
    }

    return pack_json_result(data)


@router.get(
    "/{project_id}/{heuristic_id}/heuristic-by-id",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def get_heuristic_by_heuristic_id(
    project_id: str,
    heuristic_id: str,
):
    data = sql_alchemy_to_dict(
        information_source.get_heuristic_id_with_most_recent_payload(
            project_id, heuristic_id
        )
    )
    statistics = sql_alchemy_to_dict(
        information_source.get_source_statistics(project_id, heuristic_id)
    )
    data["sourceStatistics"] = statistics
    return pack_json_result(data)


@router.get(
    "/{project_id}/{payload_id}/payload-by-id",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def get_payload_by_payload_id(
    project_id: str,
    payload_id: str,
):
    data = sql_alchemy_to_dict(get_payload_with_heuristic_type(project_id, payload_id))
    return pack_json_result(data)


@router.get(
    "/{project_id}/{heuristic_id}/lf-on-10-records",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def get_labeling_function_on_10_records(
    project_id: str,
    heuristic_id: str,
):
    data = payload_manager.get_labeling_function_on_10_records(project_id, heuristic_id)
    return pack_json_result(data, wrap_for_frontend=False)


@router.post(
    "/{project_id}/{information_source_id}/toggle-heuristic",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def toggle_heuristic(
    request: Request,
    project_id: str,
    information_source_id: str,
):
    manager.toggle_information_source(project_id, information_source_id)
    notification.send_organization_update(
        project_id, f"information_source_updated:{information_source_id}"
    )
    return get_silent_success()


@router.post(
    "/{project_id}/change-selection-state",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def set_information_sources(
    request: Request,
    project_id: str,
    value: bool,
):
    manager.set_all_information_source_selected(project_id, value)
    notification.send_organization_update(project_id, "information_source_updated:all")
    return get_silent_success()


@router.post(
    "/{project_id}/{heuristic_id}/payload",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def set_payload(
    request: Request,
    project_id: str,
    heuristic_id: str,
):
    user = auth_manager.get_user_by_info(request.state.info)
    org_id = user.organization_id
    information_source_item = information_source.get(project_id, heuristic_id)
    priority = True

    task_master_response = task_master_manager.queue_task(
        str(org_id),
        str(user.id),
        enums.TaskType.INFORMATION_SOURCE,
        {
            "project_id": str(project_id),
            "information_source_id": str(heuristic_id),
            "source_type": information_source_item.type,
        },
        priority=priority,
    )
    queue_id = None
    if task_master_response.ok:
        queue_id = task_master_response.json().get("task_id")

    return pack_json_result({"queueId": queue_id})


@router.delete(
    "/{project_id}/{heuristic_id}/delete-heuristic",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def delete_heuristic(
    request: Request,
    project_id: str,
    heuristic_id: str,
):
    manager.delete_information_source(project_id, heuristic_id)
    notification.send_organization_update(
        project_id, f"information_source_deleted:{heuristic_id}"
    )
    return get_silent_success()


@router.post(
    "/{project_id}/create-heuristic",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def create_heuristic(
    request: Request,
    project_id: str,
    body: CreateHeuristicBody = Body(...),
):
    user = auth_manager.get_user_by_info(request.state.info)
    info_source = manager.create_information_source(
        project_id,
        user.id,
        body.labeling_task_id,
        body.name,
        body.source_code,
        body.description,
        body.type,
    )
    notification.send_organization_update(
        project_id, f"information_source_created:{str(info_source.id)}"
    )
    return pack_json_result({"id": str(info_source.id)}, wrap_for_frontend=False)


@router.post(
    "/{project_id}/{heuristic_id}/update-heuristic",
    dependencies=[Depends(auth_manager.check_project_access_dep)],
)
def update_heuristic(
    request: Request,
    project_id: str,
    heuristic_id: str,
    body: UpdateHeuristicBody = Body(...),
):
    manager.update_information_source(
        project_id,
        heuristic_id,
        body.labeling_task_id,
        body.code,
        body.description,
        body.name,
    )
    user = auth_manager.get_user_by_info(request.state.info)
    notification.send_organization_update(
        project_id, f"information_source_updated:{heuristic_id}:{user.id}"
    )
    return get_silent_success()