fix: make attach/detach routes return None if sdk verion 1.0 (#6203)

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

This reverts commit bb0d10725f5889306de61e1758f061d6c1041c52.

* fix type checking

* revert

* return state for blocks and sources

* func signatures

* create memgpt_agent for cloud-e2e-tests

* Revert "create memgpt_agent for cloud-e2e-tests"

This reverts commit f279e5897b0942b1006a5f8527713dd801064c63.

* fix

---------

Co-authored-by: Ari Webb <ari@letta.com>
This commit is contained in:
Ari Webb
2025-11-17 15:50:26 -08:00
committed by Caren Thomas
parent b0f9986209
commit 65a733d5dd
4 changed files with 65 additions and 13 deletions

View File

@@ -5319,7 +5319,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Attach Folder To Agent"
}
}
}
@@ -5444,7 +5452,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Detach Folder From Agent"
}
}
}
@@ -7782,7 +7798,15 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/AgentState"
"anyOf": [
{
"$ref": "#/components/schemas/AgentState"
},
{
"type": "null"
}
],
"title": "Response Reset Messages"
}
}
}

View File

@@ -625,7 +625,7 @@ async def attach_source(
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,
@@ -649,6 +649,8 @@ 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
@@ -679,10 +681,11 @@ async def detach_source(
await server.block_manager.delete_block_async(block.id, actor)
except:
pass
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,
@@ -709,6 +712,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
@@ -1896,7 +1902,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(...),
@@ -1906,7 +1912,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

@@ -1498,8 +1498,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).
@@ -1513,9 +1513,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)
@@ -1536,7 +1537,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:
@@ -1706,6 +1712,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
@@ -1866,6 +1874,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()
# ======================================================================================================================

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