trauma/api/message/utils.py

import json

from trauma.api.data.model import EntityModel
from trauma.api.message.model import MessageModel


def transform_messages_to_openai(messages: list[MessageModel]) -> list[dict]:
    openai_messages = []
    for message in messages:
        content = message.text
        openai_messages.append({
            "role": message.author.value,
            "content": content
        })

    return openai_messages


def retrieve_empty_field_from_entity_data(entity_data: dict) -> str | None:
    for k, v in entity_data.items():
        if not v:
            return k
    return None


def prepare_user_messages_str(user_message: str, messages: list[dict]) -> str:
    user_message_str = ''
    for message in messages:
        if message['role'] == 'user':
            user_message_str += f'- {message['content']}\n'
    user_message_str += f'- {user_message}'
    return user_message_str


def prepare_final_entities_str(entities: list[EntityModel]) -> str:
    entities_list = []
    for entity in entities:
        entities_list.append(entity.model_dump(mode='json', exclude={'id', 'contactDetails'}))
    return json.dumps({"entities": entities_list})

def pick_empty_field_instructions(empty_field: str) -> str:
    if empty_field == "ageMin":
        return "Minimum age of the patient."
    elif empty_field == "ageMax":
        return "Maximum age of the patient."
    elif empty_field == "treatmentArea":
        return "The type of mental or physical illness/disorder."
    elif empty_field == "treatmentMethod":
        return "A method of treating the illness or disorder."