letta-nightly 0.6.4.dev20241213193437__py3-none-any.whl → 0.6.4.dev20241215104129__py3-none-any.whl

letta/orm/custom_columns.py

@@ -0,0 +1,152 @@
+import base64
+from typing import List, Union
+
+import numpy as np
+from sqlalchemy import JSON
+from sqlalchemy.types import BINARY, TypeDecorator
+
+from letta.schemas.embedding_config import EmbeddingConfig
+from letta.schemas.enums import ToolRuleType
+from letta.schemas.llm_config import LLMConfig
+from letta.schemas.openai.chat_completions import ToolCall, ToolCallFunction
+from letta.schemas.tool_rule import ChildToolRule, InitToolRule, TerminalToolRule
+
+
+class EmbeddingConfigColumn(TypeDecorator):
+    """Custom type for storing EmbeddingConfig as JSON."""
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(JSON())
+
+    def process_bind_param(self, value, dialect):
+        if value and isinstance(value, EmbeddingConfig):
+            return value.model_dump()
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value:
+            return EmbeddingConfig(**value)
+        return value
+
+
+class LLMConfigColumn(TypeDecorator):
+    """Custom type for storing LLMConfig as JSON."""
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(JSON())
+
+    def process_bind_param(self, value, dialect):
+        if value and isinstance(value, LLMConfig):
+            return value.model_dump()
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value:
+            return LLMConfig(**value)
+        return value
+
+
+class ToolRulesColumn(TypeDecorator):
+    """Custom type for storing a list of ToolRules as JSON"""
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(JSON())
+
+    def process_bind_param(self, value, dialect):
+        """Convert a list of ToolRules to JSON-serializable format."""
+        if value:
+            data = [rule.model_dump() for rule in value]
+            for d in data:
+                d["type"] = d["type"].value
+
+            for d in data:
+                assert not (d["type"] == "ToolRule" and "children" not in d), "ToolRule does not have children field"
+            return data
+        return value
+
+    def process_result_value(self, value, dialect) -> List[Union[ChildToolRule, InitToolRule, TerminalToolRule]]:
+        """Convert JSON back to a list of ToolRules."""
+        if value:
+            return [self.deserialize_tool_rule(rule_data) for rule_data in value]
+        return value
+
+    @staticmethod
+    def deserialize_tool_rule(data: dict) -> Union[ChildToolRule, InitToolRule, TerminalToolRule]:
+        """Deserialize a dictionary to the appropriate ToolRule subclass based on the 'type'."""
+        rule_type = ToolRuleType(data.get("type"))  # Remove 'type' field if it exists since it is a class var
+        if rule_type == ToolRuleType.run_first:
+            return InitToolRule(**data)
+        elif rule_type == ToolRuleType.exit_loop:
+            return TerminalToolRule(**data)
+        elif rule_type == ToolRuleType.constrain_child_tools:
+            rule = ChildToolRule(**data)
+            return rule
+        else:
+            raise ValueError(f"Unknown tool rule type: {rule_type}")
+
+
+class ToolCallColumn(TypeDecorator):
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(JSON())
+
+    def process_bind_param(self, value, dialect):
+        if value:
+            values = []
+            for v in value:
+                if isinstance(v, ToolCall):
+                    values.append(v.model_dump())
+                else:
+                    values.append(v)
+            return values
+
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value:
+            tools = []
+            for tool_value in value:
+                if "function" in tool_value:
+                    tool_call_function = ToolCallFunction(**tool_value["function"])
+                    del tool_value["function"]
+                else:
+                    tool_call_function = None
+                tools.append(ToolCall(function=tool_call_function, **tool_value))
+            return tools
+        return value
+
+
+class CommonVector(TypeDecorator):
+    """Common type for representing vectors in SQLite"""
+
+    impl = BINARY
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(BINARY())
+
+    def process_bind_param(self, value, dialect):
+        if value is None:
+            return value
+        if isinstance(value, list):
+            value = np.array(value, dtype=np.float32)
+        return base64.b64encode(value.tobytes())
+
+    def process_result_value(self, value, dialect):
+        if not value:
+            return value
+        if dialect.name == "sqlite":
+            value = base64.b64decode(value)
+        return np.frombuffer(value, dtype=np.float32)

