added highlights

trauma/api/data/dto.py

@@ -10,4 +10,5 @@ class ContactDetails(BaseModel):
     email: str | None = None
     website: str | None = None
     address: str | None = None
-    postalCode: str | None = None
+    postalCode: str | None = None
+

trauma/api/data/model.py

@@ -42,3 +42,9 @@ class EntityModel(MongoBaseModel):
     treatmentMethods: list[str]
     description: str = ''
     contactDetails: ContactDetails
+
+
+class EntityModelExtended(EntityModel):
+    highlightedAgeGroup: AgeGroup
+    highlightedTreatmentArea: str
+    highlightedTreatmentMethod: str

trauma/api/message/ai/engine.py

@@ -1,19 +1,24 @@
 import asyncio
 
+import numpy as np
+
+from trauma.api.chat.dto import EntityData
 from trauma.api.chat.model import ChatModel
+from trauma.api.data.model import EntityModel, EntityModelExtended
 from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
                                                   generate_next_question,
                                                   generate_search_request,
-                                                  generate_final_response)
+                                                  generate_final_response, convert_value_to_embeddings,
+                                                  choose_closest_treatment_method, choose_closest_treatment_area)
 from trauma.api.message.db_requests import (save_assistant_user_message,
                                             filter_entities_by_age,
-                                            search_semantic_entities,
-                                            update_entity_data_obj)
+                                            update_entity_data_obj, get_entity_by_index)
 from trauma.api.message.schemas import CreateMessageResponse
 from trauma.api.message.utils import (retrieve_empty_field_from_entity_data,
                                       prepare_user_messages_str,
                                       prepare_final_entities_str,
-                                      pick_empty_field_instructions)
+                                      pick_empty_field_instructions, find_matching_age_group)
+from trauma.core.config import settings
 
 
 async def search_entities(
@@ -34,9 +39,44 @@ async def search_entities(
             filter_entities_by_age(entity_data),
             generate_search_request(user_messages_str, entity_data)
         )
-        final_entities = await search_semantic_entities(search_request, possible_entity_indexes)
+        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)
 
     asyncio.create_task(save_assistant_user_message(user_message, response, chat.id))
     return CreateMessageResponse(text=response, entities=final_entities)
+
+
+async def search_semantic_entities(
+        search_request: str, entity_data: EntityData, entities_indexes: list[int]
+) -> list[EntityModelExtended]:
+    embedding = await convert_value_to_embeddings(search_request)
+    query_embedding = np.array([embedding], dtype=np.float32)
+    distances, indices = settings.SEMANTIC_INDEX.search(query_embedding, k=settings.SEMANTIC_INDEX.ntotal)
+    distances = distances[0]
+    indices = indices[0]
+    filtered_results = [
+        {"index": int(idx), "distance": float(dist)}
+        for idx, dist in zip(indices, distances)
+        if idx in entities_indexes and dist <= 1.3
+    ]
+    filtered_results = sorted(filtered_results, key=lambda x: x["distance"])[:5]
+    final_entities = await asyncio.gather(*[get_entity_by_index(i['index']) for i in filtered_results])
+    final_entities_extended = await asyncio.gather(
+        *[extended_entity_with_highlights(entity, entity_data) for entity in final_entities]
+    )
+    return final_entities_extended
+
+
+async def extended_entity_with_highlights(entity: EntityModel, entity_data: dict) -> EntityModelExtended:
+    age_group = find_matching_age_group(entity, entity_data)
+    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'])
+    )
+    return EntityModelExtended(
+        **entity.to_mongo(),
+        highlightedAgeGroup=age_group,
+        highlightedTreatmentArea=treatment_area,
+        highlightedTreatmentMethod=treatment_method
+    )

trauma/api/message/ai/openai_request.py

@@ -74,4 +74,30 @@ async def convert_value_to_embeddings(value: str) -> list[float]:
         model='text-embedding-3-large',
         dimensions=1536,
     )
