53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
from typing import Dict, Optional
|
|
|
|
from marshmallow import post_dump, pre_load
|
|
from marshmallow_sqlalchemy import SQLAlchemyAutoSchema
|
|
|
|
from letta.schemas.user import User
|
|
|
|
|
|
class BaseSchema(SQLAlchemyAutoSchema):
|
|
"""
|
|
Base schema for all SQLAlchemy models.
|
|
This ensures all schemas share the same session.
|
|
"""
|
|
|
|
__pydantic_model__ = None
|
|
|
|
def __init__(self, *args, actor: Optional[User] = None, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.actor = actor
|
|
|
|
@classmethod
|
|
def generate_id(cls) -> Optional[str]:
|
|
if cls.__pydantic_model__:
|
|
return cls.__pydantic_model__.generate_id()
|
|
|
|
return None
|
|
|
|
@post_dump
|
|
def sanitize_ids(self, data: Dict, **kwargs) -> Dict:
|
|
# delete id
|
|
del data["id"]
|
|
del data["_created_by_id"]
|
|
del data["_last_updated_by_id"]
|
|
del data["organization"]
|
|
|
|
return data
|
|
|
|
@pre_load
|
|
def regenerate_ids(self, data: Dict, **kwargs) -> Dict:
|
|
if self.Meta.model:
|
|
data["id"] = self.generate_id()
|
|
data["_created_by_id"] = self.actor.id
|
|
data["_last_updated_by_id"] = self.actor.id
|
|
data["organization"] = self.actor.organization_id
|
|
|
|
return data
|
|
|
|
class Meta:
|
|
model = None
|
|
include_relationships = True
|
|
load_instance = True
|
|
exclude = ()
|