letta/orm/message.py

@@ -1,46 +1,12 @@
 from typing import Optional
 
-from sqlalchemy import JSON, TypeDecorator
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from letta.orm.custom_columns import ToolCallColumn
 from letta.orm.mixins import AgentMixin, OrganizationMixin
 from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.message import Message as PydanticMessage
-from letta.schemas.openai.chat_completions import ToolCall, ToolCallFunction
-
-
-class ToolCallColumn(TypeDecorator):
-
-    impl = JSON
-    cache_ok = True
-
-    def load_dialect_impl(self, dialect):
-        return dialect.type_descriptor(JSON())
-
-    def process_bind_param(self, value, dialect):
-        if value:
-            values = []
-            for v in value:
-                if isinstance(v, ToolCall):
-                    values.append(v.model_dump())
-                else:
-                    values.append(v)
-            return values
-
-        return value
-
-    def process_result_value(self, value, dialect):
-        if value:
-            tools = []
-            for tool_value in value:
-                if "function" in tool_value:
-                    tool_call_function = ToolCallFunction(**tool_value["function"])
-                    del tool_value["function"]
-                else:
-                    tool_call_function = None
-                tools.append(ToolCall(function=tool_call_function, **tool_value))
-            return tools
-        return value
+from letta.schemas.openai.chat_completions import ToolCall
 
 
 class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
@@ -59,6 +25,5 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
     tool_call_id: Mapped[Optional[str]] = mapped_column(nullable=True, doc="ID of the tool call")
 
     # Relationships
-    # TODO: Add in after Agent ORM is created
-    # agent: Mapped["Agent"] = relationship("Agent", back_populates="messages", lazy="selectin")
+    agent: Mapped["Agent"] = relationship("Agent", back_populates="messages", lazy="selectin")
     organization: Mapped["Organization"] = relationship("Organization", back_populates="messages", lazy="selectin")

letta/orm/organization.py

@@ -7,6 +7,7 @@ from letta.schemas.organization import Organization as PydanticOrganization
 
 if TYPE_CHECKING:
 
+    from letta.orm.agent import Agent
     from letta.orm.file import FileMetadata
     from letta.orm.tool import Tool
     from letta.orm.user import User
@@ -25,7 +26,6 @@ class Organization(SqlalchemyBase):
     tools: Mapped[List["Tool"]] = relationship("Tool", back_populates="organization", cascade="all, delete-orphan")
     blocks: Mapped[List["Block"]] = relationship("Block", back_populates="organization", cascade="all, delete-orphan")
     sources: Mapped[List["Source"]] = relationship("Source", back_populates="organization", cascade="all, delete-orphan")
