letta-nightly 0.6.5.dev20241219104153__py3-none-any.whl → 0.6.6.dev20241220190343__py3-none-any.whl

letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.5"
+__version__ = "0.6.6"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client

letta/agent.py

@@ -1,5 +1,6 @@
 import datetime
 import inspect
+import json
 import time
 import traceback
 import warnings
@@ -371,6 +372,9 @@ class Agent(BaseAgent):
             self._append_to_messages(added_messages=init_messages_objs)
             self._validate_message_buffer_is_utc()
 
+        # Load last function response from message history
+        self.last_function_response = self.load_last_function_response()
+
         # Keep track of the total number of messages throughout all time
         self.messages_total = messages_total if messages_total is not None else (len(self._messages) - 1)  # (-system)
         self.messages_total_init = len(self._messages) - 1
@@ -389,6 +393,19 @@ class Agent(BaseAgent):
         else:
             self.supports_structured_output = True
 
+    def load_last_function_response(self):
+        """Load the last function response from message history"""
+        for i in range(len(self._messages) - 1, -1, -1):
+            msg = self._messages[i]
+            if msg.role == MessageRole.tool and msg.text:
+                try:
+                    response_json = json.loads(msg.text)
+                    if response_json.get("message"):
+                        return response_json["message"]
+                except (json.JSONDecodeError, KeyError):
+                    raise ValueError(f"Invalid JSON format in message: {msg.text}")
+        return None
+
     def update_memory_if_change(self, new_memory: Memory) -> bool:
         """
         Update internal memory object and system prompt if there have been modifications.
@@ -586,7 +603,7 @@ class Agent(BaseAgent):
     ) -> ChatCompletionResponse:
         """Get response from LLM API with robust retry mechanism."""
 
-        allowed_tool_names = self.tool_rules_solver.get_allowed_tool_names()
+        allowed_tool_names = self.tool_rules_solver.get_allowed_tool_names(last_function_response=self.last_function_response)
         agent_state_tool_jsons = [t.json_schema for t in self.agent_state.tools]
 
         allowed_functions = (
@@ -826,6 +843,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)
+                self.last_function_response = function_response
                 # TODO: truncate error message somehow
                 messages.append(
                     Message.dict_to_message(
@@ -861,6 +879,7 @@ class Agent(BaseAgent):
             )  # extend conversation with function response
             self.interface.function_message(f"Ran {function_name}({function_args})", msg_obj=messages[-1])
             self.interface.function_message(f"Success: {function_response_string}", msg_obj=messages[-1])
+            self.last_function_response = function_response
 
         else:
             # Standard non-function reply

letta/client/client.py

@@ -2909,7 +2909,7 @@ class LocalClient(AbstractClient):
         job = self.server.job_manager.create_job(pydantic_job=job, actor=self.user)
 
         # TODO: implement blocking vs. non-blocking
-        self.server.load_file_to_source(source_id=source_id, file_path=filename, job_id=job.id)
+        self.server.load_file_to_source(source_id=source_id, file_path=filename, job_id=job.id, actor=self.user)
         return job
 
     def delete_file_from_source(self, source_id: str, file_id: str):

letta/client/streaming.py

@@ -8,9 +8,9 @@ from letta.constants import OPENAI_CONTEXT_WINDOW_ERROR_SUBSTRING
 from letta.errors import LLMError
 from letta.schemas.enums import MessageStreamStatus
 from letta.schemas.letta_message import (
-    FunctionCallMessage,
-    FunctionReturn,
-    InternalMonologue,
+    ToolCallMessage,
+    ToolReturnMessage,
+    ReasoningMessage,
 )
 from letta.schemas.letta_response import LettaStreamingResponse
 from letta.schemas.usage import LettaUsageStatistics
@@ -53,12 +53,12 @@ def _sse_post(url: str, data: dict, headers: dict) -> Generator[LettaStreamingRe
                         yield MessageStreamStatus(sse.data)
                     else:
                         chunk_data = json.loads(sse.data)
-                        if "internal_monologue" in chunk_data:
-                            yield InternalMonologue(**chunk_data)
-                        elif "function_call" in chunk_data:
-                            yield FunctionCallMessage(**chunk_data)
-                        elif "function_return" in chunk_data:
-                            yield FunctionReturn(**chunk_data)
+                        if "reasoning" in chunk_data:
+                            yield ReasoningMessage(**chunk_data)
+                        elif "tool_call" in chunk_data:
+                            yield ToolCallMessage(**chunk_data)
+                        elif "tool_return" in chunk_data:
+                            yield ToolReturnMessage(**chunk_data)
                         elif "usage" in chunk_data:
                             yield LettaUsageStatistics(**chunk_data["usage"])
                         else:

letta/errors.py

@@ -131,16 +131,16 @@ class LettaMessageError(LettaError):
         return f"{error_msg}\n\n{message_json}"
 
 
-class MissingFunctionCallError(LettaMessageError):
-    """Error raised when a message is missing a function call."""
+class MissingToolCallError(LettaMessageError):
+    """Error raised when a message is missing a tool call."""
 
-    default_error_message = "The message is missing a function call."
+    default_error_message = "The message is missing a tool call."
 
 
-class InvalidFunctionCallError(LettaMessageError):
-    """Error raised when a message uses an invalid function call."""
+class InvalidToolCallError(LettaMessageError):
+    """Error raised when a message uses an invalid tool call."""
 
-    default_error_message = "The message uses an invalid function call or has improper usage of a function call."
+    default_error_message = "The message uses an invalid tool call or has improper usage of a tool call."
 
 
 class MissingInnerMonologueError(LettaMessageError):

letta/helpers/tool_rule_solver.py

@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional, Set
+import json
+from typing import List, Optional, Union
 
 from pydantic import BaseModel, Field
 
@@ -6,6 +7,7 @@ from letta.schemas.enums import ToolRuleType
 from letta.schemas.tool_rule import (
     BaseToolRule,
     ChildToolRule,
+    ConditionalToolRule,
     InitToolRule,
     TerminalToolRule,
 )
@@ -22,7 +24,7 @@ class ToolRulesSolver(BaseModel):
     init_tool_rules: List[InitToolRule] = Field(
         default_factory=list, description="Initial tool rules to be used at the start of tool execution."
     )
-    tool_rules: List[ChildToolRule] = Field(
+    tool_rules: List[Union[ChildToolRule, ConditionalToolRule]] = Field(
         default_factory=list, description="Standard tool rules for controlling execution sequence and allowed transitions."
     )
     terminal_tool_rules: List[TerminalToolRule] = Field(
@@ -35,21 +37,25 @@ class ToolRulesSolver(BaseModel):
         # Separate the provided tool rules into init, standard, and terminal categories
         for rule in tool_rules:
             if rule.type == ToolRuleType.run_first:
+                assert isinstance(rule, InitToolRule)
                 self.init_tool_rules.append(rule)
             elif rule.type == ToolRuleType.constrain_child_tools:
+                assert isinstance(rule, ChildToolRule)
+                self.tool_rules.append(rule)
+            elif rule.type == ToolRuleType.conditional:
+                assert isinstance(rule, ConditionalToolRule)
+                self.validate_conditional_tool(rule)
                 self.tool_rules.append(rule)
             elif rule.type == ToolRuleType.exit_loop:
+                assert isinstance(rule, TerminalToolRule)
                 self.terminal_tool_rules.append(rule)
 
-        # Validate the tool rules to ensure they form a DAG
-        if not self.validate_tool_rules():
-            raise ToolRuleValidationError("Tool rules contain cycles, which are not allowed in a valid configuration.")
 
     def update_tool_usage(self, tool_name: str):
         """Update the internal state to track the last tool called."""
         self.last_tool_name = tool_name
 
-    def get_allowed_tool_names(self, error_on_empty: bool = False) -> List[str]:
+    def get_allowed_tool_names(self, error_on_empty: bool = False, last_function_response: Optional[str] = None) -> List[str]:
         """Get a list of tool names allowed based on the last tool called."""
         if self.last_tool_name is None:
             # Use initial tool rules if no tool has been called yet
@@ -58,18 +64,21 @@ class ToolRulesSolver(BaseModel):
             # Find a matching ToolRule for the last tool used
             current_rule = next((rule for rule in self.tool_rules if rule.tool_name == self.last_tool_name), None)
 
-            # Return children which must exist on ToolRule
-            if current_rule:
-                return current_rule.children
-
-            # Default to empty if no rule matches
-            message = "User provided tool rules and execution state resolved to no more possible tool calls."
-            if error_on_empty:
-                raise RuntimeError(message)
-            else:
-                # warnings.warn(message)
+            if current_rule is None:
+                if error_on_empty:
+                    raise ValueError(f"No tool rule found for {self.last_tool_name}")
                 return []
 
+            # If the current rule is a conditional tool rule, use the LLM response to
+            # determine which child tool to use
+            if isinstance(current_rule, ConditionalToolRule):
+                if not last_function_response:
+                    raise ValueError("Conditional tool rule requires an LLM response to determine which child tool to use")
+                next_tool = self.evaluate_conditional_tool(current_rule, last_function_response)
+                return [next_tool] if next_tool else []
+
+            return current_rule.children if current_rule.children else []
+
     def is_terminal_tool(self, tool_name: str) -> bool:
         """Check if the tool is defined as a terminal tool in the terminal tool rules."""
         return any(rule.tool_name == tool_name for rule in self.terminal_tool_rules)
@@ -78,38 +87,60 @@ class ToolRulesSolver(BaseModel):
         """Check if the tool has children tools"""
         return any(rule.tool_name == tool_name for rule in self.tool_rules)
 
-    def validate_tool_rules(self) -> bool:
-        """
-        Validate that the tool rules define a directed acyclic graph (DAG).
-        Returns True if valid (no cycles), otherwise False.
-        """
-        # Build adjacency list for the tool graph
-        adjacency_list: Dict[str, List[str]] = {rule.tool_name: rule.children for rule in self.tool_rules}
-
-        # Track visited nodes
-        visited: Set[str] = set()
-        path_stack: Set[str] = set()
-
-        # Define DFS helper function
-        def dfs(tool_name: str) -> bool:
-            if tool_name in path_stack:
-                return False  # Cycle detected
-            if tool_name in visited:
-                return True  # Already validated
-
-            # Mark the node as visited in the current path
-            path_stack.add(tool_name)
-            for child in adjacency_list.get(tool_name, []):
-                if not dfs(child):
-                    return False  # Cycle detected in DFS
-            path_stack.remove(tool_name)  # Remove from current path
-            visited.add(tool_name)
-            return True
-
-        # Run DFS from each tool in `tool_rules`
-        for rule in self.tool_rules:
-            if rule.tool_name not in visited:
-                if not dfs(rule.tool_name):
-                    return False  # Cycle found, invalid tool rules
-
-        return True  # No cycles, valid DAG
+    def validate_conditional_tool(self, rule: ConditionalToolRule):
+        '''
+        Validate a conditional tool rule
+
+        Args:
+            rule (ConditionalToolRule): The conditional tool rule to validate
+
+        Raises:
+            ToolRuleValidationError: If the rule is invalid
+        '''
+        if len(rule.child_output_mapping) == 0:
+            raise ToolRuleValidationError("Conditional tool rule must have at least one child tool.")
+        return True
+
+    def evaluate_conditional_tool(self, tool: ConditionalToolRule, last_function_response: str) -> str:
+        '''
+        Parse function response to determine which child tool to use based on the mapping
+
+        Args:
+            tool (ConditionalToolRule): The conditional tool rule
+            last_function_response (str): The function response in JSON format
+
+        Returns:
+            str: The name of the child tool to use next
+        '''
+        json_response = json.loads(last_function_response)
+        function_output = json_response["message"]
+
+        # Try to match the function output with a mapping key
+        for key in tool.child_output_mapping:
+
+            # Convert function output to match key type for comparison
+            if isinstance(key, bool):
+                typed_output = function_output.lower() == "true"
+            elif isinstance(key, int):
+                try:
+                    typed_output = int(function_output)
+                except (ValueError, TypeError):
+                    continue
+            elif isinstance(key, float):
+                try:
+                    typed_output = float(function_output)
+                except (ValueError, TypeError):
+                    continue
+            else:  # string
+                if function_output == "True" or function_output == "False":
+                    typed_output = function_output.lower()
+                elif function_output == "None":
+                    typed_output = None
+                else:
+                    typed_output = function_output
+
+            if typed_output == key:
+                return tool.child_output_mapping[key]
+
+        # If no match found, use default
+        return tool.default_child

letta/orm/custom_columns.py

@@ -9,7 +9,7 @@ from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import ToolRuleType
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.openai.chat_completions import ToolCall, ToolCallFunction
-from letta.schemas.tool_rule import ChildToolRule, InitToolRule, TerminalToolRule
+from letta.schemas.tool_rule import ChildToolRule, ConditionalToolRule, InitToolRule, TerminalToolRule
 
 
 class EmbeddingConfigColumn(TypeDecorator):
@@ -80,7 +80,7 @@ class ToolRulesColumn(TypeDecorator):
         return value
 
     @staticmethod
-    def deserialize_tool_rule(data: dict) -> Union[ChildToolRule, InitToolRule, TerminalToolRule]:
+    def deserialize_tool_rule(data: dict) -> Union[ChildToolRule, InitToolRule, TerminalToolRule, ConditionalToolRule]:
         """Deserialize a dictionary to the appropriate ToolRule subclass based on the 'type'."""
         rule_type = ToolRuleType(data.get("type"))  # Remove 'type' field if it exists since it is a class var
         if rule_type == ToolRuleType.run_first:
@@ -90,6 +90,9 @@ class ToolRulesColumn(TypeDecorator):
         elif rule_type == ToolRuleType.constrain_child_tools:
             rule = ChildToolRule(**data)
             return rule
+        elif rule_type == ToolRuleType.conditional:
+            rule = ConditionalToolRule(**data)
+            return rule
         else:
             raise ValueError(f"Unknown tool rule type: {rule_type}")
 

letta/providers.py

@@ -482,7 +482,8 @@ class GoogleAIProvider(Provider):
         model_options = [mo[len("models/") :] if mo.startswith("models/") else mo for mo in model_options]
 
         # TODO remove manual filtering for gemini-pro
-        model_options = [mo for mo in model_options if str(mo).startswith("gemini") and "-pro" in str(mo)]
+        # Add support for all gemini models
+        model_options = [mo for mo in model_options if str(mo).startswith("gemini-")]
 
         configs = []
         for model in model_options:

letta/schemas/enums.py

@@ -45,5 +45,6 @@ class ToolRuleType(str, Enum):
     run_first = "InitToolRule"
     exit_loop = "TerminalToolRule"  # reasoning loop should exit
     continue_loop = "continue_loop"  # reasoning loop should continue
+    conditional = "conditional"
     constrain_child_tools = "ToolRule"
     require_parent_tools = "require_parent_tools"

letta/schemas/letta_message.py

@@ -9,7 +9,7 @@ from pydantic import BaseModel, Field, field_serializer, field_validator
 
 class LettaMessage(BaseModel):
     """
