Add entity handling to message saving and AI response logic

trauma/api/message/ai/engine.py

@@ -31,16 +31,10 @@ from trauma.core.config import settings
 async def search_entities(
         user_message: str, messages: list[dict], chat: ChatModel
 ) -> CreateMessageResponse:
-    related_entity, entity_data, is_valid = await asyncio.gather(
-        retrieve_semantic_answer(user_message),
+    entity_data, is_valid = await asyncio.gather(
         update_entity_data_with_ai(chat.entityData, user_message, messages[-1]['content']),
         check_is_valid_request(user_message, "\n".join([f"- [{i['role']}]: {i['content']}." for i in messages]))
     )
-    if related_entity:
-        response = await generate_searched_entity_response(user_message, related_entity[0])
-        asyncio.create_task(save_assistant_user_message(user_message, response, chat.id))
-        return CreateMessageResponse(text=response, entities=related_entity)
-
     final_entities = None
     if not is_valid:
         response = await generate_invalid_response(user_message, messages)
@@ -60,7 +54,7 @@ async def search_entities(
             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))
+    asyncio.create_task(save_assistant_user_message(user_message, final_entities, response, chat.id))
     return CreateMessageResponse(text=response, entities=final_entities)
 
 

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
+from trauma.api.data.model import EntityModel, EntityModelExtended
 from trauma.api.message.dto import Author
 from trauma.api.message.model import MessageModel
 from trauma.api.message.schemas import CreateMessageRequest
@@ -53,8 +53,10 @@ async def update_entity_data_obj(entity_data: dict, chat_id: str) -> None:
 
 
 @background_task()
-async def save_assistant_user_message(user_message: str, assistant_message: str, chat_id: str) -> None:
-    user_message = MessageModel(chatId=chat_id, author=Author.User, text=user_message)
+async def save_assistant_user_message(
+        user_message: str, final_entities: list[EntityModelExtended], assistant_message: str, chat_id: str
+) -> None:
+    user_message = MessageModel(chatId=chat_id, author=Author.User, entities=final_entities, text=user_message)
     assistant_message = MessageModel(chatId=chat_id, author=Author.Assistant, text=assistant_message)
     await settings.DB_CLIENT.messages.insert_one(user_message.to_mongo())
     await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())

trauma/api/message/model.py

@@ -2,6 +2,7 @@ from datetime import datetime
 
 from pydantic import Field
 
+from trauma.api.data.model import EntityModelExtended
 from trauma.api.message.dto import Author
 from trauma.core.database import MongoBaseModel
 
@@ -10,5 +11,6 @@ class MessageModel(MongoBaseModel):
     chatId: str
     author: Author
     text: str
+    entities: list[EntityModelExtended] | None = None
     datetimeInserted: datetime = Field(default_factory=datetime.now)
     datetimeUpdated: datetime = Field(default_factory=datetime.now)