letta-nightly 0.5.0.dev20241022104124__py3-none-any.whl → 0.5.0.dev20241023104105__py3-none-any.whl

letta/orm/sqlalchemy_base.py

@@ -0,0 +1,214 @@
+from typing import TYPE_CHECKING, List, Literal, Optional, Type, Union
+from uuid import UUID, uuid4
+
+from humps import depascalize
+from sqlalchemy import Boolean, String, select
+from sqlalchemy.orm import Mapped, mapped_column
+
+from letta.log import get_logger
+from letta.orm.base import Base, CommonSqlalchemyMetaMixins
+from letta.orm.errors import NoResultFound
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+    from sqlalchemy.orm import Session
+
+    # from letta.orm.user import User
+
+logger = get_logger(__name__)
+
+
+class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
+    __abstract__ = True
+
+    __order_by_default__ = "created_at"
+
+    _id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"{uuid4()}")
+
+    deleted: Mapped[bool] = mapped_column(Boolean, default=False, doc="Is this record deleted? Used for universal soft deletes.")
+
+    @classmethod
+    def __prefix__(cls) -> str:
+        return depascalize(cls.__name__)
+
+    @property
+    def id(self) -> Optional[str]:
+        if self._id:
+            return f"{self.__prefix__()}-{self._id}"
+
+    @id.setter
+    def id(self, value: str) -> None:
+        if not value:
+            return
+        prefix, id_ = value.split("-", 1)
+        assert prefix == self.__prefix__(), f"{prefix} is not a valid id prefix for {self.__class__.__name__}"
+        assert SqlalchemyBase.is_valid_uuid4(id_), f"{id_} is not a valid uuid4"
+        self._id = id_
+
+    @classmethod
+    def list(
+        cls, *, db_session: "Session", cursor: Optional[str] = None, limit: Optional[int] = 50, **kwargs
+    ) -> List[Type["SqlalchemyBase"]]:
+        """List records with optional cursor (for pagination) and limit."""
+        with db_session as session:
+            # Start with the base query filtered by kwargs
+            query = select(cls).filter_by(**kwargs)
+
+            # Add a cursor condition if provided
+            if cursor:
+                cursor_uuid = cls.get_uid_from_identifier(cursor)  # Assuming the cursor is an _id value
+                query = query.where(cls._id > cursor_uuid)
+
+            # Add a limit to the query if provided
+            query = query.order_by(cls._id).limit(limit)
+
+            # Handle soft deletes if the class has the 'is_deleted' attribute
+            if hasattr(cls, "is_deleted"):
+                query = query.where(cls.is_deleted == False)
+
+            # Execute the query and return the results as a list of model instances
+            return list(session.execute(query).scalars())
+
+    @classmethod
+    def get_uid_from_identifier(cls, identifier: str, indifferent: Optional[bool] = False) -> str:
+        """converts the id into a uuid object
+        Args:
+            identifier: the string identifier, such as `organization-xxxx-xx...`
+            indifferent: if True, will not enforce the prefix check
+        """
+        try:
+            uuid_string = identifier.split("-", 1)[1] if indifferent else identifier.replace(f"{cls.__prefix__()}-", "")
+            assert SqlalchemyBase.is_valid_uuid4(uuid_string)
+            return uuid_string
+        except ValueError as e:
+            raise ValueError(f"{identifier} is not a valid identifier for class {cls.__name__}") from e
+
+    @classmethod
+    def is_valid_uuid4(cls, uuid_string: str) -> bool:
+        try:
+            # Try to create a UUID object from the string
+            uuid_obj = UUID(uuid_string)
+            # Check if the UUID is version 4
+            return uuid_obj.version == 4
+        except ValueError:
+            # Raised if the string is not a valid UUID
+            return False
+
+    @classmethod
+    def read(
+        cls,
+        db_session: "Session",
+        identifier: Union[str, UUID],
+        actor: Optional["User"] = None,
+        access: Optional[List[Literal["read", "write", "admin"]]] = ["read"],
+        **kwargs,
+    ) -> Type["SqlalchemyBase"]:
+        """The primary accessor for an ORM record.
+        Args:
+            db_session: the database session to use when retrieving the record
+            identifier: the identifier of the record to read, can be the id string or the UUID object for backwards compatibility
+            actor: if specified, results will be scoped only to records the user is able to access
+            access: if actor is specified, records will be filtered to the minimum permission level for the actor
+            kwargs: additional arguments to pass to the read, used for more complex objects
+        Returns:
+            The matching object
+        Raises:
+            NoResultFound: if the object is not found
+        """
+        del kwargs  # arity for more complex reads
+        identifier = cls.get_uid_from_identifier(identifier)
+        query = select(cls).where(cls._id == identifier)
+        # if actor:
+        #     query = cls.apply_access_predicate(query, actor, access)
+        if hasattr(cls, "is_deleted"):
+            query = query.where(cls.is_deleted == False)
+        if found := db_session.execute(query).scalar():
+            return found
+        raise NoResultFound(f"{cls.__name__} with id {identifier} not found")
+
+    def create(self, db_session: "Session") -> Type["SqlalchemyBase"]:
+        # self._infer_organization(db_session)
+
+        with db_session as session:
+            session.add(self)
+            session.commit()
+            session.refresh(self)
+            return self
+
+    def delete(self, db_session: "Session") -> Type["SqlalchemyBase"]:
+        self.is_deleted = True
+        return self.update(db_session)
+
+    def update(self, db_session: "Session") -> Type["SqlalchemyBase"]:
+        with db_session as session:
+            session.add(self)
+            session.commit()
+            session.refresh(self)
+            return self
+
+    @classmethod
+    def read_or_create(cls, *, db_session: "Session", **kwargs) -> Type["SqlalchemyBase"]:
+        """get an instance by search criteria or create it if it doesn't exist"""
+        try:
+            return cls.read(db_session=db_session, identifier=kwargs.get("id", None))
+        except NoResultFound:
+            clean_kwargs = {k: v for k, v in kwargs.items() if k in cls.__table__.columns}
+            return cls(**clean_kwargs).create(db_session=db_session)
+
+    # TODO: Add back later when access predicates are actually important
+    # The idea behind this is that you can add a WHERE clause restricting the actions you can take, e.g. R/W
+    # @classmethod
+    # def apply_access_predicate(
+    #     cls,
+    #     query: "Select",
+    #     actor: "User",
+    #     access: List[Literal["read", "write", "admin"]],
+    # ) -> "Select":
+    #     """applies a WHERE clause restricting results to the given actor and access level
+    #     Args:
+    #         query: The initial sqlalchemy select statement
+    #         actor: The user acting on the query. **Note**: this is called 'actor' to identify the
+    #                person or system acting. Users can act on users, making naming very sticky otherwise.
+    #         access:
+    #             what mode of access should the query restrict to? This will be used with granular permissions,
+    #             but because of how it will impact every query we want to be explicitly calling access ahead of time.
+    #     Returns:
+    #         the sqlalchemy select statement restricted to the given access.
+    #     """
+    #     del access  # entrypoint for row-level permissions. Defaults to "same org as the actor, all permissions" at the moment
+    #     org_uid = getattr(actor, "_organization_id", getattr(actor.organization, "_id", None))
+    #     if not org_uid:
+    #         raise ValueError("object %s has no organization accessor", actor)
+    #     return query.where(cls._organization_id == org_uid, cls.is_deleted == False)
+
+    @property
+    def __pydantic_model__(self) -> Type["BaseModel"]:
+        raise NotImplementedError("Sqlalchemy models must declare a __pydantic_model__ property to be convertable.")
+
+    def to_pydantic(self) -> Type["BaseModel"]:
+        """converts to the basic pydantic model counterpart"""
+        return self.__pydantic_model__.model_validate(self)
+
+    def to_record(self) -> Type["BaseModel"]:
+        """Deprecated accessor for to_pydantic"""
+        logger.warning("to_record is deprecated, use to_pydantic instead.")
+        return self.to_pydantic()
+
+    # TODO: Look into this later and maybe add back?
+    # def _infer_organization(self, db_session: "Session") -> None:
+    #     """🪄 MAGIC ALERT! 🪄
+    #     Because so much of the original API is centered around user scopes,
+    #     this allows us to continue with that scope and then infer the org from the creating user.
+    #
+    #     IF a created_by_id is set, we will use that to infer the organization and magic set it at create time!
+    #     If not do nothing to the object. Mutates in place.
+    #     """
+    #     if self.created_by_id and hasattr(self, "_organization_id"):
+    #         try:
+    #             from letta.orm.user import User  # to avoid circular import
+    #
+    #             created_by = User.read(db_session, self.created_by_id)
+    #         except NoResultFound:
+    #             logger.warning(f"User {self.created_by_id} not found, unable to infer organization.")
+    #             return
+    #         self._organization_id = created_by._organization_id

