diff --git a/letta/orm/mcp_server.py b/letta/orm/mcp_server.py index 14888baf..a829dd53 100644 --- a/letta/orm/mcp_server.py +++ b/letta/orm/mcp_server.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING, Optional from sqlalchemy import JSON, String, Text, UniqueConstraint -from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.orm import Mapped, mapped_column, relationship from letta.functions.mcp_client.types import StdioServerConfig from letta.orm.custom_columns import MCPStdioServerConfigColumn @@ -13,7 +13,7 @@ from letta.schemas.enums import MCPServerType from letta.schemas.mcp import MCPServer if TYPE_CHECKING: - pass + from letta.orm.organization import Organization class MCPServer(SqlalchemyBase, OrganizationMixin): @@ -57,6 +57,9 @@ class MCPServer(SqlalchemyBase, OrganizationMixin): JSON, default=lambda: {}, doc="A dictionary of additional metadata for the MCP server." ) + # relationships + organization: Mapped["Organization"] = relationship("Organization", back_populates="mcp_servers") + class MCPTools(SqlalchemyBase, OrganizationMixin): """Represents a mapping of MCP server ID to tool ID""" diff --git a/letta/orm/organization.py b/letta/orm/organization.py index c24f5553..4e181e86 100644 --- a/letta/orm/organization.py +++ b/letta/orm/organization.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from letta.orm.job import Job from letta.orm.llm_batch_items import LLMBatchItem from letta.orm.llm_batch_job import LLMBatchJob + from letta.orm.mcp_server import MCPServer from letta.orm.message import Message from letta.orm.passage import ArchivalPassage, SourcePassage from letta.orm.passage_tag import PassageTag @@ -39,7 +40,7 @@ class Organization(SqlalchemyBase): # relationships users: Mapped[List["User"]] = relationship("User", back_populates="organization", cascade="all, delete-orphan") tools: Mapped[List["Tool"]] = relationship("Tool", back_populates="organization", cascade="all, delete-orphan") - # mcp_servers: Mapped[List["MCPServer"]] = relationship("MCPServer", back_populates="organization", cascade="all, delete-orphan") + mcp_servers: Mapped[List["MCPServer"]] = relationship("MCPServer", back_populates="organization", cascade="all, delete-orphan") blocks: Mapped[List["Block"]] = relationship("Block", back_populates="organization", cascade="all, delete-orphan") sandbox_configs: Mapped[List["SandboxConfig"]] = relationship( "SandboxConfig", back_populates="organization", cascade="all, delete-orphan"