-    Base class for simplified Letta message response type. This is intended to be used for developers who want the internal monologue, function calls, and function returns in a simplified format that does not include additional information other than the content and timestamp.
+    Base class for simplified Letta message response type. This is intended to be used for developers who want the internal monologue, tool calls, and tool returns in a simplified format that does not include additional information other than the content and timestamp.
 
     Attributes:
         id (str): The ID of the message
@@ -60,32 +60,32 @@ class UserMessage(LettaMessage):
     message: str
 
 
-class InternalMonologue(LettaMessage):
+class ReasoningMessage(LettaMessage):
     """
-    Representation of an agent's internal monologue.
+    Representation of an agent's internal reasoning.
 
     Attributes:
-        internal_monologue (str): The internal monologue of the agent
+        reasoning (str): The internal reasoning of the agent
         id (str): The ID of the message
         date (datetime): The date the message was created in ISO format
     """
 
-    message_type: Literal["internal_monologue"] = "internal_monologue"
-    internal_monologue: str
+    message_type: Literal["reasoning_message"] = "reasoning_message"
+    reasoning: str
 
 
-class FunctionCall(BaseModel):
+class ToolCall(BaseModel):
 
     name: str
     arguments: str
-    function_call_id: str
+    tool_call_id: str
 
 
-class FunctionCallDelta(BaseModel):
+class ToolCallDelta(BaseModel):
 
     name: Optional[str]
     arguments: Optional[str]
