add location postal code searching

trauma/api/account/db_requests.py

@@ -1,5 +1,6 @@
 import asyncio
 
+from trauma.api.account.dto import AccountType
 from trauma.api.account.model import AccountModel
 from trauma.api.account.schemas import CreateAccountRequest
 from trauma.api.common.db_requests import check_unique_fields_existence

trauma/api/chat/dto.py

@@ -13,3 +13,5 @@ class EntityData(BaseModel):
     age: int | None = None
     treatmentArea: str | None = None
     treatmentMethod: str | None = None
+    location: str | None = None
+    postalCode: str | None = None

trauma/api/data/__init__.py

@@ -1,7 +1,7 @@
 from fastapi.routing import APIRouter
 
 facility_router = APIRouter(
-    prefix="/api/facility", tags=["chat"]
+    prefix="/api/facility", tags=["facility"]
 )
 
 from . import views

trauma/api/data/model.py

@@ -1,39 +1,7 @@
-import datetime
-
 from trauma.api.data.dto import AgeGroup, ContactDetails
 from trauma.core.database import MongoBaseModel
 
 
-class DataModel(MongoBaseModel):
-    startDate: datetime.datetime
-    endDate: datetime.datetime
-    email: str | None = None
-    name: str | None = None
-    datetimeUpdated: datetime.datetime | None = None
-    organizationName: str | None = None
-    organizationLocation: str
-    youthCareRegion: str
-    postalCode: str
-    hasOrganizationMultipleLocations: bool
-    publicEmail: str
-    website: str
-    isTraumaTreatmentVisible: bool
-    injuryTreatmentForms: list[str]
-    offers: list[str]
-    intensities: list[str]
-    injuryTreatmentDuration: list[str]
-    ageGroups: list[AgeGroup]
-    traumaIndication: list[str]
-    supplier: str
-    treatmentFinancing: list[str]
-    isContractedCare: list[bool]
-    traumaTeamCharacteristics: list[str]
-    traumaDisciplines: list[str]
-    practitionersQualifications: list[str]
-    AQAs: list[str]
-    additions: list[str]
-
-
 class EntityModel(MongoBaseModel):
     id: str
     name: str

trauma/api/message/ai/engine.py

@@ -12,7 +12,7 @@ from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
                                                   choose_closest_treatment_method, choose_closest_treatment_area,
                                                   check_is_valid_request, generate_invalid_response, set_entity_score)
 from trauma.api.message.db_requests import (save_assistant_user_message,
-                                            filter_entities_by_age,
+                                            filter_entities_by_age_location,
                                             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,
@@ -41,7 +41,7 @@ async def search_entities(
         else:
             user_messages_str = prepare_user_messages_str(user_message, messages)
             possible_entity_indexes, search_request = await asyncio.gather(
-                filter_entities_by_age(entity_data),
+                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)
@@ -67,12 +67,12 @@ async def search_semantic_entities(
     ]
     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 extended_entities_with_highlights(final_entities, entity_data)
+    final_entities_extended = await extend_entities_with_highlights(final_entities, entity_data)
     final_entities_scored = await set_entities_score(final_entities_extended, search_request)
     return final_entities_scored
 
 
-async def extended_entities_with_highlights(entities: list[EntityModel], entity_data: dict) -> list[
+async def extend_entities_with_highlights(entities: list[EntityModel], entity_data: dict) -> list[
     EntityModelExtended]:
     async def choose_closest(entity_: EntityModel) -> tuple:
         treatment_area, treatment_method = await asyncio.gather(

trauma/api/message/ai/prompts.py

@@ -30,13 +30,17 @@ Je verzamelt informatie over een patiënt, hun ziekte en de behandelmethode zoda
 {
     “age”: integer,
     “treatmentArea”: “string”,
-    “treatmentMethod”: “string”
+    “treatmentMethod”: “string”,
+    "location": "string",
+    "postalCode": "string",
 }
 ```
 
-- **[age]**: leeftijd van de patiënt.
-- **[treatmentArea]**: Het type mentale of fysieke ziekte/stoornis.
-- **[treatmentMethod]**: Een methode voor het behandelen van de ziekte of stoornis.
+- **age**: leeftijd van de patiënt.
+- **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..
 
 ## Regels voor het bijwerken van Entity Data
 
@@ -139,6 +143,7 @@ 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.
 
 [/INST]"""
     generate_invalid_response = """## Taak

trauma/api/message/db_requests.py

@@ -60,15 +60,26 @@ 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_data: dict) -> list[int]:
+async def filter_entities_by_age_location(entity_data: dict) -> list[int]:
     query = {
         "ageGroups": {
             "$elemMatch": {
                 "ageMin": {"$lte": entity_data['age']},
                 "ageMax": {"$gte": entity_data['age']}
             }
-        }
+        },
     }
+    if entity_data.get('location'):
+        query["contactDetails.address"] = {
+            "$regex": f".*{entity_data['location']}.*",
+            "$options": "i"
+        }
+
+    if entity_data.get('postalCode'):
+        query["contactDetails.postalCode"] = {
+            "$regex": f".*{entity_data['postalCode']}.*",
+            "$options": "i"
+        }
     entities = await settings.DB_CLIENT.entities.find(query).to_list(length=None)
     return [entity['index'] for entity in entities]
 

trauma/api/message/utils.py

@@ -20,7 +20,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 not v:
+        if k not in ('location', 'postalCode') and not v:
             return k
     return None