letta-nightly 0.6.6.dev20241220190343__py3-none-any.whl → 0.6.6.dev20241221104005__py3-none-any.whl

letta/server/server.py

@@ -40,14 +40,14 @@ from letta.providers import (
     VLLMChatCompletionsProvider,
     VLLMCompletionsProvider,
 )
-from letta.schemas.agent import AgentState, AgentType, CreateAgent, UpdateAgent
+from letta.schemas.agent import AgentState, AgentType, CreateAgent
 from letta.schemas.block import BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
 
 # openai schemas
 from letta.schemas.enums import JobStatus
 from letta.schemas.job import Job, JobUpdate
-from letta.schemas.letta_message import ToolReturnMessage, LettaMessage
+from letta.schemas.letta_message import LettaMessage, ToolReturnMessage
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import (
     ArchivalMemorySummary,
@@ -59,7 +59,7 @@ from letta.schemas.message import Message, MessageCreate, MessageRole, MessageUp
 from letta.schemas.organization import Organization
 from letta.schemas.passage import Passage
 from letta.schemas.source import Source
-from letta.schemas.tool import Tool, ToolCreate
+from letta.schemas.tool import Tool
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User
 from letta.services.agent_manager import AgentManager
@@ -303,11 +303,6 @@ class SyncServer(Server):
             self.block_manager.add_default_blocks(actor=self.default_user)
             self.tool_manager.upsert_base_tools(actor=self.default_user)
 
-            # If there is a default org/user
-            # This logic may have to change in the future
-            if settings.load_default_external_tools:
-                self.add_default_external_tools(actor=self.default_user)
-
         # collect providers (always has Letta as a default)
         self._enabled_providers: List[Provider] = [LettaProvider()]
         if model_settings.openai_api_key:
@@ -376,25 +371,6 @@ class SyncServer(Server):
                 )
             )
 
-    def initialize_agent(self, agent_id, actor, interface: Union[AgentInterface, None] = None, initial_message_sequence=None) -> Agent:
-        """Initialize an agent from the database"""
-        agent_state = self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor)
-
-        interface = interface or self.default_interface_factory()
-        if agent_state.agent_type == AgentType.memgpt_agent:
-            agent = Agent(agent_state=agent_state, interface=interface, user=actor, initial_message_sequence=initial_message_sequence)
-        elif agent_state.agent_type == AgentType.offline_memory_agent:
-            agent = OfflineMemoryAgent(
-                agent_state=agent_state, interface=interface, user=actor, initial_message_sequence=initial_message_sequence
-            )
-        else:
-            assert initial_message_sequence is None, f"Initial message sequence is not supported for O1Agents"
-            agent = O1Agent(agent_state=agent_state, interface=interface, user=actor)
-
-        # Persist to agent
-        save_agent(agent)
-        return agent
-
     def load_agent(self, agent_id: str, actor: User, interface: Union[AgentInterface, None] = None) -> Agent:
         """Updated method to load agents from persisted storage"""
         agent_lock = self.per_agent_lock_manager.get_lock(agent_id)
@@ -413,11 +389,6 @@ class SyncServer(Server):
             else:
                 raise ValueError(f"Invalid agent type {agent_state.agent_type}")
 
-            # Rebuild the system prompt - may be linked to new blocks now
-            agent.rebuild_system_prompt()
-
-            # Persist to agent
-            save_agent(agent)
             return agent
 
     def _step(
@@ -455,9 +426,6 @@ class SyncServer(Server):
                 skip_verify=True,
             )
 
-            # save agent after step
-            save_agent(letta_agent)
-
         except Exception as e:
             logger.error(f"Error in server._step: {e}")
             print(traceback.print_exc())
@@ -790,129 +758,23 @@ class SyncServer(Server):
 
         """Create a new agent using a config"""
         # Invoke manager
-        agent_state = self.agent_manager.create_agent(
+        return self.agent_manager.create_agent(
             agent_create=request,
             actor=actor,
         )
 
-        # create the agent object
-        if request.initial_message_sequence is not None:
-            # init_messages = [Message(user_id=user_id, agent_id=agent_state.id, role=message.role, text=message.text) for message in request.initial_message_sequence]
-            init_messages = []
-            for message in request.initial_message_sequence:
-
-                if message.role == MessageRole.user:
-                    packed_message = system.package_user_message(
-                        user_message=message.text,
-                    )
-                elif message.role == MessageRole.system:
-                    packed_message = system.package_system_message(
-                        system_message=message.text,
-                    )
-                else:
-                    raise ValueError(f"Invalid message role: {message.role}")
-
-                init_messages.append(Message(role=message.role, text=packed_message, agent_id=agent_state.id))
-            # init_messages = [Message.dict_to_message(user_id=user_id, agent_id=agent_state.id, openai_message_dict=message.model_dump()) for message in request.initial_message_sequence]
-        else:
-            init_messages = None
-
-        # initialize the agent (generates initial message list with system prompt)
-        if interface is None:
-            interface = self.default_interface_factory()
-        self.initialize_agent(agent_id=agent_state.id, interface=interface, initial_message_sequence=init_messages, actor=actor)
-
-        in_memory_agent_state = self.agent_manager.get_agent_by_id(agent_state.id, actor=actor)
-        return in_memory_agent_state
-
-    # TODO: This is not good!
-    # TODO: Ideally, this should ALL be handled by the ORM
-    # TODO: The main blocker here IS the _message updates
-    def update_agent(
-        self,
-        agent_id: str,
-        request: UpdateAgent,
-        actor: User,
-    ) -> AgentState:
-        """Update the agents core memory block, return the new state"""
-        # Update agent state in the db first
-        agent_state = self.agent_manager.update_agent(agent_id=agent_id, agent_update=request, actor=actor)
-
-        # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-
-        # TODO: Everything below needs to get removed, no updating anything in memory
-        # update the system prompt
-        if request.system:
-            letta_agent.update_system_prompt(request.system)
-
-        # update in-context messages
-        if request.message_ids:
-            # This means the user is trying to change what messages are in the message buffer
-            # Internally this requires (1) pulling from recall,
-            # then (2) setting the attributes ._messages and .state.message_ids
-            letta_agent.set_message_buffer(message_ids=request.message_ids)
-
-        letta_agent.update_state()
-
-        return agent_state
-
-    def get_tools_from_agent(self, agent_id: str, user_id: Optional[str]) -> List[Tool]:
-        """Get tools from an existing agent"""
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-
-        # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        return letta_agent.agent_state.tools
-
-    def add_tool_to_agent(
-        self,
-        agent_id: str,
-        tool_id: str,
-        user_id: str,
-    ):
-        """Add tools from an existing agent"""
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-
-        agent_state = self.agent_manager.attach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
-
-        return agent_state
-
-    def remove_tool_from_agent(
-        self,
-        agent_id: str,
-        tool_id: str,
-        user_id: str,
-    ):
-        """Remove tools from an existing agent"""
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-        agent_state = self.agent_manager.detach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
-
-        return agent_state
-
     # convert name->id
 
+    # TODO: These can be moved to agent_manager
     def get_agent_memory(self, agent_id: str, actor: User) -> Memory:
         """Return the memory of an agent (core memory)"""
-        agent = self.load_agent(agent_id=agent_id, actor=actor)
-        return agent.agent_state.memory
+        return self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor).memory
 
     def get_archival_memory_summary(self, agent_id: str, actor: User) -> ArchivalMemorySummary:
-        agent = self.load_agent(agent_id=agent_id, actor=actor)
         return ArchivalMemorySummary(size=self.agent_manager.passage_size(actor=actor, agent_id=agent_id))
 
     def get_recall_memory_summary(self, agent_id: str, actor: User) -> RecallMemorySummary:
-        agent = self.load_agent(agent_id=agent_id, actor=actor)
-        return RecallMemorySummary(size=len(agent.message_manager))
-
-    def get_in_context_messages(self, agent_id: str, actor: User) -> List[Message]:
-        """Get the in-context messages in the agent's memory"""
-        # Get the agent object (loaded in memory)
-        agent = self.load_agent(agent_id=agent_id, actor=actor)
-        return agent._messages
+        return RecallMemorySummary(size=self.message_manager.size(actor=actor, agent_id=agent_id))
 
     def get_agent_archival(self, user_id: str, agent_id: str, cursor: Optional[str] = None, limit: int = 50) -> List[Passage]:
         """Paginated query of all messages in agent archival memory"""
