letta-nightly 0.5.2.dev20241112104101__py3-none-any.whl → 0.5.2.dev20241113234401__py3-none-any.whl

letta/agent.py

@@ -46,6 +46,8 @@ from letta.schemas.passage import Passage
 from letta.schemas.tool import Tool
 from letta.schemas.tool_rule import TerminalToolRule
 from letta.schemas.usage import LettaUsageStatistics
+from letta.services.source_manager import SourceManager
+from letta.services.user_manager import UserManager
 from letta.system import (
     get_heartbeat,
     get_initial_boot_messages,
@@ -1311,7 +1313,7 @@ class Agent(BaseAgent):
     def attach_source(self, source_id: str, source_connector: StorageConnector, ms: MetadataStore):
         """Attach data with name `source_name` to the agent from source_connector."""
         # TODO: eventually, adding a data source should just give access to the retriever the source table, rather than modifying archival memory
-
+        user = UserManager().get_user_by_id(self.agent_state.user_id)
         filters = {"user_id": self.agent_state.user_id, "source_id": source_id}
         size = source_connector.size(filters)
         page_size = 100
@@ -1339,7 +1341,7 @@ class Agent(BaseAgent):
         self.persistence_manager.archival_memory.storage.save()
 
         # attach to agent
-        source = ms.get_source(source_id=source_id)
+        source = SourceManager().get_source_by_id(source_id=source_id, actor=user)
         assert source is not None, f"Source {source_id} not found in metadata store"
         ms.attach_source(agent_id=self.agent_state.id, source_id=source_id, user_id=self.agent_state.user_id)
 

letta/agent_store/db.py

@@ -27,8 +27,9 @@ from tqdm import tqdm
 from letta.agent_store.storage import StorageConnector, TableType
 from letta.config import LettaConfig
 from letta.constants import MAX_EMBEDDING_DIM
-from letta.metadata import EmbeddingConfigColumn, FileMetadataModel, ToolCallColumn
+from letta.metadata import EmbeddingConfigColumn, ToolCallColumn
 from letta.orm.base import Base
+from letta.orm.file import FileMetadata as FileMetadataModel
 
 # from letta.schemas.message import Message, Passage, Record, RecordType, ToolCall
 from letta.schemas.message import Message

letta/cli/cli.py

@@ -47,6 +47,7 @@ def server(
     host: Annotated[Optional[str], typer.Option(help="Host to run the server on (default to localhost)")] = None,
     debug: Annotated[bool, typer.Option(help="Turn debugging output on")] = False,
     ade: Annotated[bool, typer.Option(help="Allows remote access")] = False,
+    secure: Annotated[bool, typer.Option(help="Adds simple security access")] = False,
 ):
     """Launch a Letta server process"""
     if type == ServerChoice.rest_api:

letta/client/client.py

@@ -238,7 +238,7 @@ class AbstractClient(object):
     def delete_file_from_source(self, source_id: str, file_id: str) -> None:
         raise NotImplementedError
 
-    def create_source(self, name: str) -> Source:
+    def create_source(self, name: str, embedding_config: Optional[EmbeddingConfig] = None) -> Source:
         raise NotImplementedError
 
     def delete_source(self, source_id: str):
@@ -1188,7 +1188,7 @@ class RESTClient(AbstractClient):
         if response.status_code not in [200, 204]:
             raise ValueError(f"Failed to delete tool: {response.text}")
 
-    def create_source(self, name: str) -> Source:
+    def create_source(self, name: str, embedding_config: Optional[EmbeddingConfig] = None) -> Source:
         """
         Create a source
 
@@ -1198,7 +1198,8 @@ class RESTClient(AbstractClient):
         Returns:
             source (Source): Created source
         """
-        payload = {"name": name}
+        source_create = SourceCreate(name=name, embedding_config=embedding_config or self._default_embedding_config)
+        payload = source_create.model_dump()
         response = requests.post(f"{self.base_url}/{self.api_prefix}/sources", json=payload, headers=self.headers)
         response_json = response.json()
         return Source(**response_json)
@@ -1253,7 +1254,7 @@ class RESTClient(AbstractClient):
         Returns:
             source (Source): Updated source
         """
-        request = SourceUpdate(id=source_id, name=name)
+        request = SourceUpdate(name=name)
         response = requests.patch(f"{self.base_url}/{self.api_prefix}/sources/{source_id}", json=request.model_dump(), headers=self.headers)
         if response.status_code != 200:
             raise ValueError(f"Failed to update source: {response.text}")
@@ -2394,7 +2395,7 @@ class LocalClient(AbstractClient):
         Args:
             id (str): ID of the tool
         """
-        return self.server.tool_manager.delete_tool_by_id(id, user_id=self.user_id)
+        return self.server.tool_manager.delete_tool_by_id(id, actor=self.user)
 
     def get_tool_id(self, name: str) -> Optional[str]:
         """
@@ -2439,7 +2440,7 @@ class LocalClient(AbstractClient):
         return job
 
     def delete_file_from_source(self, source_id: str, file_id: str):
-        self.server.delete_file_from_source(source_id, file_id, user_id=self.user_id)
+        self.server.source_manager.delete_file(file_id, actor=self.user)
 
     def get_job(self, job_id: str):
         return self.server.get_job(job_id=job_id)
@@ -2453,7 +2454,7 @@ class LocalClient(AbstractClient):
     def list_active_jobs(self):
         return self.server.list_active_jobs(user_id=self.user_id)
 
-    def create_source(self, name: str) -> Source:
+    def create_source(self, name: str, embedding_config: Optional[EmbeddingConfig] = None) -> Source:
         """
         Create a source
 
@@ -2463,8 +2464,10 @@ class LocalClient(AbstractClient):
         Returns:
             source (Source): Created source
         """
-        request = SourceCreate(name=name)
-        return self.server.create_source(request=request, user_id=self.user_id)
+        source = Source(
+            name=name, embedding_config=embedding_config or self._default_embedding_config, organization_id=self.user.organization_id
+        )
+        return self.server.source_manager.create_source(source=source, actor=self.user)
 
     def delete_source(self, source_id: str):
         """
@@ -2475,7 +2478,7 @@ class LocalClient(AbstractClient):
         """
 
         # TODO: delete source data
-        self.server.delete_source(source_id=source_id, user_id=self.user_id)
+        self.server.delete_source(source_id=source_id, actor=self.user)
 
     def get_source(self, source_id: str) -> Source:
         """
@@ -2487,7 +2490,7 @@ class LocalClient(AbstractClient):
         Returns:
             source (Source): Source
         """
-        return self.server.get_source(source_id=source_id, user_id=self.user_id)
+        return self.server.source_manager.get_source_by_id(source_id=source_id, actor=self.user)
 
     def get_source_id(self, source_name: str) -> str:
         """
@@ -2499,7 +2502,7 @@ class LocalClient(AbstractClient):
         Returns:
             source_id (str): ID of the source
         """
-        return self.server.get_source_id(source_name=source_name, user_id=self.user_id)
+        return self.server.source_manager.get_source_by_name(source_name=source_name, actor=self.user).id
 
     def attach_source_to_agent(self, agent_id: str, source_id: Optional[str] = None, source_name: Optional[str] = None):
         """
@@ -2532,7 +2535,7 @@ class LocalClient(AbstractClient):
             sources (List[Source]): List of sources
         """
 
-        return self.server.list_all_sources(user_id=self.user_id)
+        return self.server.list_all_sources(actor=self.user)
 
     def list_attached_sources(self, agent_id: str) -> List[Source]:
         """
@@ -2558,7 +2561,7 @@ class LocalClient(AbstractClient):
         Returns:
             files (List[FileMetadata]): List of files
         """
-        return self.server.list_files_from_source(source_id=source_id, limit=limit, cursor=cursor)
+        return self.server.source_manager.list_files(source_id=source_id, limit=limit, cursor=cursor, actor=self.user)
 
     def update_source(self, source_id: str, name: Optional[str] = None) -> Source:
         """
@@ -2572,8 +2575,8 @@ class LocalClient(AbstractClient):
             source (Source): Updated source
         """
         # TODO should the arg here just be "source_update: Source"?
-        request = SourceUpdate(id=source_id, name=name)
-        return self.server.update_source(request=request, user_id=self.user_id)
+        request = SourceUpdate(name=name)
+        return self.server.source_manager.update_source(source_id=source_id, source_update=request, actor=self.user)
 
     # archival memory
 

letta/data_sources/connectors.py

@@ -12,6 +12,7 @@ from letta.embeddings import embedding_model
 from letta.schemas.file import FileMetadata
 from letta.schemas.passage import Passage
 from letta.schemas.source import Source
+from letta.services.source_manager import SourceManager
 from letta.utils import create_uuid_from_string
 
 
@@ -41,13 +42,8 @@ class DataConnector:
         """
 
 
-def load_data(
-    connector: DataConnector,
-    source: Source,
-    passage_store: StorageConnector,
-    file_metadata_store: StorageConnector,
-):
-    """Load data from a connector (generates file and passages) into a specified source_id, associatedw with a user_id."""
+def load_data(connector: DataConnector, source: Source, passage_store: StorageConnector, source_manager: SourceManager, actor: "User"):
+    """Load data from a connector (generates file and passages) into a specified source_id, associated with a user_id."""
     embedding_config = source.embedding_config
 
     # embedding model
@@ -60,7 +56,7 @@ def load_data(
     file_count = 0
     for file_metadata in connector.find_files(source):
         file_count += 1
-        file_metadata_store.insert(file_metadata)
+        source_manager.create_file(file_metadata, actor)
 
         # generate passages
         for passage_text, passage_metadata in connector.generate_passages(file_metadata, chunk_size=embedding_config.embedding_chunk_size):
@@ -88,7 +84,7 @@ def load_data(
                 file_id=file_metadata.id,
                 source_id=source.id,
                 metadata_=passage_metadata,
-                user_id=source.user_id,
+                user_id=source.created_by_id,
                 embedding_config=source.embedding_config,
                 embedding=embedding,
             )
@@ -155,7 +151,6 @@ class DirectoryConnector(DataConnector):
 
         for metadata in extract_metadata_from_files(files):
             yield FileMetadata(
-                user_id=source.user_id,
                 source_id=source.id,
                 file_name=metadata.get("file_name"),
                 file_path=metadata.get("file_path"),

letta/llm_api/google_ai.py

@@ -95,10 +95,8 @@ def google_ai_get_model_list(base_url: str, api_key: str, key_in_header: bool =
 
     try:
         response = requests.get(url, headers=headers)
-        printd(f"response = {response}")
         response.raise_for_status()  # Raises HTTPError for 4XX/5XX status
         response = response.json()  # convert to dict from string
-        printd(f"response.json = {response}")
 
         # Grab the models out
         model_list = response["models"]

letta/llm_api/openai.py

@@ -126,6 +126,7 @@ def build_openai_chat_completions_request(
     openai_message_list = [
         cast_message_to_subtype(m.to_openai_dict(put_inner_thoughts_in_kwargs=llm_config.put_inner_thoughts_in_kwargs)) for m in messages
     ]
+
     if llm_config.model:
         model = llm_config.model
     else:

letta/memory.py

@@ -7,6 +7,7 @@ from letta.embeddings import embedding_model, parse_and_chunk_text, query_embedd
 from letta.llm_api.llm_api_tools import create
 from letta.prompts.gpt_summarize import SYSTEM as SUMMARY_PROMPT_SYSTEM
 from letta.schemas.agent import AgentState
+from letta.schemas.enums import MessageRole
 from letta.schemas.memory import Memory
 from letta.schemas.message import Message
 from letta.schemas.passage import Passage
@@ -50,7 +51,6 @@ def _format_summary_history(message_history: List[Message]):
 def summarize_messages(
     agent_state: AgentState,
     message_sequence_to_summarize: List[Message],
-    insert_acknowledgement_assistant_message: bool = True,
 ):
     """Summarize a message sequence using GPT"""
     # we need the context_window
@@ -70,13 +70,17 @@ def summarize_messages(
     dummy_user_id = agent_state.user_id
     dummy_agent_id = agent_state.id
     message_sequence = []
-    message_sequence.append(Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role="system", text=summary_prompt))
-    if insert_acknowledgement_assistant_message:
-        message_sequence.append(Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role="assistant", text=MESSAGE_SUMMARY_REQUEST_ACK))
-    message_sequence.append(Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role="user", text=summary_input))
+    message_sequence.append(Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role=MessageRole.system, text=summary_prompt))
+    message_sequence.append(
+        Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role=MessageRole.assistant, text=MESSAGE_SUMMARY_REQUEST_ACK)
+    )
+    message_sequence.append(Message(user_id=dummy_user_id, agent_id=dummy_agent_id, role=MessageRole.user, text=summary_input))
 
+    # TODO: We need to eventually have a separate LLM config for the summarizer LLM
+    llm_config_no_inner_thoughts = agent_state.llm_config.model_copy(deep=True)
+    llm_config_no_inner_thoughts.put_inner_thoughts_in_kwargs = False
     response = create(
-        llm_config=agent_state.llm_config,
+        llm_config=llm_config_no_inner_thoughts,
         user_id=agent_state.user_id,
         messages=message_sequence,
         stream=False,

letta/metadata.py

@@ -11,7 +11,6 @@ from sqlalchemy import (
     Column,
     DateTime,
     Index,
-    Integer,
     String,
     TypeDecorator,
 )
@@ -24,12 +23,10 @@ from letta.schemas.api_key import APIKey
 from letta.schemas.block import Block, Human, Persona
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import JobStatus
-from letta.schemas.file import FileMetadata
 from letta.schemas.job import Job
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import Memory
 from letta.schemas.openai.chat_completions import ToolCall, ToolCallFunction
-from letta.schemas.source import Source
 from letta.schemas.tool_rule import (
     BaseToolRule,
     InitToolRule,
@@ -41,41 +38,6 @@ from letta.settings import settings
 from letta.utils import enforce_types, get_utc_time, printd
 
 
-class FileMetadataModel(Base):
-    __tablename__ = "files"
-    __table_args__ = {"extend_existing": True}
-
-    id = Column(String, primary_key=True, nullable=False)
-    user_id = Column(String, nullable=False)
-    # TODO: Investigate why this breaks during table creation due to FK
-    # source_id = Column(String, ForeignKey("sources.id"), nullable=False)
-    source_id = Column(String, nullable=False)
-    file_name = Column(String, nullable=True)
-    file_path = Column(String, nullable=True)
-    file_type = Column(String, nullable=True)
-    file_size = Column(Integer, nullable=True)
-    file_creation_date = Column(String, nullable=True)
-    file_last_modified_date = Column(String, nullable=True)
-    created_at = Column(DateTime(timezone=True), server_default=func.now())
-
-    def __repr__(self):
-        return f"<FileMetadata(id='{self.id}', source_id='{self.source_id}', file_name='{self.file_name}')>"
-
-    def to_record(self):
-        return FileMetadata(
-            id=self.id,
-            user_id=self.user_id,
-            source_id=self.source_id,
-            file_name=self.file_name,
-            file_path=self.file_path,
-            file_type=self.file_type,
-            file_size=self.file_size,
-            file_creation_date=self.file_creation_date,
-            file_last_modified_date=self.file_last_modified_date,
-            created_at=self.created_at,
-        )
-
-
 class LLMConfigColumn(TypeDecorator):
     """Custom type for storing LLMConfig as JSON"""
 
@@ -292,40 +254,6 @@ class AgentModel(Base):
         return agent_state
 
 
-class SourceModel(Base):
-    """Defines data model for storing Passages (consisting of text, embedding)"""
-
-    __tablename__ = "sources"
-    __table_args__ = {"extend_existing": True}
-
-    # Assuming passage_id is the primary key
-    # id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
-    id = Column(String, primary_key=True)
-    user_id = Column(String, nullable=False)
-    name = Column(String, nullable=False)
-    created_at = Column(DateTime(timezone=True), server_default=func.now())
-    embedding_config = Column(EmbeddingConfigColumn)
-    description = Column(String)
-    metadata_ = Column(JSON)
-    Index(__tablename__ + "_idx_user", user_id),
-
-    # TODO: add num passages
-
-    def __repr__(self) -> str:
-        return f"<Source(passage_id='{self.id}', name='{self.name}')>"
-
-    def to_record(self) -> Source:
-        return Source(
-            id=self.id,
-            user_id=self.user_id,
-            name=self.name,
-            created_at=self.created_at,
-            embedding_config=self.embedding_config,
-            description=self.description,
-            metadata_=self.metadata_,
-        )
-
-
 class AgentSourceMappingModel(Base):
     """Stores mapping between agent -> source"""
 
@@ -497,14 +425,6 @@ class MetadataStore:
             session.add(AgentModel(**fields))
             session.commit()
 
-    @enforce_types
-    def create_source(self, source: Source):
-        with self.session_maker() as session:
-            if session.query(SourceModel).filter(SourceModel.name == source.name).filter(SourceModel.user_id == source.user_id).count() > 0:
-                raise ValueError(f"Source with name {source.name} already exists for user {source.user_id}")
-            session.add(SourceModel(**vars(source)))
-            session.commit()
-
     @enforce_types
     def create_block(self, block: Block):
         with self.session_maker() as session:
@@ -522,6 +442,7 @@ class MetadataStore:
             ):
 
                 raise ValueError(f"Block with name {block.template_name} already exists")
+
             session.add(BlockModel(**vars(block)))
             session.commit()
 
@@ -536,12 +457,6 @@ class MetadataStore:
             session.query(AgentModel).filter(AgentModel.id == agent.id).update(fields)
             session.commit()
 
-    @enforce_types
-    def update_source(self, source: Source):
-        with self.session_maker() as session:
-            session.query(SourceModel).filter(SourceModel.id == source.id).update(vars(source))
-            session.commit()
-
     @enforce_types
     def update_block(self, block: Block):
         with self.session_maker() as session:
@@ -558,21 +473,6 @@ class MetadataStore:
                 session.add(BlockModel(**vars(block)))
             session.commit()
 
-    @enforce_types
-    def delete_file_from_source(self, source_id: str, file_id: str, user_id: Optional[str]):
-        with self.session_maker() as session:
-            file_metadata = (
-                session.query(FileMetadataModel)
-                .filter(FileMetadataModel.source_id == source_id, FileMetadataModel.id == file_id, FileMetadataModel.user_id == user_id)
-                .first()
-            )
-
-            if file_metadata:
-                session.delete(file_metadata)
-                session.commit()
-
-            return file_metadata
-
     @enforce_types
     def delete_block(self, block_id: str):
         with self.session_maker() as session:
@@ -591,29 +491,12 @@ class MetadataStore:
 
             session.commit()
 
-    @enforce_types
-    def delete_source(self, source_id: str):
-        with self.session_maker() as session:
-            # delete from sources table
-            session.query(SourceModel).filter(SourceModel.id == source_id).delete()
-
-            # delete any mappings
-            session.query(AgentSourceMappingModel).filter(AgentSourceMappingModel.source_id == source_id).delete()
-
-            session.commit()
-
     @enforce_types
     def list_agents(self, user_id: str) -> List[AgentState]:
         with self.session_maker() as session:
             results = session.query(AgentModel).filter(AgentModel.user_id == user_id).all()
             return [r.to_record() for r in results]
 
-    @enforce_types
-    def list_sources(self, user_id: str) -> List[Source]:
-        with self.session_maker() as session:
-            results = session.query(SourceModel).filter(SourceModel.user_id == user_id).all()
-            return [r.to_record() for r in results]
-
     @enforce_types
     def get_agent(
         self, agent_id: Optional[str] = None, agent_name: Optional[str] = None, user_id: Optional[str] = None
@@ -630,21 +513,6 @@ class MetadataStore:
             assert len(results) == 1, f"Expected 1 result, got {len(results)}"  # should only be one result
             return results[0].to_record()
 
-    @enforce_types
-    def get_source(
-        self, source_id: Optional[str] = None, user_id: Optional[str] = None, source_name: Optional[str] = None
-    ) -> Optional[Source]:
-        with self.session_maker() as session:
-            if source_id:
-                results = session.query(SourceModel).filter(SourceModel.id == source_id).all()
-            else:
-                assert user_id is not None and source_name is not None
-                results = session.query(SourceModel).filter(SourceModel.name == source_name).filter(SourceModel.user_id == user_id).all()
-            if len(results) == 0:
-                return None
-            assert len(results) == 1, f"Expected 1 result, got {len(results)}"
-            return results[0].to_record()
-
     @enforce_types
     def get_block(self, block_id: str) -> Optional[Block]:
         with self.session_maker() as session:
@@ -699,19 +567,10 @@ class MetadataStore:
             session.commit()
 
     @enforce_types
-    def list_attached_sources(self, agent_id: str) -> List[Source]:
+    def list_attached_source_ids(self, agent_id: str) -> List[str]:
         with self.session_maker() as session:
             results = session.query(AgentSourceMappingModel).filter(AgentSourceMappingModel.agent_id == agent_id).all()
-
-            sources = []
-            # make sure source exists
-            for r in results:
-                source = self.get_source(source_id=r.source_id)
-                if source:
-                    sources.append(source)
-                else:
-                    printd(f"Warning: source {r.source_id} does not exist but exists in mapping database. This should never happen.")
-            return sources
+            return [r.source_id for r in results]
 
     @enforce_types
     def list_attached_agents(self, source_id: str) -> List[str]:
@@ -742,27 +601,6 @@ class MetadataStore:
             session.add(JobModel(**vars(job)))
             session.commit()
 
-    @enforce_types
-    def list_files_from_source(self, source_id: str, limit: int, cursor: Optional[str]):
-        with self.session_maker() as session:
-            # Start with the basic query filtered by source_id
-            query = session.query(FileMetadataModel).filter(FileMetadataModel.source_id == source_id)
-
-            if cursor:
-                # Assuming cursor is the ID of the last file in the previous page
-                query = query.filter(FileMetadataModel.id > cursor)
-
-            # Order by ID or other ordering criteria to ensure correct pagination
-            query = query.order_by(FileMetadataModel.id)
-
-            # Limit the number of results returned
-            results = query.limit(limit).all()
-
-            # Convert the results to the required FileMetadata objects
-            files = [r.to_record() for r in results]
-
-            return files
-
     def delete_job(self, job_id: str):
         with self.session_maker() as session:
             session.query(JobModel).filter(JobModel.id == job_id).delete()

letta/orm/__init__.py

@@ -1,4 +1,6 @@
 from letta.orm.base import Base
+from letta.orm.file import FileMetadata
 from letta.orm.organization import Organization
+from letta.orm.source import Source
 from letta.orm.tool import Tool
 from letta.orm.user import User

letta/orm/file.py

@@ -0,0 +1,29 @@
+from typing import TYPE_CHECKING, Optional
+
+from sqlalchemy import Integer, String
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from letta.orm.mixins import OrganizationMixin, SourceMixin
+from letta.orm.sqlalchemy_base import SqlalchemyBase
+from letta.schemas.file import FileMetadata as PydanticFileMetadata
+
+if TYPE_CHECKING:
+    from letta.orm.organization import Organization
+
+
+class FileMetadata(SqlalchemyBase, OrganizationMixin, SourceMixin):
+    """Represents metadata for an uploaded file."""
+
+    __tablename__ = "files"
+    __pydantic_model__ = PydanticFileMetadata
+
+    file_name: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The name of the file.")
+    file_path: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The file path on the system.")
+    file_type: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The type of the file.")
+    file_size: Mapped[Optional[int]] = mapped_column(Integer, nullable=True, doc="The size of the file in bytes.")
+    file_creation_date: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The creation date of the file.")
+    file_last_modified_date: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The last modified date of the file.")
+
+    # relationships
+    organization: Mapped["Organization"] = relationship("Organization", back_populates="files", lazy="selectin")
+    source: Mapped["Source"] = relationship("Source", back_populates="files", lazy="selectin")

letta/orm/mixins.py

@@ -29,3 +29,11 @@ class UserMixin(Base):
     __abstract__ = True
 
     user_id: Mapped[str] = mapped_column(String, ForeignKey("users.id"))
+
+
+class SourceMixin(Base):
+    """Mixin for models (e.g. file) that belong to a source."""
+
+    __abstract__ = True
+
+    source_id: Mapped[str] = mapped_column(String, ForeignKey("sources.id"))

letta/orm/organization.py

@@ -1,8 +1,8 @@
 from typing import TYPE_CHECKING, List
 
-from sqlalchemy import String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from letta.orm.file import FileMetadata
 from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.organization import Organization as PydanticOrganization
 
@@ -18,16 +18,16 @@ class Organization(SqlalchemyBase):
     __tablename__ = "organizations"
     __pydantic_model__ = PydanticOrganization
 
-    id: Mapped[str] = mapped_column(String, primary_key=True)
     name: Mapped[str] = mapped_column(doc="The display name of the organization.")
 
+    # 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")
+    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")
     # 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")
-    # sources: Mapped[List["Source"]] = relationship("Source", 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")