feat: add last_stop_reason to AgentState [LET-5911] (#5772)
* feat: add last_stop_reason to AgentState [LET-5911] * undo agent loop changes, use update_run_by_id_async * add run manager test * add integration tests * remove comment * remove duplicate test
This commit is contained in:
committed by
Caren Thomas
parent
dbad510a6e
commit
ef3df907c5
@@ -0,0 +1,31 @@
|
||||
"""add last_stop_reason to agent state
|
||||
|
||||
Revision ID: 89fd4648866b
|
||||
Revises: f6cd5a1e519d
|
||||
Create Date: 2025-10-27 16:55:54.383688
|
||||
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "89fd4648866b"
|
||||
down_revision: Union[str, None] = "f6cd5a1e519d"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.add_column("agents", sa.Column("last_stop_reason", sa.String(), nullable=True))
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_column("agents", "last_stop_reason")
|
||||
# ### end Alembic commands ###
|
||||
@@ -18981,6 +18981,17 @@
|
||||
"title": "Last Run Duration Ms",
|
||||
"description": "The duration in milliseconds of the agent's last run."
|
||||
},
|
||||
"last_stop_reason": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/StopReasonType"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "The stop reason from the agent's last run."
|
||||
},
|
||||
"timezone": {
|
||||
"anyOf": [
|
||||
{
|
||||
@@ -34653,6 +34664,17 @@
|
||||
"title": "Last Run Duration Ms",
|
||||
"description": "The duration in milliseconds of the agent's last run."
|
||||
},
|
||||
"last_stop_reason": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/StopReasonType"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "The stop reason from the agent's last run."
|
||||
},
|
||||
"timezone": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
||||
@@ -16,6 +16,7 @@ from letta.orm.sqlalchemy_base import SqlalchemyBase
|
||||
from letta.schemas.agent import AgentState as PydanticAgentState
|
||||
from letta.schemas.embedding_config import EmbeddingConfig
|
||||
from letta.schemas.enums import AgentType
|
||||
from letta.schemas.letta_stop_reason import StopReasonType
|
||||
from letta.schemas.llm_config import LLMConfig
|
||||
from letta.schemas.memory import Memory
|
||||
from letta.schemas.response_format import ResponseFormatUnion
|
||||
@@ -93,6 +94,9 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
|
||||
last_run_duration_ms: Mapped[Optional[int]] = mapped_column(
|
||||
Integer, nullable=True, doc="The duration in milliseconds of the agent's last run."
|
||||
)
|
||||
last_stop_reason: Mapped[Optional[StopReasonType]] = mapped_column(
|
||||
String, nullable=True, doc="The stop reason from the agent's last run."
|
||||
)
|
||||
|
||||
# timezone
|
||||
timezone: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The timezone of the agent (for the context window).")
|
||||
@@ -232,6 +236,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
|
||||
"response_format": self.response_format,
|
||||
"last_run_completion": self.last_run_completion,
|
||||
"last_run_duration_ms": self.last_run_duration_ms,
|
||||
"last_stop_reason": self.last_stop_reason,
|
||||
"timezone": self.timezone,
|
||||
"max_files_open": self.max_files_open,
|
||||
"per_file_view_window_char_limit": self.per_file_view_window_char_limit,
|
||||
@@ -334,6 +339,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
|
||||
"response_format": self.response_format,
|
||||
"last_run_completion": self.last_run_completion,
|
||||
"last_run_duration_ms": self.last_run_duration_ms,
|
||||
"last_stop_reason": self.last_stop_reason,
|
||||
"max_files_open": self.max_files_open,
|
||||
"per_file_view_window_char_limit": self.per_file_view_window_char_limit,
|
||||
"hidden": self.hidden,
|
||||
|
||||
@@ -14,6 +14,7 @@ from letta.schemas.file import FileStatus
|
||||
from letta.schemas.group import Group
|
||||
from letta.schemas.identity import Identity
|
||||
from letta.schemas.letta_base import OrmMetadataBase
|
||||
from letta.schemas.letta_stop_reason import StopReasonType
|
||||
from letta.schemas.llm_config import LLMConfig
|
||||
from letta.schemas.memory import Memory
|
||||
from letta.schemas.message import Message, MessageCreate
|
||||
@@ -135,6 +136,7 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
|
||||
# Run metrics
|
||||
last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.")
|
||||
last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.")
|
||||
last_stop_reason: Optional[StopReasonType] = Field(None, description="The stop reason from the agent's last run.")
|
||||
|
||||
# timezone
|
||||
timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).")
|
||||
@@ -388,6 +390,7 @@ class UpdateAgent(BaseModel):
|
||||
response_format: Optional[ResponseFormatUnion] = Field(None, description="The response format for the agent.")
|
||||
last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.")
|
||||
last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.")
|
||||
last_stop_reason: Optional[StopReasonType] = Field(None, description="The stop reason from the agent's last run.")
|
||||
timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).")
|
||||
max_files_open: Optional[int] = Field(
|
||||
None,
|
||||
|
||||
@@ -712,6 +712,7 @@ class AgentManager:
|
||||
"response_format": agent_update.response_format,
|
||||
"last_run_completion": agent_update.last_run_completion,
|
||||
"last_run_duration_ms": agent_update.last_run_duration_ms,
|
||||
"last_stop_reason": agent_update.last_stop_reason,
|
||||
"timezone": agent_update.timezone,
|
||||
"max_files_open": agent_update.max_files_open,
|
||||
"per_file_view_window_char_limit": agent_update.per_file_view_window_char_limit,
|
||||
|
||||
@@ -341,6 +341,20 @@ class RunManager:
|
||||
|
||||
await session.commit()
|
||||
|
||||
# Update agent's last_stop_reason when run completes
|
||||
# Do this after run update is committed to database
|
||||
if is_terminal_update and update.stop_reason:
|
||||
try:
|
||||
from letta.schemas.agent import UpdateAgent
|
||||
|
||||
await self.agent_manager.update_agent_async(
|
||||
agent_id=pydantic_run.agent_id,
|
||||
agent_update=UpdateAgent(last_stop_reason=update.stop_reason),
|
||||
actor=actor,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update agent's last_stop_reason for run {run_id}: {e}")
|
||||
|
||||
# update run metrics table
|
||||
num_steps = len(await self.step_manager.list_steps_async(run_id=run_id, actor=actor))
|
||||
|
||||
|
||||
@@ -1277,3 +1277,51 @@ def test_parallel_tool_calling(
|
||||
assert messages[1].message_type == "assistant_message"
|
||||
assert messages[2].message_type == "stop_reason"
|
||||
assert messages[3].message_type == "usage_statistics"
|
||||
|
||||
|
||||
def test_agent_records_last_stop_reason_after_approval_flow(
|
||||
client: Letta,
|
||||
agent: AgentState,
|
||||
) -> None:
|
||||
"""
|
||||
Test that the agent's last_stop_reason is properly updated after a human-in-the-loop flow.
|
||||
This verifies the integration between run completion and agent state updates.
|
||||
"""
|
||||
# Get initial agent state
|
||||
initial_agent = client.agents.retrieve(agent_id=agent.id)
|
||||
initial_stop_reason = initial_agent.last_stop_reason
|
||||
|
||||
# Trigger approval request
|
||||
response = client.agents.messages.create(
|
||||
agent_id=agent.id,
|
||||
messages=USER_MESSAGE_TEST_APPROVAL,
|
||||
)
|
||||
|
||||
# Verify we got an approval request
|
||||
messages = response.messages
|
||||
assert messages is not None
|
||||
assert len(messages) == 3
|
||||
assert messages[2].message_type == "approval_request_message"
|
||||
|
||||
# Check agent after approval request (run should be paused with requires_approval)
|
||||
agent_after_request = client.agents.retrieve(agent_id=agent.id)
|
||||
assert agent_after_request.last_stop_reason == "requires_approval"
|
||||
|
||||
# Approve the tool call
|
||||
approve_tool_call(client, agent.id, response.messages[2].tool_call.tool_call_id)
|
||||
|
||||
# Check agent after approval (run should complete with end_turn or similar)
|
||||
agent_after_approval = client.agents.retrieve(agent_id=agent.id)
|
||||
# After approval and run completion, stop reason should be updated (could be end_turn or other terminal reason)
|
||||
assert agent_after_approval.last_stop_reason is not None
|
||||
assert agent_after_approval.last_stop_reason != initial_stop_reason # Should be different from initial
|
||||
|
||||
# Send follow-up message to complete the flow
|
||||
response2 = client.agents.messages.create(
|
||||
agent_id=agent.id,
|
||||
messages=USER_MESSAGE_FOLLOW_UP,
|
||||
)
|
||||
|
||||
# Verify final agent state has the most recent stop reason
|
||||
final_agent = client.agents.retrieve(agent_id=agent.id)
|
||||
assert final_agent.last_stop_reason is not None
|
||||
|
||||
@@ -625,6 +625,46 @@ async def test_update_agent_file_fields(server: SyncServer, comprehensive_test_a
|
||||
assert updated_agent.per_file_view_window_char_limit == 150_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_agent_last_stop_reason(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
||||
"""Test updating last_stop_reason field on an existing agent"""
|
||||
|
||||
agent, _ = comprehensive_test_agent_fixture
|
||||
|
||||
assert agent.last_stop_reason is None
|
||||
|
||||
# Update with end_turn stop reason
|
||||
update_request = UpdateAgent(
|
||||
last_stop_reason=StopReasonType.end_turn,
|
||||
last_run_completion=datetime.now(timezone.utc),
|
||||
last_run_duration_ms=1500,
|
||||
)
|
||||
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
||||
|
||||
assert updated_agent.last_stop_reason == StopReasonType.end_turn
|
||||
assert updated_agent.last_run_completion is not None
|
||||
assert updated_agent.last_run_duration_ms == 1500
|
||||
|
||||
# Update with error stop reason
|
||||
update_request = UpdateAgent(
|
||||
last_stop_reason=StopReasonType.error,
|
||||
last_run_completion=datetime.now(timezone.utc),
|
||||
last_run_duration_ms=2500,
|
||||
)
|
||||
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
||||
|
||||
assert updated_agent.last_stop_reason == StopReasonType.error
|
||||
assert updated_agent.last_run_duration_ms == 2500
|
||||
|
||||
# Update with requires_approval stop reason
|
||||
update_request = UpdateAgent(
|
||||
last_stop_reason=StopReasonType.requires_approval,
|
||||
)
|
||||
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
||||
|
||||
assert updated_agent.last_stop_reason == StopReasonType.requires_approval
|
||||
|
||||
|
||||
# ======================================================================================================================
|
||||
# AgentManager Tests - Listing
|
||||
# ======================================================================================================================
|
||||
@@ -1086,6 +1126,7 @@ async def test_agent_state_schema_unchanged(server: SyncServer):
|
||||
# Run metrics
|
||||
"last_run_completion": (datetime, type(None)),
|
||||
"last_run_duration_ms": (int, type(None)),
|
||||
"last_stop_reason": (StopReasonType, type(None)),
|
||||
# Timezone
|
||||
"timezone": (str, type(None)),
|
||||
# File controls
|
||||
|
||||
@@ -232,6 +232,41 @@ async def test_update_run_metadata_persistence(server: SyncServer, sarah_agent,
|
||||
assert fetched_run.metadata["error"]["type"] == "llm_timeout"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_run_updates_agent_last_stop_reason(server: SyncServer, sarah_agent, default_user):
|
||||
"""Test that completing a run updates the agent's last_stop_reason."""
|
||||
|
||||
# Verify agent starts with no last_stop_reason
|
||||
agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
|
||||
initial_stop_reason = agent.last_stop_reason
|
||||
|
||||
# Create a run
|
||||
run_data = PydanticRun(agent_id=sarah_agent.id)
|
||||
created_run = await server.run_manager.create_run(pydantic_run=run_data, actor=default_user)
|
||||
|
||||
# Complete the run with end_turn stop reason
|
||||
await server.run_manager.update_run_by_id_async(
|
||||
created_run.id, RunUpdate(status=RunStatus.completed, stop_reason=StopReasonType.end_turn), actor=default_user
|
||||
)
|
||||
|
||||
# Verify agent's last_stop_reason was updated to end_turn
|
||||
updated_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
|
||||
assert updated_agent.last_stop_reason == StopReasonType.end_turn
|
||||
|
||||
# Create another run and complete with different stop reason
|
||||
run_data2 = PydanticRun(agent_id=sarah_agent.id)
|
||||
created_run2 = await server.run_manager.create_run(pydantic_run=run_data2, actor=default_user)
|
||||
|
||||
# Complete with error stop reason
|
||||
await server.run_manager.update_run_by_id_async(
|
||||
created_run2.id, RunUpdate(status=RunStatus.failed, stop_reason=StopReasonType.error), actor=default_user
|
||||
)
|
||||
|
||||
# Verify agent's last_stop_reason was updated to error
|
||||
final_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
|
||||
assert final_agent.last_stop_reason == StopReasonType.error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_run_by_id(server: SyncServer, sarah_agent, default_user):
|
||||
"""Test deleting a run by its ID."""
|
||||
|
||||
@@ -1179,3 +1179,50 @@ def test_parallel_tool_calling(
|
||||
assert messages[1].message_type == "assistant_message"
|
||||
assert messages[2].message_type == "stop_reason"
|
||||
assert messages[3].message_type == "usage_statistics"
|
||||
|
||||
|
||||
def test_agent_records_last_stop_reason_after_approval_flow(
|
||||
client: Letta,
|
||||
agent: AgentState,
|
||||
) -> None:
|
||||
"""
|
||||
Test that the agent's last_stop_reason is properly updated after a human-in-the-loop flow.
|
||||
This verifies the integration between run completion and agent state updates.
|
||||
"""
|
||||
# Get initial agent state
|
||||
initial_agent = client.agents.retrieve(agent_id=agent.id)
|
||||
initial_stop_reason = initial_agent.last_stop_reason
|
||||
|
||||
# Trigger approval request
|
||||
response = client.agents.messages.send(
|
||||
agent_id=agent.id,
|
||||
messages=USER_MESSAGE_TEST_APPROVAL,
|
||||
)
|
||||
|
||||
# Verify we got an approval request
|
||||
messages = response.messages
|
||||
assert messages is not None
|
||||
assert len(messages) == 3
|
||||
assert messages[2].message_type == "approval_request_message"
|
||||
|
||||
# Check agent after approval request (run should be paused with requires_approval)
|
||||
agent_after_request = client.agents.retrieve(agent_id=agent.id)
|
||||
assert agent_after_request.last_stop_reason == "requires_approval"
|
||||
|
||||
# Approve the tool call
|
||||
approve_tool_call(client, agent.id, response.messages[2].tool_call.tool_call_id)
|
||||
|
||||
# Check agent after approval (run should complete with end_turn or similar)
|
||||
agent_after_approval = client.agents.retrieve(agent_id=agent.id)
|
||||
assert agent_after_approval.last_stop_reason is not None
|
||||
assert agent_after_approval.last_stop_reason != initial_stop_reason
|
||||
|
||||
# Send follow-up message to complete the flow
|
||||
response2 = client.agents.messages.send(
|
||||
agent_id=agent.id,
|
||||
messages=USER_MESSAGE_FOLLOW_UP,
|
||||
)
|
||||
|
||||
# Verify final agent state has the most recent stop reason
|
||||
final_agent = client.agents.retrieve(agent_id=agent.id)
|
||||
assert final_agent.last_stop_reason is not None
|
||||
|
||||
Reference in New Issue
Block a user