letta-nightly 0.6.5.dev20241218213641__tar.gz → 0.6.5.dev20241220104040__tar.gz

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.6.5.dev20241218213641
+Version: 0.6.5.dev20241220104040
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/agent.py

@@ -1,5 +1,6 @@
 import datetime
 import inspect
+import json
 import time
 import traceback
 import warnings
@@ -20,7 +21,7 @@ from letta.constants import (
     REQ_HEARTBEAT_MESSAGE,
     STRUCTURED_OUTPUT_MODELS,
 )
-from letta.errors import LLMError
+from letta.errors import ContextWindowExceededError
 from letta.helpers import ToolRulesSolver
 from letta.interface import AgentInterface
 from letta.llm_api.helpers import is_context_overflow_error
@@ -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
@@ -1094,6 +1113,7 @@ class Agent(BaseAgent):
 
             # If we got a context alert, try trimming the messages length, then try again
             if is_context_overflow_error(e):
+                printd(f"context window exceeded with limit {self.agent_state.llm_config.context_window}, running summarizer to trim messages")
                 # A separate API call to run a summarizer
                 self.summarize_messages_inplace()
 
@@ -1169,8 +1189,13 @@ class Agent(BaseAgent):
 
         # If at this point there's nothing to summarize, throw an error
         if len(candidate_messages_to_summarize) == 0:
