feat: make attach/detach routes return None if version is 1.0 [LET-5844] (#6141)

---------

Co-authored-by: Ari Webb <ari@letta.com>
This commit is contained in:
Ari Webb
2025-11-14 17:27:44 -08:00
committed by Caren Thomas
parent 099da97e7b
commit e20feaa3e4
4 changed files with 130 additions and 27 deletions

View File

@@ -5257,7 +5257,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Attach Source To Agent"
}
}
}
@@ -5319,7 +5327,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Attach Folder To Agent"
}
}
}
@@ -5382,7 +5398,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Detach Source From Agent"
}
}
}
@@ -5444,7 +5468,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Detach Folder From Agent"
}
}
}
@@ -6414,7 +6446,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Attach Core Memory Block"
}
}
}
@@ -6476,7 +6516,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Detach Core Memory Block"
}
}
}
@@ -7782,7 +7830,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Reset Messages"
}
}
}

View File

@@ -575,7 +575,9 @@ async def modify_approval_for_tool(
return await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
@router.patch("/{agent_id}/sources/attach/{source_id}", response_model=AgentState, operation_id="attach_source_to_agent", deprecated=True)
@router.patch(
"/{agent_id}/sources/attach/{source_id}", response_model=Optional[AgentState], operation_id="attach_source_to_agent", deprecated=True
)
async def attach_source(
source_id: SourceId,
agent_id: AgentId,
@@ -599,10 +601,12 @@ async def attach_source(
source = await server.source_manager.get_source_by_id(source_id=source_id)
safe_create_task(server.sleeptime_document_ingest_async(agent_state, source, actor), label="sleeptime_document_ingest_async")
if is_1_0_sdk_version(headers):
return None
return agent_state
@router.patch("/{agent_id}/folders/attach/{folder_id}", response_model=AgentState, operation_id="attach_folder_to_agent")
@router.patch("/{agent_id}/folders/attach/{folder_id}", response_model=Optional[AgentState], operation_id="attach_folder_to_agent")
async def attach_folder_to_agent(
folder_id: SourceId,
agent_id: AgentId,
@@ -626,10 +630,14 @@ async def attach_folder_to_agent(
source = await server.source_manager.get_source_by_id(source_id=folder_id)
safe_create_task(server.sleeptime_document_ingest_async(agent_state, source, actor), label="sleeptime_document_ingest_async")
if is_1_0_sdk_version(headers):
return None
return agent_state
@router.patch("/{agent_id}/sources/detach/{source_id}", response_model=AgentState, operation_id="detach_source_from_agent", deprecated=True)
@router.patch(
"/{agent_id}/sources/detach/{source_id}", response_model=Optional[AgentState], operation_id="detach_source_from_agent", deprecated=True
)
async def detach_source(
source_id: SourceId,
agent_id: AgentId,
@@ -656,10 +664,13 @@ async def detach_source(
await server.block_manager.delete_block_async(block.id, actor)
except:
pass
if is_1_0_sdk_version(headers):
return None
return agent_state
@router.patch("/{agent_id}/folders/detach/{folder_id}", response_model=AgentState, operation_id="detach_folder_from_agent")
@router.patch("/{agent_id}/folders/detach/{folder_id}", response_model=Optional[AgentState], operation_id="detach_folder_from_agent")
async def detach_folder_from_agent(
folder_id: SourceId,
agent_id: AgentId,
@@ -686,6 +697,9 @@ async def detach_folder_from_agent(
await server.block_manager.delete_block_async(block.id, actor)
except:
pass
if is_1_0_sdk_version(headers):
return None
return agent_state
@@ -1038,7 +1052,9 @@ async def modify_block_for_agent(
return block
@router.patch("/{agent_id}/core-memory/blocks/attach/{block_id}", response_model=AgentState, operation_id="attach_core_memory_block")
@router.patch(
"/{agent_id}/core-memory/blocks/attach/{block_id}", response_model=Optional[AgentState], operation_id="attach_core_memory_block"
)
async def attach_block_to_agent(
block_id: BlockId,
agent_id: AgentId,
@@ -1049,10 +1065,14 @@ async def attach_block_to_agent(
Attach a core memory block to an agent.
"""
actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id)
return await server.agent_manager.attach_block_async(agent_id=agent_id, block_id=block_id, actor=actor)
return await server.agent_manager.attach_block_async(
agent_id=agent_id, block_id=block_id, actor=actor, needs_agent_state=not is_1_0_sdk_version(headers)
)
@router.patch("/{agent_id}/core-memory/blocks/detach/{block_id}", response_model=AgentState, operation_id="detach_core_memory_block")
@router.patch(
"/{agent_id}/core-memory/blocks/detach/{block_id}", response_model=Optional[AgentState], operation_id="detach_core_memory_block"
)
async def detach_block_from_agent(
block_id: BlockId,
agent_id: AgentId,
@@ -1063,7 +1083,9 @@ async def detach_block_from_agent(
Detach a core memory block from an agent.
"""
actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id)
return await server.agent_manager.detach_block_async(agent_id=agent_id, block_id=block_id, actor=actor)
return await server.agent_manager.detach_block_async(
agent_id=agent_id, block_id=block_id, actor=actor, needs_agent_state=not is_1_0_sdk_version(headers)
)
@router.patch("/{agent_id}/archives/attach/{archive_id}", response_model=None, operation_id="attach_archive_to_agent")
@@ -1873,7 +1895,7 @@ class ResetMessagesRequest(BaseModel):
)
@router.patch("/{agent_id}/reset-messages", response_model=AgentState, operation_id="reset_messages")
@router.patch("/{agent_id}/reset-messages", response_model=Optional[AgentState], operation_id="reset_messages")
async def reset_messages(
agent_id: AgentId,
request: ResetMessagesRequest = Body(...),
@@ -1883,7 +1905,10 @@ async def reset_messages(
"""Resets the messages for an agent"""
actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id)
return await server.agent_manager.reset_messages_async(
agent_id=agent_id, actor=actor, add_default_initial_messages=request.add_default_initial_messages
agent_id=agent_id,
actor=actor,
add_default_initial_messages=request.add_default_initial_messages,
needs_agent_state=not is_1_0_sdk_version(headers),
)

View File

@@ -1495,8 +1495,8 @@ class AgentManager:
@enforce_types
@trace_method
async def reset_messages_async(
self, agent_id: str, actor: PydanticUser, add_default_initial_messages: bool = False
) -> PydanticAgentState:
self, agent_id: str, actor: PydanticUser, add_default_initial_messages: bool = False, needs_agent_state: bool = True
) -> Optional[PydanticAgentState]:
"""
Removes all in-context messages for the specified agent except the original system message by:
1) Preserving the first message ID (original system message).
@@ -1510,9 +1510,10 @@ class AgentManager:
add_default_initial_messages: If true, adds the default initial messages after resetting.
agent_id (str): The ID of the agent whose messages will be reset.
actor (PydanticUser): The user performing this action.
needs_agent_state: If True, returns the updated agent state. If False, returns None (for performance optimization)
Returns:
PydanticAgentState: The updated agent state with only the original system message preserved.
Optional[PydanticAgentState]: The updated agent state with only the original system message preserved, or None if needs_agent_state=False.
"""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
@@ -1533,7 +1534,12 @@ class AgentManager:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
agent.message_ids = [system_message_id]
await agent.update_async(db_session=session, actor=actor)
agent_state = await agent.to_pydantic_async(include_relationships=["sources"])
# Only convert to pydantic if we need to return it or add initial messages
if add_default_initial_messages or needs_agent_state:
agent_state = await agent.to_pydantic_async(include_relationships=["sources"] if add_default_initial_messages else None)
else:
agent_state = None
# Optionally add default initial messages after the system message
if add_default_initial_messages:
@@ -1703,6 +1709,8 @@ class AgentManager:
# Commit the changes
agent = await agent.update_async(session, actor=actor)
# TODO: This refresh is expensive. If we can find out which fields are needed, we can save cost by only refreshing those fields.
# or even better, not refresh at all.
return await agent.to_pydantic_async()
@enforce_types
@@ -1863,6 +1871,8 @@ class AgentManager:
# Get agent without loading relationships for return value
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
# TODO: This refresh is expensive. If we can find out which fields are needed, we can save cost by only refreshing those fields.
# or even better, not refresh at all.
return await agent.to_pydantic_async()
# ======================================================================================================================
@@ -1919,7 +1929,9 @@ class AgentManager:
@trace_method
@raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
@raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK)
async def attach_block_async(self, agent_id: str, block_id: str, actor: PydanticUser) -> PydanticAgentState:
async def attach_block_async(
self, agent_id: str, block_id: str, actor: PydanticUser, needs_agent_state: bool = True
) -> Optional[PydanticAgentState]:
"""Attaches a block to an agent. For sleeptime agents, also attaches to paired agents in the same group."""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
@@ -1951,7 +1963,7 @@ class AgentManager:
# TODO: I have too many things rn so lets look at this later
# await session.commit()
return await agent.to_pydantic_async()
return await agent.to_pydantic_async() if needs_agent_state else None
@enforce_types
@trace_method
@@ -1960,7 +1972,8 @@ class AgentManager:
agent_id: str,
block_id: str,
actor: PydanticUser,
) -> PydanticAgentState:
needs_agent_state: bool = True,
) -> Optional[PydanticAgentState]:
"""Detaches a block from an agent."""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
@@ -1972,7 +1985,7 @@ class AgentManager:
raise NoResultFound(f"No block with id '{block_id}' found for agent '{agent_id}' with actor id: '{actor.id}'")
await agent.update_async(session, actor=actor)
return await agent.to_pydantic_async()
return await agent.to_pydantic_async() if needs_agent_state else None
# ======================================================================================================================
# Passage Management

View File

@@ -460,8 +460,17 @@ def test_reset_messages(client: LettaSDKClient):
# After reset, messages should be empty or only have default initial messages
# Messages returns SyncArrayPage, check items
assert isinstance(messages_after.items, list), "Should return list of messages"
assert isinstance(reset_agent, AgentState), "Should return updated agent state"
assert reset_agent.id == agent.id, "Should return the same agent"
# In SDK v1.0, reset-messages returns None, so we need to retrieve the agent to verify
if reset_agent is None:
# Retrieve the agent state after reset
agent_after_reset = client.agents.retrieve(agent_id=agent.id)
assert isinstance(agent_after_reset, AgentState), "Should be able to retrieve agent after reset"
assert agent_after_reset.id == agent.id, "Should be the same agent"
else:
# For older SDK versions that still return AgentState
assert isinstance(reset_agent, AgentState), "Should return updated agent state"
assert reset_agent.id == agent.id, "Should return the same agent"
finally:
# Clean up