| """ |
| ByteAstra backend — MongoDB document wrapper models. |
| Behaves like SQLAlchemy ORM classes for seamless API compatibility. |
| """ |
| import enum |
| from datetime import datetime, timezone |
|
|
|
|
| class MessageRole(str, enum.Enum): |
| user = "user" |
| assistant = "assistant" |
| system = "system" |
|
|
|
|
| class Message: |
| """Wraps a MongoDB message document.""" |
| def __init__(self, doc: dict): |
| self.id = doc["_id"] |
| self.session_id = doc["session_id"] |
| self.role = MessageRole(doc["role"]) |
| self.content = doc["content"] |
| self.citations_json = doc.get("citations_json") |
| self.is_grounded = doc.get("is_grounded") |
| self.created_at = doc.get("created_at") or datetime.now(timezone.utc) |
|
|
|
|
| class Session: |
| """Wraps a MongoDB session document and its messages.""" |
| def __init__(self, doc: dict, messages: list[Message] = None): |
| self.id = doc["_id"] |
| self.domain = doc["domain"] |
| self.title = doc["title"] |
| self.user_id = doc.get("user_id") |
| self.subscription_tier = doc.get("subscription_tier", "free") |
| self.created_at = doc.get("created_at") or datetime.now(timezone.utc) |
| self.updated_at = doc.get("updated_at") or datetime.now(timezone.utc) |
| self.messages = messages or [] |
|
|
