From 7a37557bb91614043ceb2b9ed83a177eea7df4ac Mon Sep 17 00:00:00 2001
From: cthomas <caren@letta.com>
Date: Sat, 19 Jul 2025 22:47:32 -0700
Subject: [PATCH] feat: remove organization from agents pydantic schema (#3426)

---
 letta/orm/agent.py                           |  2 +-
 letta/schemas/agent.py                       |  2 --
 letta/serialize_schemas/marshmallow_agent.py | 14 +++++++++++++-
 letta/services/agent_manager.py              |  3 ++-
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/letta/orm/agent.py b/letta/orm/agent.py
index 8713506d..4504bbc0 100644
--- a/letta/orm/agent.py
+++ b/letta/orm/agent.py
@@ -93,7 +93,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs):
     timezone: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The timezone of the agent (for the context window).")
 
     # relationships
-    organization: Mapped["Organization"] = relationship("Organization", back_populates="agents")
+    organization: Mapped["Organization"] = relationship("Organization", back_populates="agents", lazy="raise")
     tool_exec_environment_variables: Mapped[List["AgentEnvironmentVariable"]] = relationship(
         "AgentEnvironmentVariable",
         back_populates="agent",
diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py
index 2db5e2bf..2ffd4139 100644
--- a/letta/schemas/agent.py
+++ b/letta/schemas/agent.py
@@ -78,8 +78,6 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
 
     # This is an object representing the in-process state of a running `Agent`
     # Field in this object can be theoretically edited by tools, and will be persisted by the ORM
-    organization_id: Optional[str] = Field(None, description="The unique identifier of the organization associated with the agent.")
-
     description: Optional[str] = Field(None, description="The description of the agent.")
     metadata: Optional[Dict] = Field(None, description="The metadata of the agent.")
 
diff --git a/letta/serialize_schemas/marshmallow_agent.py b/letta/serialize_schemas/marshmallow_agent.py
index b1340071..4e9f68aa 100644
--- a/letta/serialize_schemas/marshmallow_agent.py
+++ b/letta/serialize_schemas/marshmallow_agent.py
@@ -86,7 +86,9 @@ class MarshmallowAgentSchema(BaseSchema):
         - Marks messages as in-context, preserving the order of the original `message_ids`
         - Removes individual message `id` fields
         """
-        data = super().sanitize_ids(data, **kwargs)
+        del data["id"]
+        del data["_created_by_id"]
+        del data["_last_updated_by_id"]
         data[self.FIELD_VERSION] = letta.__version__
 
         original_message_ids = data.pop(self.FIELD_MESSAGE_IDS, [])
@@ -107,6 +109,15 @@ class MarshmallowAgentSchema(BaseSchema):
 
         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
+
+        return data
+
     @post_dump
     def hide_tool_exec_environment_variables(self, data: Dict, **kwargs):
         """Hide the value of tool_exec_environment_variables"""
@@ -136,4 +147,5 @@ class MarshmallowAgentSchema(BaseSchema):
             "is_deleted",
             "groups",
             "batch_items",
+            "organization",
         )
diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py
index 0b872a97..cf3408d2 100644
--- a/letta/services/agent_manager.py
+++ b/letta/services/agent_manager.py
@@ -1319,6 +1319,7 @@ class AgentManager:
             schema = MarshmallowAgentSchema(session=session, actor=actor)
             agent = schema.load(serialized_agent_dict, session=session)
 
+            agent.organization_id = actor.organization_id
             if append_copy_suffix:
                 agent.name += "_copy"
             if project_id:
@@ -1698,7 +1699,7 @@ class AgentManager:
                 logger.error(
                     f"Agent {agent_id} has no message_ids. Agent details: "
                     f"name={agent.name}, created_at={agent.created_at}, "
-                    f"message_ids={agent.message_ids}, organization_id={agent.organization_id}"
+                    f"message_ids={agent.message_ids}, organization_id={actor.organization_id}"
                 )
                 raise ValueError(f"Agent {agent_id} has no message_ids - cannot preserve system message")