-    agents_tags: Mapped[List["AgentsTags"]] = relationship("AgentsTags", back_populates="organization", cascade="all, delete-orphan")
     files: Mapped[List["FileMetadata"]] = relationship("FileMetadata", back_populates="organization", cascade="all, delete-orphan")
     sandbox_configs: Mapped[List["SandboxConfig"]] = relationship(
         "SandboxConfig", back_populates="organization", cascade="all, delete-orphan"
@@ -36,10 +36,5 @@ class Organization(SqlalchemyBase):
 
     # relationships
     messages: Mapped[List["Message"]] = relationship("Message", back_populates="organization", cascade="all, delete-orphan")
+    agents: Mapped[List["Agent"]] = relationship("Agent", back_populates="organization", cascade="all, delete-orphan")
     passages: Mapped[List["Passage"]] = relationship("Passage", back_populates="organization", cascade="all, delete-orphan")
-
-    # TODO: Map these relationships later when we actually make these models
-    # below is just a suggestion
-    # agents: Mapped[List["Agent"]] = relationship("Agent", back_populates="organization", cascade="all, delete-orphan")
-    # tools: Mapped[List["Tool"]] = relationship("Tool", back_populates="organization", cascade="all, delete-orphan")
-    # documents: Mapped[List["Document"]] = relationship("Document", back_populates="organization", cascade="all, delete-orphan")

letta/orm/passage.py

@@ -1,19 +1,16 @@
 from datetime import datetime
-from typing import Optional, TYPE_CHECKING
-from sqlalchemy import Column, String, DateTime, JSON, ForeignKey
-from sqlalchemy.orm import Mapped, mapped_column, relationship
-from sqlalchemy.types import TypeDecorator, BINARY
+from typing import TYPE_CHECKING, Optional
 
-import numpy as np
-import base64
+from sqlalchemy import JSON, Column, DateTime, ForeignKey, String
+from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from letta.config import LettaConfig
+from letta.constants import MAX_EMBEDDING_DIM
+from letta.orm.custom_columns import CommonVector
+from letta.orm.mixins import FileMixin, OrganizationMixin
 from letta.orm.source import EmbeddingConfigColumn
 from letta.orm.sqlalchemy_base import SqlalchemyBase
-from letta.orm.mixins import FileMixin, OrganizationMixin
 from letta.schemas.passage import Passage as PydanticPassage
-
-from letta.config import LettaConfig
-from letta.constants import MAX_EMBEDDING_DIM
 from letta.settings import settings
 
 config = LettaConfig()
@@ -21,32 +18,12 @@ config = LettaConfig()
 if TYPE_CHECKING:
     from letta.orm.organization import Organization
 
-class CommonVector(TypeDecorator):
-    """Common type for representing vectors in SQLite"""
-    impl = BINARY
-    cache_ok = True
-
-    def load_dialect_impl(self, dialect):
-        return dialect.type_descriptor(BINARY())
-
-    def process_bind_param(self, value, dialect):
-        if value is None:
-            return value
-        if isinstance(value, list):
-            value = np.array(value, dtype=np.float32)
-        return base64.b64encode(value.tobytes())
 
-    def process_result_value(self, value, dialect):
-        if not value:
-            return value
-        if dialect.name == "sqlite":
-            value = base64.b64decode(value)
-        return np.frombuffer(value, dtype=np.float32)
-
-# TODO: After migration to Passage, will need to manually delete passages where files 
+# TODO: After migration to Passage, will need to manually delete passages where files
 #       are deleted on web
 class Passage(SqlalchemyBase, OrganizationMixin, FileMixin):
     """Defines data model for storing Passages"""
+
     __tablename__ = "passages"
     __table_args__ = {"extend_existing": True}
     __pydantic_model__ = PydanticPassage
@@ -59,6 +36,7 @@ class Passage(SqlalchemyBase, OrganizationMixin, FileMixin):
     created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=datetime.utcnow)
     if settings.letta_pg_uri_no_default:
         from pgvector.sqlalchemy import Vector
+
         embedding = mapped_column(Vector(MAX_EMBEDDING_DIM))
     else:
         embedding = Column(CommonVector)

letta/orm/source.py

@@ -1,8 +1,10 @@
 from typing import TYPE_CHECKING, List, Optional
 
-from sqlalchemy import JSON, TypeDecorator
+from sqlalchemy import JSON
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from letta.orm import FileMetadata
+from letta.orm.custom_columns import EmbeddingConfigColumn
 from letta.orm.mixins import OrganizationMixin
 from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.embedding_config import EmbeddingConfig
@@ -12,28 +14,6 @@ if TYPE_CHECKING:
     from letta.orm.organization import Organization
 
 
-class EmbeddingConfigColumn(TypeDecorator):
-    """Custom type for storing EmbeddingConfig as JSON"""
-
-    impl = JSON
-    cache_ok = True
-
-    def load_dialect_impl(self, dialect):
-        return dialect.type_descriptor(JSON())
-
-    def process_bind_param(self, value, dialect):
-        if value:
-            # return vars(value)
-            if isinstance(value, EmbeddingConfig):
-                return value.model_dump()
-        return value
-
-    def process_result_value(self, value, dialect):
-        if value:
-            return EmbeddingConfig(**value)
-        return value
-
-
 class Source(SqlalchemyBase, OrganizationMixin):
     """A source represents an embedded text passage"""
 
@@ -47,5 +27,5 @@ class Source(SqlalchemyBase, OrganizationMixin):
 
     # relationships
     organization: Mapped["Organization"] = relationship("Organization", back_populates="sources")
