diff --git a/alembic/versions/5fb8bba2c373_add_step_metrics.py b/alembic/versions/5fb8bba2c373_add_step_metrics.py new file mode 100644 index 00000000..137b20db --- /dev/null +++ b/alembic/versions/5fb8bba2c373_add_step_metrics.py @@ -0,0 +1,55 @@ +"""add_step_metrics + +Revision ID: 5fb8bba2c373 +Revises: f7f757414d20 +Create Date: 2025-08-07 17:40:11.923402 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "5fb8bba2c373" +down_revision: Union[str, None] = "f7f757414d20" +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.create_table( + "step_metrics", + sa.Column("id", sa.String(), nullable=False), + sa.Column("organization_id", sa.String(), nullable=True), + sa.Column("provider_id", sa.String(), nullable=True), + sa.Column("job_id", sa.String(), nullable=True), + sa.Column("llm_request_ns", sa.BigInteger(), nullable=True), + sa.Column("tool_execution_ns", sa.BigInteger(), nullable=True), + sa.Column("step_ns", sa.BigInteger(), nullable=True), + sa.Column("base_template_id", sa.String(), nullable=True), + sa.Column("template_id", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True), + sa.Column("is_deleted", sa.Boolean(), server_default=sa.text("FALSE"), nullable=False), + sa.Column("_created_by_id", sa.String(), nullable=True), + sa.Column("_last_updated_by_id", sa.String(), nullable=True), + sa.Column("project_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=False), + sa.ForeignKeyConstraint(["agent_id"], ["agents.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["id"], ["steps.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["job_id"], ["jobs.id"], ondelete="SET NULL"), + sa.ForeignKeyConstraint(["organization_id"], ["organizations.id"], ondelete="RESTRICT"), + sa.ForeignKeyConstraint(["provider_id"], ["providers.id"], ondelete="RESTRICT"), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("step_metrics") + # ### end Alembic commands ### diff --git a/letta/orm/__init__.py b/letta/orm/__init__.py index b84923a3..e5dcd47e 100644 --- a/letta/orm/__init__.py +++ b/letta/orm/__init__.py @@ -29,6 +29,7 @@ from letta.orm.sandbox_config import AgentEnvironmentVariable, SandboxConfig, Sa from letta.orm.source import Source from letta.orm.sources_agents import SourcesAgents from letta.orm.step import Step +from letta.orm.step_metrics import StepMetrics from letta.orm.tool import Tool from letta.orm.tools_agents import ToolsAgents from letta.orm.user import User diff --git a/letta/orm/step.py b/letta/orm/step.py index d616b85b..a70e831f 100644 --- a/letta/orm/step.py +++ b/letta/orm/step.py @@ -12,7 +12,10 @@ from letta.schemas.step import Step as PydanticStep if TYPE_CHECKING: from letta.orm.job import Job + from letta.orm.message import Message + from letta.orm.organization import Organization from letta.orm.provider import Provider + from letta.orm.step_metrics import StepMetrics class Step(SqlalchemyBase, ProjectMixin): @@ -70,3 +73,6 @@ class Step(SqlalchemyBase, ProjectMixin): # Relationships (backrefs) messages: Mapped[List["Message"]] = relationship("Message", back_populates="step", cascade="save-update", lazy="noload") + metrics: Mapped[Optional["StepMetrics"]] = relationship( + "StepMetrics", back_populates="step", cascade="all, delete-orphan", lazy="noload", uselist=False + ) diff --git a/letta/orm/step_metrics.py b/letta/orm/step_metrics.py new file mode 100644 index 00000000..c85607cd --- /dev/null +++ b/letta/orm/step_metrics.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING, Optional + +from sqlalchemy import BigInteger, ForeignKey, String +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from letta.orm.mixins import AgentMixin, ProjectMixin +from letta.orm.sqlalchemy_base import SqlalchemyBase +from letta.schemas.step_metrics import StepMetrics as PydanticStepMetrics + +if TYPE_CHECKING: + from letta.orm.agent import Agent + from letta.orm.job import Job + from letta.orm.step import Step + + +class StepMetrics(SqlalchemyBase, ProjectMixin, AgentMixin): + """Tracks performance metrics for agent steps.""" + + __tablename__ = "step_metrics" + __pydantic_model__ = PydanticStepMetrics + + id: Mapped[str] = mapped_column( + ForeignKey("steps.id", ondelete="CASCADE"), + primary_key=True, + doc="The unique identifier of the step this metric belongs to (also serves as PK)", + ) + organization_id: Mapped[str] = mapped_column( + ForeignKey("organizations.id", ondelete="RESTRICT"), + nullable=True, + doc="The unique identifier of the organization", + ) + provider_id: Mapped[Optional[str]] = mapped_column( + ForeignKey("providers.id", ondelete="RESTRICT"), + nullable=True, + doc="The unique identifier of the provider", + ) + job_id: Mapped[Optional[str]] = mapped_column( + ForeignKey("jobs.id", ondelete="SET NULL"), + nullable=True, + doc="The unique identifier of the job", + ) + llm_request_ns: Mapped[Optional[int]] = mapped_column( + BigInteger, + nullable=True, + doc="Time spent on the LLM request in nanoseconds", + ) + tool_execution_ns: Mapped[Optional[int]] = mapped_column( + BigInteger, + nullable=True, + doc="Time spent on tool execution in nanoseconds", + ) + step_ns: Mapped[Optional[int]] = mapped_column( + BigInteger, + nullable=True, + doc="Total time for the step in nanoseconds", + ) + base_template_id: Mapped[Optional[str]] = mapped_column( + String, + nullable=True, + doc="The base template ID for the step", + ) + template_id: Mapped[Optional[str]] = mapped_column( + String, + nullable=True, + doc="The template ID for the step", + ) + + # Relationships (foreign keys) + step: Mapped["Step"] = relationship("Step", back_populates="metrics", uselist=False) + job: Mapped[Optional["Job"]] = relationship("Job") + agent: Mapped[Optional["Agent"]] = relationship("Agent") diff --git a/letta/schemas/step_metrics.py b/letta/schemas/step_metrics.py new file mode 100644 index 00000000..9a5ea8ea --- /dev/null +++ b/letta/schemas/step_metrics.py @@ -0,0 +1,23 @@ +from typing import Optional + +from pydantic import Field + +from letta.schemas.letta_base import LettaBase + + +class StepMetricsBase(LettaBase): + __id_prefix__ = "step" + + +class StepMetrics(StepMetricsBase): + id: str = Field(..., description="The id of the step this metric belongs to (matches steps.id).") + organization_id: Optional[str] = Field(None, description="The unique identifier of the organization.") + provider_id: Optional[str] = Field(None, description="The unique identifier of the provider.") + job_id: Optional[str] = Field(None, description="The unique identifier of the job.") + agent_id: Optional[str] = Field(None, description="The unique identifier of the agent.") + llm_request_ns: Optional[int] = Field(None, description="Time spent on LLM requests in nanoseconds.") + tool_execution_ns: Optional[int] = Field(None, description="Time spent on tool execution in nanoseconds.") + step_ns: Optional[int] = Field(None, description="Total time for the step in nanoseconds.") + base_template_id: Optional[str] = Field(None, description="The base template ID that the step belongs to (cloud only).") + template_id: Optional[str] = Field(None, description="The template ID that the step belongs to (cloud only).") + project_id: Optional[str] = Field(None, description="The project that the step belongs to (cloud only).")