@@ -947,24 +809,17 @@ class SyncServer(Server):
 
     def insert_archival_memory(self, agent_id: str, memory_contents: str, actor: User) -> List[Passage]:
         # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-
+        agent_state = self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor)
         # Insert into archival memory
-        passages = self.passage_manager.insert_passage(
-            agent_state=letta_agent.agent_state, agent_id=agent_id, text=memory_contents, actor=actor
-        )
-
-        save_agent(letta_agent)
+        # TODO: @mindy look at moving this to agent_manager to avoid above extra call
+        passages = self.passage_manager.insert_passage(agent_state=agent_state, agent_id=agent_id, text=memory_contents, actor=actor)
 
         return passages
 
-    def delete_archival_memory(self, agent_id: str, memory_id: str, actor: User):
-        # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-
-        # Delete by ID
+    def delete_archival_memory(self, memory_id: str, actor: User):
         # TODO check if it exists first, and throw error if not
-        letta_agent.passage_manager.delete_passage_by_id(passage_id=memory_id, actor=actor)
+        # TODO: @mindy make this return the deleted passage instead
+        self.passage_manager.delete_passage_by_id(passage_id=memory_id, actor=actor)
 
         # TODO: return archival memory
 
@@ -1042,9 +897,8 @@ class SyncServer(Server):
         # update the block
         self.block_manager.update_block(block_id=block.id, block_update=BlockUpdate(value=value), actor=actor)
 
-        # load agent
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        return letta_agent.agent_state.memory
+        # rebuild system prompt for agent, potentially changed
+        return self.agent_manager.rebuild_system_prompt(agent_id=agent_id, actor=actor).memory
 
     def delete_source(self, source_id: str, actor: User):
         """Delete a data source"""
@@ -1082,11 +936,10 @@ class SyncServer(Server):
         agent_states = self.source_manager.list_attached_agents(source_id=source_id, actor=actor)
         for agent_state in agent_states:
             agent_id = agent_state.id
-            agent = self.load_agent(agent_id=agent_id, actor=actor)
 
             # Attach source to agent
             curr_passage_size = self.agent_manager.passage_size(actor=actor, agent_id=agent_id)
-            agent.attach_source(user=actor, source_id=source_id, source_manager=self.source_manager, agent_manager=self.agent_manager)
+            self.agent_manager.attach_source(agent_id=agent_state.id, source_id=source_id, actor=actor)
             new_passage_size = self.agent_manager.passage_size(actor=actor, agent_id=agent_id)
             assert new_passage_size >= curr_passage_size  # in case empty files are added
 
@@ -1111,56 +964,6 @@ class SyncServer(Server):
         passage_count, document_count = load_data(connector, source, self.passage_manager, self.source_manager, actor=user)
         return passage_count, document_count
 
-    def attach_source_to_agent(
-        self,
-        user_id: str,
-        agent_id: str,
-        source_id: Optional[str] = None,
-        source_name: Optional[str] = None,
-    ) -> Source:
-        # attach a data source to an agent
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-        if source_id:
-            data_source = self.source_manager.get_source_by_id(source_id=source_id, actor=actor)
-        elif source_name:
-            data_source = self.source_manager.get_source_by_name(source_name=source_name, actor=actor)
-        else:
-            raise ValueError(f"Need to provide at least source_id or source_name to find the source.")
-
-        assert data_source, f"Data source with id={source_id} or name={source_name} does not exist"
-
-        # load agent
-        agent = self.load_agent(agent_id=agent_id, actor=actor)
-
-        # attach source to agent
-        agent.attach_source(user=actor, source_id=data_source.id, source_manager=self.source_manager, agent_manager=self.agent_manager)
-
-        return data_source
-
-    def detach_source_from_agent(
-        self,
-        user_id: str,
-        agent_id: str,
-        source_id: Optional[str] = None,
-        source_name: Optional[str] = None,
-    ) -> Source:
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-        if source_id:
-            source = self.source_manager.get_source_by_id(source_id=source_id, actor=actor)
-        elif source_name:
-            source = self.source_manager.get_source_by_name(source_name=source_name, actor=actor)
-            source_id = source.id
-        else:
-            raise ValueError(f"Need to provide at least source_id or source_name to find the source.")
-
-        # delete agent-source mapping
-        self.agent_manager.detach_source(agent_id=agent_id, source_id=source_id, actor=actor)
-
-        # return back source data
-        return source
-
     def list_data_source_passages(self, user_id: str, source_id: str) -> List[Passage]:
         warnings.warn("list_data_source_passages is not yet implemented, returning empty list.", category=UserWarning)
         return []
@@ -1198,52 +1001,11 @@ class SyncServer(Server):
 
         return sources_with_metadata
 
-    def add_default_external_tools(self, actor: User) -> bool:
-        """Add default langchain tools. Return true if successful, false otherwise."""
-        success = True
-        tool_creates = ToolCreate.load_default_langchain_tools()
-        if tool_settings.composio_api_key:
-            tool_creates += ToolCreate.load_default_composio_tools()
-        for tool_create in tool_creates:
-            try:
-                self.tool_manager.create_or_update_tool(Tool(**tool_create.model_dump()), actor=actor)
-            except Exception as e:
-                warnings.warn(f"An error occurred while creating tool {tool_create}: {e}")
-                warnings.warn(traceback.format_exc())
-                success = False
-
-        return success
-
-    def update_agent_message(self, agent_id: str, message_id: str, request: MessageUpdate, actor: User) -> Message:
+    def update_agent_message(self, message_id: str, request: MessageUpdate, actor: User) -> Message:
         """Update the details of a message associated with an agent"""
 
         # Get the current message
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        response = letta_agent.update_message(message_id=message_id, request=request)
-        save_agent(letta_agent)
-        return response
-
-    def rewrite_agent_message(self, agent_id: str, new_text: str, actor: User) -> Message:
-
-        # Get the current message
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        response = letta_agent.rewrite_message(new_text=new_text)
-        save_agent(letta_agent)
-        return response
-
-    def rethink_agent_message(self, agent_id: str, new_thought: str, actor: User) -> Message:
-        # Get the current message
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        response = letta_agent.rethink_message(new_thought=new_thought)
-        save_agent(letta_agent)
-        return response
-
-    def retry_agent_message(self, agent_id: str, actor: User) -> List[Message]:
-        # Get the current message
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        response = letta_agent.retry_message()
-        save_agent(letta_agent)
-        return response
+        return self.message_manager.update_message_by_id(message_id=message_id, message_update=request, actor=actor)
 
     def get_organization_or_default(self, org_id: Optional[str]) -> Organization:
         """Get the organization object for org_id if it exists, otherwise return the default organization object"""
@@ -1331,15 +1093,7 @@ class SyncServer(Server):
     def add_embedding_model(self, request: EmbeddingConfig) -> EmbeddingConfig:
         """Add a new embedding model"""
 
-    def get_agent_context_window(
-        self,
-        user_id: str,
-        agent_id: str,
-    ) -> ContextWindowOverview:
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-
-        # Get the current message
+    def get_agent_context_window(self, agent_id: str, actor: User) -> ContextWindowOverview:
         letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
         return letta_agent.get_context_window()
 

letta/services/agent_manager.py

@@ -20,6 +20,7 @@ from letta.schemas.agent import AgentType, CreateAgent, UpdateAgent
 from letta.schemas.block import Block as PydanticBlock
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.llm_config import LLMConfig
+from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.passage import Passage as PydanticPassage
 from letta.schemas.source import Source as PydanticSource
 from letta.schemas.tool_rule import ToolRule as PydanticToolRule
@@ -28,12 +29,17 @@ from letta.services.block_manager import BlockManager
 from letta.services.helpers.agent_manager_helper import (
     _process_relationship,
     _process_tags,
+    check_supports_structured_output,
+    compile_system_message,
     derive_system_message,
+    initialize_message_sequence,
+    package_initial_message_sequence,
 )
+from letta.services.message_manager import MessageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
 from letta.settings import settings
-from letta.utils import enforce_types
+from letta.utils import enforce_types, get_utc_time, united_diff
 
 logger = get_logger(__name__)
 
@@ -49,6 +55,7 @@ class AgentManager:
         self.block_manager = BlockManager()
         self.tool_manager = ToolManager()
         self.source_manager = SourceManager()
+        self.message_manager = MessageManager()
 
     # ======================================================================================================================
     # Basic CRUD operations
@@ -64,6 +71,10 @@ class AgentManager:
         if not agent_create.llm_config or not agent_create.embedding_config:
             raise ValueError("llm_config and embedding_config are required")
 
+        # Check tool rules are valid
+        if agent_create.tool_rules:
+            check_supports_structured_output(model=agent_create.llm_config.model, tool_rules=agent_create.tool_rules)
+
         # create blocks (note: cannot be linked into the agent_id is created)
         block_ids = list(agent_create.block_ids or [])  # Create a local copy to avoid modifying the original
         for create_block in agent_create.memory_blocks:
@@ -88,7 +99,8 @@ class AgentManager:
         # Remove duplicates
         tool_ids = list(set(tool_ids))
 
-        return self._create_agent(
+        # Create the agent
+        agent_state = self._create_agent(
             name=agent_create.name,
             system=system,
             agent_type=agent_create.agent_type,
@@ -104,6 +116,35 @@ class AgentManager:
             actor=actor,
         )
 
+        # TODO: See if we can merge this into the above SQL create call for performance reasons
+        # Generate a sequence of initial messages to put in the buffer
+        init_messages = initialize_message_sequence(
+            agent_state=agent_state, memory_edit_timestamp=get_utc_time(), include_initial_boot_message=True
+        )
+
+        if agent_create.initial_message_sequence is not None:
+            # We always need the system prompt up front
+            system_message_obj = PydanticMessage.dict_to_message(
+                agent_id=agent_state.id,
+                user_id=agent_state.created_by_id,
+                model=agent_state.llm_config.model,
+                openai_message_dict=init_messages[0],
+            )
+            # Don't use anything else in the pregen sequence, instead use the provided sequence
+            init_messages = [system_message_obj]
+            init_messages.extend(
+                package_initial_message_sequence(agent_state.id, agent_create.initial_message_sequence, agent_state.llm_config.model, actor)
+            )
+        else:
+            init_messages = [
+                PydanticMessage.dict_to_message(
+                    agent_id=agent_state.id, user_id=agent_state.created_by_id, model=agent_state.llm_config.model, openai_message_dict=msg
+                )
+                for msg in init_messages
+            ]
+
+        return self.append_to_in_context_messages(init_messages, agent_id=agent_state.id, actor=actor)
+
     @enforce_types
     def _create_agent(
         self,
@@ -149,6 +190,16 @@ class AgentManager:
 
     @enforce_types
     def update_agent(self, agent_id: str, agent_update: UpdateAgent, actor: PydanticUser) -> PydanticAgentState:
+        agent_state = self._update_agent(agent_id=agent_id, agent_update=agent_update, actor=actor)
+
+        # Rebuild the system prompt if it's different
+        if agent_update.system and agent_update.system != agent_state.system:
+            agent_state = self.rebuild_system_prompt(agent_id=agent_state.id, actor=actor, force=True, update_timestamp=False)
+
+        return agent_state
+
+    @enforce_types
+    def _update_agent(self, agent_id: str, agent_update: UpdateAgent, actor: PydanticUser) -> PydanticAgentState:
         """
         Update an existing agent.
 
@@ -247,6 +298,105 @@ class AgentManager:
             agent.hard_delete(session)
             return agent_state
 
+    # ======================================================================================================================
+    # In Context Messages Management
+    # ======================================================================================================================
+    # TODO: There are several assumptions here that are not explicitly checked
+    # TODO: 1) These message ids are valid
+    # TODO: 2) These messages are ordered from oldest to newest
+    # TODO: This can be fixed by having an actual relationship in the ORM for message_ids
+    # TODO: This can also be made more efficient, instead of getting, setting, we can do it all in one db session for one query.
+    @enforce_types
+    def get_in_context_messages(self, agent_id: str, actor: PydanticUser) -> List[PydanticMessage]:
+        message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
+        return self.message_manager.get_messages_by_ids(message_ids=message_ids, actor=actor)
+
+    @enforce_types
+    def get_system_message(self, agent_id: str, actor: PydanticUser) -> PydanticMessage:
+        message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
+        return self.message_manager.get_message_by_id(message_id=message_ids[0], actor=actor)
+
+    @enforce_types
+    def rebuild_system_prompt(self, agent_id: str, actor: PydanticUser, force=False, update_timestamp=True) -> PydanticAgentState:
+        """Rebuilds the system message with the latest memory object and any shared memory block updates
+
+        Updates to core memory blocks should trigger a "rebuild", which itself will create a new message object
+
+        Updates to the memory header should *not* trigger a rebuild, since that will simply flood recall storage with excess messages
+        """
+        agent_state = self.get_agent_by_id(agent_id=agent_id, actor=actor)
+
+        curr_system_message = self.get_system_message(
+            agent_id=agent_id, actor=actor
+        )  # this is the system + memory bank, not just the system prompt
+        curr_system_message_openai = curr_system_message.to_openai_dict()
+
+        # note: we only update the system prompt if the core memory is changed
+        # this means that the archival/recall memory statistics may be someout out of date
+        curr_memory_str = agent_state.memory.compile()
+        if curr_memory_str in curr_system_message_openai["content"] and not force:
+            # NOTE: could this cause issues if a block is removed? (substring match would still work)
+            logger.info(
+                f"Memory hasn't changed for agent id={agent_id} and actor=({actor.id}, {actor.name}), skipping system prompt rebuild"
+            )
+            return agent_state
+
+        # If the memory didn't update, we probably don't want to update the timestamp inside
+        # For example, if we're doing a system prompt swap, this should probably be False
+        if update_timestamp:
+            memory_edit_timestamp = get_utc_time()
+        else:
+            # NOTE: a bit of a hack - we pull the timestamp from the message created_by
+            memory_edit_timestamp = curr_system_message.created_at
+
+        # update memory (TODO: potentially update recall/archival stats separately)
+        new_system_message_str = compile_system_message(
+            system_prompt=agent_state.system,
+            in_context_memory=agent_state.memory,
+            in_context_memory_last_edit=memory_edit_timestamp,
+        )
+
+        diff = united_diff(curr_system_message_openai["content"], new_system_message_str)
+        if len(diff) > 0:  # there was a diff
+            logger.info(f"Rebuilding system with new memory...\nDiff:\n{diff}")
+
+            # Swap the system message out (only if there is a diff)
+            message = PydanticMessage.dict_to_message(
+                agent_id=agent_id,
+                user_id=actor.id,
+                model=agent_state.llm_config.model,
+                openai_message_dict={"role": "system", "content": new_system_message_str},
+            )
+            message = self.message_manager.create_message(message, actor=actor)
+            message_ids = [message.id] + agent_state.message_ids[1:]  # swap index 0 (system)
+            return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+        else:
+            return agent_state
+
+    @enforce_types
+    def set_in_context_messages(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState:
+        return self.update_agent(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor)
+
+    @enforce_types
+    def trim_older_in_context_messages(self, num: int, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
+        message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
+        new_messages = [message_ids[0]] + message_ids[num:]  # 0 is system message
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=new_messages, actor=actor)
+
+    @enforce_types
+    def prepend_to_in_context_messages(self, messages: List[PydanticMessage], agent_id: str, actor: PydanticUser) -> PydanticAgentState:
+        message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
+        new_messages = self.message_manager.create_many_messages(messages, actor=actor)
+        message_ids = [message_ids[0]] + [m.id for m in new_messages] + message_ids[1:]
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+
+    @enforce_types
+    def append_to_in_context_messages(self, messages: List[PydanticMessage], agent_id: str, actor: PydanticUser) -> PydanticAgentState:
+        messages = self.message_manager.create_many_messages(messages, actor=actor)
+        message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids or []
+        message_ids += [m.id for m in messages]
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+
     # ======================================================================================================================
     # Source Management
     # ======================================================================================================================