-    files: Mapped[List["Source"]] = relationship("FileMetadata", back_populates="source", cascade="all, delete-orphan")
-    # agents: Mapped[List["Agent"]] = relationship("Agent", secondary="sources_agents", back_populates="sources")
+    files: Mapped[List["FileMetadata"]] = relationship("FileMetadata", back_populates="source", cascade="all, delete-orphan")
+    agents: Mapped[List["Agent"]] = relationship("Agent", secondary="sources_agents", back_populates="sources")

letta/orm/sources_agents.py

@@ -0,0 +1,13 @@
+from sqlalchemy import ForeignKey, String
+from sqlalchemy.orm import Mapped, mapped_column
+
+from letta.orm.base import Base
+
+
+class SourcesAgents(Base):
+    """Agents can have zero to many sources"""
+
+    __tablename__ = "sources_agents"
+
+    agent_id: Mapped[String] = mapped_column(String, ForeignKey("agents.id"), primary_key=True)
+    source_id: Mapped[String] = mapped_column(String, ForeignKey("sources.id"), primary_key=True)

letta/orm/sqlalchemy_base.py

@@ -1,7 +1,6 @@
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, List, Literal, Optional, Type
-import sqlite3
+from typing import TYPE_CHECKING, List, Literal, Optional
 
 from sqlalchemy import String, desc, func, or_, select
 from sqlalchemy.exc import DBAPIError
@@ -9,12 +8,12 @@ from sqlalchemy.orm import Mapped, Session, mapped_column
 
 from letta.log import get_logger
 from letta.orm.base import Base, CommonSqlalchemyMetaMixins
-from letta.orm.sqlite_functions import adapt_array, convert_array, cosine_distance
 from letta.orm.errors import (
     ForeignKeyConstraintViolationError,
     NoResultFound,
     UniqueConstraintViolationError,
 )
+from letta.orm.sqlite_functions import adapt_array
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
@@ -64,11 +63,26 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         query_text: Optional[str] = None,
         query_embedding: Optional[List[float]] = None,
         ascending: bool = True,
+        tags: Optional[List[str]] = None,
+        match_all_tags: bool = False,
         **kwargs,
-    ) -> List[Type["SqlalchemyBase"]]:
+    ) -> List["SqlalchemyBase"]:
         """
         List records with cursor-based pagination, ordering by created_at.
         Cursor is an ID, but pagination is based on the cursor object's created_at value.
+
+        Args:
+            db_session: SQLAlchemy session
+            cursor: ID of the last item seen (for pagination)
+            start_date: Filter items after this date
+            end_date: Filter items before this date
+            limit: Maximum number of items to return
+            query_text: Text to search for
+            query_embedding: Vector to search for similar embeddings
+            ascending: Sort direction
+            tags: List of tags to filter by
+            match_all_tags: If True, return items matching all tags. If False, match any tag.
+            **kwargs: Additional filters to apply
         """
         if start_date and end_date and start_date > end_date:
             raise ValueError("start_date must be earlier than or equal to end_date")
@@ -84,7 +98,25 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
 
             query = select(cls)
 
-            # Apply filtering logic
+            # Handle tag filtering if the model has tags
+            if tags and hasattr(cls, "tags"):
+                query = select(cls)
+
+                if match_all_tags:
+                    # Match ALL tags - use subqueries
+                    for tag in tags:
+                        subquery = select(cls.tags.property.mapper.class_.agent_id).where(cls.tags.property.mapper.class_.tag == tag)
+                        query = query.filter(cls.id.in_(subquery))
+                else:
+                    # Match ANY tag - use join and filter
+                    query = (
+                        query.join(cls.tags).filter(cls.tags.property.mapper.class_.tag.in_(tags)).group_by(cls.id)  # Deduplicate results
+                    )
+
+                # Group by primary key and all necessary columns to avoid JSON comparison
+                query = query.group_by(cls.id)
+
+            # Apply filtering logic from kwargs
             for key, value in kwargs.items():
                 column = getattr(cls, key)
                 if isinstance(value, (list, tuple, set)):
@@ -98,9 +130,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             if end_date:
                 query = query.filter(cls.created_at < end_date)
 
