37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
from marshmallow import fields
|
|
|
|
from letta.orm import Agent
|
|
from letta.serialize_schemas.base import BaseSchema
|
|
from letta.serialize_schemas.custom_fields import EmbeddingConfigField, LLMConfigField, ToolRulesField
|
|
from letta.serialize_schemas.message import SerializedMessageSchema
|
|
|
|
|
|
class SerializedAgentSchema(BaseSchema):
|
|
"""
|
|
Marshmallow schema for serializing/deserializing Agent objects.
|
|
Excludes relational fields.
|
|
"""
|
|
|
|
llm_config = LLMConfigField()
|
|
embedding_config = EmbeddingConfigField()
|
|
tool_rules = ToolRulesField()
|
|
|
|
messages = fields.List(fields.Nested(SerializedMessageSchema))
|
|
|
|
def __init__(self, *args, session=None, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
if session:
|
|
self.session = session
|
|
|
|
# propagate session to nested schemas
|
|
for field_name, field_obj in self.fields.items():
|
|
if isinstance(field_obj, fields.List) and hasattr(field_obj.inner, "schema"):
|
|
field_obj.inner.schema.session = session
|
|
elif hasattr(field_obj, "schema"):
|
|
field_obj.schema.session = session
|
|
|
|
class Meta(BaseSchema.Meta):
|
|
model = Agent
|
|
# TODO: Serialize these as well...
|
|
exclude = ("tools", "sources", "core_memory", "tags", "source_passages", "agent_passages", "organization")
|