-    function_call_id: Optional[str]
+    tool_call_id: Optional[str]
 
     # NOTE: this is a workaround to exclude None values from the JSON dump,
     # since the OpenAI style of returning chunks doesn't include keys with null values
@@ -97,50 +97,84 @@ class FunctionCallDelta(BaseModel):
         return json.dumps(self.model_dump(exclude_none=True), *args, **kwargs)
 
 
-class FunctionCallMessage(LettaMessage):
+class ToolCallMessage(LettaMessage):
     """
-    A message representing a request to call a function (generated by the LLM to trigger function execution).
+    A message representing a request to call a tool (generated by the LLM to trigger tool execution).
 
     Attributes:
-        function_call (Union[FunctionCall, FunctionCallDelta]): The function call
+        tool_call (Union[ToolCall, ToolCallDelta]): The tool call
         id (str): The ID of the message
         date (datetime): The date the message was created in ISO format
     """
 
-    message_type: Literal["function_call"] = "function_call"
-    function_call: Union[FunctionCall, FunctionCallDelta]
+    message_type: Literal["tool_call_message"] = "tool_call_message"
+    tool_call: Union[ToolCall, ToolCallDelta]
 
-    # NOTE: this is required for the FunctionCallDelta exclude_none to work correctly
+    # NOTE: this is required for the ToolCallDelta exclude_none to work correctly
     def model_dump(self, *args, **kwargs):
         kwargs["exclude_none"] = True
         data = super().model_dump(*args, **kwargs)
