mirror of
https://github.com/cpacker/MemGPT.git
synced 2025-06-03 04:30:22 +00:00
58 lines
3.1 KiB
Python
58 lines
3.1 KiB
Python
import uuid
|
|
from typing import TYPE_CHECKING, Dict, List, Optional
|
|
|
|
from sqlalchemy import JSON, ForeignKey, String
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from letta.orm.sqlalchemy_base import SqlalchemyBase
|
|
from letta.schemas.step import Step as PydanticStep
|
|
|
|
if TYPE_CHECKING:
|
|
from letta.orm.job import Job
|
|
from letta.orm.provider import Provider
|
|
|
|
|
|
class Step(SqlalchemyBase):
|
|
"""Tracks all metadata for agent step."""
|
|
|
|
__tablename__ = "steps"
|
|
__pydantic_model__ = PydanticStep
|
|
|
|
id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"step-{uuid.uuid4()}")
|
|
origin: Mapped[Optional[str]] = mapped_column(nullable=True, doc="The surface that this agent step was initiated from.")
|
|
organization_id: Mapped[str] = mapped_column(
|
|
ForeignKey("organizations.id", ondelete="RESTRICT"),
|
|
nullable=True,
|
|
doc="The unique identifier of the organization that this step ran for",
|
|
)
|
|
provider_id: Mapped[Optional[str]] = mapped_column(
|
|
ForeignKey("providers.id", ondelete="RESTRICT"),
|
|
nullable=True,
|
|
doc="The unique identifier of the provider that was configured for this step",
|
|
)
|
|
job_id: Mapped[Optional[str]] = mapped_column(
|
|
ForeignKey("jobs.id", ondelete="SET NULL"), nullable=True, doc="The unique identified of the job run that triggered this step"
|
|
)
|
|
agent_id: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the model used for this step.")
|
|
provider_name: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the provider used for this step.")
|
|
model: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the model used for this step.")
|
|
model_endpoint: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The model endpoint url used for this step.")
|
|
context_window_limit: Mapped[Optional[int]] = mapped_column(
|
|
None, nullable=True, doc="The context window limit configured for this step."
|
|
)
|
|
completion_tokens: Mapped[int] = mapped_column(default=0, doc="Number of tokens generated by the agent")
|
|
prompt_tokens: Mapped[int] = mapped_column(default=0, doc="Number of tokens in the prompt")
|
|
total_tokens: Mapped[int] = mapped_column(default=0, doc="Total number of tokens processed by the agent")
|
|
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")
|
|
provider: Mapped[Optional["Provider"]] = relationship("Provider")
|
|
job: Mapped[Optional["Job"]] = relationship("Job", back_populates="steps")
|
|
|
|
# Relationships (backrefs)
|
|
messages: Mapped[List["Message"]] = relationship("Message", back_populates="step", cascade="save-update", lazy="noload")
|