From 1a477ccbfe45e93cf78b505a14c58f4ce374219a Mon Sep 17 00:00:00 2001 From: cthomas Date: Thu, 19 Jun 2025 16:04:52 -0700 Subject: [PATCH] chore: move agent step state def (#2923) --- letta/agents/letta_agent_batch.py | 4 ++-- letta/helpers/converters.py | 2 +- letta/orm/llm_batch_items.py | 2 +- letta/schemas/agent.py | 6 ------ letta/schemas/llm_batch_job.py | 9 +++++++-- letta/services/llm_batch_manager.py | 2 +- tests/integration_test_batch_api_cron_jobs.py | 3 ++- tests/test_letta_agent_batch.py | 3 ++- tests/test_managers.py | 4 ++-- 9 files changed, 18 insertions(+), 17 deletions(-) diff --git a/letta/agents/letta_agent_batch.py b/letta/agents/letta_agent_batch.py index ce4ee643..629714f2 100644 --- a/letta/agents/letta_agent_batch.py +++ b/letta/agents/letta_agent_batch.py @@ -18,14 +18,14 @@ from letta.local_llm.constants import INNER_THOUGHTS_KWARG from letta.log import get_logger from letta.orm.enums import ToolType from letta.otel.tracing import log_event, trace_method -from letta.schemas.agent import AgentState, AgentStepState +from letta.schemas.agent import AgentState from letta.schemas.enums import AgentStepStatus, JobStatus, MessageStreamStatus, ProviderType from letta.schemas.job import JobUpdate from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage from letta.schemas.letta_message_content import OmittedReasoningContent, ReasoningContent, RedactedReasoningContent, TextContent from letta.schemas.letta_request import LettaBatchRequest from letta.schemas.letta_response import LettaBatchResponse, LettaResponse -from letta.schemas.llm_batch_job import LLMBatchItem +from letta.schemas.llm_batch_job import AgentStepState, LLMBatchItem from letta.schemas.message import Message, MessageCreate from letta.schemas.openai.chat_completion_response import ToolCall as OpenAIToolCall from letta.schemas.sandbox_config import SandboxConfig, SandboxType diff --git a/letta/helpers/converters.py b/letta/helpers/converters.py index 78c8b45c..9a8bceda 100644 --- a/letta/helpers/converters.py +++ b/letta/helpers/converters.py @@ -8,7 +8,6 @@ from openai.types.chat.chat_completion_message_tool_call import Function as Open from sqlalchemy import Dialect from letta.functions.mcp_client.types import StdioServerConfig -from letta.schemas.agent import AgentStepState from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.enums import ProviderType, ToolRuleType from letta.schemas.letta_message_content import ( @@ -23,6 +22,7 @@ from letta.schemas.letta_message_content import ( ToolCallContent, ToolReturnContent, ) +from letta.schemas.llm_batch_job import AgentStepState from letta.schemas.llm_config import LLMConfig from letta.schemas.message import ToolReturn from letta.schemas.response_format import ( diff --git a/letta/orm/llm_batch_items.py b/letta/orm/llm_batch_items.py index d9d97d1a..ee696905 100644 --- a/letta/orm/llm_batch_items.py +++ b/letta/orm/llm_batch_items.py @@ -8,8 +8,8 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from letta.orm.custom_columns import AgentStepStateColumn, BatchRequestResultColumn, LLMConfigColumn from letta.orm.mixins import AgentMixin, OrganizationMixin from letta.orm.sqlalchemy_base import SqlalchemyBase -from letta.schemas.agent import AgentStepState from letta.schemas.enums import AgentStepStatus, JobStatus +from letta.schemas.llm_batch_job import AgentStepState from letta.schemas.llm_batch_job import LLMBatchItem as PydanticLLMBatchItem from letta.schemas.llm_config import LLMConfig diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index 0220e10d..8b5d7dd4 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -5,7 +5,6 @@ from typing import Dict, List, Optional from pydantic import BaseModel, Field, field_validator, model_validator from letta.constants import CORE_MEMORY_LINE_NUMBER_WARNING, DEFAULT_EMBEDDING_CHUNK_SIZE -from letta.helpers import ToolRulesSolver from letta.schemas.block import CreateBlock from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.environment_variables import AgentEnvironmentVariable @@ -301,11 +300,6 @@ class AgentStepResponse(BaseModel): usage: UsageStatistics = Field(..., description="Usage statistics of the LLM call during the agent's step.") -class AgentStepState(BaseModel): - step_number: int = Field(..., description="The current step number in the agent loop") - tool_rules_solver: ToolRulesSolver = Field(..., description="The current state of the ToolRulesSolver") - - def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): # Sleeptime agents use the MemGPT v2 memory tools (line numbers) diff --git a/letta/schemas/llm_batch_job.py b/letta/schemas/llm_batch_job.py index a6e537f0..e07e148e 100644 --- a/letta/schemas/llm_batch_job.py +++ b/letta/schemas/llm_batch_job.py @@ -2,14 +2,19 @@ from datetime import datetime from typing import Optional, Union from anthropic.types.beta.messages import BetaMessageBatch, BetaMessageBatchIndividualResponse -from pydantic import Field +from pydantic import BaseModel, Field -from letta.schemas.agent import AgentStepState +from letta.helpers import ToolRulesSolver from letta.schemas.enums import AgentStepStatus, JobStatus, ProviderType from letta.schemas.letta_base import OrmMetadataBase from letta.schemas.llm_config import LLMConfig +class AgentStepState(BaseModel): + step_number: int = Field(..., description="The current step number in the agent loop") + tool_rules_solver: ToolRulesSolver = Field(..., description="The current state of the ToolRulesSolver") + + class LLMBatchItemBase(OrmMetadataBase, validate_assignment=True): __id_prefix__ = "batch_item" diff --git a/letta/services/llm_batch_manager.py b/letta/services/llm_batch_manager.py index 4700d146..b2241e44 100644 --- a/letta/services/llm_batch_manager.py +++ b/letta/services/llm_batch_manager.py @@ -10,8 +10,8 @@ from letta.orm import Message as MessageModel from letta.orm.llm_batch_items import LLMBatchItem from letta.orm.llm_batch_job import LLMBatchJob from letta.otel.tracing import trace_method -from letta.schemas.agent import AgentStepState from letta.schemas.enums import AgentStepStatus, JobStatus, ProviderType +from letta.schemas.llm_batch_job import AgentStepState from letta.schemas.llm_batch_job import LLMBatchItem as PydanticLLMBatchItem from letta.schemas.llm_batch_job import LLMBatchJob as PydanticLLMBatchJob from letta.schemas.llm_config import LLMConfig diff --git a/tests/integration_test_batch_api_cron_jobs.py b/tests/integration_test_batch_api_cron_jobs.py index 8a07b9a4..5d9b6739 100644 --- a/tests/integration_test_batch_api_cron_jobs.py +++ b/tests/integration_test_batch_api_cron_jobs.py @@ -21,10 +21,11 @@ from letta.config import LettaConfig from letta.helpers import ToolRulesSolver from letta.jobs.llm_batch_job_polling import poll_running_llm_batches from letta.orm import Base -from letta.schemas.agent import AgentStepState, CreateAgent +from letta.schemas.agent import CreateAgent from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.enums import JobStatus, ProviderType from letta.schemas.job import BatchJob +from letta.schemas.llm_batch_job import AgentStepState from letta.schemas.llm_config import LLMConfig from letta.schemas.tool_rule import InitToolRule from letta.server.db import db_context diff --git a/tests/test_letta_agent_batch.py b/tests/test_letta_agent_batch.py index e5eeaa49..8e798cb6 100644 --- a/tests/test_letta_agent_batch.py +++ b/tests/test_letta_agent_batch.py @@ -19,11 +19,12 @@ from letta.config import LettaConfig from letta.helpers import ToolRulesSolver from letta.jobs.llm_batch_job_polling import poll_running_llm_batches from letta.orm import Base -from letta.schemas.agent import AgentState, AgentStepState, CreateAgent +from letta.schemas.agent import AgentState, CreateAgent from letta.schemas.enums import AgentStepStatus, JobStatus, MessageRole, ProviderType from letta.schemas.job import BatchJob from letta.schemas.letta_message_content import TextContent from letta.schemas.letta_request import LettaBatchRequest +from letta.schemas.llm_batch_job import AgentStepState from letta.schemas.message import MessageCreate from letta.schemas.tool_rule import InitToolRule from letta.server.db import db_context diff --git a/tests/test_managers.py b/tests/test_managers.py index d3375f8a..411f18ea 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -50,7 +50,7 @@ from letta.orm.enums import ToolType from letta.orm.errors import NoResultFound, UniqueConstraintViolationError from letta.orm.file import FileContent as FileContentModel from letta.orm.file import FileMetadata as FileMetadataModel -from letta.schemas.agent import AgentStepState, CreateAgent, UpdateAgent +from letta.schemas.agent import CreateAgent, UpdateAgent from letta.schemas.block import Block as PydanticBlock from letta.schemas.block import BlockUpdate, CreateBlock from letta.schemas.embedding_config import EmbeddingConfig @@ -64,7 +64,7 @@ from letta.schemas.job import Job as PydanticJob from letta.schemas.job import JobUpdate, LettaRequestConfig from letta.schemas.letta_message import UpdateAssistantMessage, UpdateReasoningMessage, UpdateSystemMessage, UpdateUserMessage from letta.schemas.letta_message_content import TextContent -from letta.schemas.llm_batch_job import LLMBatchItem +from letta.schemas.llm_batch_job import AgentStepState, LLMBatchItem from letta.schemas.llm_config import LLMConfig from letta.schemas.message import Message as PydanticMessage from letta.schemas.message import MessageCreate, MessageUpdate