Optimize entity retrieval and limit final entity processing.

Replaced individual entity fetches with bulk retrieval to improve performance. Increased the number of filtered results but reduced output entities for better efficiency. Adjusted returned fields and introduced shortening in utility functions for streamlined processing.

trauma/api/message/ai/engine.py

@@ -17,7 +17,7 @@ from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
                                                   set_entity_score)
 from trauma.api.message.db_requests import (save_assistant_user_message,
                                             filter_entities_by_age_location,
-                                            update_entity_data_obj, get_entity_by_index)
+                                            update_entity_data_obj, get_entities_bulk)
 from trauma.api.message.dto import Author
 from trauma.api.message.model import MessageModel
 from trauma.api.message.schemas import CreateMessageResponse
@@ -86,8 +86,8 @@ async def search_semantic_entities(
         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])
+    filtered_results = sorted(filtered_results, key=lambda x: x["distance"])[:50]
+    final_entities = await get_entities_bulk([i['index'] for i in filtered_results])
     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

trauma/api/message/db_requests.py

@@ -5,7 +5,7 @@ from fastapi import HTTPException
 from trauma.api.account.dto import AccountType
 from trauma.api.account.model import AccountModel
 from trauma.api.chat.model import ChatModel
-from trauma.api.data.model import EntityModel, EntityModelExtended
+from trauma.api.data.model import EntityModel
 from trauma.api.message.dto import Author, Feedback
 from trauma.api.message.model import MessageModel
 from trauma.api.message.schemas import CreateMessageRequest
@@ -80,7 +80,7 @@ async def filter_entities_by_age_location(entity_data: dict) -> list[int]:
             "$regex": f".*{entity_data['postalCode']}.*",
             "$options": "i"
         }
-    entities = await settings.DB_CLIENT.entities.find(query, {"embedding": 0}).to_list(length=None)
+    entities = await settings.DB_CLIENT.entities.find(query, {"index": 1, "_id": 0}).to_list(length=None)
     return [entity['index'] for entity in entities]
 
 
@@ -89,6 +89,12 @@ async def get_entity_by_index(index: int) -> EntityModel:
     return EntityModel.from_mongo(entity)
 
 
+async def get_entities_bulk(indices: list[int]) -> list[EntityModel]:
+    entities = await settings.DB_CLIENT.entities.find({"index": {"$in": indices}},
+                                                       {"embedding": 0}).to_list(length=None)
+    return [EntityModel.from_mongo(entity) for entity in entities]
+
+
 async def update_message_feedback_obj(message_id: str, feedback_data: Feedback) -> MessageModel:
     message = await settings.DB_CLIENT.messages.find_one({"id": message_id})
     if not message:

trauma/api/message/utils.py

@@ -44,8 +44,9 @@ def prepare_user_messages_str(user_message: str, messages: list[MessageModel]) -
 
 
 def prepare_final_entities_str(entities: list[EntityModel]) -> str:
+    shortened_entities = entities[:3]
     entities_list = []
-    for entity in entities:
+    for entity in shortened_entities:
         entities_list.append(entity.model_dump(mode='json', exclude={
             'id', 'contactDetails', "highlightedAgeGroup", "highlightedTreatmentArea", "highlightedTreatmentMethod"
         }))