letta-nightly 0.6.48.dev20250407104216__py3-none-any.whl → 0.6.49.dev20250408030511__py3-none-any.whl

letta/services/tool_manager.py

@@ -2,7 +2,15 @@ import importlib
 import warnings
 from typing import List, Optional
 
-from letta.constants import BASE_FUNCTION_RETURN_CHAR_LIMIT, BASE_MEMORY_TOOLS, BASE_TOOLS, MCP_TOOL_TAG_NAME_PREFIX, MULTI_AGENT_TOOLS
+from letta.constants import (
+    BASE_FUNCTION_RETURN_CHAR_LIMIT,
+    BASE_MEMORY_TOOLS,
+    BASE_SLEEPTIME_TOOLS,
+    BASE_TOOLS,
+    LETTA_TOOL_SET,
+    MCP_TOOL_TAG_NAME_PREFIX,
+    MULTI_AGENT_TOOLS,
+)
 from letta.functions.functions import derive_openai_json_schema, load_function_set
 from letta.log import get_logger
 from letta.orm.enums import ToolType
@@ -194,7 +202,7 @@ class ToolManager:
         # create tool in db
         tools = []
         for name, schema in functions_to_schema.items():
-            if name in BASE_TOOLS + BASE_MEMORY_TOOLS + MULTI_AGENT_TOOLS:
+            if name in LETTA_TOOL_SET:
                 if name in BASE_TOOLS:
                     tool_type = ToolType.LETTA_CORE
                     tags = [tool_type.value]
@@ -204,9 +212,12 @@ class ToolManager:
                 elif name in MULTI_AGENT_TOOLS:
                     tool_type = ToolType.LETTA_MULTI_AGENT_CORE
                     tags = [tool_type.value]
+                elif name in BASE_SLEEPTIME_TOOLS:
+                    tool_type = ToolType.LETTA_SLEEPTIME_CORE
+                    tags = [tool_type.value]
                 else:
                     raise ValueError(
-                        f"Tool name {name} is not in the list of base tool names: {BASE_TOOLS + BASE_MEMORY_TOOLS + MULTI_AGENT_TOOLS}"
+                        f"Tool name {name} is not in the list of base tool names: {BASE_TOOLS + BASE_MEMORY_TOOLS + MULTI_AGENT_TOOLS + BASE_SLEEPTIME_TOOLS}"
                     )
 
                 # create to tool

letta/services/tool_sandbox/base.py

