From b93b75ac79f6555482cd08ed4b952dafc0a7909e Mon Sep 17 00:00:00 2001
From: cthomas <caren@letta.com>
Date: Thu, 27 Feb 2025 14:25:22 -0800
Subject: [PATCH] feat: add tracing to steps table (#1056)

---
 ...183663c6769_add_trace_id_to_steps_table.py | 31 +++++++++++++++++++
 letta/orm/step.py                             |  1 +
 letta/schemas/step.py                         |  2 +-
 letta/services/step_manager.py                |  2 ++
 4 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 alembic/versions/b183663c6769_add_trace_id_to_steps_table.py

diff --git a/alembic/versions/b183663c6769_add_trace_id_to_steps_table.py b/alembic/versions/b183663c6769_add_trace_id_to_steps_table.py
new file mode 100644
index 00000000..5f69f492
--- /dev/null
+++ b/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 ###
diff --git a/letta/orm/step.py b/letta/orm/step.py
index e5c33347..f13fac6e 100644
--- a/letta/orm/step.py
+++ b/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")
diff --git a/letta/schemas/step.py b/letta/schemas/step.py
index 98bc51c7..f0e7f080 100644
--- a/letta/schemas/step.py
+++ b/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.")
diff --git a/letta/services/step_manager.py b/letta/services/step_manager.py
index 278dd292..dbaf9f90 100644
--- a/letta/services/step_manager.py
+++ b/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: