diff --git a/alembic/versions/60ed28ee7138_add_project_id_to_step_model.py b/alembic/versions/60ed28ee7138_add_project_id_to_step_model.py new file mode 100644 index 00000000..a6e23444 --- /dev/null +++ b/alembic/versions/60ed28ee7138_add_project_id_to_step_model.py @@ -0,0 +1,41 @@ +"""add project id to step model + +Revision ID: 60ed28ee7138 +Revises: 46699adc71a7 +Create Date: 2025-07-01 13:12:44.485233 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "60ed28ee7138" +down_revision: Union[str, None] = "46699adc71a7" +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("steps", sa.Column("project_id", sa.String(), nullable=True)) + op.execute( + """ + UPDATE steps + SET project_id = agents.project_id + FROM agents + WHERE steps.agent_id = agents.id + AND steps.agent_id IS NOT NULL + AND agents.project_id IS NOT NULL + """ + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("steps", "project_id") + # ### end Alembic commands ### diff --git a/letta/agent.py b/letta/agent.py index 90f5494f..87b083c5 100644 --- a/letta/agent.py +++ b/letta/agent.py @@ -990,6 +990,7 @@ class Agent(BaseAgent): ), job_id=job_id, step_id=step_id, + project_id=self.agent_state.project_id, ) for message in all_new_messages: message.step_id = step.id diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py index 721f915a..d13546ca 100644 --- a/letta/agents/letta_agent.py +++ b/letta/agents/letta_agent.py @@ -1018,6 +1018,7 @@ class LettaAgent(BaseAgent): provider_id=None, job_id=run_id, step_id=step_id, + project_id=agent_state.project_id, ) tool_call_messages = create_letta_messages_from_llm_response( diff --git a/letta/orm/step.py b/letta/orm/step.py index 752c492e..85d2afae 100644 --- a/letta/orm/step.py +++ b/letta/orm/step.py @@ -51,6 +51,9 @@ class Step(SqlalchemyBase): feedback: Mapped[Optional[str]] = mapped_column( None, nullable=True, doc="The feedback for this step. Must be either 'positive' or 'negative'." ) + project_id: Mapped[Optional[str]] = mapped_column( + None, nullable=True, doc="The project that the agent that executed this step belongs to (cloud only)." + ) # Relationships (foreign keys) organization: Mapped[Optional["Organization"]] = relationship("Organization") diff --git a/letta/schemas/step.py b/letta/schemas/step.py index 398199b5..37153a56 100644 --- a/letta/schemas/step.py +++ b/letta/schemas/step.py @@ -35,3 +35,4 @@ class Step(StepBase): feedback: Optional[Literal["positive", "negative"]] = Field( None, description="The feedback for this step. Must be either 'positive' or 'negative'." ) + project_id: Optional[str] = Field(None, description="The project that the agent that executed this step belongs to (cloud only).") diff --git a/letta/server/rest_api/routers/v1/steps.py b/letta/server/rest_api/routers/v1/steps.py index cdb401ed..07d513dd 100644 --- a/letta/server/rest_api/routers/v1/steps.py +++ b/letta/server/rest_api/routers/v1/steps.py @@ -26,8 +26,12 @@ async def list_steps( feedback: Optional[Literal["positive", "negative"]] = Query(None, description="Filter by feedback"), has_feedback: Optional[bool] = Query(None, description="Filter by whether steps have feedback (true) or not (false)"), tags: Optional[list[str]] = Query(None, description="Filter by tags"), + project_id: Optional[str] = Query(None, description="Filter by the project ID that is associated with the step (cloud only)."), server: SyncServer = Depends(get_letta_server), actor_id: Optional[str] = Header(None, alias="user_id"), + x_project: Optional[str] = Header( + None, alias="X-Project", description="Filter by project slug to associate with the group (cloud only)." + ), # Only handled by next js middleware ): """ List steps with optional pagination and date filters. @@ -53,6 +57,7 @@ async def list_steps( feedback=feedback, has_feedback=has_feedback, tags=tags, + project_id=project_id, ) diff --git a/letta/services/step_manager.py b/letta/services/step_manager.py index 138071ea..b95de67d 100644 --- a/letta/services/step_manager.py +++ b/letta/services/step_manager.py @@ -42,6 +42,7 @@ class StepManager: trace_ids: Optional[list[str]] = None, feedback: Optional[Literal["positive", "negative"]] = None, has_feedback: Optional[bool] = None, + project_id: Optional[str] = None, ) -> List[PydanticStep]: """List all jobs with optional pagination and status filter.""" async with db_registry.async_session() as session: @@ -54,6 +55,8 @@ class StepManager: filter_kwargs["trace_id"] = trace_ids if feedback: filter_kwargs["feedback"] = feedback + if project_id: + filter_kwargs["project_id"] = project_id steps = await StepModel.list_async( db_session=session, before=before, @@ -82,6 +85,7 @@ class StepManager: provider_id: Optional[str] = None, job_id: Optional[str] = None, step_id: Optional[str] = None, + project_id: Optional[str] = None, ) -> PydanticStep: step_data = { "origin": None, @@ -100,6 +104,7 @@ class StepManager: "tags": [], "tid": None, "trace_id": get_trace_id(), # Get the current trace ID + "project_id": project_id, } if step_id: step_data["id"] = step_id @@ -125,6 +130,7 @@ class StepManager: provider_id: Optional[str] = None, job_id: Optional[str] = None, step_id: Optional[str] = None, + project_id: Optional[str] = None, ) -> PydanticStep: step_data = { "origin": None, @@ -143,6 +149,7 @@ class StepManager: "tags": [], "tid": None, "trace_id": get_trace_id(), # Get the current trace ID + "project_id": project_id, } if step_id: step_data["id"] = step_id @@ -280,6 +287,7 @@ class NoopStepManager(StepManager): provider_id: Optional[str] = None, job_id: Optional[str] = None, step_id: Optional[str] = None, + project_id: Optional[str] = None, ) -> PydanticStep: return @@ -298,5 +306,6 @@ class NoopStepManager(StepManager): provider_id: Optional[str] = None, job_id: Optional[str] = None, step_id: Optional[str] = None, + project_id: Optional[str] = None, ) -> PydanticStep: return diff --git a/tests/test_managers.py b/tests/test_managers.py index cdd9c0a4..0503fa94 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -6160,6 +6160,7 @@ async def test_job_usage_stats_add_and_get(server: SyncServer, sarah_agent, defa total_tokens=150, ), actor=default_user, + project_id=sarah_agent.project_id, ) # Get usage statistics @@ -6213,6 +6214,7 @@ async def test_job_usage_stats_add_multiple(server: SyncServer, sarah_agent, def total_tokens=150, ), actor=default_user, + project_id=sarah_agent.project_id, ) # Add second usage statistics entry @@ -6230,6 +6232,7 @@ async def test_job_usage_stats_add_multiple(server: SyncServer, sarah_agent, def total_tokens=300, ), actor=default_user, + project_id=sarah_agent.project_id, ) # Get usage statistics (should return the latest entry) @@ -6291,6 +6294,7 @@ async def test_job_usage_stats_add_nonexistent_job(server: SyncServer, sarah_age total_tokens=150, ), actor=default_user, + project_id=sarah_agent.project_id, )