File size: 6,010 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | --- a/letta/agent.py
+++ b/letta/agent.py
@@ -1,3 +1,4 @@
+# Modified.
import json
import time
import traceback
@@ -70,6 +71,33 @@
logger = get_logger(__name__)
+def _fit_context_window(sequence: List[Message], new_messages: List[Message], limit: Optional[int]) -> List[Message]:
+ if not limit or len(sequence) <= 2 or not new_messages:
+ return sequence
+ counts = get_token_counts_for_messages(sequence)
+ if sum(counts) <= limit:
+ return sequence
+ n_new = len(new_messages)
+ room = limit - counts[0] - sum(counts[-n_new:])
+ middle = list(zip(sequence[1:-n_new], counts[1:-n_new]))
+ units, i = [], 0
+ while i < len(middle):
+ j = i + 1
+ if getattr(middle[i][0], "role", None) == "assistant" and middle[i][0].tool_calls:
+ while j < len(middle) and getattr(middle[j][0], "role", None) == "tool":
+ j += 1
+ units.append(middle[i:j])
+ i = j
+ kept = []
+ for unit in reversed(units):
+ size = sum(n for _, n in unit)
+ if size > room:
+ break
+ kept[:0] = [m for m, _ in unit]
+ room -= size
+ return [sequence[0], *kept, *new_messages]
+
+
class BaseAgent(ABC):
"""
Abstract class for all agents.
@@ -745,6 +773,23 @@
elif max_chaining_steps is not None and counter > max_chaining_steps:
self.logger.info(f"Hit max chaining steps, stopping after {counter} steps")
break
+ elif (max_chaining_steps is not None and counter == max_chaining_steps
+ and (heartbeat_request or function_failed or token_warning)
+ and __import__("os").environ.get("MEMEBENCH_LETTA_FINAL_REPLY_AT_CAP") == "1"):
+ assert self.agent_state.created_by_id is not None
+ next_input_messages = [
+ Message.dict_to_message(
+ agent_id=self.agent_state.id,
+ model=self.model,
+ openai_message_dict={
+ "role": "user",
+ "content": get_heartbeat(
+ "You have reached the maximum number of steps. Call send_message now "
+ "with your final response."),
+ },
+ )
+ ]
+ continue
# Chain handlers
elif token_warning and summarizer_settings.send_memory_warning_message:
assert self.agent_state.created_by_id is not None
@@ -828,6 +873,8 @@
in_context_messages = self.agent_manager.get_in_context_messages(agent_id=self.agent_state.id, actor=self.user)
input_message_sequence = in_context_messages + messages
+ input_message_sequence = _fit_context_window(
+ input_message_sequence, messages, self.agent_state.llm_config.context_window)
if (
len(input_message_sequence) > 1
--- a/letta/constants.py
+++ b/letta/constants.py
@@ -1,7 +1,8 @@
+# Modified.
import os
from logging import CRITICAL, DEBUG, ERROR, INFO, NOTSET, WARN, WARNING
-LETTA_DIR = os.path.join(os.path.expanduser("~"), ".letta")
+LETTA_DIR = os.environ.get("LETTA_DIR") or os.path.join(os.path.expanduser("~"), ".letta")
LETTA_TOOL_EXECUTION_DIR = os.path.join(LETTA_DIR, "tool_execution_dir")
ADMIN_PREFIX = "/v1/admin"
--- a/letta/functions/function_sets/base.py
+++ b/letta/functions/function_sets/base.py
@@ -1,3 +1,4 @@
+# Modified.
from typing import Optional
from letta.agent import Agent
@@ -43,7 +44,7 @@
page = int(page)
except:
raise ValueError(f"'page' argument must be an integer")
- count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
+ count = (self.agent_state.metadata or {}).get("retrieve_num", RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE)
# TODO: add paging by page number. currently cursor only works with strings.
# original: start=page * count
messages = self.message_manager.list_user_messages_for_agent(
@@ -104,7 +105,7 @@
page = int(page)
except:
raise ValueError(f"'page' argument must be an integer")
- count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
+ count = (self.agent_state.metadata or {}).get("retrieve_num", RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE)
try:
# Get results using passage manager
--- a/letta/schemas/openai/chat_completion_request.py
+++ b/letta/schemas/openai/chat_completion_request.py
@@ -1,3 +1,4 @@
+# Modified.
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field, field_validator
@@ -124,6 +125,7 @@
logprobs: Optional[bool] = False
top_logprobs: Optional[int] = None
max_tokens: Optional[int] = None
+ max_completion_tokens: Optional[int] = None
n: Optional[int] = 1
presence_penalty: Optional[float] = 0
response_format: Optional[ResponseFormat] = None
--- a/letta/services/message_manager.py
+++ b/letta/services/message_manager.py
@@ -1,7 +1,8 @@
+# Modified.
import json
from typing import List, Optional, Sequence
-from sqlalchemy import delete, exists, func, select, text
+from sqlalchemy import String, cast, delete, exists, func, select, text
from letta.log import get_logger
from letta.orm.agent import Agent as AgentModel
@@ -304,7 +305,9 @@
query = query.filter(MessageModel.group_id == group_id)
# If query_text is provided, filter messages using subquery + json_array_elements.
- if query_text:
+ if query_text and session.bind is not None and session.bind.dialect.name == "sqlite":
+ query = query.filter(func.lower(cast(MessageModel.content, String)).like(f"%{query_text.lower()}%"))
+ elif query_text:
content_element = func.json_array_elements(MessageModel.content).alias("content_element")
query = query.filter(
exists(
|