letta-nightly 0.5.5.dev20241122170833__tar.gz → 0.6.0.dev20241204054116__tar.gz

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.5.5.dev20241122170833
+Version: 0.6.0.dev20241204054116
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: all
 Provides-Extra: autogen
+Provides-Extra: cloud-tool-sandbox
 Provides-Extra: dev
 Provides-Extra: external-tools
 Provides-Extra: milvus
@@ -25,13 +26,14 @@ Requires-Dist: alembic (>=1.13.3,<2.0.0)
 Requires-Dist: autoflake (>=2.3.0,<3.0.0) ; extra == "dev" or extra == "all"
 Requires-Dist: black[jupyter] (>=24.2.0,<25.0.0) ; extra == "dev" or extra == "all"
 Requires-Dist: chromadb (>=0.4.24,<0.5.0)
-Requires-Dist: composio-core (>=0.5.34,<0.6.0) ; extra == "external-tools" or extra == "all"
-Requires-Dist: composio-langchain (>=0.5.28,<0.6.0) ; extra == "external-tools" or extra == "all"
+Requires-Dist: composio-core (>=0.5.34,<0.6.0)
+Requires-Dist: composio-langchain (>=0.5.28,<0.6.0)
 Requires-Dist: datasets (>=2.14.6,<3.0.0) ; extra == "dev" or extra == "all"
 Requires-Dist: demjson3 (>=3.0.6,<4.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0) ; extra == "external-tools" or extra == "all"
 Requires-Dist: docstring-parser (>=0.16,<0.17)
 Requires-Dist: docx2txt (>=0.8,<0.9)
+Requires-Dist: e2b-code-interpreter (>=1.0.1,<2.0.0) ; extra == "cloud-tool-sandbox"
 Requires-Dist: fastapi (>=0.104.1,<0.105.0) ; extra == "server" or extra == "all"
 Requires-Dist: html2text (>=2020.1.16,<2021.0.0)
 Requires-Dist: httpx (>=0.27.2,<0.28.0)