-    return embeddings.data[0].embedding
+    return embeddings.data[0].embedding
+
+
+@openai_wrapper(is_json=True, return_='result')
+async def choose_closest_treatment_area(treatment_areas: list[str], treatment_area: str):
+    messages = [
+        {
+            "role": "system",
+            "content": TraumaPrompts.choose_closest_treatment_area
+            .replace("{treatment_areas}", ", ".join(treatment_areas))
+            .replace("{treatment_area}", treatment_area)
+        }
+    ]
+    return messages
+
+
+@openai_wrapper(is_json=True, return_='result')
+async def choose_closest_treatment_method(treatment_methods: list[str], treatment_method: str):
+    messages = [
+        {
+            "role": "system",
+            "content": TraumaPrompts.choose_closest_treatment_method
+            .replace("{treatment_methods}", ", ".join(treatment_methods))
+            .replace("{treatment_method}", treatment_method)
+        }
+    ]
+    return messages

trauma/api/message/ai/prompts.py

@@ -137,3 +137,60 @@ Je bent verplicht om een beschrijving voor een kliniek te genereren op basis van
 - De beschrijving moet beknopt en bondig zijn.
 
 [/INST]"""
+
+    choose_closest_treatment_area = """
+
+## Task
+
+You must determine the most semantically similar disorder or disease from the list of [treatment areas] to the requested disease [requested treatment area]. The most similar disease should be returned in the [result] field of the JSON.
+
+## Data
+
+**treatment areas**:
+```
+{treatment_areas}
+```
+
+**requested treatment area**:
+```
+{treatment_area}
+```
+
+## JSON Response format
+
+```json
+{
+    "result": "string"
+}
+```
+
+## Instructions for filling JSON
+
+- [result]: The item from the [treatment areas] list that is most semantically similar to the requested disease. The disease name in the result field must exactly match the name as it appears in the [treatment areas] list."""
+    choose_closest_treatment_method = """## Task
+
+You must determine the most semantically similar treatment method from the list of [treatment methods] to the requested treatment method [requested treatment method]. The most similar treatment method should be returned in the [result] field of the JSON.
+
+## Data
+
+**treatment methods**:
+```
+{treatment_methods}
+```
+
+**requested treatment method**:
+```
+{treatment_method}
+```
+
+## JSON Response format
+
+```json
+{
+    "result": "string"
+}
+```
+
+## Instructions for filling JSON
+
+- [result]: The item from the [treatment methods] list that is most semantically similar to the requested treatment method. The treatment method name in the result field must exactly match the name as it appears in the [treatment methods] list."""

trauma/api/message/db_requests.py

@@ -1,11 +1,9 @@
 import asyncio
 
-import numpy as np
 from fastapi import HTTPException
 
 from trauma.api.chat.model import ChatModel
 from trauma.api.data.model import EntityModel
-from trauma.api.message.ai.openai_request import convert_value_to_embeddings
 from trauma.api.message.dto import Author
 from trauma.api.message.model import MessageModel
 from trauma.api.message.schemas import CreateMessageRequest
@@ -58,12 +56,12 @@ async def save_assistant_user_message(user_message: str, assistant_message: str,
     await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())
 
 
-async def filter_entities_by_age(entity: dict) -> list[int]:
+async def filter_entities_by_age(entity_data: dict) -> list[int]:
     query = {
         "ageGroups": {
             "$elemMatch": {
-                "ageMin": {"$lte": entity['ageMax']},
-                "ageMax": {"$gte": entity['ageMin']}
+                "ageMin": {"$lte": entity_data['ageMax']},
+                "ageMax": {"$gte": entity_data['ageMin']}
             }
         }
     }
@@ -74,18 +72,3 @@ async def filter_entities_by_age(entity: dict) -> list[int]:
 async def get_entity_by_index(index: int) -> EntityModel:
     entity = await settings.DB_CLIENT.entities.find_one({"index": index})
     return EntityModel.from_mongo(entity)
-
-async def search_semantic_entities(search_request: str, entities_indexes: list[int]) -> list[EntityModel]:
-    embedding = await convert_value_to_embeddings(search_request)
-    query_embedding = np.array([embedding], dtype=np.float32)
-    distances, indices = settings.SEMANTIC_INDEX.search(query_embedding, k=settings.SEMANTIC_INDEX.ntotal)
-    distances = distances[0]
-    indices = indices[0]
-    filtered_results = [
-        {"index": int(idx), "distance": float(dist)}
-        for idx, dist in zip(indices, distances)
-        if idx in entities_indexes and dist <= 1.3
-    ]
-    filtered_results = sorted(filtered_results, key=lambda x: x["distance"])[:5]
-    final_entities = await asyncio.gather(*[get_entity_by_index(i['index']) for i in filtered_results])
-    return final_entities

trauma/api/message/schemas.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel
 
 from trauma.api.common.dto import Paging
-from trauma.api.data.model import EntityModel
+from trauma.api.data.model import EntityModelExtended
 from trauma.api.message.model import MessageModel
 from trauma.core.wrappers import TraumaResponseWrapper
 
@@ -24,4 +24,4 @@ class AllMessageWrapper(TraumaResponseWrapper[AllMessageResponse]):
 
 class CreateMessageResponse(BaseModel):
     text: str
-    entities: list[EntityModel] | None = None
+    entities: list[EntityModelExtended] | None = None

trauma/api/message/utils.py

@@ -1,5 +1,7 @@
 import json
 
+from trauma.api.chat.dto import EntityData
+from trauma.api.data.dto import AgeGroup
 from trauma.api.data.model import EntityModel
 from trauma.api.message.model import MessageModel
 
@@ -35,7 +37,9 @@ def prepare_user_messages_str(user_message: str, messages: list[dict]) -> 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'}))
+        entities_list.append(entity.model_dump(mode='json', exclude={
+            'id', 'contactDetails', "highlightedAgeGroup", "highlightedTreatmentArea", "highlightedTreatmentMethod"
+        }))
     return json.dumps({"klinieken": entities_list})
 
 
@@ -48,3 +52,15 @@ def pick_empty_field_instructions(empty_field: str) -> str:
         return "Het type psychische of lichamelijke ziekte / stoornis."
     elif empty_field == "treatmentMethod":
         return "Een methode om de ziekte of stoornis te behandelen."
+
+
+def find_matching_age_group(entity: EntityModel, entity_data: dict) -> AgeGroup:
+    age_groups = entity.ageGroups
+    best_match = None
+    for age_group in age_groups:
+        if age_group.ageMin <= entity_data['ageMax'] and age_group.ageMax >= entity_data['ageMin']:
+            return age_group
+        if age_group.ageMax < entity_data['ageMin']:
+            if best_match is None or age_group.ageMax > best_match.ageMax:
+                best_match = age_group
+    return best_match

trauma/core/config.py

@@ -16,7 +16,7 @@ class BaseConfig:
     DB_CLIENT = motor.motor_asyncio.AsyncIOMotorClient(os.getenv("MONGO_DB_URL")).AtlasCluster
     OPENAI_CLIENT = AsyncClient(api_key=os.getenv('OPENAI_API_KEY'))
     SEMANTIC_INDEX = faiss.read_index(str(pathlib.Path(__file__).parent.parent.parent / 'indexes' / 'entities.index'))
-    INTRO_MESSAGE = """Hello! I am an AI assistant here to help find the perfect clinic for every patient. Please share the patient’s age restrictions."""
+    INTRO_MESSAGE = """Hallo! Ik ben een AI-assistent hier om te helpen bij het vinden van de perfecte kliniek voor elke patiënt. Deel de gegevens van de patiënt."""
 
 class DevelopmentConfig(BaseConfig):
     Issuer = "http://localhost:8000"