-            # Cursor-based pagination using created_at
-            # TODO: There is a really nasty race condition issue here with Sqlite
-            # TODO: If they have the same created_at timestamp, this query does NOT match for whatever reason
+            # Cursor-based pagination
             if cursor_obj:
                 if ascending:
                     query = query.where(cls.created_at >= cursor_obj.created_at).where(
@@ -111,40 +141,34 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
                         or_(cls.created_at < cursor_obj.created_at, cls.id < cursor_obj.id)
                     )
 
-            # Apply text search
+            # Text search
             if query_text:
-                from sqlalchemy import func
                 query = query.filter(func.lower(cls.text).contains(func.lower(query_text)))
 
-            # Apply embedding search (Passages)
+            # Embedding search (for Passages)
             is_ordered = False
             if query_embedding:
-                # check if embedding column exists. should only exist for passages
                 if not hasattr(cls, "embedding"):
                     raise ValueError(f"Class {cls.__name__} does not have an embedding column")
-                
+
                 from letta.settings import settings
+
                 if settings.letta_pg_uri_no_default:
                     # PostgreSQL with pgvector
-                    from pgvector.sqlalchemy import Vector
                     query = query.order_by(cls.embedding.cosine_distance(query_embedding).asc())
                 else:
                     # SQLite with custom vector type
-                    from sqlalchemy import func
-
                     query_embedding_binary = adapt_array(query_embedding)
                     query = query.order_by(
-                        func.cosine_distance(cls.embedding, query_embedding_binary).asc(),
-                        cls.created_at.asc(),
-                        cls.id.asc()
+                        func.cosine_distance(cls.embedding, query_embedding_binary).asc(), cls.created_at.asc(), cls.id.asc()
                     )
                     is_ordered = True
 
-            # Handle ordering and soft deletes
+            # Handle soft deletes
             if hasattr(cls, "is_deleted"):
                 query = query.where(cls.is_deleted == False)
-            
-            # Apply ordering by created_at
+
+            # Apply ordering
             if not is_ordered:
                 if ascending:
                     query = query.order_by(cls.created_at, cls.id)
@@ -164,7 +188,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         access: Optional[List[Literal["read", "write", "admin"]]] = ["read"],
         access_type: AccessType = AccessType.ORGANIZATION,
         **kwargs,
