feat: add tracing to steps table (#1056)

alembic/versions/b183663c6769_add_trace_id_to_steps_table.py

@@ -0,0 +1,31 @@
+"""add trace id to steps table
+
+Revision ID: b183663c6769
+Revises: fdcdafdb11cf
+Create Date: 2025-02-26 14:38:06.095556
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "b183663c6769"
+down_revision: Union[str, None] = "fdcdafdb11cf"
+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("trace_id", sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("steps", "trace_id")
+    # ### end Alembic commands ###

letta/orm/step.py

@@ -45,6 +45,7 @@ class Step(SqlalchemyBase):
     completion_tokens_details: Mapped[Optional[Dict]] = mapped_column(JSON, nullable=True, doc="metadata for the agent.")
     tags: Mapped[Optional[List]] = mapped_column(JSON, doc="Metadata tags.")
     tid: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="Transaction ID that processed the step.")
+    trace_id: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The trace id of the agent step.")
 
     # Relationships (foreign keys)
     organization: Mapped[Optional["Organization"]] = relationship("Organization")

letta/schemas/step.py

@@ -26,7 +26,7 @@ class Step(StepBase):
     prompt_tokens: Optional[int] = Field(None, description="The number of tokens in the prompt during this step.")
     total_tokens: Optional[int] = Field(None, description="The total number of tokens processed by the agent during this step.")
     completion_tokens_details: Optional[Dict] = Field(None, description="Metadata for the agent.")
-
     tags: List[str] = Field([], description="Metadata tags.")
     tid: Optional[str] = Field(None, description="The unique identifier of the transaction that processed this step.")
+    trace_id: Optional[str] = Field(None, description="The trace id of the agent step.")
     messages: List[Message] = Field([], description="The messages generated during this step.")

letta/services/step_manager.py

@@ -11,6 +11,7 @@ from letta.orm.step import Step as StepModel
 from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.step import Step as PydanticStep
 from letta.schemas.user import User as PydanticUser
+from letta.tracing import get_trace_id
 from letta.utils import enforce_types
 
 
@@ -75,6 +76,7 @@ class StepManager:
             "job_id": job_id,
             "tags": [],
             "tid": None,
+            "trace_id": get_trace_id(),  # Get the current trace ID
         }
         with self.session_maker() as session:
             if job_id: