letta-nightly 0.5.4.dev20241126104249__py3-none-any.whl → 0.5.4.dev20241128000451__py3-none-any.whl

letta/server/server.py

@@ -3,6 +3,7 @@ import os
 import traceback
 import warnings
 from abc import abstractmethod
+from asyncio import Lock
 from datetime import datetime
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
@@ -17,24 +18,10 @@ from letta.agent_store.storage import StorageConnector, TableType
 from letta.credentials import LettaCredentials
 from letta.data_sources.connectors import DataConnector, load_data
 
-# from letta.data_types import (
-#    AgentState,
-#    EmbeddingConfig,
-#    LLMConfig,
-#    Message,
-#    Preset,
-#    Source,
-#    Token,
-#    User,
-# )
-from letta.functions.functions import generate_schema, parse_source_code
-from letta.functions.schema_generator import generate_schema
-
 # TODO use custom interface
 from letta.interface import AgentInterface  # abstract
 from letta.interface import CLIInterface  # for printing to terminal
 from letta.log import get_logger
-from letta.memory import get_memory_functions
 from letta.metadata import MetadataStore
 from letta.o1_agent import O1Agent
 from letta.orm import Base
@@ -53,8 +40,15 @@ from letta.providers import (
     VLLMChatCompletionsProvider,
     VLLMCompletionsProvider,
 )
-from letta.schemas.agent import AgentState, AgentType, CreateAgent, UpdateAgentState
+from letta.schemas.agent import (
+    AgentState,
+    AgentType,
+    CreateAgent,
+    PersistedAgentState,
+    UpdateAgentState,
+)
 from letta.schemas.api_key import APIKey, APIKeyCreate
+from letta.schemas.block import Block, BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
 
 # openai schemas
@@ -79,15 +73,13 @@ from letta.services.agents_tags_manager import AgentsTagsManager
 from letta.services.block_manager import BlockManager
 from letta.services.blocks_agents_manager import BlocksAgentsManager
 from letta.services.organization_manager import OrganizationManager
+from letta.services.per_agent_lock_manager import PerAgentLockManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
 from letta.services.user_manager import UserManager
 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
-
-
 logger = get_logger(__name__)
 
 
@@ -127,10 +119,11 @@ class Server(object):
     @abstractmethod
     def create_agent(
         self,
-        user_id: str,
-        agent_config: Union[dict, AgentState],
-        interface: Union[AgentInterface, None],
-    ) -> str:
+        request: CreateAgent,
+        actor: User,
+        # interface
+        interface: Union[AgentInterface, None] = None,
+    ) -> AgentState:
         """Create a new agent using a config"""
         raise NotImplementedError
 
@@ -231,6 +224,9 @@ class SyncServer(Server):
 
         self.credentials = LettaCredentials.load()
 
+        # Locks
+        self.send_message_lock = Lock()
+
         # Initialize the metadata store
         config = LettaConfig.load()
         if settings.letta_pg_uri_no_default:
@@ -249,8 +245,11 @@ class SyncServer(Server):
         self.block_manager = BlockManager()
         self.source_manager = SourceManager()
         self.agents_tags_manager = AgentsTagsManager()
-        self.blocks_agents_manager = BlocksAgentsManager()
         self.sandbox_config_manager = SandboxConfigManager(tool_settings)
+        self.blocks_agents_manager = BlocksAgentsManager()
+
+        # Managers that interface with parallelism
+        self.per_agent_lock_manager = PerAgentLockManager()
 
         # Make default user and org
         if init_with_default_org_and_user:
@@ -364,92 +363,23 @@ class SyncServer(Server):
             }
         )
 
-    def _load_agent(self, agent_id: str, actor: User, interface: Union[AgentInterface, None] = None) -> Agent:
-        """Loads a saved agent into memory (if it doesn't exist, throw an error)"""
-        assert isinstance(agent_id, str), agent_id
-        user_id = actor.id
-
-        # If an interface isn't specified, use the default
-        if interface is None:
-            interface = self.default_interface_factory()
-
-        try:
-            logger.debug(f"Grabbing agent user_id={user_id} agent_id={agent_id} from database")
-            agent_state = self.ms.get_agent(agent_id=agent_id, user_id=user_id)
-            if not agent_state:
-                logger.exception(f"agent_id {agent_id} does not exist")
-                raise ValueError(f"agent_id {agent_id} does not exist")
-
-            # Instantiate an agent object using the state retrieved
-            logger.debug(f"Creating an agent object")
-            tool_objs = []
-            for name in agent_state.tools:
-                # TODO: This should be a hard failure, but for migration reasons, we patch it for now
-                tool_obj = self.tool_manager.get_tool_by_name(tool_name=name, actor=actor)
-                if tool_obj:
-                    tool_obj = self.tool_manager.get_tool_by_name(tool_name=name, actor=actor)
-                    tool_objs.append(tool_obj)
-                else:
-                    warnings.warn(f"Tried to retrieve a tool with name {name} from the agent_state, but does not exist in tool db.")
-
-            # set agent_state tools to only the names of the available tools
-            agent_state.tools = [t.name for t in tool_objs]
+    def load_agent(self, agent_id: str, interface: Union[AgentInterface, None] = None) -> Agent:
+        """Updated method to load agents from persisted storage"""
+        agent_state = self.get_agent(agent_id=agent_id)
+        actor = self.user_manager.get_user_by_id(user_id=agent_state.user_id)
 
-            # Make sure the memory is a memory object
-            assert isinstance(agent_state.memory, Memory)
-
-            if agent_state.agent_type == AgentType.memgpt_agent:
-                letta_agent = Agent(agent_state=agent_state, interface=interface, tools=tool_objs, user=actor)
-            elif agent_state.agent_type == AgentType.o1_agent:
-                letta_agent = O1Agent(agent_state=agent_state, interface=interface, tools=tool_objs, user=actor)
-            else:
-                raise NotImplementedError("Not a supported agent type")
-
-            # Add the agent to the in-memory store and return its reference
-            logger.debug(f"Adding agent to the agent cache: user_id={user_id}, agent_id={agent_id}")
-            self._add_agent(user_id=user_id, agent_id=agent_id, agent_obj=letta_agent)
-            return letta_agent
-
-        except Exception as e:
-            logger.exception(f"Error occurred while trying to get agent {agent_id}:\n{e}")
-            raise
-
-    def _get_or_load_agent(self, agent_id: str, caching: bool = True) -> Agent:
-        """Check if the agent is in-memory, then load"""
-
-        # Gets the agent state
-        agent_state = self.ms.get_agent(agent_id=agent_id)
-        if not agent_state:
-            raise ValueError(f"Agent does not exist")
-        user_id = agent_state.user_id
-        actor = self.user_manager.get_user_by_id(user_id)
-
-        logger.debug(f"Checking for agent user_id={user_id} agent_id={agent_id}")
-        if caching:
-            # TODO: consider disabling loading cached agents due to potential concurrency issues
-            letta_agent = self._get_agent(user_id=user_id, agent_id=agent_id)
-            if not letta_agent:
-                logger.debug(f"Agent not loaded, loading agent user_id={user_id} agent_id={agent_id}")
-                letta_agent = self._load_agent(agent_id=agent_id, actor=actor)
+        interface = interface or self.default_interface_factory()
+        if agent_state.agent_type == AgentType.memgpt_agent:
+            return Agent(agent_state=agent_state, interface=interface, user=actor)
         else:
-            # This breaks unit tests in test_local_client.py
-            letta_agent = self._load_agent(agent_id=agent_id, actor=actor)
-
-        # letta_agent = self._get_agent(user_id=user_id, agent_id=agent_id)
-        # if not letta_agent:
-        # logger.debug(f"Agent not loaded, loading agent user_id={user_id} agent_id={agent_id}")
-
-        # NOTE: no longer caching, always forcing a lot from the database
-        # Loads the agent objects
-        # letta_agent = self._load_agent(agent_id=agent_id, actor=actor)
-
-        return letta_agent
+            return O1Agent(agent_state=agent_state, interface=interface, user=actor)
 
     def _step(
         self,
         user_id: str,
         agent_id: str,
         input_messages: Union[Message, List[Message]],
+        interface: Union[AgentInterface, None] = None,  # needed to getting responses
         # timestamp: Optional[datetime],
     ) -> LettaUsageStatistics:
         """Send the input message through the agent"""
@@ -465,7 +395,8 @@ class SyncServer(Server):
         try:
 
             # Get the agent object (loaded in memory)
-            letta_agent = self._get_or_load_agent(agent_id=agent_id)
+            # letta_agent = self._get_or_load_agent(agent_id=agent_id)
+            letta_agent = self.load_agent(agent_id=agent_id, interface=interface)
             if letta_agent is None:
                 raise KeyError(f"Agent (user={user_id}, agent={agent_id}) is not loaded")
 
@@ -482,6 +413,9 @@ class SyncServer(Server):
                 skip_verify=True,
             )
 
+            # save agent after step
+            save_agent(letta_agent, self.ms)
+
         except Exception as e:
             logger.error(f"Error in server._step: {e}")
             print(traceback.print_exc())
@@ -499,7 +433,7 @@ class SyncServer(Server):
         logger.debug(f"Got command: {command}")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
         usage = None
 
         if command.lower() == "exit":
@@ -536,7 +470,7 @@ class SyncServer(Server):
         elif command.lower() == "memory":
             ret_str = (
                 f"\nDumping memory contents:\n"
-                + f"\n{str(letta_agent.memory)}"
+                + f"\n{str(letta_agent.agent_state.memory)}"
                 + f"\n{str(letta_agent.persistence_manager.archival_memory)}"
                 + f"\n{str(letta_agent.persistence_manager.recall_memory)}"
             )
@@ -728,6 +662,7 @@ class SyncServer(Server):
         # whether or not to wrap user and system message as MemGPT-style stringified JSON
         wrap_user_message: bool = True,
         wrap_system_message: bool = True,
+        interface: Union[AgentInterface, None] = None,  # needed to getting responses
     ) -> LettaUsageStatistics:
         """Send a list of messages to the agent
 
@@ -780,7 +715,7 @@ class SyncServer(Server):
             raise ValueError(f"All messages must be of type Message or MessageCreate, got {[type(message) for message in messages]}")
 
         # Run the agent state forward
-        return self._step(user_id=user_id, agent_id=agent_id, input_messages=message_objects)
+        return self._step(user_id=user_id, agent_id=agent_id, input_messages=message_objects, interface=interface)
 
     # @LockingServer.agent_lock_decorator
     def run_command(self, user_id: str, agent_id: str, command: str) -> LettaUsageStatistics:
@@ -828,129 +763,109 @@ class SyncServer(Server):
             else:
                 raise ValueError(f"Invalid agent type: {request.agent_type}")
 
+        # create blocks (note: cannot be linked into the agent_id is created)
+        blocks = []
+        for create_block in request.memory_blocks:
+            block = self.block_manager.create_or_update_block(Block(**create_block.model_dump()), actor=actor)
+            blocks.append(block)
+
+        # get tools + only add if they exist
+        tool_objs = []
+        if request.tools:
+            for tool_name in request.tools:
+                tool_obj = self.tool_manager.get_tool_by_name(tool_name=tool_name, actor=actor)
+                if tool_obj:
+                    tool_objs.append(tool_obj)
+                else:
+                    warnings.warn(f"Attempted to add a nonexistent tool {tool_name} to agent {request.name}, skipping.")
+        # reset the request.tools to only valid tools
+        request.tools = [t.name for t in tool_objs]
+
+        # get the user
         logger.debug(f"Attempting to find user: {user_id}")
         user = self.user_manager.get_user_by_id(user_id=user_id)
         if not user:
             raise ValueError(f"cannot find user with associated client id: {user_id}")
 
-        try:
-            # model configuration
-            llm_config = request.llm_config
-            embedding_config = request.embedding_config
-
-            # get tools + only add if they exist
-            tool_objs = []
-            if request.tools:
-                for tool_name in request.tools:
-                    tool_obj = self.tool_manager.get_tool_by_name(tool_name=tool_name, actor=actor)
-                    if tool_obj:
-                        tool_objs.append(tool_obj)
-                    else:
-                        warnings.warn(f"Attempted to add a nonexistent tool {tool_name} to agent {request.name}, skipping.")
-            # reset the request.tools to only valid tools
-            request.tools = [t.name for t in tool_objs]
-
-            assert request.memory is not None
-            memory_functions = get_memory_functions(request.memory)
-            for func_name, func in memory_functions.items():
-
-                if request.tools and func_name in request.tools:
-                    # tool already added
-                    continue
-                source_code = parse_source_code(func)
-                # memory functions are not terminal
-                json_schema = generate_schema(func, name=func_name)
-                source_type = "python"
-                tags = ["memory", "memgpt-base"]
-                tool = self.tool_manager.create_or_update_tool(
-                    Tool(
-                        source_code=source_code,
-                        source_type=source_type,
-                        tags=tags,
-                        json_schema=json_schema,
-                    ),
-                    actor=actor,
-                )
-                tool_objs.append(tool)
-                if not request.tools:
-                    request.tools = []
-                request.tools.append(tool.name)
-
-            # TODO: save the agent state
-            agent_state = AgentState(
-                name=request.name,
-                user_id=user_id,
-                tools=request.tools if request.tools else [],
-                tool_rules=request.tool_rules if request.tool_rules else [],
-                agent_type=request.agent_type or AgentType.memgpt_agent,
-                llm_config=llm_config,
-                embedding_config=embedding_config,
-                system=request.system,
-                memory=request.memory,
-                description=request.description,
-                metadata_=request.metadata_,
-                tags=request.tags,
-            )
-            if request.agent_type == AgentType.memgpt_agent:
-                agent = Agent(
-                    interface=interface,
-                    agent_state=agent_state,
-                    tools=tool_objs,
-                    # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now
-                    first_message_verify_mono=(
-                        True if (llm_config and llm_config.model is not None and "gpt-4" in llm_config.model) else False
-                    ),
-                    user=actor,
-                    initial_message_sequence=request.initial_message_sequence,
-                )
-            elif request.agent_type == AgentType.o1_agent:
-                agent = O1Agent(
-                    interface=interface,
-                    agent_state=agent_state,
-                    tools=tool_objs,
-                    # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now
-                    first_message_verify_mono=(
-                        True if (llm_config and llm_config.model is not None and "gpt-4" in llm_config.model) else False
-                    ),
-                    user=actor,
-                )
-            # rebuilding agent memory on agent create in case shared memory blocks
-            # were specified in the new agent's memory config. we're doing this for two reasons:
-            # 1. if only the ID of the shared memory block was specified, we can fetch its most recent value
-            # 2. if the shared block state changed since this agent initialization started, we can be sure to have the latest value
-            agent.rebuild_memory(force=True, ms=self.ms)
-            # FIXME: this is a hacky way to get the system prompts injected into agent into the DB
-            # self.ms.update_agent(agent.agent_state)
-        except Exception as e:
-            logger.exception(e)
-            try:
-                if agent:
-                    self.ms.delete_agent(agent_id=agent.agent_state.id)
-            except Exception as delete_e:
-                logger.exception(f"Failed to delete_agent:\n{delete_e}")
-            raise e
+        # TODO: create the message objects (NOTE: do this after we migrate to `CreateMessage`)
 
-        # save agent
-        save_agent(agent, self.ms)
-        logger.debug(f"Created new agent from config: {agent}")
+        # created and persist the agent state in the DB
+        agent_state = PersistedAgentState(
+            name=request.name,
+            user_id=user_id,
+            tool_names=request.tools if request.tools else [],
+            tool_rules=request.tool_rules,
+            agent_type=request.agent_type or AgentType.memgpt_agent,
+            llm_config=request.llm_config,
+            embedding_config=request.embedding_config,
+            system=request.system,
+            # other metadata
+            description=request.description,
+            metadata_=request.metadata_,
+        )
+        # TODO: move this to agent ORM
+        # this saves the agent ID and state into the DB
+        self.ms.create_agent(agent_state)
 
-        # TODO: move this into save_agent. save_agent should be moved to server.py
+        # Note: mappings (e.g. tags, blocks) are created after the agent is persisted
+        # TODO: add source mappings here as well
+
+        # create the tags
         if request.tags:
             for tag in request.tags:
-                self.agents_tags_manager.add_tag_to_agent(agent_id=agent.agent_state.id, tag=tag, actor=actor)
+                self.agents_tags_manager.add_tag_to_agent(agent_id=agent_state.id, tag=tag, actor=actor)
+
+        # create block mappins (now that agent is persisted)
+        for block in blocks:
+            # this links the created block to the agent
+            self.blocks_agents_manager.add_block_to_agent(block_id=block.id, agent_id=agent_state.id, block_label=block.label)
+
+        in_memory_agent_state = self.get_agent(agent_state.id)
+        return in_memory_agent_state
+
+    def get_agent(self, agent_id: str) -> AgentState:
+        """
+        Retrieve the full agent state from the DB.
+        This gathers data accross multiple tables to provide the full state of an agent, which is passed into the `Agent` object for creation.
+        """
+
+        # get data persisted from the DB
+        agent_state = self.ms.get_agent(agent_id=agent_id)
+        if agent_state is None:
+            # agent does not exist
+            return None
+        user = self.user_manager.get_user_by_id(user_id=agent_state.user_id)
 
-        assert isinstance(agent.agent_state.memory, Memory), f"Invalid memory type: {type(agent_state.memory)}"
+        # construct the in-memory, full agent state - this gather data stored in different tables but that needs to be passed to `Agent`
+        # we also return this data to the user to provide all the state related to an agent
 
-        # TODO: remove (hacky)
-        agent.agent_state.tags = self.agents_tags_manager.get_tags_for_agent(agent_id=agent.agent_state.id, actor=actor)
+        # get `Memory` object by getting the linked block IDs and fetching the blocks, then putting that into a `Memory` object
+        # this is the "in memory" representation of the in-context memory
+        block_ids = self.blocks_agents_manager.list_block_ids_for_agent(agent_id=agent_id)
+        blocks = []
+        for block_id in block_ids:
+            block = self.block_manager.get_block_by_id(block_id=block_id, actor=user)
+            assert block, f"Block with ID {block_id} does not exist"
+            blocks.append(block)
+        memory = Memory(blocks=blocks)
 
-        return agent.agent_state
+        # get `Tool` objects
+        tools = [self.tool_manager.get_tool_by_name(tool_name=tool_name, actor=user) for tool_name in agent_state.tool_names]
+
+        # get `Source` objects
+        sources = self.list_attached_sources(agent_id=agent_id)
+
+        # get the tags
+        tags = self.agents_tags_manager.get_tags_for_agent(agent_id=agent_id, actor=user)
+
+        # return the full agent state - this contains all data needed to recreate the agent
+        return AgentState(**agent_state.model_dump(), memory=memory, tools=tools, sources=sources, tags=tags)
 
     def update_agent(
         self,
         request: UpdateAgentState,
         actor: User,
-    ):
+    ) -> AgentState:
         """Update the agents core memory block, return the new state"""
         try:
             self.user_manager.get_user_by_id(user_id=actor.id)
@@ -961,13 +876,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={request.id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=request.id)
-
-        # update the core memory of the agent
-        if request.memory:
-            assert isinstance(request.memory, Memory), type(request.memory)
-            new_memory_contents = request.memory.to_flat_dict()
-            _ = self.update_agent_core_memory(user_id=actor.id, agent_id=request.id, new_memory_contents=new_memory_contents)
+        letta_agent = self.load_agent(agent_id=request.id)
 
         # update the system prompt
         if request.system:
@@ -981,13 +890,13 @@ class SyncServer(Server):
             letta_agent.set_message_buffer(message_ids=request.message_ids)
 
         # tools
-        if request.tools:
+        if request.tool_names:
             # Replace tools and also re-link
 
             # (1) get tools + make sure they exist
             # Current and target tools as sets of tool names
-            current_tools = set(letta_agent.agent_state.tools)
-            target_tools = set(request.tools)
+            current_tools = set(letta_agent.agent_state.tool_names)
+            target_tools = set(request.tool_names)
 
             # Calculate tools to add and remove
             tools_to_add = target_tools - current_tools
@@ -1004,7 +913,7 @@ class SyncServer(Server):
                 self.add_tool_to_agent(agent_id=request.id, tool_id=tool.id, user_id=actor.id)
 
             # reload agent
-            letta_agent = self._get_or_load_agent(agent_id=request.id)
+            letta_agent = self.load_agent(agent_id=request.id)
 
         # configs
         if request.llm_config:
@@ -1032,7 +941,6 @@ class SyncServer(Server):
                 self.agents_tags_manager.delete_tag_from_agent(agent_id=letta_agent.agent_state.id, tag=tag, actor=actor)
 
         # save the agent
-        assert isinstance(letta_agent.memory, Memory)
         save_agent(letta_agent, self.ms)
         # TODO: probably reload the agent somehow?
         return letta_agent.agent_state
@@ -1045,8 +953,8 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent.tools
+        letta_agent = self.load_agent(agent_id=agent_id)
+        return letta_agent.agent_state.tools
 
     def add_tool_to_agent(
         self,
@@ -1064,7 +972,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # Get all the tool objects from the request
         tool_objs = []
@@ -1072,7 +980,7 @@ class SyncServer(Server):
         assert tool_obj, f"Tool with id={tool_id} does not exist"
         tool_objs.append(tool_obj)
 
-        for tool in letta_agent.tools:
+        for tool in letta_agent.agent_state.tools:
             tool_obj = self.tool_manager.get_tool_by_id(tool_id=tool.id, actor=user)
             assert tool_obj, f"Tool with id={tool.id} does not exist"
 
@@ -1081,7 +989,7 @@ class SyncServer(Server):
                 tool_objs.append(tool_obj)
 
         # replace the list of tool names ("ids") inside the agent state
-        letta_agent.agent_state.tools = [tool.name for tool in tool_objs]
+        letta_agent.agent_state.tool_names = [tool.name for tool in tool_objs]
 
         # then attempt to link the tools modules
         letta_agent.link_tools(tool_objs)
@@ -1106,11 +1014,11 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # Get all the tool_objs
         tool_objs = []
-        for tool in letta_agent.tools:
+        for tool in letta_agent.agent_state.tools:
             tool_obj = self.tool_manager.get_tool_by_id(tool_id=tool.id, actor=user)
             assert tool_obj, f"Tool with id={tool.id} does not exist"
 
@@ -1119,7 +1027,7 @@ class SyncServer(Server):
                 tool_objs.append(tool_obj)
 
         # replace the list of tool names ("ids") inside the agent state
-        letta_agent.agent_state.tools = [tool.name for tool in tool_objs]
+        letta_agent.agent_state.tool_names = [tool.name for tool in tool_objs]
 
         # then attempt to link the tools modules
         letta_agent.link_tools(tool_objs)
@@ -1128,18 +1036,9 @@ class SyncServer(Server):
         save_agent(letta_agent, self.ms)
         return letta_agent.agent_state
 
-    def _agent_state_to_config(self, agent_state: AgentState) -> dict:
-        """Convert AgentState to a dict for a JSON response"""
-        assert agent_state is not None
-
-        agent_config = {
-            "id": agent_state.id,
-            "name": agent_state.name,
-            "human": agent_state._metadata.get("human", None),
-            "persona": agent_state._metadata.get("persona", None),
-            "created_at": agent_state.created_at.isoformat(),
-        }
-        return agent_config
+    def get_agent_state(self, user_id: str, agent_id: str) -> AgentState:
+        # TODO: duplicate, remove
+        return self.get_agent(agent_id=agent_id)
 
     def list_agents(self, user_id: str, tags: Optional[List[str]] = None) -> List[AgentState]:
         """List all available agents to a user"""
@@ -1147,13 +1046,13 @@ class SyncServer(Server):
 
         if tags is None:
             agents_states = self.ms.list_agents(user_id=user_id)
-            return agents_states
+            agent_ids = [agent.id for agent in agents_states]
         else:
             agent_ids = []
             for tag in tags:
                 agent_ids += self.agents_tags_manager.get_agents_by_tag(tag=tag, actor=user)
 
-            return [self.get_agent_state(user_id=user.id, agent_id=agent_id) for agent_id in agent_ids]
+        return [self.get_agent(agent_id=agent_id) for agent_id in agent_ids]
 
     # convert name->id
 
@@ -1177,34 +1076,34 @@ class SyncServer(Server):
 
     def get_agent_memory(self, agent_id: str) -> Memory:
         """Return the memory of an agent (core memory)"""
-        agent = self._get_or_load_agent(agent_id=agent_id)
-        return agent.memory
+        agent = self.load_agent(agent_id=agent_id)
+        return agent.agent_state.memory
 
     def get_archival_memory_summary(self, agent_id: str) -> ArchivalMemorySummary:
-        agent = self._get_or_load_agent(agent_id=agent_id)
+        agent = self.load_agent(agent_id=agent_id)
         return ArchivalMemorySummary(size=len(agent.persistence_manager.archival_memory))
 
     def get_recall_memory_summary(self, agent_id: str) -> RecallMemorySummary:
-        agent = self._get_or_load_agent(agent_id=agent_id)
+        agent = self.load_agent(agent_id=agent_id)
         return RecallMemorySummary(size=len(agent.persistence_manager.recall_memory))
 
     def get_in_context_message_ids(self, agent_id: str) -> List[str]:
         """Get the message ids of the in-context messages in the agent's memory"""
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return [m.id for m in letta_agent._messages]
+        agent = self.load_agent(agent_id=agent_id)
+        return [m.id for m in agent._messages]
 
     def get_in_context_messages(self, agent_id: str) -> List[Message]:
         """Get the in-context messages in the agent's memory"""
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent._messages
+        agent = self.load_agent(agent_id=agent_id)
+        return agent._messages
 
     def get_agent_message(self, agent_id: str, message_id: str) -> Message:
         """Get a single message from the agent's memory"""
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        message = letta_agent.persistence_manager.recall_memory.storage.get(id=message_id)
+        agent = self.load_agent(agent_id=agent_id)
+        message = agent.persistence_manager.recall_memory.storage.get(id=message_id)
         return message
 
     def get_agent_messages(
@@ -1212,11 +1111,10 @@ class SyncServer(Server):
         agent_id: str,
         start: int,
         count: int,
-        return_message_object: bool = True,
     ) -> Union[List[Message], List[LettaMessage]]:
         """Paginated query of all messages in agent message queue"""
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         if start < 0 or count < 0:
             raise ValueError("Start and count values should be non-negative")
@@ -1234,10 +1132,6 @@ class SyncServer(Server):
             # Slice the list for pagination
             messages = reversed_messages[start:end_index]
 
-            ## Convert to json
-            ## Add a tag indicating in-context or not
-            # json_messages = [{**record.to_json(), "in_context": True} for record in messages]
-
         else:
             # need to access persistence manager for additional messages
             db_iterator = letta_agent.persistence_manager.recall_memory.storage.get_all_paginated(page_size=count, offset=start)
@@ -1257,9 +1151,6 @@ class SyncServer(Server):
             # for d in json_messages:
             #    d["in_context"] = True if str(d["id"]) in in_context_message_ids else False
 
-        if not return_message_object:
-            messages = [msg for m in messages for msg in m.to_letta_message()]
-
         return messages
 
     def get_agent_archival(self, user_id: str, agent_id: str, start: int, count: int) -> List[Passage]:
@@ -1270,7 +1161,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # iterate over records
         db_iterator = letta_agent.persistence_manager.archival_memory.storage.get_all_paginated(page_size=count, offset=start)
@@ -1295,7 +1186,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # iterate over recorde
         cursor, records = letta_agent.persistence_manager.archival_memory.storage.get_all_cursor(
@@ -1310,11 +1201,15 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # Insert into archival memory
         passage_ids = letta_agent.persistence_manager.archival_memory.insert(memory_string=memory_contents, return_ids=True)
 
+        # Update the agent
+        # TODO: should this update the system prompt?
+        save_agent(letta_agent, self.ms)
+
         # TODO: this is gross, fix
         return [letta_agent.persistence_manager.archival_memory.storage.get(id=passage_id) for passage_id in passage_ids]
 
@@ -1327,7 +1222,7 @@ class SyncServer(Server):
         # TODO: should return a passage
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # Delete by ID
         # TODO check if it exists first, and throw error if not
@@ -1346,9 +1241,8 @@ class SyncServer(Server):
         order: Optional[str] = "asc",
         reverse: Optional[bool] = False,
         return_message_object: bool = True,
-        use_assistant_message: bool = False,
-        assistant_message_function_name: str = constants.DEFAULT_MESSAGE_TOOL,
-        assistant_message_function_kwarg: str = constants.DEFAULT_MESSAGE_TOOL_KWARG,
+        assistant_message_tool_name: str = constants.DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = constants.DEFAULT_MESSAGE_TOOL_KWARG,
     ) -> Union[List[Message], List[LettaMessage]]:
         if self.user_manager.get_user_by_id(user_id=user_id) is None:
             raise ValueError(f"User user_id={user_id} does not exist")
@@ -1356,7 +1250,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         # iterate over records
         cursor, records = letta_agent.persistence_manager.recall_memory.storage.get_all_cursor(
@@ -1367,50 +1261,19 @@ class SyncServer(Server):
 
         if not return_message_object:
             # If we're GETing messages in reverse, we need to reverse the inner list (generated by to_letta_message)
-            if reverse:
-                records = [
-                    msg
-                    for m in records
-                    for msg in m.to_letta_message(
-                        assistant_message=use_assistant_message,
-                        assistant_message_function_name=assistant_message_function_name,
-                        assistant_message_function_kwarg=assistant_message_function_kwarg,
-                    )[::-1]
-                ]
-            else:
-                records = [
-                    msg
-                    for m in records
-                    for msg in m.to_letta_message(
-                        assistant_message=use_assistant_message,
-                        assistant_message_function_name=assistant_message_function_name,
-                        assistant_message_function_kwarg=assistant_message_function_kwarg,
-                    )
-                ]
-
-        return records
-
-    def get_agent_state(self, user_id: str, agent_id: Optional[str], agent_name: Optional[str] = None) -> Optional[AgentState]:
-        """Return the config of an agent"""
-        user = self.user_manager.get_user_by_id(user_id=user_id)
-        if agent_id:
-            if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None:
-                return None
-        else:
-            agent_state = self.ms.get_agent(agent_name=agent_name, user_id=user_id)
-            if agent_state is None:
-                raise ValueError(f"Agent agent_name={agent_name} does not exist")
-            agent_id = agent_state.id
+            records = [
+                msg
+                for m in records
+                for msg in m.to_letta_message(
+                    assistant_message_tool_name=assistant_message_tool_name,
+                    assistant_message_tool_kwarg=assistant_message_tool_kwarg,
+                )
+            ]
 
-        # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        assert isinstance(letta_agent.memory, Memory)
+        if reverse:
+            records = records[::-1]
 
-        letta_agent.update_memory_blocks_from_db()
-        agent_state = letta_agent.agent_state.model_copy(deep=True)
-        # Load the tags in for the agent_state
-        agent_state.tags = self.agents_tags_manager.get_tags_for_agent(agent_id=agent_id, actor=user)
-        return agent_state
+        return records
 
     def get_server_config(self, include_defaults: bool = False) -> dict:
         """Return the base config"""
@@ -1435,39 +1298,23 @@ class SyncServer(Server):
 
         return response
 
-    def update_agent_core_memory(self, user_id: str, agent_id: str, new_memory_contents: dict) -> Memory:
-        """Update the agents core memory block, return the new state"""
-        if self.user_manager.get_user_by_id(user_id=user_id) is None:
-            raise ValueError(f"User user_id={user_id} does not exist")
-        if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None:
-            raise ValueError(f"Agent agent_id={agent_id} does not exist")
+    def update_agent_core_memory(self, user_id: str, agent_id: str, label: str, value: str) -> Memory:
+        """Update the value of a block in the agent's memory"""
 
-        # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-
-        # old_core_memory = self.get_agent_memory(agent_id=agent_id)
-
-        modified = False
-        for key, value in new_memory_contents.items():
-            if letta_agent.memory.get_block(key) is None:
-                # raise ValueError(f"Key {key} not found in agent memory {list(letta_agent.memory.list_block_names())}")
-                raise ValueError(f"Key {key} not found in agent memory {str(letta_agent.memory.memory)}")
-            if value is None:
-                continue
-            if letta_agent.memory.get_block(key) != value:
-                letta_agent.memory.update_block_value(label=key, value=value)  # update agent memory
-                modified = True
-
-        # If we modified the memory contents, we need to rebuild the memory block inside the system message
-        if modified:
-            letta_agent.rebuild_memory()
-            # letta_agent.rebuild_memory(force=True, ms=self.ms)  # This breaks unit tests in test_local_client.py
-            # save agent
-            save_agent(letta_agent, self.ms)
+        # get the block id
+        block = self.get_agent_block_by_label(user_id=user_id, agent_id=agent_id, label=label)
+        block_id = block.id
+
+        # update the block
+        self.block_manager.update_block(
+            block_id=block_id, block_update=BlockUpdate(value=value), actor=self.user_manager.get_user_by_id(user_id=user_id)
+        )
 
-        return self.ms.get_agent(agent_id=agent_id).memory
+        # load agent
+        letta_agent = self.load_agent(agent_id=agent_id)
+        return letta_agent.agent_state.memory
 
-    def rename_agent(self, user_id: str, agent_id: str, new_agent_name: str) -> AgentState:
+    def rename_agent(self, user_id: str, agent_id: str, new_agent_name: str) -> PersistedAgentState:
         """Update the name of the agent in the database"""
         if self.user_manager.get_user_by_id(user_id=user_id) is None:
             raise ValueError(f"User user_id={user_id} does not exist")
@@ -1475,7 +1322,7 @@ class SyncServer(Server):
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
 
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
 
         current_name = letta_agent.agent_state.name
         if current_name == new_agent_name:
@@ -1497,6 +1344,7 @@ class SyncServer(Server):
         # TODO: REMOVE THIS ONCE WE MIGRATE AGENTMODEL TO ORM MODEL
         # TODO: EVENTUALLY WE GET AUTO-DELETES WHEN WE SPECIFY RELATIONSHIPS IN THE ORM
         self.agents_tags_manager.delete_all_tags_from_agent(agent_id=agent_id, actor=actor)
+        self.blocks_agents_manager.remove_all_agent_blocks(agent_id=agent_id)
 
         if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None:
             raise ValueError(f"Agent agent_id={agent_id} does not exist")
@@ -1522,7 +1370,7 @@ class SyncServer(Server):
 
         # Next, attempt to delete it from the actual database
         try:
-            self.ms.delete_agent(agent_id=agent_id)
+            self.ms.delete_agent(agent_id=agent_id, per_agent_lock_manager=self.per_agent_lock_manager)
         except Exception as e:
             logger.exception(f"Failed to delete agent {agent_id} via ID with:\n{str(e)}")
             raise ValueError(f"Failed to delete agent {agent_id} in database")
@@ -1664,9 +1512,10 @@ class SyncServer(Server):
             raise ValueError(f"Need to provide at least source_id or source_name to find the source.")
         # get connection to data source storage
         source_connector = StorageConnector.get_storage_connector(TableType.PASSAGES, self.config, user_id=user_id)
+        assert data_source, f"Data source with id={source_id} or name={source_name} does not exist"
 
         # load agent
-        agent = self._get_or_load_agent(agent_id=agent_id)
+        agent = self.load_agent(agent_id=agent_id)
 
         # attach source to agent
         agent.attach_source(data_source.id, source_connector, self.ms)
@@ -1691,7 +1540,7 @@ class SyncServer(Server):
         source_id = source.id
 
         # delete all Passage objects with source_id==source_id from agent's archival memory
-        agent = self._get_or_load_agent(agent_id=agent_id)
+        agent = self.load_agent(agent_id=agent_id)
         archival_memory = agent.persistence_manager.archival_memory
         archival_memory.storage.delete({"source_id": source_id})
 
@@ -1770,34 +1619,43 @@ class SyncServer(Server):
     def get_agent_message(self, agent_id: str, message_id: str) -> Optional[Message]:
         """Get a single message from the agent's memory"""
         # Get the agent object (loaded in memory)
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
         message = letta_agent.persistence_manager.recall_memory.storage.get(id=message_id)
+        save_agent(letta_agent, self.ms)
         return message
 
     def update_agent_message(self, agent_id: str, request: UpdateMessage) -> Message:
         """Update the details of a message associated with an agent"""
 
         # Get the current message
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent.update_message(request=request)
+        letta_agent = self.load_agent(agent_id=agent_id)
+        response = letta_agent.update_message(request=request)
+        save_agent(letta_agent, self.ms)
+        return response
 
     def rewrite_agent_message(self, agent_id: str, new_text: str) -> Message:
 
         # Get the current message
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent.rewrite_message(new_text=new_text)
+        letta_agent = self.load_agent(agent_id=agent_id)
+        response = letta_agent.rewrite_message(new_text=new_text)
+        save_agent(letta_agent, self.ms)
+        return response
 
     def rethink_agent_message(self, agent_id: str, new_thought: str) -> Message:
 
         # Get the current message
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent.rethink_message(new_thought=new_thought)
+        letta_agent = self.load_agent(agent_id=agent_id)
+        response = letta_agent.rethink_message(new_thought=new_thought)
+        save_agent(letta_agent, self.ms)
+        return response
 
     def retry_agent_message(self, agent_id: str) -> List[Message]:
 
         # Get the current message
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        return letta_agent.retry_message()
+        letta_agent = self.load_agent(agent_id=agent_id)
+        response = letta_agent.retry_message()
+        save_agent(letta_agent, self.ms)
+        return response
 
     def get_user_or_default(self, user_id: Optional[str]) -> User:
         """Get the user object for user_id if it exists, otherwise return the default user object"""
@@ -1846,121 +1704,49 @@ class SyncServer(Server):
         agent_id: str,
     ) -> ContextWindowOverview:
         # Get the current message
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        letta_agent = self.load_agent(agent_id=agent_id)
         return letta_agent.get_context_window()
 
-    def update_agent_memory_label(self, user_id: str, agent_id: str, current_block_label: str, new_block_label: str) -> Memory:
-        """Update the label of a block in an agent's memory"""
-
-        # Get the user
-        user = self.user_manager.get_user_by_id(user_id=user_id)
-
-        # Link a block to an agent's memory
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        letta_agent.memory.update_block_label(current_label=current_block_label, new_label=new_block_label)
-        assert new_block_label in letta_agent.memory.list_block_labels()
-        self.block_manager.create_or_update_block(block=letta_agent.memory.get_block(new_block_label), actor=user)
-
-        # check that the block was updated
-        updated_block = self.block_manager.get_block_by_id(block_id=letta_agent.memory.get_block(new_block_label).id, actor=user)
-
-        # Recompile the agent memory
-        letta_agent.rebuild_memory(force=True, ms=self.ms)
-
-        # save agent
-        save_agent(letta_agent, self.ms)
-
-        updated_agent = self.ms.get_agent(agent_id=agent_id)
-        if updated_agent is None:
-            raise ValueError(f"Agent with id {agent_id} not found after linking block")
-        assert new_block_label in updated_agent.memory.list_block_labels()
-        assert current_block_label not in updated_agent.memory.list_block_labels()
-        return updated_agent.memory
-
     def link_block_to_agent_memory(self, user_id: str, agent_id: str, block_id: str) -> Memory:
         """Link a block to an agent's memory"""
-
-        # Get the user
-        user = self.user_manager.get_user_by_id(user_id=user_id)
-
-        # Get the block first
-        block = self.block_manager.get_block_by_id(block_id=block_id, actor=user)
+        block = self.block_manager.get_block_by_id(block_id=block_id, actor=self.user_manager.get_user_by_id(user_id=user_id))
         if block is None:
             raise ValueError(f"Block with id {block_id} not found")
+        self.blocks_agents_manager.add_block_to_agent(agent_id, block_id, block_label=block.label)
 
-        # Link a block to an agent's memory
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        letta_agent.memory.link_block(block=block)
-        assert block.label in letta_agent.memory.list_block_labels()
-
-        # Recompile the agent memory
-        letta_agent.rebuild_memory(force=True, ms=self.ms)
-
-        # save agent
-        save_agent(letta_agent, self.ms)
-
-        updated_agent = self.ms.get_agent(agent_id=agent_id)
-        if updated_agent is None:
-            raise ValueError(f"Agent with id {agent_id} not found after linking block")
-        assert block.label in updated_agent.memory.list_block_labels()
-
-        return updated_agent.memory
+        # get agent memory
+        memory = self.load_agent(agent_id=agent_id).agent_state.memory
+        return memory
 
     def unlink_block_from_agent_memory(self, user_id: str, agent_id: str, block_label: str, delete_if_no_ref: bool = True) -> Memory:
         """Unlink a block from an agent's memory. If the block is not linked to any agent, delete it."""
+        self.blocks_agents_manager.remove_block_with_label_from_agent(agent_id=agent_id, block_label=block_label)
 
-        # Get the user
-        user = self.user_manager.get_user_by_id(user_id=user_id)
-
-        # Link a block to an agent's memory
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        unlinked_block = letta_agent.memory.unlink_block(block_label=block_label)
-        assert unlinked_block.label not in letta_agent.memory.list_block_labels()
-
-        # Check if the block is linked to any other agent
-        # TODO needs reference counting GC to handle loose blocks
-        # block = self.block_manager.get_block_by_id(block_id=unlinked_block.id, actor=user)
-        # if block is None:
-        # raise ValueError(f"Block with id {block_id} not found")
-
-        # Recompile the agent memory
-        letta_agent.rebuild_memory(force=True, ms=self.ms)
-
-        # save agent
-        save_agent(letta_agent, self.ms)
-
-        updated_agent = self.ms.get_agent(agent_id=agent_id)
-        if updated_agent is None:
-            raise ValueError(f"Agent with id {agent_id} not found after linking block")
-        assert unlinked_block.label not in updated_agent.memory.list_block_labels()
-        return updated_agent.memory
+        # get agent memory
+        memory = self.load_agent(agent_id=agent_id).agent_state.memory
+        return memory
 
     def update_agent_memory_limit(self, user_id: str, agent_id: str, block_label: str, limit: int) -> Memory:
         """Update the limit of a block in an agent's memory"""
+        block = self.get_agent_block_by_label(user_id=user_id, agent_id=agent_id, label=block_label)
+        self.block_manager.update_block(
+            block_id=block.id, block_update=BlockUpdate(limit=limit), actor=self.user_manager.get_user_by_id(user_id=user_id)
+        )
+        # get agent memory
+        memory = self.load_agent(agent_id=agent_id).agent_state.memory
+        return memory
+
+    def upate_block(self, user_id: str, block_id: str, block_update: BlockUpdate) -> Block:
+        """Update a block"""
+        return self.block_manager.update_block(
+            block_id=block_id, block_update=block_update, actor=self.user_manager.get_user_by_id(user_id=user_id)
+        )
 
-        # Get the user
-        user = self.user_manager.get_user_by_id(user_id=user_id)
-
-        # Link a block to an agent's memory
-        letta_agent = self._get_or_load_agent(agent_id=agent_id)
-        letta_agent.memory.update_block_limit(label=block_label, limit=limit)
-        assert block_label in letta_agent.memory.list_block_labels()
-
-        # write out the update the database
-        self.block_manager.create_or_update_block(block=letta_agent.memory.get_block(block_label), actor=user)
-
-        # check that the block was updated
-        updated_block = self.block_manager.get_block_by_id(block_id=letta_agent.memory.get_block(block_label).id, actor=user)
-        assert updated_block and updated_block.limit == limit
-
-        # Recompile the agent memory
-        letta_agent.rebuild_memory(force=True, ms=self.ms)
-
-        # save agent
-        save_agent(letta_agent, self.ms)
-
-        updated_agent = self.ms.get_agent(agent_id=agent_id)
-        if updated_agent is None:
-            raise ValueError(f"Agent with id {agent_id} not found after linking block")
-        assert updated_agent.memory.get_block(label=block_label).limit == limit
-        return updated_agent.memory
+    def get_agent_block_by_label(self, user_id: str, agent_id: str, label: str) -> Block:
+        """Get a block by label"""
+        # TODO: implement at ORM?
+        for block_id in self.blocks_agents_manager.list_block_ids_for_agent(agent_id=agent_id):
+            block = self.block_manager.get_block_by_id(block_id=block_id, actor=self.user_manager.get_user_by_id(user_id=user_id))
+            if block.label == label:
+                return block
+        return None