@@ -127,24 +129,24 @@ The two main ways to install Letta are through **pypi** (`pip`) or via **Docker*
 
 ### Step 1 - Install Letta using `pip`
 ```sh
-$ pip install -U letta
+pip install -U letta
 ```
 
 ### Step 2 - Set your environment variables for your chosen LLM / embedding providers
 ```sh
-$ export OPENAI_API_KEY=sk-...
+export OPENAI_API_KEY=sk-...
 ```
 
 For Ollama (see our full [documentation](https://docs.letta.com/install) for examples of how to set up various providers):
 ```sh
-$ export OLLAMA_BASE_URL=http://localhost:11434
+export OLLAMA_BASE_URL=http://localhost:11434
 ```
 
 ### Step 3 - Run the Letta CLI
 
 You can create agents and chat with them via the Letta CLI tool (`letta run`):
 ```sh
-$ letta run
+letta run
 ```
 ```
 🧬 Creating new agent...
@@ -177,7 +179,7 @@ Hit enter to begin (will request first Letta message)
 
 You can start the Letta API server with `letta server` (see the full API reference [here](https://docs.letta.com/api-reference)):
 ```sh
-$ letta server
+letta server
 ```
 ```
 Initializing database...

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/README.md

@@ -44,24 +44,24 @@ The two main ways to install Letta are through **pypi** (`pip`) or via **Docker*
 
 ### Step 1 - Install Letta using `pip`
 ```sh
-$ pip install -U letta
+pip install -U letta
 ```
 
 ### Step 2 - Set your environment variables for your chosen LLM / embedding providers
 ```sh
-$ export OPENAI_API_KEY=sk-...
+export OPENAI_API_KEY=sk-...
 ```
 
 For Ollama (see our full [documentation](https://docs.letta.com/install) for examples of how to set up various providers):
 ```sh
-$ export OLLAMA_BASE_URL=http://localhost:11434
+export OLLAMA_BASE_URL=http://localhost:11434
 ```
 
 ### Step 3 - Run the Letta CLI
 
 You can create agents and chat with them via the Letta CLI tool (`letta run`):
 ```sh
-$ letta run
+letta run
 ```
 ```
 🧬 Creating new agent...
@@ -94,7 +94,7 @@ Hit enter to begin (will request first Letta message)
 
 You can start the Letta API server with `letta server` (see the full API reference [here](https://docs.letta.com/api-reference)):
 ```sh
-$ letta server
+letta server
 ```
 ```
 Initializing database...

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/letta/__init__.py

@@ -1,10 +1,10 @@
-__version__ = "0.5.5"
+__version__ = "0.6.0"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client
 
 # imports for easier access
-from letta.schemas.agent import AgentState
+from letta.schemas.agent import AgentState, PersistedAgentState
 from letta.schemas.block import Block
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import JobStatus

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/letta/agent.py

@@ -1,5 +1,6 @@
 import datetime
 import inspect
+import time
 import traceback
 import warnings
 from abc import ABC, abstractmethod
@@ -9,6 +10,7 @@ from tqdm import tqdm
 
 from letta.agent_store.storage import StorageConnector
 from letta.constants import (
+    BASE_TOOLS,
     CLI_WARNING_PREFIX,
     FIRST_MESSAGE_ATTEMPTS,
     FUNC_FAILED_HEARTBEAT_MESSAGE,
@@ -30,7 +32,7 @@ from letta.metadata import MetadataStore
 from letta.orm import User
 from letta.persistence_manager import LocalStateManager
 from letta.schemas.agent import AgentState, AgentStepResponse
-from letta.schemas.block import Block
+from letta.schemas.block import BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import MessageRole
 from letta.schemas.memory import ContextWindowOverview, Memory
@@ -49,6 +51,7 @@ from letta.schemas.tool_rule import TerminalToolRule
 from letta.schemas.usage import LettaUsageStatistics
 from letta.services.block_manager import BlockManager
 from letta.services.source_manager import SourceManager
+from letta.services.tool_execution_sandbox import ToolExecutionSandbox
 from letta.services.user_manager import UserManager
 from letta.streaming_interface import StreamingRefreshCLIInterface
 from letta.system import (
@@ -233,11 +236,8 @@ class Agent(BaseAgent):
     def __init__(
         self,
         interface: Optional[Union[AgentInterface, StreamingRefreshCLIInterface]],
-        # agents can be created from providing agent_state
-        agent_state: AgentState,
-        tools: List[Tool],
+        agent_state: AgentState,  # in-memory representation of the agent state (read from multiple tables)
         user: User,
-        # memory: Memory,
         # extras
         messages_total: Optional[int] = None,  # TODO remove?
         first_message_verify_mono: bool = True,  # TODO move to config?
@@ -251,7 +251,7 @@ class Agent(BaseAgent):
         self.user = user
 
         # link tools
-        self.link_tools(tools)
+        self.link_tools(agent_state.tools)
 
         # initialize a tool rules solver
         if agent_state.tool_rules:
@@ -263,26 +263,14 @@ class Agent(BaseAgent):
         # add default rule for having send_message be a terminal tool
         if agent_state.tool_rules is None:
             agent_state.tool_rules = []
-        # Define the rule to add
-        send_message_terminal_rule = TerminalToolRule(tool_name="send_message")
-        # Check if an equivalent rule is already present
-        if not any(
-            isinstance(rule, TerminalToolRule) and rule.tool_name == send_message_terminal_rule.tool_name for rule in agent_state.tool_rules
-        ):
-            agent_state.tool_rules.append(send_message_terminal_rule)
 
         self.tool_rules_solver = ToolRulesSolver(tool_rules=agent_state.tool_rules)
 
         # gpt-4, gpt-3.5-turbo, ...
         self.model = self.agent_state.llm_config.model
 
-        # Store the system instructions (used to rebuild memory)
-        self.system = self.agent_state.system
-
-        # Initialize the memory object
-        self.memory = self.agent_state.memory
-        assert isinstance(self.memory, Memory), f"Memory object is not of type Memory: {type(self.memory)}"
-        printd("Initialized memory object", self.memory.compile())
+        # state managers
+        self.block_manager = BlockManager()
 
         # Interface must implement:
         # - internal_monologue
@@ -320,8 +308,8 @@ class Agent(BaseAgent):
             # Generate a sequence of initial messages to put in the buffer
             init_messages = initialize_message_sequence(
                 model=self.model,
-                system=self.system,
-                memory=self.memory,
+                system=self.agent_state.system,
+                memory=self.agent_state.memory,
                 archival_memory=None,
                 recall_memory=None,
                 memory_edit_timestamp=get_utc_time(),
@@ -343,8 +331,8 @@ class Agent(BaseAgent):
                 # Basic "more human than human" initial message sequence
                 init_messages = initialize_message_sequence(
                     model=self.model,
-                    system=self.system,
-                    memory=self.memory,
+                    system=self.agent_state.system,
+                    memory=self.agent_state.memory,
                     archival_memory=None,
                     recall_memory=None,
                     memory_edit_timestamp=get_utc_time(),
@@ -378,6 +366,76 @@ class Agent(BaseAgent):
         # Create the agent in the DB
         self.update_state()
 
+    def update_memory_if_change(self, new_memory: Memory) -> bool:
+        """
+        Update internal memory object and system prompt if there have been modifications.
+
+        Args:
+            new_memory (Memory): the new memory object to compare to the current memory object
+
+        Returns:
+            modified (bool): whether the memory was updated
+        """
+        if self.agent_state.memory.compile() != new_memory.compile():
+            # update the blocks (LRW) in the DB
+            for label in self.agent_state.memory.list_block_labels():
+                updated_value = new_memory.get_block(label).value
+                if updated_value != self.agent_state.memory.get_block(label).value:
+                    # update the block if it's changed
+                    block_id = self.agent_state.memory.get_block(label).id
+                    block = self.block_manager.update_block(
+                        block_id=block_id, block_update=BlockUpdate(value=updated_value), actor=self.user
+                    )
+
+            # refresh memory from DB (using block ids)
+            self.agent_state.memory = Memory(
+                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()]
+            )
+
+            # NOTE: don't do this since re-buildin the memory is handled at the start of the step
+            # rebuild memory - this records the last edited timestamp of the memory
+            # TODO: pass in update timestamp from block edit time
+            self.rebuild_system_prompt()
+
+            return True
+        return False
+
+    def execute_tool_and_persist_state(self, function_name, function_to_call, function_args):
+        """
+        Execute tool modifications and persist the state of the agent.
+        Note: only some agent state modifications will be persisted, such as data in the AgentState ORM and block data
+        """
+        # TODO: add agent manager here
+        orig_memory_str = self.agent_state.memory.compile()
+
+        # TODO: need to have an AgentState object that actually has full access to the block data
+        # this is because the sandbox tools need to be able to access block.value to edit this data
+        try:
+            if function_name in BASE_TOOLS:
+                # base tools are allowed to access the `Agent` object and run on the database
+                function_args["self"] = self  # need to attach self to arg since it's dynamically linked
+                function_response = function_to_call(**function_args)
+            else:
+                # execute tool in a sandbox
+                # TODO: allow agent_state to specify which sandbox to execute tools in
+                sandbox_run_result = ToolExecutionSandbox(function_name, function_args, self.agent_state.user_id).run(
+                    agent_state=self.agent_state.__deepcopy__()
+                )
+                function_response, updated_agent_state = sandbox_run_result.func_return, sandbox_run_result.agent_state
+                assert orig_memory_str == self.agent_state.memory.compile(), "Memory should not be modified in a sandbox tool"
+                self.update_memory_if_change(updated_agent_state.memory)
+        except Exception as e:
+            # Need to catch error here, or else trunction wont happen
+            # TODO: modify to function execution error
+            from letta.constants import MAX_ERROR_MESSAGE_CHAR_LIMIT
+
+            error_msg = f"Error executing tool {function_name}: {e}"
+            if len(error_msg) > MAX_ERROR_MESSAGE_CHAR_LIMIT:
+                error_msg = error_msg[:MAX_ERROR_MESSAGE_CHAR_LIMIT]
+            raise ValueError(error_msg)
+
+        return function_response
+
     @property
     def messages(self) -> List[dict]:
         """Getter method that converts the internal Message list into OpenAI-style dicts"""
@@ -390,16 +448,6 @@ class Agent(BaseAgent):
     def link_tools(self, tools: List[Tool]):
         """Bind a tool object (schema + python function) to the agent object"""
 
-        # tools
-        for tool in tools:
-            assert tool, f"Tool is None - must be error in querying tool from DB"
-            assert tool.name in self.agent_state.tools, f"Tool {tool} not found in agent_state.tools"
-        for tool_name in self.agent_state.tools:
-            assert tool_name in [tool.name for tool in tools], f"Tool name {tool_name} not included in agent tool list"
-
-        # Update tools
-        self.tools = tools
-
         # Store the functions schemas (this is passed as an argument to ChatCompletion)
         self.functions = []
         self.functions_python = {}
@@ -414,9 +462,8 @@ class Agent(BaseAgent):
                     exec(tool.source_code, env)
                 self.functions_python[tool.json_schema["name"]] = env[tool.json_schema["name"]]
                 self.functions.append(tool.json_schema)
-            except Exception as e:
+            except Exception:
                 warnings.warn(f"WARNING: tool {tool.name} failed to link")
-                print(e)
         assert all([callable(f) for k, f in self.functions_python.items()]), self.functions_python
 
     def _load_messages_from_recall(self, message_ids: List[str]) -> List[Message]:
@@ -520,60 +567,60 @@ class Agent(BaseAgent):
         self,
         message_sequence: List[Message],
         function_call: str = "auto",
-        first_message: bool = False,  # hint
+        first_message: bool = False,
         stream: bool = False,  # TODO move to config?
-        fail_on_empty_response: bool = False,
         empty_response_retry_limit: int = 3,
+        backoff_factor: float = 0.5,  # delay multiplier for exponential backoff
+        max_delay: float = 10.0,  # max delay between retries
     ) -> ChatCompletionResponse:
-        """Get response from LLM API"""
-        # Get the allowed tools based on the ToolRulesSolver state
+        """Get response from LLM API with robust retry mechanism."""
+
         allowed_tool_names = self.tool_rules_solver.get_allowed_tool_names()
+        allowed_functions = (
+            self.functions if not allowed_tool_names else [func for func in self.functions if func["name"] in allowed_tool_names]
+        )
 
-        if not allowed_tool_names:
-            # if it's empty, any available tools are fair game
-            allowed_functions = self.functions
-        else:
-            allowed_functions = [func for func in self.functions if func["name"] in allowed_tool_names]
+        for attempt in range(1, empty_response_retry_limit + 1):
+            try:
+                response = create(
+                    llm_config=self.agent_state.llm_config,
+                    messages=message_sequence,
+                    user_id=self.agent_state.user_id,
+                    functions=allowed_functions,
+                    functions_python=self.functions_python,
+                    function_call=function_call,
+                    first_message=first_message,
+                    stream=stream,
+                    stream_interface=self.interface,
+                )
 
-        try:
-            response = create(
-                # agent_state=self.agent_state,
-                llm_config=self.agent_state.llm_config,
-                messages=message_sequence,
-                user_id=self.agent_state.user_id,
-                functions=allowed_functions,
-                functions_python=self.functions_python,
-                function_call=function_call,
-                # hint
-                first_message=first_message,
-                # streaming
-                stream=stream,
-                stream_interface=self.interface,
-            )
+                # These bottom two are retryable
+                if len(response.choices) == 0 or response.choices[0] is None:
+                    raise ValueError(f"API call returned an empty message: {response}")
 
-            if len(response.choices) == 0 or response.choices[0] is None:
-                empty_api_err_message = f"API call didn't return a message: {response}"
-                if fail_on_empty_response or empty_response_retry_limit == 0:
-                    raise Exception(empty_api_err_message)
-                else:
-                    # Decrement retry limit and try again
-                    warnings.warn(empty_api_err_message)
-                    return self._get_ai_reply(
-                        message_sequence, function_call, first_message, stream, fail_on_empty_response, empty_response_retry_limit - 1
-                    )
+                if response.choices[0].finish_reason not in ["stop", "function_call", "tool_calls"]:
+                    if response.choices[0].finish_reason == "length":
+                        # This is not retryable, hence RuntimeError v.s. ValueError
+                        raise RuntimeError("Finish reason was length (maximum context length)")
+                    else:
+                        raise ValueError(f"Bad finish reason from API: {response.choices[0].finish_reason}")
 
-            # special case for 'length'
-            if response.choices[0].finish_reason == "length":
-                raise Exception("Finish reason was length (maximum context length)")
+                return response
 
-            # catches for soft errors
-            if response.choices[0].finish_reason not in ["stop", "function_call", "tool_calls"]:
-                raise Exception(f"API call finish with bad finish reason: {response}")
+            except ValueError as ve:
+                if attempt >= empty_response_retry_limit:
+                    warnings.warn(f"Retry limit reached. Final error: {ve}")
+                    break
+                else:
+                    delay = min(backoff_factor * (2 ** (attempt - 1)), max_delay)
+                    warnings.warn(f"Attempt {attempt} failed: {ve}. Retrying in {delay} seconds...")
+                    time.sleep(delay)
 
-            # unpack with response.choices[0].message.content
-            return response
-        except Exception as e:
-            raise e
+            except Exception as e:
+                # For non-retryable errors, exit immediately
+                raise e
+
+        raise Exception("Retries exhausted and no valid response received.")
 
     def _handle_ai_response(
         self,
@@ -725,9 +772,10 @@ class Agent(BaseAgent):
                     if isinstance(function_args[name], dict):
                         function_args[name] = spec[name](**function_args[name])
 
-                function_args["self"] = self  # need to attach self to arg since it's dynamically linked
+                # handle tool execution (sandbox) and state updates
+                function_response = self.execute_tool_and_persist_state(function_name, function_to_call, function_args)
 
-                function_response = function_to_call(**function_args)
+                # handle trunction
                 if function_name in ["conversation_search", "conversation_search_date", "archival_memory_search"]:
                     # with certain functions we rely on the paging mechanism to handle overflow
                     truncate = False
@@ -747,6 +795,7 @@ class Agent(BaseAgent):
                 error_msg_user = f"{error_msg}\n{traceback.format_exc()}"
                 printd(error_msg_user)
                 function_response = package_function_response(False, error_msg)
+                # TODO: truncate error message somehow
                 messages.append(
                     Message.dict_to_message(
                         agent_id=self.agent_state.id,
@@ -799,7 +848,7 @@ class Agent(BaseAgent):
 
         # rebuild memory
         # TODO: @charles please check this
-        self.rebuild_memory()
+        self.rebuild_system_prompt()
 
         # Update ToolRulesSolver state with last called function
         self.tool_rules_solver.update_tool_usage(function_name)
@@ -915,17 +964,10 @@ class Agent(BaseAgent):
 
             # Step 0: update core memory
             # only pulling latest block data if shared memory is being used
-            # TODO: ensure we're passing in metadata store from all surfaces
-            if ms is not None:
-                should_update = False
-                for block in self.agent_state.memory.to_dict()["memory"].values():
-                    if not block.get("template", False):
-                        should_update = True
-                if should_update:
-                    # TODO: the force=True can be optimized away
-                    # once we ensure we're correctly comparing whether in-memory core
-                    # data is different than persisted core data.
-                    self.rebuild_memory(force=True, ms=ms)
+            current_persisted_memory = Memory(
+                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()]
+            )  # read blocks from DB
+            self.update_memory_if_change(current_persisted_memory)
 
             # Step 1: add user message
             if isinstance(messages, Message):
@@ -1208,43 +1250,10 @@ class Agent(BaseAgent):
         new_messages = [new_system_message_obj] + self._messages[1:]  # swap index 0 (system)
         self._messages = new_messages
 
-    def update_memory_blocks_from_db(self):
-        for block in self.memory.to_dict()["memory"].values():
-            if block.get("templates", False):
-                # we don't expect to update shared memory blocks that
-                # are templates. this is something we could update in the
-                # future if we expect templates to change often.
-                continue
-            block_id = block.get("id")
-
-            # TODO: This is really hacky and we should probably figure out how to
-            db_block = BlockManager().get_block_by_id(block_id=block_id, actor=self.user)
-            if db_block is None:
-                # this case covers if someone has deleted a shared block by interacting
-                # with some other agent.
-                # in that case we should remove this shared block from the agent currently being
-                # evaluated.
-                printd(f"removing block: {block_id=}")
-                continue
-            if not isinstance(db_block.value, str):
-                printd(f"skipping block update, unexpected value: {block_id=}")
-                continue
-            # TODO: we may want to update which columns we're updating from shared memory e.g. the limit
-            self.memory.update_block_value(label=block.get("label", ""), value=db_block.value)
-
-    def rebuild_memory(self, force=False, update_timestamp=True, ms: Optional[MetadataStore] = None):
+    def rebuild_system_prompt(self, force=False, update_timestamp=True):
         """Rebuilds the system message with the latest memory object and any shared memory block updates"""
         curr_system_message = self.messages[0]  # this is the system + memory bank, not just the system prompt
 
-        # NOTE: This is a hacky way to check if the memory has changed
-        memory_repr = self.memory.compile()
-        if not force and memory_repr == curr_system_message["content"][-(len(memory_repr)) :]:
-            printd(f"Memory has not changed, not rebuilding system")
-            return
-
-        if ms:
-            self.update_memory_blocks_from_db()
-
         # 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:
@@ -1255,8 +1264,8 @@ class Agent(BaseAgent):
 
         # update memory (TODO: potentially update recall/archival stats seperately)
         new_system_message_str = compile_system_message(
-            system_prompt=self.system,
-            in_context_memory=self.memory,
+            system_prompt=self.agent_state.system,
+            in_context_memory=self.agent_state.memory,
             in_context_memory_last_edit=memory_edit_timestamp,
             archival_memory=self.persistence_manager.archival_memory,
             recall_memory=self.persistence_manager.recall_memory,
@@ -1283,14 +1292,13 @@ class Agent(BaseAgent):
         """Update the system prompt of the agent (requires rebuilding the memory block if there's a difference)"""
         assert isinstance(new_system_prompt, str)
 
-        if new_system_prompt == self.system:
-            input("same???")
+        if new_system_prompt == self.agent_state.system:
             return
 
-        self.system = new_system_prompt
+        self.agent_state.system = new_system_prompt
 
         # updating the system prompt requires rebuilding the memory block inside the compiled system message
-        self.rebuild_memory(force=True, update_timestamp=False)
+        self.rebuild_system_prompt(force=True, update_timestamp=False)
 
         # make sure to persist the change
         _ = self.update_state()
@@ -1304,13 +1312,16 @@ class Agent(BaseAgent):
         raise NotImplementedError
 
     def update_state(self) -> AgentState:
+        # TODO: this should be removed and self._messages should be moved into self.agent_state.in_context_messages
         message_ids = [msg.id for msg in self._messages]
-        assert isinstance(self.memory, Memory), f"Memory is not a Memory object: {type(self.memory)}"
+
+        # Assert that these are all strings
+        if any(not isinstance(m_id, str) for m_id in message_ids):
+            warnings.warn(f"Non-string message IDs found in agent state: {message_ids}")
+            message_ids = [m_id for m_id in message_ids if isinstance(m_id, str)]
 
         # override any fields that may have been updated
         self.agent_state.message_ids = message_ids
-        self.agent_state.memory = self.memory
-        self.agent_state.system = self.system
 
         return self.agent_state
 
@@ -1511,7 +1522,7 @@ class Agent(BaseAgent):
 
         system_prompt = self.agent_state.system  # TODO is this the current system or the initial system?
         num_tokens_system = count_tokens(system_prompt)
-        core_memory = self.memory.compile()
+        core_memory = self.agent_state.memory.compile()
         num_tokens_core_memory = count_tokens(core_memory)
 
         # conversion of messages to OpenAI dict format, which is passed to the token counter
@@ -1603,37 +1614,15 @@ def save_agent(agent: Agent, ms: MetadataStore):
 
     agent.update_state()
     agent_state = agent.agent_state
-    agent_id = agent_state.id
     assert isinstance(agent_state.memory, Memory), f"Memory is not a Memory object: {type(agent_state.memory)}"
 
-    # NOTE: we're saving agent memory before persisting the agent to ensure
-    # that allocated block_ids for each memory block are present in the agent model
-    save_agent_memory(agent=agent)
-
-    if ms.get_agent(agent_id=agent.agent_state.id):
-        ms.update_agent(agent_state)
+    # TODO: move this to agent manager
+    # convert to persisted model
+    persisted_agent_state = agent.agent_state.to_persisted_agent_state()
+    if ms.get_agent(agent_id=persisted_agent_state.id):
+        ms.update_agent(persisted_agent_state)
     else:
-        ms.create_agent(agent_state)
-
-    agent.agent_state = ms.get_agent(agent_id=agent_id)
-    assert isinstance(agent.agent_state.memory, Memory), f"Memory is not a Memory object: {type(agent_state.memory)}"
-
-
-def save_agent_memory(agent: Agent):
-    """
-    Save agent memory to metadata store. Memory is a collection of blocks and each block is persisted to the block table.
-
-    NOTE: we are assuming agent.update_state has already been called.
-    """
-
-    for block_dict in agent.memory.to_dict()["memory"].values():
-        # TODO: block creation should happen in one place to enforce these sort of constraints consistently.
-        block = Block(**block_dict)
-        # FIXME: should we expect for block values to be None? If not, we need to figure out why that is
-        # the case in some tests, if so we should relax the DB constraint.
-        if block.value is None:
-            block.value = ""
-        BlockManager().create_or_update_block(block, actor=agent.user)
+        ms.create_agent(persisted_agent_state)
 
 
 def strip_name_field_from_user_message(user_message_text: str) -> Tuple[str, Optional[str]]:

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/letta/agent_store/chroma.py

@@ -125,6 +125,8 @@ class ChromaStorageConnector(StorageConnector):
         ids, filters = self.get_filters(filters)
         if self.collection.count() == 0:
             return []
+        if ids == []:
+            ids = None
         if limit:
             results = self.collection.get(ids=ids, include=self.include, where=filters, limit=limit)
         else:

{letta_nightly-0.5.5.dev20241122170833 → letta_nightly-0.6.0.dev20241204054116}/letta/agent_store/db.py

@@ -433,7 +433,7 @@ class PostgresStorageConnector(SQLStorageConnector):
                 else:
                     db_record = self.db_model(**record.dict())
                     session.add(db_record)
-                    print(f"Added record with id {record.id}")
+                    # print(f"Added record with id {record.id}")
                     session.commit()
 
                 added_ids.append(record.id)