Spaces:
Sleeping
Sleeping
File size: 4,026 Bytes
50553ea c6cc0f2 50553ea 030cdaf f000a3c 50553ea 9150f8e 50553ea 9150f8e c6cc0f2 50553ea c6cc0f2 9150f8e 50553ea 3ec35ef 9150f8e 3a56394 198c726 3a56394 9150f8e 50553ea 211d915 9150f8e 82c856f 9150f8e 211d915 9150f8e 211d915 030cdaf 9150f8e 50553ea 3ec35ef f000a3c 030cdaf f000a3c 198c726 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import asyncio
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.message.dto import Author, Feedback
from trauma.api.message.model import MessageModel
from trauma.api.message.schemas import CreateMessageRequest
from trauma.core.config import settings
from trauma.core.wrappers import background_task
async def create_message_obj(
chat_id: str, message_data: CreateMessageRequest
) -> tuple[MessageModel, ChatModel]:
chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
if not chat:
raise HTTPException(status_code=404, detail="Chat not found.")
message = MessageModel(**message_data.model_dump(), chatId=chat_id, author=Author.User)
await settings.DB_CLIENT.messages.insert_one(message.to_mongo())
return message, chat
async def get_all_chat_messages_obj(chat_id: str, account: AccountModel) -> tuple[list[MessageModel], ChatModel]:
messages, chat = await asyncio.gather(
settings.DB_CLIENT.messages.find({"chatId": chat_id}).to_list(length=None),
settings.DB_CLIENT.chats.find_one({"id": chat_id})
)
messages = [MessageModel.from_mongo(message) for message in messages]
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
chat = ChatModel.from_mongo(chat)
if chat.account.id != account.id and account.accountType != AccountType.Admin:
raise HTTPException(status_code=404, detail="Not Authorized.")
return messages, chat
@background_task()
async def update_entity_data_obj(entity_data: dict, chat_id: str) -> None:
await settings.DB_CLIENT.chats.update_one(
{"id": chat_id},
{"$set": {
"entityData": entity_data,
}}
)
@background_task()
async def save_assistant_user_message(
user_message: MessageModel, assistant_message: MessageModel
) -> None:
await settings.DB_CLIENT.messages.insert_one(user_message.to_mongo())
await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())
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, {"index": 1, "_id": 0}).to_list(length=None)
return [entity['index'] for entity in entities]
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 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:
raise HTTPException(status_code=404, detail="Message not found")
message = MessageModel.from_mongo(message)
message.feedback = feedback_data
await settings.DB_CLIENT.messages.update_one({"id": message_id},
{"$set": {"feedback": feedback_data.model_dump(mode='json')}})
return message
|