-        if isinstance(data["function_call"], dict):
-            data["function_call"] = {k: v for k, v in data["function_call"].items() if v is not None}
+        if isinstance(data["tool_call"], dict):
+            data["tool_call"] = {k: v for k, v in data["tool_call"].items() if v is not None}
         return data
 
     class Config:
         json_encoders = {
-            FunctionCallDelta: lambda v: v.model_dump(exclude_none=True),
-            FunctionCall: lambda v: v.model_dump(exclude_none=True),
+            ToolCallDelta: lambda v: v.model_dump(exclude_none=True),
+            ToolCall: lambda v: v.model_dump(exclude_none=True),
         }
 
-    # NOTE: this is required to cast dicts into FunctionCallMessage objects
+    # NOTE: this is required to cast dicts into ToolCallMessage objects
     # Without this extra validator, Pydantic will throw an error if 'name' or 'arguments' are None
-    # (instead of properly casting to FunctionCallDelta instead of FunctionCall)
-    @field_validator("function_call", mode="before")
+    # (instead of properly casting to ToolCallDelta instead of ToolCall)
+    @field_validator("tool_call", mode="before")
     @classmethod
-    def validate_function_call(cls, v):
+    def validate_tool_call(cls, v):
         if isinstance(v, dict):
-            if "name" in v and "arguments" in v and "function_call_id" in v:
-                return FunctionCall(name=v["name"], arguments=v["arguments"], function_call_id=v["function_call_id"])
-            elif "name" in v or "arguments" in v or "function_call_id" in v:
-                return FunctionCallDelta(name=v.get("name"), arguments=v.get("arguments"), function_call_id=v.get("function_call_id"))
+            if "name" in v and "arguments" in v and "tool_call_id" in v:
+                return ToolCall(name=v["name"], arguments=v["arguments"], tool_call_id=v["tool_call_id"])
+            elif "name" in v or "arguments" in v or "tool_call_id" in v:
+                return ToolCallDelta(name=v.get("name"), arguments=v.get("arguments"), tool_call_id=v.get("tool_call_id"))
             else:
-                raise ValueError("function_call must contain either 'name' or 'arguments'")
+                raise ValueError("tool_call must contain either 'name' or 'arguments'")
         return v
 
 
-class FunctionReturn(LettaMessage):
+class ToolReturnMessage(LettaMessage):
+    """
+    A message representing the return value of a tool call (generated by Letta executing the requested tool).
+
+    Attributes:
+        tool_return (str): The return value of the tool
+        status (Literal["success", "error"]): The status of the tool call
+        id (str): The ID of the message
+        date (datetime): The date the message was created in ISO format
+        tool_call_id (str): A unique identifier for the tool call that generated this message
+        stdout (Optional[List(str)]): Captured stdout (e.g. prints, logs) from the tool invocation
+        stderr (Optional[List(str)]): Captured stderr from the tool invocation
+    """
+
+    message_type: Literal["tool_return_message"] = "tool_return_message"
+    tool_return: str
+    status: Literal["success", "error"]
+    tool_call_id: str
+    stdout: Optional[List[str]] = None
+    stderr: Optional[List[str]] = None
+
+
+# Legacy Letta API had an additional type "assistant_message" and the "function_call" was a formatted string
+
+
+class AssistantMessage(LettaMessage):
+    message_type: Literal["assistant_message"] = "assistant_message"
+    assistant_message: str
+
+
+class LegacyFunctionCallMessage(LettaMessage):
+    function_call: str
+
+
+class LegacyFunctionReturn(LettaMessage):
     """
     A message representing the return value of a function call (generated by Letta executing the requested function).
 
@@ -162,22 +196,24 @@ class FunctionReturn(LettaMessage):
     stderr: Optional[List[str]] = None
 
 
-# Legacy Letta API had an additional type "assistant_message" and the "function_call" was a formatted string
-
-
-class AssistantMessage(LettaMessage):
-    message_type: Literal["assistant_message"] = "assistant_message"
-    assistant_message: str
+class LegacyInternalMonologue(LettaMessage):
+    """
+    Representation of an agent's internal monologue.
 