letta/schemas/organization.py

@@ -7,13 +7,13 @@ from letta.schemas.letta_base import LettaBase
 
 
 class OrganizationBase(LettaBase):
-    __id_prefix__ = "org"
+    __id_prefix__ = "organization"
 
 
 class Organization(OrganizationBase):
-    id: str = OrganizationBase.generate_id_field()
+    id: str = Field(..., description="The id of the organization.")
     name: str = Field(..., description="The name of the organization.")
-    created_at: datetime = Field(default_factory=datetime.utcnow, description="The creation date of the user.")
+    created_at: datetime = Field(default_factory=datetime.utcnow, description="The creation date of the organization.")
 
 
 class OrganizationCreate(OrganizationBase):

letta/server/rest_api/interface.py

@@ -312,11 +312,20 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
         # Two buffers used to make sure that the 'name' comes after the inner thoughts stream (if inner_thoughts_in_kwargs)
         self.function_name_buffer = None
         self.function_args_buffer = None
+        self.function_id_buffer = None
 
         # extra prints
         self.debug = False
         self.timeout = 30
 
+    def _reset_inner_thoughts_json_reader(self):
+        # A buffer for accumulating function arguments (we want to buffer keys and run checks on each one)
+        self.function_args_reader = JSONInnerThoughtsExtractor(inner_thoughts_key=self.inner_thoughts_kwarg, wait_for_first_key=True)
+        # Two buffers used to make sure that the 'name' comes after the inner thoughts stream (if inner_thoughts_in_kwargs)
+        self.function_name_buffer = None
+        self.function_args_buffer = None
+        self.function_id_buffer = None
+
     async def _create_generator(self) -> AsyncGenerator[Union[LettaMessage, LegacyLettaMessage, MessageStreamStatus], None]:
         """An asynchronous generator that yields chunks as they become available."""
         while self._active:
@@ -376,6 +385,9 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
         if not self.streaming_chat_completion_mode and not self.nonstreaming_legacy_mode:
             self._push_to_buffer(self.multi_step_gen_indicator)
 
+        # Wipe the inner thoughts buffers
+        self._reset_inner_thoughts_json_reader()
+
     def step_complete(self):
         """Signal from the agent that one 'step' finished (step = LLM response + tool execution)"""
         if not self.multi_step:
@@ -386,6 +398,9 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
             # signal that a new step has started in the stream
             self._push_to_buffer(self.multi_step_indicator)
 
+        # Wipe the inner thoughts buffers
+        self._reset_inner_thoughts_json_reader()
+
     def step_yield(self):
         """If multi_step, this is the true 'stream_end' function."""
         self._active = False
@@ -498,6 +513,13 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                     else:
                         self.function_name_buffer += tool_call.function.name
 
+                if tool_call.id:
+                    # Buffer until next time
+                    if self.function_id_buffer is None:
+                        self.function_id_buffer = tool_call.id
+                    else:
+                        self.function_id_buffer += tool_call.id
+
                 if tool_call.function.arguments:
                     updates_main_json, updates_inner_thoughts = self.function_args_reader.process_fragment(tool_call.function.arguments)
 
@@ -518,6 +540,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
 
                     # If we have main_json, we should output a FunctionCallMessage
                     elif updates_main_json:
+
                         # If there's something in the function_name buffer, we should release it first
                         # NOTE: we could output it as part of a chunk that has both name and args,
                         #       however the frontend may expect name first, then args, so to be
@@ -526,18 +549,23 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                             processed_chunk = FunctionCallMessage(
                                 id=message_id,
                                 date=message_date,
-                                function_call=FunctionCallDelta(name=self.function_name_buffer, arguments=None),
+                                function_call=FunctionCallDelta(
+                                    name=self.function_name_buffer,
+                                    arguments=None,
+                                    function_call_id=self.function_id_buffer,
+                                ),
                             )
                             # Clear the buffer
                             self.function_name_buffer = None
+                            self.function_id_buffer = None
                             # Since we're clearing the name buffer, we should store
                             # any updates to the arguments inside a separate buffer
-                            if updates_main_json:
-                                # Add any main_json updates to the arguments buffer
-                                if self.function_args_buffer is None:
-                                    self.function_args_buffer = updates_main_json
-                                else:
-                                    self.function_args_buffer += updates_main_json
+
+                            # Add any main_json updates to the arguments buffer
+                            if self.function_args_buffer is None:
+                                self.function_args_buffer = updates_main_json
+                            else:
+                                self.function_args_buffer += updates_main_json
 
                         # If there was nothing in the name buffer, we can proceed to
                         # output the arguments chunk as a FunctionCallMessage
@@ -550,17 +578,27 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                                 processed_chunk = FunctionCallMessage(
                                     id=message_id,
                                     date=message_date,
-                                    function_call=FunctionCallDelta(name=None, arguments=combined_chunk),
+                                    function_call=FunctionCallDelta(
+                                        name=None,
+                                        arguments=combined_chunk,
+                                        function_call_id=self.function_id_buffer,
+                                    ),
                                 )
                                 # clear buffer
                                 self.function_args_buffer = None
+                                self.function_id_buffer = None
                             else:
                                 # If there's no buffer to clear, just output a new chunk with new data
                                 processed_chunk = FunctionCallMessage(
                                     id=message_id,
                                     date=message_date,
-                                    function_call=FunctionCallDelta(name=None, arguments=updates_main_json),
+                                    function_call=FunctionCallDelta(
+                                        name=None,
+                                        arguments=updates_main_json,
+                                        function_call_id=self.function_id_buffer,
+                                    ),
                                 )
+                                self.function_id_buffer = None
 
                         # # If there's something in the main_json buffer, we should add if to the arguments and release it together
                         # tool_call_delta = {}

letta/server/rest_api/routers/v1/organizations.py

@@ -22,7 +22,7 @@ def get_all_orgs(
     Get a list of all orgs in the database
     """
     try:
-        next_cursor, orgs = server.ms.list_organizations(cursor=cursor, limit=limit)
+        next_cursor, orgs = server.organization_manager.list_organizations(cursor=cursor, limit=limit)
     except HTTPException:
         raise
     except Exception as e:
@@ -38,8 +38,7 @@ def create_org(
     """
     Create a new org in the database
     """
-
-    org = server.create_organization(request)
+    org = server.organization_manager.create_organization(request)
     return org
 
 
@@ -50,10 +49,10 @@ def delete_org(
 ):
     # TODO make a soft deletion, instead of a hard deletion
     try:
-        org = server.ms.get_organization(org_id=org_id)
+        org = server.organization_manager.get_organization_by_id(org_id=org_id)
         if org is None:
             raise HTTPException(status_code=404, detail=f"Organization does not exist")
-        server.ms.delete_organization(org_id=org_id)
+        server.organization_manager.delete_organization(org_id=org_id)
     except HTTPException:
         raise
     except Exception as e:

letta/server/server.py

@@ -44,6 +44,7 @@ from letta.log import get_logger
 from letta.memory import get_memory_functions
 from letta.metadata import Base, MetadataStore
 from letta.o1_agent import O1Agent
+from letta.orm.errors import NoResultFound
 from letta.prompts import gpt_system
 from letta.providers import (
     AnthropicProvider,
@@ -80,12 +81,12 @@ from letta.schemas.memory import (
     RecallMemorySummary,
 )
 from letta.schemas.message import Message, MessageCreate, MessageRole, UpdateMessage
-from letta.schemas.organization import Organization, OrganizationCreate
 from letta.schemas.passage import Passage
 from letta.schemas.source import Source, SourceCreate, SourceUpdate
 from letta.schemas.tool import Tool, ToolCreate, ToolUpdate
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User, UserCreate
+from letta.services.organization_manager import OrganizationManager
 from letta.utils import create_random_username, json_dumps, json_loads
 
 # from letta.llm_api_tools import openai_get_model_list, azure_openai_get_model_list, smart_urljoin
@@ -245,6 +246,9 @@ class SyncServer(Server):
         self.config = config
         self.ms = MetadataStore(self.config)
 
+        # Managers that interface with data models
+        self.organization_manager = OrganizationManager()
+
         # TODO: this should be removed
         # add global default tools (for admin)
         self.add_default_tools(module_name="base")
@@ -773,20 +777,6 @@ class SyncServer(Server):
 
         return user
 
-    def create_organization(self, request: OrganizationCreate) -> Organization:
-        """Create a new org using a config"""
-        if not request.name:
-            # auto-generate a name
-            request.name = create_random_username()
-        org = Organization(name=request.name)
-        self.ms.create_organization(org)
-        logger.info(f"Created new org from config: {org}")
-
-        # add default for the org
-        # TODO: add default data
-
-        return org
-
     def create_agent(
         self,
         request: CreateAgent,
@@ -2125,18 +2115,13 @@ class SyncServer(Server):
 
     def get_default_user(self) -> User:
 
-        from letta.constants import (
-            DEFAULT_ORG_ID,
-            DEFAULT_ORG_NAME,
-            DEFAULT_USER_ID,
-            DEFAULT_USER_NAME,
-        )
+        from letta.constants import DEFAULT_ORG_ID, DEFAULT_USER_ID, DEFAULT_USER_NAME
 
         # check if default org exists
-        default_org = self.ms.get_organization(DEFAULT_ORG_ID)
-        if not default_org:
-            org = Organization(name=DEFAULT_ORG_NAME, id=DEFAULT_ORG_ID)
-            self.ms.create_organization(org)
+        try:
+            self.organization_manager.get_organization_by_id(DEFAULT_ORG_ID)
+        except NoResultFound:
+            self.organization_manager.create_default_organization()
 
         # check if default user exists
         try:

letta/services/organization_manager.py

@@ -0,0 +1,66 @@
+from typing import List, Optional
+
+from sqlalchemy.exc import NoResultFound
+
+from letta.constants import DEFAULT_ORG_ID, DEFAULT_ORG_NAME
+from letta.orm.organization import Organization
+from letta.schemas.organization import Organization as PydanticOrganization
+from letta.utils import create_random_username
+
+
+class OrganizationManager:
+    """Manager class to handle business logic related to Organizations."""
+
+    def __init__(self):
+        # This is probably horrible but we reuse this technique from metadata.py
+        # TODO: Please refactor this out
+        # I am currently working on a ORM refactor and would like to make a more minimal set of changes
+        # - Matt
+        from letta.server.server import db_context
+
+        self.session_maker = db_context
+
+    def get_organization_by_id(self, org_id: str) -> PydanticOrganization:
+        """Fetch an organization by ID."""
+        with self.session_maker() as session:
+            try:
+                organization = Organization.read(db_session=session, identifier=org_id)
+                return organization.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Organization with id {org_id} not found.")
+
+    def create_organization(self, name: Optional[str] = None) -> PydanticOrganization:
+        """Create a new organization. If a name is provided, it is used, otherwise, a random one is generated."""
+        with self.session_maker() as session:
+            org = Organization(name=name if name else create_random_username())
+            org.create(session)
+            return org.to_pydantic()
+
+    def create_default_organization(self) -> PydanticOrganization:
+        """Create the default organization."""
+        with self.session_maker() as session:
+            org = Organization(name=DEFAULT_ORG_NAME)
+            org.id = DEFAULT_ORG_ID
+            org.create(session)
+            return org.to_pydantic()
+
+    def update_organization_name_using_id(self, org_id: str, name: Optional[str] = None) -> PydanticOrganization:
+        """Update an organization."""
+        with self.session_maker() as session:
+            organization = Organization.read(db_session=session, identifier=org_id)
+            if name:
+                organization.name = name
+            organization.update(session)
+            return organization.to_pydantic()
+
+    def delete_organization(self, org_id: str):
+        """Delete an organization by marking it as deleted."""
+        with self.session_maker() as session:
+            organization = Organization.read(db_session=session, identifier=org_id)
+            organization.delete(session)
+
+    def list_organizations(self, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[PydanticOrganization]:
+        """List organizations with pagination based on cursor (org_id) and limit."""
+        with self.session_maker() as session:
+            results = Organization.list(db_session=session, cursor=cursor, limit=limit)
+            return [org.to_pydantic() for org in results]

{letta_nightly-0.5.0.dev20241022104124.dist-info → letta_nightly-0.5.0.dev20241023104105.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.5.0.dev20241022104124
+Version: 0.5.0.dev20241023104105
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team
@@ -55,6 +55,7 @@ Requires-Dist: prettytable (>=3.9.0,<4.0.0)
 Requires-Dist: pyautogen (==0.2.22) ; extra == "autogen"
 Requires-Dist: pydantic (>=2.7.4,<3.0.0)
 Requires-Dist: pydantic-settings (>=2.2.1,<3.0.0)
+Requires-Dist: pyhumps (>=3.8.0,<4.0.0)
 Requires-Dist: pymilvus (>=2.4.3,<3.0.0) ; extra == "milvus"
 Requires-Dist: pyright (>=1.1.347,<2.0.0) ; extra == "dev"
 Requires-Dist: pytest-asyncio (>=0.23.2,<0.24.0) ; extra == "dev"