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:
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
# ======================================================================================================================
|
||||
|
||||
@@ -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