-    ) -> Type["SqlalchemyBase"]:
+    ) -> "SqlalchemyBase":
         """The primary accessor for an ORM record.
         Args:
             db_session: the database session to use when retrieving the record
@@ -207,7 +231,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         conditions_str = ", ".join(query_conditions) if query_conditions else "no specific conditions"
         raise NoResultFound(f"{cls.__name__} not found with {conditions_str}")
 
-    def create(self, db_session: "Session", actor: Optional["User"] = None) -> Type["SqlalchemyBase"]:
+    def create(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Creating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
 
         if actor:
@@ -221,7 +245,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         except DBAPIError as e:
             self._handle_dbapi_error(e)
 
-    def delete(self, db_session: "Session", actor: Optional["User"] = None) -> Type["SqlalchemyBase"]:
+    def delete(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Soft deleting {self.__class__.__name__} with ID: {self.id} with actor={actor}")
 
         if actor:
@@ -245,7 +269,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             else:
                 logger.debug(f"{self.__class__.__name__} with ID {self.id} successfully hard deleted")
 
-    def update(self, db_session: "Session", actor: Optional["User"] = None) -> Type["SqlalchemyBase"]:
+    def update(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Updating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
         if actor:
             self._set_created_and_updated_by_fields(actor.id)
@@ -388,14 +412,14 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         raise
 
     @property
-    def __pydantic_model__(self) -> Type["BaseModel"]:
+    def __pydantic_model__(self) -> "BaseModel":
         raise NotImplementedError("Sqlalchemy models must declare a __pydantic_model__ property to be convertable.")
 
-    def to_pydantic(self) -> Type["BaseModel"]:
+    def to_pydantic(self) -> "BaseModel":
         """converts to the basic pydantic model counterpart"""
         return self.__pydantic_model__.model_validate(self)
 
-    def to_record(self) -> Type["BaseModel"]:
+    def to_record(self) -> "BaseModel":
         """Deprecated accessor for to_pydantic"""
         logger.warning("to_record is deprecated, use to_pydantic instead.")
-        return self.to_pydantic()
+        return self.to_pydantic()

letta/orm/tool.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, List, Optional
 
-from sqlalchemy import JSON, String, UniqueConstraint, event
+from sqlalchemy import JSON, String, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 # TODO everything in functions should live in this model
@@ -11,7 +11,6 @@ from letta.schemas.tool import Tool as PydanticTool
 
 if TYPE_CHECKING:
     from letta.orm.organization import Organization
-    from letta.orm.tools_agents import ToolsAgents
 
 
 class Tool(SqlalchemyBase, OrganizationMixin):
@@ -42,20 +41,3 @@ class Tool(SqlalchemyBase, OrganizationMixin):
 
     # relationships
     organization: Mapped["Organization"] = relationship("Organization", back_populates="tools", lazy="selectin")
-    tools_agents: Mapped[List["ToolsAgents"]] = relationship("ToolsAgents", back_populates="tool", cascade="all, delete-orphan")
-
-
-# Add event listener to update tool_name in ToolsAgents when Tool name changes
-@event.listens_for(Tool, "before_update")
-def update_tool_name_in_tools_agents(mapper, connection, target):
-    """Update tool_name in ToolsAgents when Tool name changes."""
-    state = target._sa_instance_state
-    history = state.get_history("name", passive=True)
-    if not history.has_changes():
-        return
-
-    # Get the new name and update all associated ToolsAgents records
-    new_name = target.name
-    from letta.orm.tools_agents import ToolsAgents
-
-    connection.execute(ToolsAgents.__table__.update().where(ToolsAgents.tool_id == target.id).values(tool_name=new_name))

letta/orm/tools_agents.py

@@ -1,32 +1,15 @@
-from sqlalchemy import ForeignKey, ForeignKeyConstraint, String, UniqueConstraint
-from sqlalchemy.orm import Mapped, mapped_column, relationship
+from sqlalchemy import ForeignKey, String, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column
 
-from letta.orm.sqlalchemy_base import SqlalchemyBase
-from letta.schemas.tools_agents import ToolsAgents as PydanticToolsAgents
+from letta.orm import Base
 
 
-class ToolsAgents(SqlalchemyBase):
+class ToolsAgents(Base):
     """Agents can have one or many tools associated with them."""
 
     __tablename__ = "tools_agents"
-    __pydantic_model__ = PydanticToolsAgents
-    __table_args__ = (
-        UniqueConstraint(
-            "agent_id",
-            "tool_name",
-            name="unique_tool_per_agent",
-        ),
-        ForeignKeyConstraint(
-            ["tool_id"],
-            ["tools.id"],
-            name="fk_tool_id",
-        ),
-    )
+    __table_args__ = (UniqueConstraint("agent_id", "tool_id", name="unique_agent_tool"),)
 
     # Each agent must have unique tool names
-    agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id"), primary_key=True)
-    tool_id: Mapped[str] = mapped_column(String, primary_key=True)
-    tool_name: Mapped[str] = mapped_column(String, primary_key=True)
-
-    # relationships
-    tool: Mapped["Tool"] = relationship("Tool", back_populates="tools_agents")    # agent: Mapped["Agent"] = relationship("Agent", back_populates="tools_agents")
+    agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True)
+    tool_id: Mapped[str] = mapped_column(String, ForeignKey("tools.id", ondelete="CASCADE"), primary_key=True)

letta/orm/user.py

@@ -20,10 +20,9 @@ class User(SqlalchemyBase, OrganizationMixin):
 
     # relationships
     organization: Mapped["Organization"] = relationship("Organization", back_populates="users")
-    jobs: Mapped[List["Job"]] = relationship("Job", back_populates="user", doc="the jobs associated with this user.", cascade="all, delete-orphan")
+    jobs: Mapped[List["Job"]] = relationship(
+        "Job", back_populates="user", doc="the jobs associated with this user.", cascade="all, delete-orphan"
+    )
 
     # TODO: Add this back later potentially
-    # agents: Mapped[List["Agent"]] = relationship(
-    #     "Agent", secondary="users_agents", back_populates="users", doc="the agents associated with this user."
-    # )
     # tokens: Mapped[List["Token"]] = relationship("Token", back_populates="user", doc="the tokens associated with this user.")