letta-nightly 0.6.34.dev20250302104001__py3-none-any.whl → 0.6.34.dev20250303230404__py3-none-any.whl

letta/services/agent_manager.py

@@ -50,6 +50,7 @@ from letta.services.message_manager import MessageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
 from letta.settings import settings
+from letta.tracing import trace_method
 from letta.utils import enforce_types, united_diff
 
 logger = get_logger(__name__)
@@ -72,6 +73,7 @@ class AgentManager:
     # ======================================================================================================================
     # Basic CRUD operations
     # ======================================================================================================================
+    @trace_method
     @enforce_types
     def create_agent(
         self,
@@ -368,6 +370,7 @@ class AgentManager:
             agent = AgentModel.read(db_session=session, name=agent_name, actor=actor)
             return agent.to_pydantic()
 
+    @trace_method
     @enforce_types
     def delete_agent(self, agent_id: str, actor: PydanticUser) -> None:
         """
@@ -529,42 +532,43 @@ class AgentManager:
                 model=agent_state.llm_config.model,
                 openai_message_dict={"role": "system", "content": new_system_message_str},
             )
+            # TODO: This seems kind of silly, why not just update the message?
             message = self.message_manager.create_message(message, actor=actor)
             message_ids = [message.id] + agent_state.message_ids[1:]  # swap index 0 (system)
-            return self._set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+            return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
         else:
             return agent_state
 
     @enforce_types
-    def _set_in_context_messages(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState:
+    def set_in_context_messages(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState:
         return self.update_agent(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor)
 
     @enforce_types
     def trim_older_in_context_messages(self, num: int, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
         new_messages = [message_ids[0]] + message_ids[num:]  # 0 is system message
-        return self._set_in_context_messages(agent_id=agent_id, message_ids=new_messages, actor=actor)
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=new_messages, actor=actor)
 
     @enforce_types
     def trim_all_in_context_messages_except_system(self, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
         # TODO: How do we know this?
         new_messages = [message_ids[0]]  # 0 is system message
-        return self._set_in_context_messages(agent_id=agent_id, message_ids=new_messages, actor=actor)
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=new_messages, actor=actor)
 
     @enforce_types
     def prepend_to_in_context_messages(self, messages: List[PydanticMessage], agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
         new_messages = self.message_manager.create_many_messages(messages, actor=actor)
         message_ids = [message_ids[0]] + [m.id for m in new_messages] + message_ids[1:]
-        return self._set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
 
     @enforce_types
     def append_to_in_context_messages(self, messages: List[PydanticMessage], agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         messages = self.message_manager.create_many_messages(messages, actor=actor)
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids or []
         message_ids += [m.id for m in messages]
-        return self._set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
+        return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor)
 
     @enforce_types
     def reset_messages(self, agent_id: str, actor: PydanticUser, add_default_initial_messages: bool = False) -> PydanticAgentState:

letta/services/helpers/agent_manager_helper.py

@@ -91,8 +91,6 @@ def derive_system_message(agent_type: AgentType, system: Optional[str] = None):
             system = gpt_system.get_system_text("memgpt_chat")
         elif agent_type == AgentType.offline_memory_agent:
             system = gpt_system.get_system_text("memgpt_offline_memory")
-        elif agent_type == AgentType.chat_only_agent:
-            system = gpt_system.get_system_text("memgpt_convo_only")
         else:
             raise ValueError(f"Invalid agent type: {agent_type}")
 

letta/services/helpers/tool_execution_helper.py

@@ -4,6 +4,10 @@ import subprocess
 import venv
 from typing import Dict, Optional
 
+from datamodel_code_generator import DataModelType, PythonVersion
+from datamodel_code_generator.model import get_data_model_types
+from datamodel_code_generator.parser.jsonschema import JsonSchemaParser
+
 from letta.log import get_logger
 from letta.schemas.sandbox_config import LocalSandboxConfig
 
@@ -153,3 +157,17 @@ def create_venv_for_local_sandbox(sandbox_dir_path: str, venv_path: str, env: Di
     except subprocess.CalledProcessError as e:
         logger.error(f"Error while setting up the virtual environment: {e}")
         raise RuntimeError(f"Failed to set up the virtual environment: {e}")
+
+
+def add_imports_and_pydantic_schemas_for_args(args_json_schema: dict) -> str:
+    data_model_types = get_data_model_types(DataModelType.PydanticV2BaseModel, target_python_version=PythonVersion.PY_311)
+    parser = JsonSchemaParser(
+        str(args_json_schema),
+        data_model_type=data_model_types.data_model,
+        data_model_root_type=data_model_types.root_model,
+        data_model_field_type=data_model_types.field_model,
+        data_type_manager_type=data_model_types.data_type_manager,
+        dump_resolve_reference_action=data_model_types.dump_resolve_reference_action,
+    )
+    result = parser.parse()
+    return result

letta/services/job_manager.py

@@ -13,12 +13,14 @@ from letta.orm.job_messages import JobMessage
 from letta.orm.message import Message as MessageModel
 from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step
+from letta.orm.step import Step as StepModel
 from letta.schemas.enums import JobStatus, MessageRole
 from letta.schemas.job import Job as PydanticJob
 from letta.schemas.job import JobUpdate, LettaRequestConfig
 from letta.schemas.letta_message import LettaMessage
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.run import Run as PydanticRun
+from letta.schemas.step import Step as PydanticStep
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User as PydanticUser
 from letta.utils import enforce_types
@@ -161,6 +163,51 @@ class JobManager:
 
         return [message.to_pydantic() for message in messages]
 
+    @enforce_types
+    def get_job_steps(
+        self,
+        job_id: str,
+        actor: PydanticUser,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 100,
+        ascending: bool = True,
+    ) -> List[PydanticStep]:
+        """
+        Get all steps associated with a job.
+
+        Args:
+            job_id: The ID of the job to get steps for
+            actor: The user making the request
+            before: Cursor for pagination
+            after: Cursor for pagination
+            limit: Maximum number of steps to return
+            ascending: Optional flag to sort in ascending order
+
+        Returns:
+            List of steps associated with the job
+
+        Raises:
+            NoResultFound: If the job does not exist or user does not have access
+        """
+        with self.session_maker() as session:
+            # Build filters
+            filters = {}
+            filters["job_id"] = job_id
+
+            # Get steps
+            steps = StepModel.list(
+                db_session=session,
+                before=before,
+                after=after,
+                ascending=ascending,
+                limit=limit,
+                actor=actor,
+                **filters,
+            )
+
+        return [step.to_pydantic() for step in steps]
+
     @enforce_types
     def add_message_to_job(self, job_id: str, message_id: str, actor: PydanticUser) -> None:
         """
@@ -312,6 +359,57 @@ class JobManager:
 
         return messages
 
+    @enforce_types
+    def get_step_messages(
+        self,
+        run_id: str,
+        actor: PydanticUser,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 100,
+        role: Optional[MessageRole] = None,
+        ascending: bool = True,
+    ) -> List[LettaMessage]:
+        """
+        Get steps associated with a job using cursor-based pagination.
+        This is a wrapper around get_job_messages that provides cursor-based pagination.
+
+        Args:
+            run_id: The ID of the run to get steps for
+            actor: The user making the request
+            before: Message ID to get messages after
+            after: Message ID to get messages before
+            limit: Maximum number of messages to return
+            ascending: Whether to return messages in ascending order
+            role: Optional role filter
+
+        Returns:
+            List of Steps associated with the job
+
+        Raises:
+            NoResultFound: If the job does not exist or user does not have access
+        """
+        messages = self.get_job_messages(
+            job_id=run_id,
+            actor=actor,
+            before=before,
+            after=after,
+            limit=limit,
+            role=role,
+            ascending=ascending,
+        )
+
+        request_config = self._get_run_request_config(run_id)
+
+        messages = PydanticMessage.to_letta_messages_from_list(
+            messages=messages,
+            use_assistant_message=request_config["use_assistant_message"],
+            assistant_message_tool_name=request_config["assistant_message_tool_name"],
+            assistant_message_tool_kwarg=request_config["assistant_message_tool_kwarg"],
+        )
+
+        return messages
+
     def _verify_job_access(
         self,
         session: Session,

letta/services/step_manager.py

@@ -11,6 +11,7 @@ from letta.orm.step import Step as StepModel
 from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.step import Step as PydanticStep
 from letta.schemas.user import User as PydanticUser
+from letta.tracing import get_trace_id
 from letta.utils import enforce_types
 
 
@@ -75,6 +76,7 @@ class StepManager:
             "job_id": job_id,
             "tags": [],
             "tid": None,
+            "trace_id": get_trace_id(),  # Get the current trace ID
         }
         with self.session_maker() as session:
             if job_id:

letta/services/summarizer/enums.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class SummarizationMode(str, Enum):
+    """
+    Represents possible modes of summarization for conversation trimming.
+    """
+
+    STATIC_MESSAGE_BUFFER = "static_message_buffer_mode"

letta/services/summarizer/summarizer.py

@@ -0,0 +1,102 @@
+import json
+from json import JSONDecodeError
+from typing import List, Tuple
+
+from letta.agents.base_agent import BaseAgent
+from letta.schemas.enums import MessageRole
+from letta.schemas.message import Message
+from letta.schemas.openai.chat_completion_request import UserMessage
+from letta.services.summarizer.enums import SummarizationMode
+
+
+class Summarizer:
+    """
+    Handles summarization or trimming of conversation messages based on
+    the specified SummarizationMode. For now, we demonstrate a simple
+    static buffer approach but leave room for more advanced strategies.
+    """
+
+    def __init__(self, mode: SummarizationMode, summarizer_agent: BaseAgent, message_buffer_limit: int = 10, message_buffer_min: int = 3):
+        self.mode = mode
+
+        # Need to do validation on this
+        self.message_buffer_limit = message_buffer_limit
+        self.message_buffer_min = message_buffer_min
+        self.summarizer_agent = summarizer_agent
+        # TODO: Move this to config
+        self.summary_prefix = "Out of context message summarization:\n"
+
+    async def summarize(
+        self, in_context_messages: List[Message], new_letta_messages: List[Message], previous_summary: str
+    ) -> Tuple[List[Message], str, bool]:
+        """
+        Summarizes or trims in_context_messages according to the chosen mode,
+        and returns the updated messages plus any optional "summary message".
+
+        Args:
+            in_context_messages: The existing messages in the conversation's context.
+            new_letta_messages: The newly added Letta messages (just appended).
+            previous_summary: The previous summary string.
+
+        Returns:
+            (updated_messages, summary_message)
+            updated_messages: The new context after trimming/summary
+            summary_message: Optional summarization message that was created
+                             (could be appended to the conversation if desired)
+        """
+        if self.mode == SummarizationMode.STATIC_MESSAGE_BUFFER:
+            return await self._static_buffer_summarization(in_context_messages, new_letta_messages, previous_summary)
+        else:
+            # Fallback or future logic
+            return in_context_messages, "", False
+
+    async def _static_buffer_summarization(
+        self, in_context_messages: List[Message], new_letta_messages: List[Message], previous_summary: str
+    ) -> Tuple[List[Message], str, bool]:
+        previous_summary = previous_summary[: len(self.summary_prefix)]
+        all_in_context_messages = in_context_messages + new_letta_messages
+
+        # Only summarize if we exceed `message_buffer_limit`
+        if len(all_in_context_messages) <= self.message_buffer_limit:
+            return all_in_context_messages, previous_summary, False
+
+        # Aim to trim down to `message_buffer_min`
+        target_trim_index = len(all_in_context_messages) - self.message_buffer_min + 1
+
+        # Move the trim index forward until it's at a `MessageRole.user`
+        while target_trim_index < len(all_in_context_messages) and all_in_context_messages[target_trim_index].role != MessageRole.user:
+            target_trim_index += 1
+
+        # TODO: Assuming system message is always at index 0
+        updated_in_context_messages = [all_in_context_messages[0]] + all_in_context_messages[target_trim_index:]
+        out_of_context_messages = all_in_context_messages[:target_trim_index]
+
+        formatted_messages = []
+        for m in out_of_context_messages:
+            if m.content:
+                try:
+                    message = json.loads(m.content[0].text).get("message")
+                except JSONDecodeError:
+                    continue
+                if message:
+                    formatted_messages.append(f"{m.role.value}: {message}")
+
+        # If we didn't trim any messages, return as-is
+        if not formatted_messages:
+            return all_in_context_messages, previous_summary, False
+
+        # Generate summarization request
+        summary_request_text = (
+            "These are messages that are soon to be removed from the context window:\n"
+            f"{formatted_messages}\n\n"
+            "This is the current memory:\n"
+            f"{previous_summary}\n\n"
+            "Your task is to integrate any relevant updates from the messages into the memory."
+            "It should be in note-taking format in natural English. You are to return the new, updated memory only."
+        )
+
+        messages = await self.summarizer_agent.step(UserMessage(content=summary_request_text))
+        current_summary = "\n".join([m.text for m in messages])
+        current_summary = f"{self.summary_prefix}{current_summary}"
+
+        return updated_in_context_messages, current_summary, True

letta/services/tool_execution_sandbox.py

@@ -11,12 +11,14 @@ import traceback
 import uuid
 from typing import Any, Dict, Optional
 
+from letta.functions.helpers import generate_model_from_args_json_schema
 from letta.log import get_logger
 from letta.schemas.agent import AgentState
 from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult, SandboxType
 from letta.schemas.tool import Tool
 from letta.schemas.user import User
 from letta.services.helpers.tool_execution_helper import (
+    add_imports_and_pydantic_schemas_for_args,
     create_venv_for_local_sandbox,
     find_python_executable,
     install_pip_requirements_for_sandbox,
@@ -408,20 +410,35 @@ class ToolExecutionSandbox:
         code += "import sys\n"
         code += "import base64\n"
 
-        # Load the agent state data into the program
+        # imports to support agent state
         if agent_state:
             code += "import letta\n"
             code += "from letta import * \n"
             import pickle
 
+        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 agent_state:
             agent_state_pickle = pickle.dumps(agent_state)
             code += f"agent_state = pickle.loads({agent_state_pickle})\n"
         else:
             # agent state is None
             code += "agent_state = None\n"
 
-        for param in self.args:
-            code += self.initialize_param(param, self.args[param])
+        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])
 
         if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
             inject_agent_state = True

letta/services/tool_manager.py

@@ -42,7 +42,7 @@ class ToolManager:
         tool = self.get_tool_by_name(tool_name=pydantic_tool.name, actor=actor)
         if tool:
             # Put to dict and remove fields that should not be reset
-            update_data = pydantic_tool.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
+            update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True)
 
             # If there's anything to update
             if update_data:

letta/settings.py

@@ -50,6 +50,8 @@ class ModelSettings(BaseSettings):
 
     model_config = SettingsConfigDict(env_file=".env", extra="ignore")
 
+    global_max_context_window_limit: int = 32000
+
     # env_prefix='my_prefix_'
 
     # when we use /completions APIs (instead of /chat/completions), we need to specify a model wrapper