+    Attributes:
+        internal_monologue (str): The internal monologue of the agent
+        id (str): The ID of the message
+        date (datetime): The date the message was created in ISO format
+    """
 
-class LegacyFunctionCallMessage(LettaMessage):
-    function_call: str
+    message_type: Literal["internal_monologue"] = "internal_monologue"
+    internal_monologue: str
 
 
-LegacyLettaMessage = Union[InternalMonologue, AssistantMessage, LegacyFunctionCallMessage, FunctionReturn]
+LegacyLettaMessage = Union[LegacyInternalMonologue, AssistantMessage, LegacyFunctionCallMessage, LegacyFunctionReturn]
 
 
 LettaMessageUnion = Annotated[
-    Union[SystemMessage, UserMessage, InternalMonologue, FunctionCallMessage, FunctionReturn, AssistantMessage],
+    Union[SystemMessage, UserMessage, ReasoningMessage, ToolCallMessage, ToolReturnMessage, AssistantMessage],
     Field(discriminator="message_type"),
 ]

letta/schemas/letta_response.py

@@ -40,14 +40,22 @@ class LettaResponse(BaseModel):
         def get_formatted_content(msg):
             if msg.message_type == "internal_monologue":
                 return f'<div class="content"><span class="internal-monologue">{html.escape(msg.internal_monologue)}</span></div>'
+            if msg.message_type == "reasoning_message":
+                return f'<div class="content"><span class="internal-monologue">{html.escape(msg.reasoning)}</span></div>'
             elif msg.message_type == "function_call":
                 args = format_json(msg.function_call.arguments)
                 return f'<div class="content"><span class="function-name">{html.escape(msg.function_call.name)}</span>({args})</div>'
+            elif msg.message_type == "tool_call_message":
+                args = format_json(msg.tool_call.arguments)
+                return f'<div class="content"><span class="function-name">{html.escape(msg.function_call.name)}</span>({args})</div>'
             elif msg.message_type == "function_return":
-
                 return_value = format_json(msg.function_return)
                 # return f'<div class="status-line">Status: {html.escape(msg.status)}</div><div class="content">{return_value}</div>'
                 return f'<div class="content">{return_value}</div>'
+            elif msg.message_type == "tool_return_message":
+                return_value = format_json(msg.tool_return)
+                # return f'<div class="status-line">Status: {html.escape(msg.status)}</div><div class="content">{return_value}</div>'
+                return f'<div class="content">{return_value}</div>'
             elif msg.message_type == "user_message":
                 if is_json(msg.message):
                     return f'<div class="content">{format_json(msg.message)}</div>'

letta/schemas/message.py

@@ -16,10 +16,10 @@ from letta.schemas.enums import MessageRole
 from letta.schemas.letta_base import OrmMetadataBase
 from letta.schemas.letta_message import (
     AssistantMessage,
-    FunctionCall,
-    FunctionCallMessage,
-    FunctionReturn,
-    InternalMonologue,
+    ToolCall as LettaToolCall,
+    ToolCallMessage,
+    ToolReturnMessage,
+    ReasoningMessage,
     LettaMessage,
     SystemMessage,
     UserMessage,
@@ -145,10 +145,10 @@ class Message(BaseMessage):
             if self.text is not None:
                 # This is type InnerThoughts
                 messages.append(
-                    InternalMonologue(
+                    ReasoningMessage(
                         id=self.id,
                         date=self.created_at,
-                        internal_monologue=self.text,
+                        reasoning=self.text,
                     )
                 )
             if self.tool_calls is not None:
@@ -172,18 +172,18 @@ class Message(BaseMessage):
                         )
                     else:
                         messages.append(
-                            FunctionCallMessage(
+                            ToolCallMessage(
                                 id=self.id,
                                 date=self.created_at,
-                                function_call=FunctionCall(
+                                tool_call=LettaToolCall(
                                     name=tool_call.function.name,
                                     arguments=tool_call.function.arguments,
-                                    function_call_id=tool_call.id,
+                                    tool_call_id=tool_call.id,
                                 ),
                             )
                         )
         elif self.role == MessageRole.tool:
-            # This is type FunctionReturn
+            # This is type ToolReturnMessage
             # Try to interpret the function return, recall that this is how we packaged:
             # def package_function_response(was_success, response_string, timestamp=None):
             #     formatted_time = get_local_time() if timestamp is None else timestamp
@@ -208,12 +208,12 @@ class Message(BaseMessage):
             messages.append(
                 # TODO make sure this is what the API returns
                 # function_return may not match exactly...
-                FunctionReturn(
+                ToolReturnMessage(
                     id=self.id,
                     date=self.created_at,
-                    function_return=self.text,
+                    tool_return=self.text,
                     status=status_enum,
-                    function_call_id=self.tool_call_id,
+                    tool_call_id=self.tool_call_id,
                 )
             )
         elif self.role == MessageRole.user: