fix

trauma/api/data/db_requests.py

@@ -1,6 +1,10 @@
+import asyncio
+import re
+
 from fastapi import HTTPException
 
 from trauma.api.data.model import EntityModel
+from trauma.api.data.schemas import SearchRequest
 from trauma.core.config import settings
 
 
@@ -15,3 +19,19 @@ async def get_all_model_obj() -> list[EntityModel]:
     sort_v = -1
     objects = await settings.DB_CLIENT.entities.find({}, {"embedding": 0}).sort("_id", sort_v).to_list(length=None)
     return objects
+
+
+async def search_facilities_obj(data: SearchRequest) -> tuple[list[EntityModel], int]:
+    skip = data.pageSize * data.pageIndex
+    sort_v = -1
+    regex_filter = {"name": {"$regex": f"^{re.escape(data.name)}", "$options": "i"}}
+    objects, total_count = await asyncio.gather(
+        settings.DB_CLIENT.entities
+        .find(regex_filter, {"embeddings": 0})
+        .sort("_id", sort_v)
+        .skip(skip)
+        .limit(data.pageSize)
+        .to_list(length=data.pageSize),
+        settings.DB_CLIENT.entities.count_documents(regex_filter)
+    )
+    return [EntityModel.from_mongo(ent) for ent in objects], total_count

trauma/api/data/dto.py

@@ -11,4 +11,3 @@ class ContactDetails(BaseModel):
     website: str | None = None
     address: str | None = None
     postalCode: str | None = None
-

trauma/api/data/schemas.py

@@ -9,5 +9,12 @@ class AllFacilitiesResponse(BaseModel):
     paging: Paging
     data: list[EntityModel]
 
+
 class AllFacilitiesWrapper(TraumaResponseWrapper[AllFacilitiesResponse]):
-    pass
+    pass
+
+
+class SearchRequest(BaseModel):
+    name: str
+    pageSize: int
+    pageIndex: int

trauma/api/data/views.py

@@ -4,9 +4,9 @@ from trauma.api.account.dto import AccountType
 from trauma.api.account.model import AccountModel
 from trauma.api.common.dto import Paging
 from trauma.api.data import facility_router
-from trauma.api.data.db_requests import get_facility_by_id, get_all_model_obj
+from trauma.api.data.db_requests import get_facility_by_id, get_all_model_obj, search_facilities_obj
 from trauma.api.data.model import EntityModel
-from trauma.api.data.schemas import AllFacilitiesWrapper, AllFacilitiesResponse
+from trauma.api.data.schemas import AllFacilitiesWrapper, AllFacilitiesResponse, SearchRequest
 from trauma.core.security import PermissionDependency
 from trauma.core.wrappers import TraumaResponseWrapper
 
@@ -23,6 +23,18 @@ async def get_all_countries(
     return AllFacilitiesWrapper(data=response)
 
 
+@facility_router.post('/search')
+async def search_facilities(
+        data: SearchRequest,
+) -> AllFacilitiesWrapper:
+    countries, total_count = await search_facilities_obj(data)
+    response = AllFacilitiesResponse(
+        paging=Paging(pageSize=data.pageSize, pageIndex=data.pageIndex, totalCount=total_count),
+        data=countries
+    )
+    return AllFacilitiesWrapper(data=response)
+
+
 @facility_router.get('/{facilityId}')
 async def get_country(
         facilityId: str,

trauma/api/message/ai/engine.py

@@ -8,10 +8,15 @@ 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, convert_value_to_embeddings,
-                                                  choose_closest_treatment_method, choose_closest_treatment_area,
-                                                  check_is_valid_request, generate_invalid_response, set_entity_score,
-                                                  retrieve_semantic_answer, generate_searched_entity_response)
+                                                  generate_final_response,
+                                                  convert_value_to_embeddings,
+                                                  choose_closest_treatment_method,
+                                                  choose_closest_treatment_area,
+                                                  check_is_valid_request,
+                                                  generate_invalid_response,
+                                                  set_entity_score,
+                                                  retrieve_semantic_answer,
+                                                  generate_searched_entity_response)
 from trauma.api.message.db_requests import (save_assistant_user_message,
                                             filter_entities_by_age_location,
                                             update_entity_data_obj, get_entity_by_index)