Refactor message handling and AI prompts for clarity.

Streamlined AI prompt generation by replacing list-based message history with a structured string representation. Simplified assistant-user message saving and updated related utility methods for better readability. Enhanced context in invalid and recommendation responses.

trauma/api/message/ai/engine.py

@@ -18,8 +18,12 @@ from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
 from trauma.api.message.db_requests import (save_assistant_user_message,
                                             filter_entities_by_age_location,
                                             update_entity_data_obj, get_entity_by_index)
+from trauma.api.message.dto import Author
+from trauma.api.message.model import MessageModel
 from trauma.api.message.schemas import CreateMessageResponse
-from trauma.api.message.utils import (retrieve_empty_field_from_entity_data,
+from trauma.api.message.utils import (decode_treatment_letters, 
+                                      prepare_message_history_str,
+                                      retrieve_empty_field_from_entity_data,
                                       prepare_user_messages_str,
                                       prepare_final_entities_str,
                                       pick_empty_field_instructions, find_matching_age_group)
@@ -27,33 +31,44 @@ from trauma.core.config import settings
 
 
 async def search_entities(
-        user_message: str, messages: list[dict], chat: ChatModel
+        user_message: str, messages: list[MessageModel], chat: ChatModel
 ) -> CreateMessageResponse:
+    decoded_message = decode_treatment_letters(user_message)
+    message_history_str = prepare_message_history_str(messages, decoded_message)
+
     entity_data, is_valid = await asyncio.gather(
-        update_entity_data_with_ai(chat.entityData, user_message, messages[-1]['content']),
-        check_is_valid_request(user_message, "\n".join([f"- [{i['role']}]: {i['content']}." for i in messages]))
+        update_entity_data_with_ai(chat.entityData, decoded_message, messages[-1].text),
+        check_is_valid_request(decoded_message, message_history_str)
     )
     final_entities = None
+    
     if not is_valid:
-        response = await generate_invalid_response(user_message, messages)
+        empty_field = retrieve_empty_field_from_entity_data(chat.entityData.model_dump(mode='json'))
+        response = await generate_invalid_response(decoded_message, message_history_str, empty_field)
+        final_entities = messages[-1].entities if messages else None
+
     else:
         asyncio.create_task(update_entity_data_obj(entity_data, chat.id))
         empty_field = retrieve_empty_field_from_entity_data(entity_data)
+     
         if empty_field:
             empty_field_instructions = pick_empty_field_instructions(empty_field)
-            response = await generate_next_question(empty_field, empty_field_instructions, user_message, messages)
+            response = await generate_next_question(empty_field, empty_field_instructions, message_history_str)
+      
         else:
-            user_messages_str = prepare_user_messages_str(user_message, messages)
+            user_messages_str = prepare_user_messages_str(decoded_message, messages)
             possible_entity_indexes, search_request = await asyncio.gather(
                 filter_entities_by_age_location(entity_data),
                 generate_search_request(user_messages_str, entity_data)
             )
             final_entities = await search_semantic_entities(search_request, entity_data, possible_entity_indexes)
             final_entities_str = prepare_final_entities_str(final_entities)
-            response = await generate_final_response(final_entities_str, user_message, messages)
+            response = await generate_final_response(final_entities_str, decoded_message, message_history_str)
 
-    asyncio.create_task(save_assistant_user_message(user_message, final_entities, response, chat.id))
-    return CreateMessageResponse(text=response, entities=final_entities)
+    user_message = MessageModel(chatId=chat.id, author=Author.User, text=decoded_message)
+    assistant_message = MessageModel(chatId=chat.id, author=Author.Assistant, text=response, entities=final_entities)
+    asyncio.create_task(save_assistant_user_message(user_message, assistant_message))
+    return assistant_message
 
 
 async def search_semantic_entities(

trauma/api/message/ai/openai_request.py

@@ -23,18 +23,14 @@ async def update_entity_data_with_ai(entity_data: EntityData, user_message: str,
 
 
 @openai_wrapper(temperature=0.8)
-async def generate_next_question(empty_field: str, instructions: str, user_message: str, message_history: list[dict]):
+async def generate_next_question(empty_field: str, instructions: str, message_history_str: str):
     messages = [
         {
             "role": "system",
             "content": TraumaPrompts.generate_next_question
             .replace("{empty_field}", empty_field)
             .replace("{instructions}", instructions)
-        },
-        *message_history,
-        {
-            "role": "user",
-            "content": user_message
+            .replace("{message_history}", message_history_str)
         }
     ]
     return messages
@@ -54,7 +50,7 @@ async def generate_search_request(user_messages_str: str, entity_data: dict):
 
 
 @openai_wrapper(temperature=0.8)
-async def generate_final_response(final_entities: str, user_message: str, message_history: list[dict]):
+async def generate_final_response(final_entities: str, user_message: str, message_history_str: str):
     if not json.loads(final_entities)['klinieken']:
         prompt = TraumaPrompts.generate_empty_recommendations
     else:
@@ -63,12 +59,8 @@ async def generate_final_response(final_entities: str, user_message: str, messag
         {
             "role": "system",
             "content": prompt
-
-        },
-        *message_history,
-        {
-            "role": "user",
-            "content": user_message
+            .replace("{message_history}", message_history_str)
+            .replace("{user_message}", user_message)
         }
     ]
     return messages
@@ -123,16 +115,20 @@ async def check_is_valid_request(user_message: str, message_history: str):
 
 
 @openai_wrapper(temperature=0.9)
-async def generate_invalid_response(user_message: str, message_history: list[dict]):
+async def generate_invalid_response(user_message: str, message_history_str: str, empty_field: str | None):
+    from trauma.api.message.utils import pick_empty_field_instructions
+
+    if empty_field:
+        empty_field_instructions = pick_empty_field_instructions(empty_field)
+        prompt = TraumaPrompts.generate_invalid_response.replace("{empty_field}", empty_field_instructions)
+    else:
+        prompt = TraumaPrompts.generate_invalid_response_with_recs
     messages = [
         {
             "role": "system",
-            "content": TraumaPrompts.generate_invalid_response
-        },
-        *message_history,
-        {
-            "role": "user",
-            "content": user_message
+            "content": prompt
+            .replace("{message_history}", message_history_str)
+            .replace("{user_message}", user_message)
         }
     ]
     return messages

trauma/api/message/ai/prompts.py

@@ -28,9 +28,9 @@ Je verzamelt informatie over een patiënt, hun ziekte en de behandelmethode zoda
 
 ```json
 {
-    “age”: integer,
-    “treatmentArea”: “string”,
-    “treatmentMethod”: “string”,
+    "age": integer,
+    "treatmentArea": "string",
+    "treatmentMethod": "string",
     "location": "string",
     "postalCode": "string",
 }
@@ -40,7 +40,7 @@ Je verzamelt informatie over een patiënt, hun ziekte en de behandelmethode zoda
 - **treatmentArea**: Het type mentale of fysieke ziekte/stoornis.
 - **treatmentMethod**: Een methode voor het behandelen van de ziekte of stoornis.
 - **location**: Stad of adres waar de facility zich bevindt  
-- **postalCode**: Postcode van de facility..
+- **postalCode**: Postcode van de facility.
 
 ## Regels voor het bijwerken van Entity Data
 
@@ -49,11 +49,12 @@ Je verzamelt informatie over een patiënt, hun ziekte en de behandelmethode zoda
 - Werk alleen de velden bij die expliciet door de gebruiker zijn genoemd."""
     generate_next_question = """## Taak
 
-Stel een korte, gerichte vraag aan de gebruiker om informatie te verkrijgen over het veld: [Gemist veld].
+Stel de gebruiker vriendelijk een gerichte vraag over het gemiste veld (`Gemist veld`) om de noodzakelijke informatie te verkrijgen voor het aanbevelen van een medische instelling.
 
 ## Context
 
-Je bent een informatieanalist die gegevens verzamelt van een verpleegkundige over een patiënt. Deze gegevens worden gebruikt door het systeem om een passende kliniek te adviseren. Stel vragen die specifiek en gericht zijn op relevante kenmerken, met bijzondere aandacht voor het [Gemist veld].
+Je bent een informaticus die gegevens verzamelt over de patiënt en de gewenste kliniek. Deze gegevens worden door het systeem gebruikt om een geschikte kliniek aan te bevelen.
+
 ## Gegevens
 
 **Gemist veld**:
@@ -61,12 +62,16 @@ Je bent een informatieanalist die gegevens verzamelt van een verpleegkundige ove
 {empty_field}
 ```
 
-- {instructions}
+- {instructies}
+
+**Gespreksgeschiedenis**:
+```
+{message_history}
+```
 
-## Belangrijke Opmerking
+## Belangrijke opmerking
 
-- De gegenereerde vraag moet kort, to the point, en uitnodigend zijn.
-- Gebruik een empathische toon, bijvoorbeeld: "Kun je me iets meer vertellen over [Gemist veld]? Dit helpt me om je beter te helpen."""
+- De geformuleerde vraag moet beknopt zijn en een empathische toon hebben."""
     generate_search_request = """## Taak
 
 Je moet een beknopte zoekopdracht genereren op basis van de berichten van de gebruiker [berichtgeschiedenis] en de verzamelde patiëntgegevens [patiëntgegevens].
@@ -113,7 +118,7 @@ Je moet een antwoord genereren aan de gebruiker waarin je aangeeft dat je geschi
 - Geweldig nieuws! Ik heb klinieken gevonden die perfect passen bij de aangegeven behoeften. Deze klinieken zijn gespecialiseerd in [treatmentMethod] en werken met patiënten in de leeftijd van [ageMin] tot [ageMax]."""
     decide_is_valid_request = """## Task
 
-You must determine whether the user's request is valid. Provide your answer in the [is_valid] field in JSON format.
+You must determine whether the user's request pertains to the topic of medical facility recommendations or not. Use `User Request` to determine and `Message History` as context for making the decision.
 
 ## Data
 
@@ -139,21 +144,57 @@ You must determine whether the user's request is valid. Provide your answer in t
 
 ## Instructions for Filling JSON
 
-The field is considered valid (`is_valid = true`) if:
-- The user describes the patient, their data, illness, treatment method, etc.
-- The user's request relates to a medical topic.
-- The user's request is a valid response to the assistant's question.
-- The user's request describes desired facility.
+`User Request` pertains to the topic of medical facility recommendations (`is_valid = true`), if:
+- The user describes the patient, for example, age, illness, treatment method, etc.
+- The user describes the medical facility, for example, its location, postal code, treatment methods, supported age groups, etc.
 
 [/INST]"""
     generate_invalid_response = """## Taak
 
-Je moet een antwoord genereren voor de gebruiker waarin je aangeeft dat hun verzoek buiten jouw specificatie valt voor het verzamelen van informatie en het geven van aanbevelingen.
+Je moet de gebruiker vriendelijk laten weten dat hun verzoek (`Gebruikersquery`) niet betrekking heeft op het onderwerp van medische instellingen aanbevelingen, en empathisch doorverwijzen zodat ze de verplichte informatie (`Required field`) kunnen delen. Gebruik (`Gesprekshistorie`) om de toon en stijl van communicatie te behouden.
+
+## Gegevens
+
+**Gebruikersquery**:
+```
+{user_message}
+```
+
+**Gesprekshistorie**:
+```
+{message_history}
+```
+
+**Verplichte veld**:
+```
+{empty_field}
+```
 
 ## Belangrijke opmerkingen
 
-- Je antwoord moet kort en bondig zijn, bestaande uit twee zinnen.  
-- Gebruik een empathische toon en moedig de gebruiker aan om opnieuw informatie te delen, bijvoorbeeld: "Ik begrijp het niet helemaal. Kun je me wat meer vertellen over de patiënt, zodat ik je beter kan helpen?"""
+- Je antwoord moet beknopt, empathisch en vriendelijk zijn.
+- Je antwoord moet in twee zinnen worden gepresenteerd."""
+    generate_invalid_response_with_recs = """## Taak
+
+Je moet op een beleefde manier verschillende medische instellingen voorstellen op basis van de informatie die de gebruiker heeft verstrekt, en ook toevoegen dat als er aanvullende informatie is, de gebruiker altijd kan terugkomen om zijn behoeften verder te verduidelijken en nieuwe aanbevelingen te krijgen.
+
+## Gegevens
+
+**Gebruikersquery**:
+```
+{user_message}
+```
+
+**Gesprekshistorie**:
+```
+{message_history}
+```
+
+## Belangrijke notities
+
+- Je antwoord moet beknopt, empathisch en vriendelijk zijn.
+- Gebruik de `Gesprekshistorie` om de toon en stijl van de communicatie te behouden.
+- Je antwoord moet in twee zinnen zijn."""
     generate_empty_recommendations = """## Taak
 
 Je moet de gebruiker op een empathische en ondersteunende manier informeren dat er geen geschikte klinieken voor de patiënt zijn gevonden. Moedig de gebruiker aan om de informatie opnieuw in te voeren met aanvullende details om betere zoekresultaten te genereren. Geef hierbij een voorbeeld van een geschikte invoer, zoals: 'jongen van 16 jaar met LVB waarvoor ik EMDR zoek'.
@@ -162,6 +203,18 @@ Je moet de gebruiker op een empathische en ondersteunende manier informeren dat
 
 De gebruiker zoekt naar een geschikte kliniek voor een patiënt en deelt hierbij details zoals ziekte, leeftijd en behandelmethoden. Ondanks deze informatie heeft het systeem geen geschikte kliniek kunnen vinden die voldoet aan de eisen. Het is jouw taak om dit op een vriendelijke manier te communiceren en suggesties te geven voor het herzien of aanvullen van de ingevoerde gegevens.
 
+## Data
+
+**Gebruikersquery**:
+```
+{user_message}
+```
+
+**Gesprekshistorie**:
+```
+{message_history}
+```
+
 ## Belangrijke opmerkingen
 
 - Gebruik een vriendelijke en geruststellende toon, bijvoorbeeld: "Ik heb op basis van de ingevoerde gegevens geen kliniek kunnen vinden. Kunnen we samen kijken of we de informatie iets kunnen aanpassen om betere resultaten te krijgen?"

trauma/api/message/db_requests.py

@@ -54,10 +54,8 @@ async def update_entity_data_obj(entity_data: dict, chat_id: str) -> None:
 
 @background_task()
 async def save_assistant_user_message(
-        user_message: str, final_entities: list[EntityModelExtended], assistant_message: str, chat_id: str
+        user_message: MessageModel, assistant_message: MessageModel
 ) -> None:
-    user_message = MessageModel(chatId=chat_id, author=Author.User, entities=final_entities, text=user_message)
-    assistant_message = MessageModel(chatId=chat_id, author=Author.Assistant, text=assistant_message)
     await settings.DB_CLIENT.messages.insert_one(user_message.to_mongo())
     await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())
 
@@ -99,4 +97,4 @@ async def update_message_feedback_obj(message_id: str, feedback_data: Feedback)
     message.feedback = feedback_data
     await settings.DB_CLIENT.messages.update_one({"id": message_id},
                                                  {"$set": {"feedback": feedback_data.model_dump(mode='json')}})
-    return message
+    return message

trauma/api/message/utils.py

@@ -1,7 +1,9 @@
 import json
+import re
 
 from trauma.api.data.dto import AgeGroup
 from trauma.api.data.model import EntityModel
+from trauma.api.message.dto import Author
 from trauma.api.message.model import MessageModel
 
 
@@ -19,16 +21,24 @@ def transform_messages_to_openai(messages: list[MessageModel]) -> list[dict]:
 
 def retrieve_empty_field_from_entity_data(entity_data: dict) -> str | None:
     for k, v in entity_data.items():
-        if k not in ('location', 'postalCode') and not v:
+        if k not in ('treatmentArea', 'postalCode') and not v:
             return k
     return None
 
+def prepare_message_history_str(messages: list[MessageModel], user_message: str) -> str:
+    results = ""
+    shorted_message_history = messages[-8:]
+    for message in shorted_message_history:
+        results += f"[{message.author.value}] {message.text}\n"
+    results += f"[User] {user_message}"
+    return results
 
-def prepare_user_messages_str(user_message: str, messages: list[dict]) -> str:
+
+def prepare_user_messages_str(user_message: str, messages: list[MessageModel]) -> str:
     user_message_str = ''
     for message in messages:
-        if message['role'] == 'user':
-            user_message_str += f"- {message['content']}\n"
+        if message.author == Author.User:
+            user_message_str += f"- {message.text}\n"
     user_message_str += f'- {user_message}'
     return user_message_str
 
@@ -45,10 +55,10 @@ def prepare_final_entities_str(entities: list[EntityModel]) -> str:
 def pick_empty_field_instructions(empty_field: str) -> str:
     if empty_field == "age":
         return "Leeftijd van de patiënt."
-    elif empty_field == "treatmentArea":
-        return "Het type psychische of lichamelijke ziekte / stoornis."
     elif empty_field == "treatmentMethod":
         return "Een methode om de ziekte of stoornis te behandelen."
+    elif empty_field == "location":
+        return "Stad of adres waar de facility zich bevindt."
 
 
 def find_matching_age_group(entity: EntityModel, entity_data: dict) -> AgeGroup:
@@ -61,3 +71,32 @@ def find_matching_age_group(entity: EntityModel, entity_data: dict) -> AgeGroup:
             if best_match is None or age_group.ageMax > best_match.ageMax:
                 best_match = age_group
     return best_match
+
+def decode_treatment_letters(user_message: str) -> str:
+    replacements = {
+        "lvb": "Licht Verstandelijke Beperking",
+        "emdr": "Eye Movement Desensitization and Reprocessing",
+        "ptss": "Posttraumatische Stress Stoornis",
+        "cgt": "Cognitieve Gedragstherapie",
+        "ass": "Autisme Spectrum Stoornis",
+        "adhd": "Attention Deficit Hyperactivity Disorder",
+        "add": "Attention Deficit Disorder",
+        "bps": "Borderline Persoonlijkheidsstoornis",
+        "gad": "Gegeneraliseerde Angststoornis",
+        "ocd": "Obsessieve-Compulsieve Stoornis",
+        "an": "Anorexia Nervosa",
+        "bed": "Eetbuistoornis",
+        "solk": "Somatisch Onvoldoende verklaarde Lichamelijke Klachten",
+        "dis": "Dissociatieve Identiteitsstoornis",
+        "nah": "Niet Aangeboren Hersenletsel",
+        "sas": "Stemmings-, angst- en somatoforme stoornissen",
+        "epa": "Ernstige Psychiatrische Aandoeningen",
+        "net": "Narratieve Exposure Therapie",
+        "bepp": "Beknopte Eclectische Psychotherapie voor PTSS",
+        "dbt": "Dialectische Gedragstherapie",
+        "act": "Acceptance and Commitment Therapy",
+        "iht": "Intensive Home Treatment"
+    }
+    for short, full in replacements.items():
+        user_message = re.sub(rf'\b{short}\b', full, user_message, flags=re.IGNORECASE)
+    return user_message

trauma/api/message/views.py

@@ -35,10 +35,9 @@ async def create_message(
         chatId: str,
         message_data: CreateMessageRequest,
         account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
-) -> TraumaResponseWrapper[CreateMessageResponse]:
+) -> TraumaResponseWrapper[MessageModel]:
     messages, chat = await get_all_chat_messages_obj(chatId, account)
-    message_history = transform_messages_to_openai(messages)
-    response = await search_entities(message_data.text, message_history, chat)
+    response = await search_entities(message_data.text, messages, chat)
     return TraumaResponseWrapper(data=response)