-            raise LLMError(
-                f"Summarize error: tried to run summarize, but couldn't find enough messages to compress [len={len(self.messages)}, preserve_N={MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST}]"
+            raise ContextWindowExceededError(
+                "Not enough messages to compress for summarization",
+                details={
+                    "num_candidate_messages": len(candidate_messages_to_summarize),
+                    "num_total_messages": len(self.messages),
+                    "preserve_N": MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST,
+                },
             )
 
         # Walk down the message buffer (front-to-back) until we hit the target token count
@@ -1204,8 +1229,13 @@ class Agent(BaseAgent):
         message_sequence_to_summarize = self._messages[1:cutoff]  # do NOT get rid of the system message
         if len(message_sequence_to_summarize) <= 1:
             # This prevents a potential infinite loop of summarizing the same message over and over
-            raise LLMError(
-                f"Summarize error: tried to run summarize, but couldn't find enough messages to compress [len={len(message_sequence_to_summarize)} <= 1]"
+            raise ContextWindowExceededError(
+                "Not enough messages to compress for summarization after determining cutoff",
+                details={
+                    "num_candidate_messages": len(message_sequence_to_summarize),
+                    "num_total_messages": len(self.messages),
+                    "preserve_N": MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST,
+                },
             )
         else:
             printd(f"Attempting to summarize {len(message_sequence_to_summarize)} messages [1:{cutoff}] of {len(self._messages)}")
@@ -1218,6 +1248,7 @@ class Agent(BaseAgent):
             self.agent_state.llm_config.context_window = (
                 LLM_MAX_TOKENS[self.model] if (self.model is not None and self.model in LLM_MAX_TOKENS) else LLM_MAX_TOKENS["DEFAULT"]
             )
+
         summary = summarize_messages(agent_state=self.agent_state, message_sequence_to_summarize=message_sequence_to_summarize)
         printd(f"Got summary: {summary}")
 

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/client/client.py

@@ -233,7 +233,7 @@ class AbstractClient(object):
     def get_tool_id(self, name: str) -> Optional[str]:
         raise NotImplementedError
 
-    def add_base_tools(self) -> List[Tool]:
+    def upsert_base_tools(self) -> List[Tool]:
         raise NotImplementedError
 
     def load_data(self, connector: DataConnector, source_name: str):
@@ -1466,7 +1466,7 @@ class RESTClient(AbstractClient):
             raise ValueError(f"Failed to get tool: {response.text}")
         return response.json()
 
-    def add_base_tools(self) -> List[Tool]:
+    def upsert_base_tools(self) -> List[Tool]:
         response = requests.post(f"{self.base_url}/{self.api_prefix}/tools/add-base-tools/", headers=self.headers)
         if response.status_code != 200:
             raise ValueError(f"Failed to add base tools: {response.text}")

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/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_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/errors.py

@@ -1,4 +1,5 @@
 import json
+from enum import Enum
 from typing import TYPE_CHECKING, List, Optional, Union
 
 # Avoid circular imports
@@ -6,9 +7,31 @@ if TYPE_CHECKING:
     from letta.schemas.message import Message
 
 
+class ErrorCode(Enum):
+    """Enum for error codes used by client."""
+
+    INTERNAL_SERVER_ERROR = "INTERNAL_SERVER_ERROR"
+    CONTEXT_WINDOW_EXCEEDED = "CONTEXT_WINDOW_EXCEEDED"
+    RATE_LIMIT_EXCEEDED = "RATE_LIMIT_EXCEEDED"
+
+
 class LettaError(Exception):
     """Base class for all Letta related errors."""
 
+    def __init__(self, message: str, code: Optional[ErrorCode] = None, details: dict = {}):
+        self.message = message
+        self.code = code
+        self.details = details
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        if self.code:
+            return f"{self.code.value}: {self.message}"
+        return self.message
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(message='{self.message}', code='{self.code}', details={self.details})"
+
 
 class LettaToolCreateError(LettaError):
     """Error raised when a tool cannot be created."""
@@ -16,10 +39,7 @@ class LettaToolCreateError(LettaError):
     default_error_message = "Error creating tool."
 
     def __init__(self, message=None):
-        if message is None:
-            message = self.default_error_message
-        self.message = message
-        super().__init__(self.message)
+        super().__init__(message=message or self.default_error_message)
 
 
 class LettaConfigurationError(LettaError):
@@ -27,23 +47,17 @@ class LettaConfigurationError(LettaError):
 
     def __init__(self, message: str, missing_fields: Optional[List[str]] = None):
         self.missing_fields = missing_fields or []
-        super().__init__(message)
+        super().__init__(message=message, details={"missing_fields": self.missing_fields})
 
 
 class LettaAgentNotFoundError(LettaError):
     """Error raised when an agent is not found."""
-
-    def __init__(self, message: str):
-        self.message = message
-        super().__init__(self.message)
+    pass
 
 
 class LettaUserNotFoundError(LettaError):
     """Error raised when a user is not found."""
-
-    def __init__(self, message: str):
-        self.message = message
-        super().__init__(self.message)
+    pass
 
 
 class LLMError(LettaError):
@@ -54,24 +68,45 @@ class LLMJSONParsingError(LettaError):
     """Exception raised for errors in the JSON parsing process."""
 
     def __init__(self, message="Error parsing JSON generated by LLM"):
-        self.message = message
-        super().__init__(self.message)
+        super().__init__(message=message)
 
 
 class LocalLLMError(LettaError):
     """Generic catch-all error for local LLM problems"""
 
     def __init__(self, message="Encountered an error while running local LLM"):
-        self.message = message
-        super().__init__(self.message)
+        super().__init__(message=message)
 
 
 class LocalLLMConnectionError(LettaError):
     """Error for when local LLM cannot be reached with provided IP/port"""
 
     def __init__(self, message="Could not connect to local LLM"):
-        self.message = message
-        super().__init__(self.message)
+        super().__init__(message=message)
+
+
+class ContextWindowExceededError(LettaError):
+    """Error raised when the context window is exceeded but further summarization fails."""
+
+    def __init__(self, message: str, details: dict = {}):
+        error_message = f"{message} ({details})"
+        super().__init__(
+            message=error_message,
+            code=ErrorCode.CONTEXT_WINDOW_EXCEEDED,
+            details=details,
+        )
+
+
+class RateLimitExceededError(LettaError):
+    """Error raised when the llm rate limiter throttles api requests."""
+
+    def __init__(self, message: str, max_retries: int):
+        error_message = f"{message} ({max_retries})"
+        super().__init__(
+            message=error_message,
+            code=ErrorCode.RATE_LIMIT_EXCEEDED,
+            details={"max_retries": max_retries},
+        )
 
 
 class LettaMessageError(LettaError):
@@ -96,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_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/functions/function_sets/base.py

@@ -61,60 +61,6 @@ def conversation_search(self: "Agent", query: str, page: Optional[int] = 0) -> O
     return results_str
 
 
-def conversation_search_date(self: "Agent", start_date: str, end_date: str, page: Optional[int] = 0) -> Optional[str]:
-    """
-    Search prior conversation history using a date range.
-
-    Args:
-        start_date (str): The start of the date range to search, in the format 'YYYY-MM-DD'.
-        end_date (str): The end of the date range to search, in the format 'YYYY-MM-DD'.
-        page (int): Allows you to page through results. Only use on a follow-up query. Defaults to 0 (first page).
-
-    Returns:
-        str: Query result string
-    """
-    import math
-    from datetime import datetime
-
-    from letta.constants import RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
-    from letta.utils import json_dumps
-
-    if page is None or (isinstance(page, str) and page.lower().strip() == "none"):
-        page = 0
-    try:
-        page = int(page)
-        if page < 0:
-            raise ValueError
-    except:
-        raise ValueError(f"'page' argument must be an integer")
-
-    # Convert date strings to datetime objects
-    try:
-        start_datetime = datetime.strptime(start_date, "%Y-%m-%d").replace(hour=0, minute=0, second=0, microsecond=0)
-        end_datetime = datetime.strptime(end_date, "%Y-%m-%d").replace(hour=23, minute=59, second=59, microsecond=999999)
-    except ValueError:
-        raise ValueError("Dates must be in the format 'YYYY-MM-DD'")
-
-    count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
-    results = self.message_manager.list_user_messages_for_agent(
-        # TODO: add paging by page number. currently cursor only works with strings.
-        agent_id=self.agent_state.id,
-        actor=self.user,
-        start_date=start_datetime,
-        end_date=end_datetime,
-        limit=count,
-    )
-    total = len(results)
-    num_pages = math.ceil(total / count) - 1  # 0 index
-    if len(results) == 0:
-        results_str = f"No results found."
-    else:
-        results_pref = f"Showing {len(results)} of {total} results (page {page}/{num_pages}):"
-        results_formatted = [f"timestamp: {d['timestamp']}, {d['message']['role']} - {d['message']['content']}" for d in results]
-        results_str = f"{results_pref} {json_dumps(results_formatted)}"
-    return results_str
-
-
 def archival_memory_insert(self: "Agent", content: str) -> Optional[str]:
     """
     Add to archival memory. Make sure to phrase the memory contents such that it can be easily queried later.

letta_nightly-0.6.5.dev20241220104040/letta/helpers/tool_rule_solver.py

@@ -0,0 +1,146 @@
+import json
+from typing import List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from letta.schemas.enums import ToolRuleType
+from letta.schemas.tool_rule import (
+    BaseToolRule,
+    ChildToolRule,
+    ConditionalToolRule,
+    InitToolRule,
+    TerminalToolRule,
+)
+
+
+class ToolRuleValidationError(Exception):
+    """Custom exception for tool rule validation errors in ToolRulesSolver."""
+
+    def __init__(self, message: str):
+        super().__init__(f"ToolRuleValidationError: {message}")
+
+
+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[Union[ChildToolRule, ConditionalToolRule]] = Field(
+        default_factory=list, description="Standard tool rules for controlling execution sequence and allowed transitions."
+    )
+    terminal_tool_rules: List[TerminalToolRule] = Field(
+        default_factory=list, description="Terminal tool rules that end the agent loop if called."
+    )
+    last_tool_name: Optional[str] = Field(None, description="The most recent tool used, updated with each tool call.")
+
+    def __init__(self, tool_rules: List[BaseToolRule], **kwargs):
+        super().__init__(**kwargs)
+        # 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)
+
+
+    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, 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
+            return [rule.tool_name for rule in self.init_tool_rules]
+        else:
+            # 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)
+
+            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)
+
+    def has_children_tools(self, tool_name):
+        """Check if the tool has children tools"""
+        return any(rule.tool_name == tool_name for rule in self.tool_rules)
+
+    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_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/llm_api/llm_api_tools.py

@@ -5,7 +5,7 @@ from typing import List, Optional, Union
 import requests
 
 from letta.constants import CLI_WARNING_PREFIX
-from letta.errors import LettaConfigurationError
+from letta.errors import LettaConfigurationError, RateLimitExceededError
 from letta.llm_api.anthropic import anthropic_chat_completions_request
 from letta.llm_api.azure_openai import azure_openai_chat_completions_request
 from letta.llm_api.google_ai import (
@@ -80,7 +80,7 @@ def retry_with_exponential_backoff(
 
                     # Check if max retries has been reached
                     if num_retries > max_retries:
-                        raise Exception(f"Maximum number of retries ({max_retries}) exceeded.")
+                        raise RateLimitExceededError("Maximum number of retries exceeded", max_retries=max_retries)
 
                     # Increment the delay
                     delay *= exponential_base * (1 + jitter * random.random())

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/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_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/orm/message.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from sqlalchemy import Index
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.custom_columns import ToolCallColumn
@@ -13,7 +14,7 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
     """Defines data model for storing Message objects"""
 
     __tablename__ = "messages"
-    __table_args__ = {"extend_existing": True}
+    __table_args__ = (Index("ix_messages_agent_created_at", "agent_id", "created_at"),)
     __pydantic_model__ = PydanticMessage
 
     id: Mapped[str] = mapped_column(primary_key=True, doc="Unique message identifier")

{letta_nightly-0.6.5.dev20241218213641 → letta_nightly-0.6.5.dev20241220104040}/letta/orm/passage.py

@@ -1,26 +1,26 @@
 from typing import TYPE_CHECKING
-from sqlalchemy import Column, JSON, Index
-from sqlalchemy.orm import Mapped, mapped_column, relationship, declared_attr
 
-from letta.orm.mixins import FileMixin, OrganizationMixin
+from sqlalchemy import JSON, Column, Index
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
+
+from letta.config import LettaConfig
+from letta.constants import MAX_EMBEDDING_DIM
 from letta.orm.custom_columns import CommonVector, EmbeddingConfigColumn
-from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.orm.mixins import AgentMixin, FileMixin, OrganizationMixin, SourceMixin
+from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.passage import Passage as PydanticPassage
 from letta.settings import settings
 
-from letta.config import LettaConfig
-from letta.constants import MAX_EMBEDDING_DIM
-
 config = LettaConfig()
 
 if TYPE_CHECKING:
-    from letta.orm.organization import Organization
     from letta.orm.agent import Agent
+    from letta.orm.organization import Organization
 
 
 class BasePassage(SqlalchemyBase, OrganizationMixin):
     """Base class for all passage types with common fields"""
+
     __abstract__ = True
     __pydantic_model__ = PydanticPassage
 
@@ -45,17 +45,15 @@ class BasePassage(SqlalchemyBase, OrganizationMixin):
     @declared_attr
     def __table_args__(cls):
         if settings.letta_pg_uri_no_default:
-            return (
-                Index(f'{cls.__tablename__}_org_idx', 'organization_id'),
-                {"extend_existing": True}
-            )
+            return (Index(f"{cls.__tablename__}_org_idx", "organization_id"), {"extend_existing": True})
         return ({"extend_existing": True},)
 
 
 class SourcePassage(BasePassage, FileMixin, SourceMixin):
     """Passages derived from external files/sources"""
+
     __tablename__ = "source_passages"
-    
+
     @declared_attr
     def file(cls) -> Mapped["FileMetadata"]:
         """Relationship to file"""
@@ -64,7 +62,7 @@ class SourcePassage(BasePassage, FileMixin, SourceMixin):
     @declared_attr
     def organization(cls) -> Mapped["Organization"]:
         return relationship("Organization", back_populates="source_passages", lazy="selectin")
-        
+
     @declared_attr
     def source(cls) -> Mapped["Source"]:
         """Relationship to source"""
@@ -73,8 +71,9 @@ class SourcePassage(BasePassage, FileMixin, SourceMixin):
 
 class AgentPassage(BasePassage, AgentMixin):
     """Passages created by agents as archival memories"""
+
     __tablename__ = "agent_passages"
-    
+
     @declared_attr
     def organization(cls) -> Mapped["Organization"]:
         return relationship("Organization", back_populates="agent_passages", lazy="selectin")