Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: add project id to step model (#3121)

Browse Source
This commit is contained in:
cthomas
2025-07-02 12:03:55 -07:00
committed by GitHub
parent cf75cb67a0
commit 87b23c8a29
8 changed files with 65 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
alembic/versions/60ed28ee7138_add_project_id_to_step_model.py Normal file
View File

@@ -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 ###

1
letta/agent.py
View File

@@ -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

1
letta/agents/letta_agent.py
View File

@@ -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(

3
letta/orm/step.py
View File

@@ -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")

1
letta/schemas/step.py
View File

@@ -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).")

5
letta/server/rest_api/routers/v1/steps.py
View File

@@ -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,
)

9
letta/services/step_manager.py
View File

@@ -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

4
tests/test_managers.py
View File

@@ -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,
)