Refactor message handling to support optional fields.

Added support for handling missing or optional fields when generating prompts, responding to user inputs, and selecting treatment areas or methods. Updated utilities and enhanced prompt templates to improve flexibility and user interaction.

trauma/api/message/ai/engine.py

@@ -50,11 +50,10 @@ async def search_entities(
     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)
+        empty_field_instructions = pick_empty_field_instructions(empty_field)
+
+        if empty_field == 'age':
             response = await generate_next_question(empty_field, empty_field_instructions, message_history_str)
-      
         else:
             user_messages_str = prepare_user_messages_str(decoded_message, messages)
             possible_entity_indexes, search_request = await asyncio.gather(
@@ -63,7 +62,9 @@ async def search_entities(
             )
             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, decoded_message, message_history_str)
+            response = await generate_final_response(
+                final_entities_str, decoded_message, message_history_str, empty_field_instructions, empty_field
+            )
 
     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)
@@ -95,8 +96,8 @@ async def extend_entities_with_highlights(entities: list[EntityModel], entity_da
     EntityModelExtended]:
     async def choose_closest(entity_: EntityModel) -> tuple:
         treatment_area, treatment_method = await asyncio.gather(
-            choose_closest_treatment_area(entity_.treatmentAreas, entity_data['treatmentArea']),
-            choose_closest_treatment_method(entity_.treatmentMethods, entity_data['treatmentMethod'])
+            choose_closest_treatment_area(entity_.treatmentAreas, entity_data.get('treatmentArea')),
+            choose_closest_treatment_method(entity_.treatmentMethods, entity_data.get('treatmentMethod'))
         )
         return treatment_area, treatment_method
 

trauma/api/message/ai/openai_request.py

@@ -50,11 +50,18 @@ 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_str: str):
-    if not json.loads(final_entities)['klinieken']:
-        prompt = TraumaPrompts.generate_empty_recommendations
+async def generate_final_response(
+    final_entities: str, user_message: str, message_history_str: str, empty_field_instructions: str, empty_field: str
+):
+    if empty_field_instructions:
+        prompt = (TraumaPrompts.generate_not_fully_recommendations
+                  .replace("{instructions}", empty_field_instructions)
+                  .replace("{empty_field}", empty_field))
+    elif json.loads(final_entities)['klinieken']:
+        prompt = (TraumaPrompts.generate_recommendation_decision
+                  .replace("{final_entities}", final_entities))
     else:
-        prompt = TraumaPrompts.generate_recommendation_decision.replace("{final_entities}", final_entities)
+        prompt = TraumaPrompts.generate_empty_recommendations
     messages = [
         {
             "role": "system",
@@ -76,7 +83,9 @@ async def convert_value_to_embeddings(value: str, dimensions: int = 1536) -> lis
 
 
 @openai_wrapper(is_json=True, return_='result')
-async def choose_closest_treatment_area(treatment_areas: list[str], treatment_area: str):
+async def choose_closest_treatment_area(treatment_areas: list[str], treatment_area: str | None):
+    if not treatment_area:
+        return None
     messages = [
         {
             "role": "system",
@@ -89,7 +98,9 @@ async def choose_closest_treatment_area(treatment_areas: list[str], treatment_ar
 
 
 @openai_wrapper(is_json=True, return_='result')
-async def choose_closest_treatment_method(treatment_methods: list[str], treatment_method: str):
+async def choose_closest_treatment_method(treatment_methods: list[str], treatment_method: str | None):
+    if not treatment_method:
+        return None
     messages = [
         {
             "role": "system",

trauma/api/message/ai/prompts.py

@@ -62,7 +62,7 @@ Je bent een informaticus die gegevens verzamelt over de patiënt en de gewenste
 {empty_field}
 ```
 
-- {instructies}
+- {instructions}
 
 **Gespreksgeschiedenis**:
 ```
@@ -72,6 +72,31 @@ Je bent een informaticus die gegevens verzamelt over de patiënt en de gewenste
 ## Belangrijke opmerking
 
 - De geformuleerde vraag moet beknopt zijn en een empathische toon hebben."""
+    generate_not_fully_recommendations = """## Taak
+
+Vertel de gebruiker vriendelijk dat hij de aanbevelingen van de medische instelling kan bekijken op basis van de informatie die hij heeft verstrekt, en stel vervolgens de exacte vraag over het ontbrekende veld (`Ontbrekend veld`), met de uitleg dat dit de zoekresultaten zal verbeteren.
+
+## Context
+
+Je bent een informatieve assistent die gegevens verzamelt over de patiënt en de gewenste kliniek. Deze gegevens worden gebruikt door het systeem om een geschikte kliniek aan te bevelen.
+
+## Gegevens
+
+**Ontbrekend veld**:
+```
+{empty_field}
+```
+
+- {instructions}
+
+**Berichtenhistorie**:
+```
+{message_history}
+```
+
+## Belangrijke opmerking
+
+- Verwijs naar de specifieke informatie die de gebruiker al heeft verstrekt (bijvoorbeeld leeftijd en locatie, enz.)."""
     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].

trauma/api/message/utils.py

@@ -21,7 +21,7 @@ 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 ('treatmentArea', 'postalCode') and not v:
+        if k != 'postalCode' and not v:
             return k
     return None
 
@@ -59,7 +59,9 @@ def pick_empty_field_instructions(empty_field: str) -> str:
         return "Een methode om de ziekte of stoornis te behandelen."
     elif empty_field == "location":
         return "Stad of adres waar de facility zich bevindt."
-
+    elif empty_field == "treatmentArea":
+        return "Een gebied waar de facility zich bevindt."
+    return None
 
 def find_matching_age_group(entity: EntityModel, entity_data: dict) -> AgeGroup:
     age_groups = entity.ageGroups

trauma/core/wrappers.py

@@ -53,6 +53,8 @@ def openai_wrapper(
         @wraps(func)
         async def wrapper(*args, **kwargs) -> str:
             messages = await func(*args, **kwargs)
+            if not isinstance(messages, list):
+                return messages
             completion = await settings.OPENAI_CLIENT.chat.completions.create(
                 messages=messages,
                 temperature=temperature,