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:
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user