f5c4ab50f4
* auto fixes * auto fix pt2 and transitive deps and undefined var checking locals() * manual fixes (ignored or letta-code fixed) * fix circular import * remove all ignores, add FastAPI rules and Ruff rules * add ty and precommit * ruff stuff * ty check fixes * ty check fixes pt 2 * error on invalid
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
from typing import Dict
|
|
|
|
from marshmallow import post_dump, pre_load
|
|
|
|
from letta.orm.message import Message
|
|
from letta.schemas.message import Message as PydanticMessage
|
|
from letta.serialize_schemas.marshmallow_base import BaseSchema
|
|
from letta.serialize_schemas.marshmallow_custom_fields import ToolCallField
|
|
|
|
|
|
class SerializedMessageSchema(BaseSchema):
|
|
"""
|
|
Marshmallow schema for serializing/deserializing Message objects.
|
|
"""
|
|
|
|
__pydantic_model__ = PydanticMessage
|
|
|
|
tool_calls = ToolCallField()
|
|
|
|
@post_dump
|
|
def sanitize_ids(self, data: Dict, **kwargs) -> Dict:
|
|
# keep id for remapping later on agent dump
|
|
# agent dump will then get rid of message ids
|
|
del data["_created_by_id"]
|
|
del data["_last_updated_by_id"]
|
|
|
|
return data
|
|
|
|
@pre_load
|
|
def regenerate_ids(self, data: Dict, **kwargs) -> Dict:
|
|
if self.Meta.model:
|
|
# Skip regenerating ID, as agent dump will do it
|
|
data["_created_by_id"] = self.actor.id
|
|
data["_last_updated_by_id"] = self.actor.id
|
|
|
|
return data
|
|
|
|
class Meta(BaseSchema.Meta):
|
|
model = Message
|
|
exclude = (*BaseSchema.Meta.exclude, "step", "job_message", "otid", "is_deleted", "organization")
|