@@ -0,0 +1,188 @@
+import ast
+import base64
+import pickle
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Tuple
+
+from letta.functions.helpers import generate_model_from_args_json_schema
+from letta.schemas.agent import AgentState
+from letta.schemas.sandbox_config import SandboxRunResult
+from letta.services.helpers.tool_execution_helper import add_imports_and_pydantic_schemas_for_args
+from letta.services.organization_manager import OrganizationManager
+from letta.services.sandbox_config_manager import SandboxConfigManager
+from letta.services.tool_manager import ToolManager
+
+
+class AsyncToolSandboxBase(ABC):
+    NAMESPACE = uuid.NAMESPACE_DNS
+    LOCAL_SANDBOX_RESULT_START_MARKER = str(uuid.uuid5(NAMESPACE, "local-sandbox-result-start-marker"))
+    LOCAL_SANDBOX_RESULT_END_MARKER = str(uuid.uuid5(NAMESPACE, "local-sandbox-result-end-marker"))
+    LOCAL_SANDBOX_RESULT_VAR_NAME = "result_ZQqiequkcFwRwwGQMqkt"
+
+    def __init__(self, tool_name: str, args: dict, user, tool_object=None):
+        self.tool_name = tool_name
+        self.args = args
+        self.user = user
+        self.organization = OrganizationManager().get_organization_by_id(self.user.organization_id)
+        self.privileged_tools = self.organization.privileged_tools
+
+        self.tool = tool_object or ToolManager().get_tool_by_name(tool_name=tool_name, actor=self.user)
+        if self.tool is None:
+            raise ValueError(
+                f"Agent attempted to invoke tool {self.tool_name} that does not exist for organization {self.user.organization_id}"
+            )
+
+        self.sandbox_config_manager = SandboxConfigManager()
+
+        # See if we should inject agent_state or not based on the presence of the "agent_state" arg
+        if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
+            self.inject_agent_state = True
+        else:
+            self.inject_agent_state = False
+
+    @abstractmethod
+    async def run(
+        self,
+        agent_state: Optional[AgentState] = None,
+        additional_env_vars: Optional[Dict] = None,
+    ) -> SandboxRunResult:
+        """
+        Run the tool in a sandbox environment asynchronously.
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def generate_execution_script(self, agent_state: Optional[AgentState], wrap_print_with_markers: bool = False) -> str:
+        """
+        Generate code to run inside of execution sandbox.
+        Serialize the agent state and arguments, call the tool,
+        then base64-encode/pickle the result.
+        """
+        code = "from typing import *\n"
+        code += "import pickle\n"
+        code += "import sys\n"
+        code += "import base64\n"
+
+        # Additional imports to support agent state
+        if self.inject_agent_state:
+            code += "import letta\n"
+            code += "from letta import * \n"
+
+        # Add schema code if available
+        if self.tool.args_json_schema:
+            schema_code = add_imports_and_pydantic_schemas_for_args(self.tool.args_json_schema)
+            if "from __future__ import annotations" in schema_code:
+                schema_code = schema_code.replace("from __future__ import annotations", "").lstrip()
+                code = "from __future__ import annotations\n\n" + code
+            code += schema_code + "\n"
+
+        # Load the agent state
+        if self.inject_agent_state:
+            agent_state_pickle = pickle.dumps(agent_state)
+            code += f"agent_state = pickle.loads({agent_state_pickle})\n"
+        else:
+            code += "agent_state = None\n"
+
+        # Initialize arguments
+        if self.tool.args_json_schema:
+            args_schema = generate_model_from_args_json_schema(self.tool.args_json_schema)
+            code += f"args_object = {args_schema.__name__}(**{self.args})\n"
+            for param in self.args:
+                code += f"{param} = args_object.{param}\n"
+        else:
+            for param in self.args:
+                code += self.initialize_param(param, self.args[param])
+
+        # Insert the tool's source code
+        code += "\n" + self.tool.source_code + "\n"
+
+        # Invoke the function and store the result in a global variable
+        code += (
+            f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME}" + ' = {"results": ' + self.invoke_function_call() + ', "agent_state": agent_state}\n'
+        )
+        code += (
+            f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME} = base64.b64encode("
+            f"pickle.dumps({self.LOCAL_SANDBOX_RESULT_VAR_NAME})"
+            ").decode('utf-8')\n"
+        )
+
+        if wrap_print_with_markers:
+            code += f"sys.stdout.write('{self.LOCAL_SANDBOX_RESULT_START_MARKER}')\n"
+            code += f"sys.stdout.write(str({self.LOCAL_SANDBOX_RESULT_VAR_NAME}))\n"
+            code += f"sys.stdout.write('{self.LOCAL_SANDBOX_RESULT_END_MARKER}')\n"
+        else:
+            code += f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME}\n"
+
+        return code
+
+    def _convert_param_to_value(self, param_type: str, raw_value: str) -> str:
+        """
+        Convert parameter to Python code representation based on JSON schema type.
+        """
+        if param_type == "string":
+            # Safely inject a Python string via pickle
+            value = "pickle.loads(" + str(pickle.dumps(raw_value)) + ")"
+        elif param_type in ["integer", "boolean", "number", "array", "object"]:
+            # This is simplistic. In real usage, ensure correct type-casting or sanitization.
+            value = raw_value
+        else:
+            raise TypeError(f"Unsupported type: {param_type}, raw_value={raw_value}")
+
+        return str(value)
+
+    def initialize_param(self, name: str, raw_value: str) -> str:
+        """
+        Produce code for initializing a single parameter in the generated script.
+        """
+        params = self.tool.json_schema["parameters"]["properties"]
+        spec = params.get(name)
+        if spec is None:
+            # Possibly an extra param like 'self' that we ignore
+            return ""
+
+        param_type = spec.get("type")
+        if param_type is None and spec.get("parameters"):
+            param_type = spec["parameters"].get("type")
+
+        value = self._convert_param_to_value(param_type, raw_value)
+        return f"{name} = {value}\n"
+
+    def invoke_function_call(self) -> str:
+        """
+        Generate the function call code string with the appropriate arguments.
+        """
+        kwargs = []
+        for name in self.args:
+            if name in self.tool.json_schema["parameters"]["properties"]:
+                kwargs.append(name)
+
+        param_list = [f"{arg}={arg}" for arg in kwargs]
+        if self.inject_agent_state:
+            param_list.append("agent_state=agent_state")
+
+        params = ", ".join(param_list)
+        func_call_str = self.tool.name + "(" + params + ")"
+        return func_call_str
+
+    def parse_best_effort(self, text: str) -> Tuple[Any, Optional[AgentState]]:
+        """
+        Decode and unpickle the result from the function execution if possible.
+        Returns (function_return_value, agent_state).
+        """
+        if not text:
+            return None, None
+
+        result = pickle.loads(base64.b64decode(text))
+        agent_state = result["agent_state"]
+        return result["results"], agent_state
+
+    def parse_function_arguments(self, source_code: str, tool_name: str):
+        """Get arguments of a function from its source code"""
+        tree = ast.parse(source_code)
+        args = []
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef) and node.name == tool_name:
+                for arg in node.args.args:
+                    args.append(arg.arg)
+        return args

letta/services/tool_sandbox/e2b_sandbox.py

@@ -0,0 +1,116 @@
+from typing import Dict, Optional
+
+from letta.log import get_logger
+from letta.schemas.agent import AgentState
+from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult, SandboxType
+from letta.services.tool_sandbox.base import AsyncToolSandboxBase
+from letta.utils import get_friendly_error_msg
+
+logger = get_logger(__name__)
+
+
+class AsyncToolSandboxE2B(AsyncToolSandboxBase):
+    METADATA_CONFIG_STATE_KEY = "config_state"
+
+    def __init__(self, tool_name: str, args: dict, user, force_recreate=True, tool_object=None):
+        super().__init__(tool_name, args, user, tool_object)
+        self.force_recreate = force_recreate
+
+    async def run(
+        self,
+        agent_state: Optional[AgentState] = None,
+        additional_env_vars: Optional[Dict] = None,
+    ) -> SandboxRunResult:
+        """
+        Run the tool in a sandbox environment asynchronously,
+        *always* using a subprocess for execution.
+        """
+        result = await self.run_e2b_sandbox(agent_state=agent_state, additional_env_vars=additional_env_vars)
+
+        # Simple console logging for demonstration
+        for log_line in (result.stdout or []) + (result.stderr or []):
+            print(f"Tool execution log: {log_line}")
+
+        return result
+
+    async def run_e2b_sandbox(
+        self, agent_state: Optional[AgentState] = None, additional_env_vars: Optional[Dict] = None
+    ) -> SandboxRunResult:
+        sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=self.user)
+        # TODO: So this defaults to force recreating always
+        # TODO: Eventually, provision one sandbox PER agent, and that agent re-uses that one specifically
+        e2b_sandbox = await self.create_e2b_sandbox_with_metadata_hash(sandbox_config=sbx_config)
+
+        logger.info(f"E2B Sandbox configurations: {sbx_config}")
+        logger.info(f"E2B Sandbox ID: {e2b_sandbox.sandbox_id}")
+
+        # TODO: This only makes sense if we re-use sandboxes
+        # # Since this sandbox was used, we extend its lifecycle by the timeout
+        # await sbx.set_timeout(sbx_config.get_e2b_config().timeout)
+
+        # Get environment variables for the sandbox
+        # TODO: We set limit to 100 here, but maybe we want it uncapped? Realistically this should be fine.
+        env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100)
+        # Get environment variables for this agent specifically
+        if agent_state:
+            env_vars.update(agent_state.get_agent_env_vars_as_dict())
+
+        # Finally, get any that are passed explicitly into the `run` function call
+        if additional_env_vars:
+            env_vars.update(additional_env_vars)
+        code = self.generate_execution_script(agent_state=agent_state)
+
+        execution = await e2b_sandbox.run_code(code, envs=env_vars)
+
+        if execution.results:
+            func_return, agent_state = self.parse_best_effort(execution.results[0].text)
+        elif execution.error:
+            logger.error(f"Executing tool {self.tool_name} raised a {execution.error.name} with message: \n{execution.error.value}")
+            logger.error(f"Traceback from e2b sandbox: \n{execution.error.traceback}")
+            func_return = get_friendly_error_msg(
+                function_name=self.tool_name, exception_name=execution.error.name, exception_message=execution.error.value
+            )
+            execution.logs.stderr.append(execution.error.traceback)
+        else:
+            raise ValueError(f"Tool {self.tool_name} returned execution with None")
+
+        return SandboxRunResult(
+            func_return=func_return,
+            agent_state=agent_state,
+            stdout=execution.logs.stdout,
+            stderr=execution.logs.stderr,
+            status="error" if execution.error else "success",
+            sandbox_config_fingerprint=sbx_config.fingerprint(),
+        )
+
+    def parse_exception_from_e2b_execution(self, e2b_execution: "Execution") -> Exception:
+        builtins_dict = __builtins__ if isinstance(__builtins__, dict) else vars(__builtins__)
+        # Dynamically fetch the exception class from builtins, defaulting to Exception if not found
+        exception_class = builtins_dict.get(e2b_execution.error.name, Exception)
+        return exception_class(e2b_execution.error.value)
+
+    async def create_e2b_sandbox_with_metadata_hash(self, sandbox_config: SandboxConfig) -> "Sandbox":
+        from e2b_code_interpreter import AsyncSandbox
+
+        state_hash = sandbox_config.fingerprint()
+        e2b_config = sandbox_config.get_e2b_config()
+
+        if e2b_config.template:
+            sbx = await AsyncSandbox.create(sandbox_config.get_e2b_config().template, metadata={self.METADATA_CONFIG_STATE_KEY: state_hash})
+        else:
+            # no template
+            sbx = await AsyncSandbox.create(
+                metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(exclude={"pip_requirements"})
+            )
+
+        # install pip requirements
+        if e2b_config.pip_requirements:
+            for package in e2b_config.pip_requirements:
+                await sbx.commands.run(f"pip install {package}")
+        return sbx
+
+    async def list_running_e2b_sandboxes(self):
+        from e2b_code_interpreter import AsyncSandbox
+
+        # List running sandboxes and access metadata.
+        return await AsyncSandbox.list()

letta/services/tool_sandbox/local_sandbox.py

@@ -0,0 +1,221 @@
+import asyncio
+import os
+import sys
+import tempfile
+from typing import Dict, Optional, Tuple
+
+from letta.schemas.agent import AgentState
+from letta.schemas.sandbox_config import SandboxRunResult, SandboxType
+from letta.services.helpers.tool_execution_helper import (
+    create_venv_for_local_sandbox,
+    find_python_executable,
+    install_pip_requirements_for_sandbox,
+)
+from letta.services.tool_sandbox.base import AsyncToolSandboxBase
+from letta.tracing import log_event, trace_method
+from letta.utils import get_friendly_error_msg
+
+
+class AsyncToolSandboxLocal(AsyncToolSandboxBase):
+    METADATA_CONFIG_STATE_KEY = "config_state"
+    REQUIREMENT_TXT_NAME = "requirements.txt"
+
+    def __init__(self, tool_name: str, args: dict, user, force_recreate_venv=False, tool_object=None):
+        super().__init__(tool_name, args, user, tool_object)
+        self.force_recreate_venv = force_recreate_venv
+
+    async def run(
+        self,
+        agent_state: Optional[AgentState] = None,
+        additional_env_vars: Optional[Dict] = None,
+    ) -> SandboxRunResult:
+        """
+        Run the tool in a sandbox environment asynchronously,
+        *always* using a subprocess for execution.
+        """
+        result = await self.run_local_dir_sandbox(agent_state=agent_state, additional_env_vars=additional_env_vars)
+
+        # Simple console logging for demonstration
+        for log_line in (result.stdout or []) + (result.stderr or []):
+            print(f"Tool execution log: {log_line}")
+
+        return result
+
+    @trace_method
+    async def run_local_dir_sandbox(self, agent_state: Optional[AgentState], additional_env_vars: Optional[Dict]) -> SandboxRunResult:
+        """
+        Unified asynchronougit pus method to run the tool in a local sandbox environment,
+        always via subprocess for multi-core parallelism.
+        """
+        # Get sandbox configuration
+        sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.LOCAL, actor=self.user)
+        local_configs = sbx_config.get_local_config()
+        use_venv = local_configs.use_venv
+
+        # Prepare environment variables
+        env = os.environ.copy()
+        env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100)
+        env.update(env_vars)
+
+        if agent_state:
+            env.update(agent_state.get_agent_env_vars_as_dict())
+
+        if additional_env_vars:
+            env.update(additional_env_vars)
+
+        # Make sure sandbox directory exists
+        sandbox_dir = os.path.expanduser(local_configs.sandbox_dir)
+        if not os.path.exists(sandbox_dir) or not os.path.isdir(sandbox_dir):
+            os.makedirs(sandbox_dir)
+
+        # If using a virtual environment, ensure it's prepared in parallel
+        venv_preparation_task = None
+        if use_venv:
+            venv_path = str(os.path.join(sandbox_dir, local_configs.venv_name))
+            venv_preparation_task = asyncio.create_task(self._prepare_venv(local_configs, venv_path, env))
+
+        # Generate and write execution script (always with markers, since we rely on stdout)
+        with tempfile.NamedTemporaryFile(mode="w", dir=sandbox_dir, suffix=".py", delete=False) as temp_file:
+            code = self.generate_execution_script(agent_state=agent_state, wrap_print_with_markers=True)
+            temp_file.write(code)
+            temp_file.flush()
+            temp_file_path = temp_file.name
+
+        try:
+            # If we started a venv preparation task, wait for it to complete
+            if venv_preparation_task:
+                await venv_preparation_task
+
+            # Determine the python executable and environment for the subprocess
+            exec_env = env.copy()
+            if use_venv:
+                venv_path = str(os.path.join(sandbox_dir, local_configs.venv_name))
+                python_executable = find_python_executable(local_configs)
+                exec_env["VIRTUAL_ENV"] = venv_path
+                exec_env["PATH"] = os.path.join(venv_path, "bin") + ":" + exec_env["PATH"]
+            else:
+                # If not using venv, use whatever Python we are running on
+                python_executable = sys.executable
+
+            exec_env["PYTHONWARNINGS"] = "ignore"
+
+            # Execute in subprocess
+            return await self._execute_tool_subprocess(
+                sbx_config=sbx_config,
+                python_executable=python_executable,
+                temp_file_path=temp_file_path,
+                env=exec_env,
+                cwd=sandbox_dir,
+            )
+
+        except Exception as e:
+            print(f"Executing tool {self.tool_name} has an unexpected error: {e}")
+            print(f"Auto-generated code for debugging:\n\n{code}")
+            raise e
+        finally:
+            # Clean up the temp file
+            os.remove(temp_file_path)
+
+    async def _prepare_venv(self, local_configs, venv_path: str, env: Dict[str, str]):
+        """
+        Prepare virtual environment asynchronously (in a background thread).
+        """
+        if self.force_recreate_venv or not os.path.isdir(venv_path):
+            sandbox_dir = os.path.expanduser(local_configs.sandbox_dir)
+            log_event(name="start create_venv_for_local_sandbox", attributes={"venv_path": venv_path})
+            await asyncio.to_thread(
+                create_venv_for_local_sandbox,
+                sandbox_dir_path=sandbox_dir,
+                venv_path=venv_path,
+                env=env,
+                force_recreate=self.force_recreate_venv,
+            )
+            log_event(name="finish create_venv_for_local_sandbox")
+
+        log_event(name="start install_pip_requirements_for_sandbox", attributes={"local_configs": local_configs.model_dump_json()})
+        await asyncio.to_thread(install_pip_requirements_for_sandbox, local_configs, upgrade=True, user_install_if_no_venv=False, env=env)
+        log_event(name="finish install_pip_requirements_for_sandbox", attributes={"local_configs": local_configs.model_dump_json()})
+
+    @trace_method
+    async def _execute_tool_subprocess(
+        self, sbx_config, python_executable: str, temp_file_path: str, env: Dict[str, str], cwd: str
+    ) -> SandboxRunResult:
+        """
+        Execute user code in a subprocess, always capturing stdout and stderr.
+        We parse special markers to extract the pickled result string.
+        """
+        try:
+            log_event(name="start subprocess")
+
+            process = await asyncio.create_subprocess_exec(
+                python_executable, temp_file_path, env=env, cwd=cwd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+
+            try:
+                stdout_bytes, stderr_bytes = await asyncio.wait_for(process.communicate(), timeout=60)
+            except asyncio.TimeoutError:
+                # Terminate the process on timeout
+                if process.returncode is None:
+                    process.terminate()
+                    try:
+                        await asyncio.wait_for(process.wait(), timeout=5)
+                    except asyncio.TimeoutError:
+                        process.kill()
+
+                raise TimeoutError(f"Executing tool {self.tool_name} timed out after 60 seconds.")
+
+            stdout = stdout_bytes.decode("utf-8") if stdout_bytes else ""
+            stderr = stderr_bytes.decode("utf-8") if stderr_bytes else ""
+            log_event(name="finish subprocess")
+
+            # Parse markers to isolate the function result
+            func_result, stdout_text = self.parse_out_function_results_markers(stdout)
+            func_return, agent_state = self.parse_best_effort(func_result)
+
+            return SandboxRunResult(
+                func_return=func_return,
+                agent_state=agent_state,
+                stdout=[stdout_text] if stdout_text else [],
+                stderr=[stderr] if stderr else [],
+                status="success" if process.returncode == 0 else "error",
+                sandbox_config_fingerprint=sbx_config.fingerprint(),
+            )
+
+        except (TimeoutError, Exception) as e:
+            # Distinguish between timeouts and other exceptions for clarity
+            if isinstance(e, TimeoutError):
+                raise e
+
+            print(f"Subprocess execution for tool {self.tool_name} encountered an error: {e}")
+            func_return = get_friendly_error_msg(
+                function_name=self.tool_name,
+                exception_name=type(e).__name__,
+                exception_message=str(e),
+            )
+            return SandboxRunResult(
+                func_return=func_return,
+                agent_state=None,
+                stdout=[],
+                stderr=[str(e)],
+                status="error",
+                sandbox_config_fingerprint=sbx_config.fingerprint(),
+            )
+
+    def parse_out_function_results_markers(self, text: str) -> Tuple[str, str]:
+        """
+        Parse the function results out of the stdout using special markers.
+        Returns (function_result_str, stripped_stdout).
+        """
+        if self.LOCAL_SANDBOX_RESULT_START_MARKER not in text:
+            # No markers found, so nothing to parse
+            return "", text
+
+        marker_len = len(self.LOCAL_SANDBOX_RESULT_START_MARKER)
+        start_index = text.index(self.LOCAL_SANDBOX_RESULT_START_MARKER) + marker_len
+        end_index = text.index(self.LOCAL_SANDBOX_RESULT_END_MARKER)
+
+        # The actual pickled base64 is between start_index and end_index
+        results_str = text[start_index:end_index]
+        # The rest of stdout (minus the markers)
+        remainder = text[: start_index - marker_len] + text[end_index + marker_len :]
+        return results_str, remainder

letta/sleeptime_agent.py

@@ -0,0 +1,61 @@
+from typing import List, Optional, Union
+
+from letta.agent import Agent, AgentState, save_agent
+from letta.interface import AgentInterface
+from letta.orm import User
+from letta.schemas.message import Message
+from letta.schemas.openai.chat_completion_response import UsageStatistics
+from letta.schemas.usage import LettaUsageStatistics
+
+
+class SleeptimeAgent(Agent):
+    def __init__(
+        self,
+        interface: AgentInterface,
+        agent_state: AgentState,
+        user: User = None,
+        # extras
+        first_message_verify_mono: bool = False,
+        max_memory_rethinks: int = 10,
+    ):
+        super().__init__(interface, agent_state, user)
+        self.first_message_verify_mono = first_message_verify_mono
+        self.max_memory_rethinks = max_memory_rethinks
+
+    def step(
+        self,
+        messages: Union[Message, List[Message]],
+        chaining: bool = True,
+        max_chaining_steps: Optional[int] = None,
+        **kwargs,
+    ) -> LettaUsageStatistics:
+        """Go through what is currently in memory core memory and integrate information."""
+        next_input_message = messages if isinstance(messages, list) else [messages]
+        counter = 0
+        total_usage = UsageStatistics()
+        step_count = 0
+
+        while counter < self.max_memory_rethinks:
+            if counter > 0:
+                next_input_message = []
+            kwargs["first_message"] = False
+            step_response = self.inner_step(
+                messages=next_input_message,
+                **kwargs,
+            )
+            for message in step_response.messages:
+                if message.tool_calls:
+                    for tool_call in message.tool_calls:
+                        # check if the function name is "finish_rethinking_memory"
+                        if tool_call.function.name == "finish_rethinking_memory":
+                            counter = self.max_memory_rethinks
+                            break
+            usage = step_response.usage
+            step_count += 1
+            total_usage += usage
+            counter += 1
+            self.interface.step_complete()
+
+            save_agent(self)
+
+        return LettaUsageStatistics(**total_usage.model_dump(), step_count=step_count)