letta-nightly 0.8.5.dev20250625104328__py3-none-any.whl → 0.8.6.dev20250625222533__py3-none-any.whl

letta/services/helpers/agent_manager_helper.py

@@ -1,15 +1,15 @@
-import datetime
+from datetime import datetime
 from typing import List, Literal, Optional
 
 import numpy as np
-from sqlalchemy import Select, and_, asc, desc, func, literal, or_, select, union_all
+from sqlalchemy import Select, and_, asc, desc, func, literal, nulls_last, or_, select, union_all
 from sqlalchemy.sql.expression import exists
 
 from letta import system
 from letta.constants import IN_CONTEXT_MEMORY_KEYWORD, MAX_EMBEDDING_DIM, STRUCTURED_OUTPUT_MODELS
 from letta.embeddings import embedding_model
 from letta.helpers import ToolRulesSolver
-from letta.helpers.datetime_helpers import get_local_time, get_local_time_fast
+from letta.helpers.datetime_helpers import format_datetime, get_local_time, get_local_time_fast
 from letta.orm import AgentPassage, SourcePassage, SourcesAgents
 from letta.orm.agent import Agent as AgentModel
 from letta.orm.agents_tags import AgentsTags
@@ -178,18 +178,19 @@ def derive_system_message(agent_type: AgentType, enable_sleeptime: Optional[bool
 
 # TODO: This code is kind of wonky and deserves a rewrite
 def compile_memory_metadata_block(
-    memory_edit_timestamp: datetime.datetime,
+    memory_edit_timestamp: datetime,
+    timezone: str,
     previous_message_count: int = 0,
     archival_memory_size: int = 0,
 ) -> str:
     # Put the timestamp in the local timezone (mimicking get_local_time())
-    timestamp_str = memory_edit_timestamp.astimezone().strftime("%Y-%m-%d %I:%M:%S %p %Z%z").strip()
+    timestamp_str = format_datetime(memory_edit_timestamp, timezone)
 
     # Create a metadata block of info so the agent knows about the metadata of out-of-context memories
     memory_metadata_block = "\n".join(
         [
             "<memory_metadata>",
-            f"- The current time is: {get_local_time_fast()}",
+            f"- The current time is: {get_local_time_fast(timezone)}",
             f"- Memory blocks were last modified: {timestamp_str}",
             f"- {previous_message_count} previous messages between you and the user are stored in recall memory (use tools to access them)",
             f"- {archival_memory_size} total memories you created are stored in archival memory (use tools to access them)",
@@ -223,7 +224,8 @@ def safe_format(template: str, variables: dict) -> str:
 def compile_system_message(
     system_prompt: str,
     in_context_memory: Memory,
-    in_context_memory_last_edit: datetime.datetime,  # TODO move this inside of BaseMemory?
+    in_context_memory_last_edit: datetime,  # TODO move this inside of BaseMemory?
+    timezone: str,
     user_defined_variables: Optional[dict] = None,
     append_icm_if_missing: bool = True,
     template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
@@ -239,10 +241,9 @@ def compile_system_message(
       - CORE_MEMORY: the in-context memory of the LLM
     """
     # Add tool rule constraints if available
+    tool_constraint_block = None
     if tool_rules_solver is not None:
         tool_constraint_block = tool_rules_solver.compile_tool_rule_prompts()
-        if tool_constraint_block:  # There may not be any depending on if there are tool rules attached
-            in_context_memory.blocks.append(tool_constraint_block)
 
     if user_defined_variables is not None:
         # TODO eventually support the user defining their own variables to inject
@@ -259,8 +260,9 @@ def compile_system_message(
             memory_edit_timestamp=in_context_memory_last_edit,
             previous_message_count=previous_message_count,
             archival_memory_size=archival_memory_size,
+            timezone=timezone,
         )
-        full_memory_string = in_context_memory.compile() + "\n\n" + memory_metadata_string
+        full_memory_string = in_context_memory.compile(tool_usage_rules=tool_constraint_block) + "\n\n" + memory_metadata_string
 
         # Add to the variables list to inject
         variables[IN_CONTEXT_MEMORY_KEYWORD] = full_memory_string
@@ -292,7 +294,7 @@ def compile_system_message(
 
 def initialize_message_sequence(
     agent_state: AgentState,
-    memory_edit_timestamp: Optional[datetime.datetime] = None,
+    memory_edit_timestamp: Optional[datetime] = None,
     include_initial_boot_message: bool = True,
     previous_message_count: int = 0,
     archival_memory_size: int = 0,
@@ -304,20 +306,21 @@ def initialize_message_sequence(
         system_prompt=agent_state.system,
         in_context_memory=agent_state.memory,
         in_context_memory_last_edit=memory_edit_timestamp,
+        timezone=agent_state.timezone,
         user_defined_variables=None,
         append_icm_if_missing=True,
         previous_message_count=previous_message_count,
         archival_memory_size=archival_memory_size,
     )
-    first_user_message = get_login_event()  # event letting Letta know the user just logged in
+    first_user_message = get_login_event(agent_state.timezone)  # event letting Letta know the user just logged in
 
     if include_initial_boot_message:
         if agent_state.agent_type == AgentType.sleeptime_agent:
             initial_boot_messages = []
         elif agent_state.llm_config.model is not None and "gpt-3.5" in agent_state.llm_config.model:
-            initial_boot_messages = get_initial_boot_messages("startup_with_send_message_gpt35")
+            initial_boot_messages = get_initial_boot_messages("startup_with_send_message_gpt35", agent_state.timezone)
         else:
-            initial_boot_messages = get_initial_boot_messages("startup_with_send_message")
+            initial_boot_messages = get_initial_boot_messages("startup_with_send_message", agent_state.timezone)
         messages = (
             [
                 {"role": "system", "content": full_system_message},
@@ -338,7 +341,7 @@ def initialize_message_sequence(
 
 
 def package_initial_message_sequence(
-    agent_id: str, initial_message_sequence: List[MessageCreate], model: str, actor: User
+    agent_id: str, initial_message_sequence: List[MessageCreate], model: str, timezone: str, actor: User
 ) -> List[Message]:
     # create the agent object
     init_messages = []
@@ -347,6 +350,7 @@ def package_initial_message_sequence(
         if message_create.role == MessageRole.user:
             packed_message = system.package_user_message(
                 user_message=message_create.content,
+                timezone=timezone,
             )
             init_messages.append(
                 Message(
@@ -361,6 +365,7 @@ def package_initial_message_sequence(
         elif message_create.role == MessageRole.system:
             packed_message = system.package_system_message(
                 system_message=message_create.content,
+                timezone=timezone,
             )
             init_messages.append(
                 Message(
@@ -402,7 +407,7 @@ def package_initial_message_sequence(
             )
 
             # add tool return
-            function_response = package_function_response(True, "None")
+            function_response = package_function_response(True, "None", timezone)
             init_messages.append(
                 Message(
                     role=MessageRole.tool,
@@ -430,23 +435,47 @@ def check_supports_structured_output(model: str, tool_rules: List[ToolRule]) ->
         return True
 
 
-def _cursor_filter(created_at_col, id_col, ref_created_at, ref_id, forward: bool):
+def _cursor_filter(sort_col, id_col, ref_sort_col, ref_id, forward: bool, nulls_last: bool = False):
     """
     Returns a SQLAlchemy filter expression for cursor-based pagination.
 
     If `forward` is True, returns records after the reference.
     If `forward` is False, returns records before the reference.
+
+    Handles NULL values in the sort column properly when nulls_last is True.
     """
-    if forward:
-        return or_(
-            created_at_col > ref_created_at,
-            and_(created_at_col == ref_created_at, id_col > ref_id),
-        )
+    if not nulls_last:
+        # Simple case: no special NULL handling needed
+        if forward:
+            return or_(
+                sort_col > ref_sort_col,
+                and_(sort_col == ref_sort_col, id_col > ref_id),
+            )
+        else:
+            return or_(
+                sort_col < ref_sort_col,
+                and_(sort_col == ref_sort_col, id_col < ref_id),
+            )
+
+    # Handle nulls_last case
+    # TODO: add tests to check if this works for ascending order but nulls are stil last?
+    if ref_sort_col is None:
+        # Reference cursor is at a NULL value
+        if forward:
+            # Moving forward (e.g. previous) from NULL: either other NULLs with greater IDs or non-NULLs
+            return or_(and_(sort_col.is_(None), id_col > ref_id), sort_col.isnot(None))
+        else:
+            # Moving backward (e.g. next) from NULL: NULLs with smaller IDs
+            return and_(sort_col.is_(None), id_col < ref_id)
     else:
-        return or_(
-            created_at_col < ref_created_at,
-            and_(created_at_col == ref_created_at, id_col < ref_id),
-        )
+        # Reference cursor is at a non-NULL value
+        if forward:
+            # Moving forward (e.g. previous) from non-NULL: only greater non-NULL values
+            # (NULLs are at the end, so we don't include them when moving forward from non-NULL)
+            return and_(sort_col.isnot(None), or_(sort_col > ref_sort_col, and_(sort_col == ref_sort_col, id_col > ref_id)))
+        else:
+            # Moving backward (e.g. next) from non-NULL: smaller non-NULL values or NULLs
+            return or_(sort_col.is_(None), or_(sort_col < ref_sort_col, and_(sort_col == ref_sort_col, id_col < ref_id)))
 
 
 def _apply_pagination(
@@ -455,30 +484,30 @@ def _apply_pagination(
     # Determine the sort column
     if sort_by == "last_run_completion":
         sort_column = AgentModel.last_run_completion
+        sort_nulls_last = True  # TODO: handle this as a query param eventually
     else:
         sort_column = AgentModel.created_at
+        sort_nulls_last = False
 
     if after:
-        if sort_by == "last_run_completion":
-            result = session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == after)).first()
-        else:
-            result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after)).first()
+        result = session.execute(select(sort_column, AgentModel.id).where(AgentModel.id == after)).first()
         if result:
             after_sort_value, after_id = result
-            query = query.where(_cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending))
+            query = query.where(
+                _cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending, nulls_last=sort_nulls_last)
+            )
 
     if before:
-        if sort_by == "last_run_completion":
-            result = session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == before)).first()
-        else:
-            result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before)).first()
+        result = session.execute(select(sort_column, AgentModel.id).where(AgentModel.id == before)).first()
         if result:
             before_sort_value, before_id = result
-            query = query.where(_cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending))
+            query = query.where(
+                _cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending, nulls_last=sort_nulls_last)
+            )
 
     # Apply ordering
     order_fn = asc if ascending else desc
-    query = query.order_by(order_fn(sort_column), order_fn(AgentModel.id))
+    query = query.order_by(nulls_last(order_fn(sort_column)) if sort_nulls_last else order_fn(sort_column), order_fn(AgentModel.id))
     return query
 
 
@@ -488,30 +517,30 @@ async def _apply_pagination_async(
     # Determine the sort column
     if sort_by == "last_run_completion":
         sort_column = AgentModel.last_run_completion
+        sort_nulls_last = True  # TODO: handle this as a query param eventually
     else:
         sort_column = AgentModel.created_at
+        sort_nulls_last = False
 
     if after:
-        if sort_by == "last_run_completion":
-            result = (await session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == after))).first()
-        else:
-            result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after))).first()
+        result = (await session.execute(select(sort_column, AgentModel.id).where(AgentModel.id == after))).first()
         if result:
             after_sort_value, after_id = result
-            query = query.where(_cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending))
+            query = query.where(
+                _cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending, nulls_last=sort_nulls_last)
+            )
 
     if before:
-        if sort_by == "last_run_completion":
-            result = (await session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == before))).first()
-        else:
-            result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before))).first()
+        result = (await session.execute(select(sort_column, AgentModel.id).where(AgentModel.id == before))).first()
         if result:
             before_sort_value, before_id = result
-            query = query.where(_cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending))
+            query = query.where(
+                _cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending, nulls_last=sort_nulls_last)
+            )
 
     # Apply ordering
     order_fn = asc if ascending else desc
-    query = query.order_by(order_fn(sort_column), order_fn(AgentModel.id))
+    query = query.order_by(nulls_last(order_fn(sort_column)) if sort_nulls_last else order_fn(sort_column), order_fn(AgentModel.id))
     return query
 
 

letta/services/helpers/tool_parser_helper.py

@@ -28,7 +28,8 @@ def parse_function_arguments(source_code: str, tool_name: str):
     tree = ast.parse(source_code)
     args = []
     for node in ast.walk(tree):
-        if isinstance(node, ast.FunctionDef) and node.name == tool_name:
+        # Handle both sync and async functions
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == tool_name:
             for arg in node.args.args:
                 args.append(arg.arg)
     return args

letta/services/job_manager.py

@@ -7,7 +7,6 @@ from sqlalchemy.orm import Session
 
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.log import get_logger
-from letta.orm.enums import JobType
 from letta.orm.errors import NoResultFound
 from letta.orm.job import Job as JobModel
 from letta.orm.job_messages import JobMessage
@@ -16,7 +15,7 @@ from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step
 from letta.orm.step import Step as StepModel
 from letta.otel.tracing import trace_method
-from letta.schemas.enums import JobStatus, MessageRole
+from letta.schemas.enums import JobStatus, JobType, MessageRole
 from letta.schemas.job import BatchJob as PydanticBatchJob
 from letta.schemas.job import Job as PydanticJob
 from letta.schemas.job import JobUpdate, LettaRequestConfig
@@ -342,6 +341,33 @@ class JobManager:
             session.add(job_message)
             session.commit()
 
+    @enforce_types
+    @trace_method
+    async def add_messages_to_job_async(self, job_id: str, message_ids: List[str], actor: PydanticUser) -> None:
+        """
+        Associate a message with a job by creating a JobMessage record.
+        Each message can only be associated with one job.
+
+        Args:
+            job_id: The ID of the job
+            message_id: The ID of the message to associate
+            actor: The user making the request
+
+        Raises:
+            NoResultFound: If the job does not exist or user does not have access
+        """
+        if not message_ids:
+            return
+
+        async with db_registry.async_session() as session:
+            # First verify job exists and user has access
+            await self._verify_job_access_async(session, job_id, actor, access=["write"])
+
+            # Create new JobMessage associations
+            job_messages = [JobMessage(job_id=job_id, message_id=message_id) for message_id in message_ids]
+            session.add_all(job_messages)
+            await session.commit()
+
     @enforce_types
     @trace_method
     def get_job_usage(self, job_id: str, actor: PydanticUser) -> LettaUsageStatistics:
@@ -463,14 +489,19 @@ class JobManager:
         )
 
         request_config = self._get_run_request_config(run_id)
+        print("request_config", request_config)
 
         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"],
+            reverse=not ascending,
         )
 
+        if request_config["include_return_message_types"]:
+            messages = [msg for msg in messages if msg.message_type in request_config["include_return_message_types"]]
+
         return messages
 
     @enforce_types

letta/services/llm_batch_manager.py

@@ -10,8 +10,8 @@ from letta.orm import Message as MessageModel
 from letta.orm.llm_batch_items import LLMBatchItem
 from letta.orm.llm_batch_job import LLMBatchJob
 from letta.otel.tracing import trace_method
-from letta.schemas.agent import AgentStepState
 from letta.schemas.enums import AgentStepStatus, JobStatus, ProviderType
+from letta.schemas.llm_batch_job import AgentStepState
 from letta.schemas.llm_batch_job import LLMBatchItem as PydanticLLMBatchItem
 from letta.schemas.llm_batch_job import LLMBatchJob as PydanticLLMBatchJob
 from letta.schemas.llm_config import LLMConfig

letta/services/provider_manager.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 from letta.orm.provider import Provider as ProviderModel
 from letta.otel.tracing import trace_method
@@ -196,10 +196,12 @@ class ProviderManager:
 
     @enforce_types
     @trace_method
-    async def get_bedrock_credentials_async(self, provider_name: Union[str, None], actor: PydanticUser) -> Optional[str]:
+    async def get_bedrock_credentials_async(
+        self, provider_name: Union[str, None], actor: PydanticUser
+    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
         providers = await self.list_providers_async(name=provider_name, actor=actor)
-        access_key = providers[0].api_key if providers else None
-        secret_key = providers[0].api_secret if providers else None
+        access_key = providers[0].access_key if providers else None
+        secret_key = providers[0].api_key if providers else None
         region = providers[0].region if providers else None
         return access_key, secret_key, region
 

letta/services/tool_executor/core_tool_executor.py

@@ -143,7 +143,7 @@ class LettaCoreToolExecutor(ToolExecutor):
 
         try:
             # Get results using passage manager
-            all_results = await AgentManager().list_passages_async(
+            all_results = await AgentManager().list_agent_passages_async(
                 actor=actor,
                 agent_id=agent_state.id,
                 query_text=query,

letta/services/tool_executor/files_tool_executor.py

@@ -3,8 +3,8 @@ import re
 from typing import Any, Dict, List, Optional, Tuple
 
 from letta.log import get_logger
+from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState
-from letta.schemas.file import FileMetadata
 from letta.schemas.sandbox_config import SandboxConfig
 from letta.schemas.tool import Tool
 from letta.schemas.tool_execution_result import ToolExecutionResult
@@ -14,6 +14,7 @@ from letta.services.block_manager import BlockManager
 from letta.services.file_manager import FileManager
 from letta.services.file_processor.chunker.line_chunker import LineChunker
 from letta.services.files_agents_manager import FileAgentManager
+from letta.services.job_manager import JobManager
 from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
 from letta.services.source_manager import SourceManager
@@ -38,6 +39,7 @@ class LettaFileToolExecutor(ToolExecutor):
         message_manager: MessageManager,
         agent_manager: AgentManager,
         block_manager: BlockManager,
+        job_manager: JobManager,
         passage_manager: PassageManager,
         actor: User,
     ):
@@ -45,6 +47,7 @@ class LettaFileToolExecutor(ToolExecutor):
             message_manager=message_manager,
             agent_manager=agent_manager,
             block_manager=block_manager,
+            job_manager=job_manager,
             passage_manager=passage_manager,
             actor=actor,
         )
@@ -94,6 +97,7 @@ class LettaFileToolExecutor(ToolExecutor):
                 stderr=[get_friendly_error_msg(function_name=function_name, exception_name=type(e).__name__, exception_message=str(e))],
             )
 
+    @trace_method
     async def open_file(self, agent_state: AgentState, file_name: str, view_range: Optional[Tuple[int, int]] = None) -> str:
         """Stub for open_file tool."""
         start, end = None, None
@@ -120,14 +124,23 @@ class LettaFileToolExecutor(ToolExecutor):
         # TODO: Inefficient, maybe we can pre-compute this
         # TODO: This is also not the best way to split things - would be cool to have "content aware" splitting
         # TODO: Split code differently from large text blurbs
-        content_lines = LineChunker().chunk_text(text=file.content, file_metadata=file, start=start, end=end)
+        content_lines = LineChunker().chunk_text(file_metadata=file, start=start, end=end)
         visible_content = "\n".join(content_lines)
 
-        await self.files_agents_manager.update_file_agent_by_id(
-            agent_id=agent_state.id, file_id=file_id, actor=self.actor, is_open=True, visible_content=visible_content
+        # Efficiently handle LRU eviction and file opening in a single transaction
+        closed_files, was_already_open = await self.files_agents_manager.enforce_max_open_files_and_open(
+            agent_id=agent_state.id, file_id=file_id, file_name=file_name, actor=self.actor, visible_content=visible_content
         )
-        return f"Successfully opened file {file_name}, lines {start} to {end} are now visible in memory block <{file_name}>"
 
+        success_msg = f"Successfully opened file {file_name}, lines {start} to {end} are now visible in memory block <{file_name}>"
+        if closed_files:
+            success_msg += (
+                f"\nNote: Closed {len(closed_files)} least recently used file(s) due to open file limit: {', '.join(closed_files)}"
+            )
+
+        return success_msg
+
+    @trace_method
     async def close_file(self, agent_state: AgentState, file_name: str) -> str:
         """Stub for close_file tool."""
         await self.files_agents_manager.update_file_agent_by_name(
@@ -146,32 +159,52 @@ class LettaFileToolExecutor(ToolExecutor):
         except re.error as e:
             raise ValueError(f"Invalid regex pattern: {e}")
 
-    def _get_context_lines(self, text: str, file_metadata: FileMetadata, match_line_idx: int, total_lines: int) -> List[str]:
-        """Get context lines around a match using LineChunker."""
-        start_idx = max(0, match_line_idx - self.MAX_CONTEXT_LINES)
-        end_idx = min(total_lines, match_line_idx + self.MAX_CONTEXT_LINES + 1)
+    def _get_context_lines(
+        self,
+        formatted_lines: List[str],
+        match_line_num: int,
+        context_lines: int,
+    ) -> List[str]:
+        """Get context lines around a match from already-chunked lines.
 
-        # Use LineChunker to get formatted lines with numbers
-        chunker = LineChunker()
-        context_lines = chunker.chunk_text(text, file_metadata=file_metadata, start=start_idx, end=end_idx, add_metadata=False)
+        Args:
+            formatted_lines: Already chunked lines from LineChunker (format: "line_num: content")
+            match_line_num: The 1-based line number of the match
+            context_lines: Number of context lines before and after
+        """
+        if not formatted_lines or context_lines < 0:
+            return []
 
-        # Add match indicator
-        formatted_lines = []
-        for line in context_lines:
+        # Find the index of the matching line in the formatted_lines list
+        match_formatted_idx = None
+        for i, line in enumerate(formatted_lines):
             if line and ":" in line:
-                line_num_str = line.split(":")[0].strip()
                 try:
-                    line_num = int(line_num_str)
-                    prefix = ">" if line_num == match_line_idx + 1 else " "
-                    formatted_lines.append(f"{prefix} {line}")
+                    line_num = int(line.split(":", 1)[0].strip())
+                    if line_num == match_line_num:
+                        match_formatted_idx = i
+                        break
                 except ValueError:
-                    formatted_lines.append(f"  {line}")
-            else:
-                formatted_lines.append(f"  {line}")
+                    continue
 
-        return formatted_lines
+        if match_formatted_idx is None:
+            return []
 
-    async def grep(self, agent_state: AgentState, pattern: str, include: Optional[str] = None) -> str:
+        # Calculate context range with bounds checking
+        start_idx = max(0, match_formatted_idx - context_lines)
+        end_idx = min(len(formatted_lines), match_formatted_idx + context_lines + 1)
+
+        # Extract context lines and add match indicator
+        context_lines_with_indicator = []
+        for i in range(start_idx, end_idx):
+            line = formatted_lines[i]
+            prefix = ">" if i == match_formatted_idx else " "
+            context_lines_with_indicator.append(f"{prefix} {line}")
+
+        return context_lines_with_indicator
+
+    @trace_method
+    async def grep(self, agent_state: AgentState, pattern: str, include: Optional[str] = None, context_lines: Optional[int] = 3) -> str:
         """
         Search for pattern in all attached files and return matches with context.
 
@@ -179,6 +212,8 @@ class LettaFileToolExecutor(ToolExecutor):
             agent_state: Current agent state
             pattern: Regular expression pattern to search for
             include: Optional pattern to filter filenames to include in the search
+            context_lines (Optional[int]): Number of lines of context to show before and after each match.
+                                       Equivalent to `-C` in grep. Defaults to 3.
 
         Returns:
             Formatted string with search results, file names, line numbers, and context
@@ -229,10 +264,11 @@ class LettaFileToolExecutor(ToolExecutor):
         total_content_size = 0
         files_processed = 0
         files_skipped = 0
+        files_with_matches = set()  # Track files that had matches for LRU policy
 
         # Use asyncio timeout to prevent hanging
         async def _search_files():
-            nonlocal results, total_matches, total_content_size, files_processed, files_skipped
+            nonlocal results, total_matches, total_content_size, files_processed, files_skipped, files_with_matches
 
             for file_agent in file_agents:
                 # Load file content
@@ -268,12 +304,27 @@ class LettaFileToolExecutor(ToolExecutor):
 
                 # Use LineChunker to get all lines with proper formatting
                 chunker = LineChunker()
-                formatted_lines = chunker.chunk_text(file.content, file_metadata=file)
+                formatted_lines = chunker.chunk_text(file_metadata=file)
 
                 # Remove metadata header
                 if formatted_lines and formatted_lines[0].startswith("[Viewing"):
                     formatted_lines = formatted_lines[1:]
 
+                # Convert 0-based line numbers to 1-based for grep compatibility
+                corrected_lines = []
+                for line in formatted_lines:
+                    if line and ":" in line:
+                        try:
+                            line_parts = line.split(":", 1)
+                            line_num = int(line_parts[0].strip())
+                            line_content = line_parts[1] if len(line_parts) > 1 else ""
+                            corrected_lines.append(f"{line_num + 1}:{line_content}")
+                        except (ValueError, IndexError):
+                            corrected_lines.append(line)
+                    else:
+                        corrected_lines.append(line)
+                formatted_lines = corrected_lines
+
                 # Search for matches in formatted lines
                 for formatted_line in formatted_lines:
                     if total_matches >= self.MAX_TOTAL_MATCHES:
@@ -294,12 +345,13 @@ class LettaFileToolExecutor(ToolExecutor):
                             continue
 
                         if pattern_regex.search(line_content):
-                            # Get context around the match (convert back to 0-based indexing)
-                            context_lines = self._get_context_lines(file.content, file, line_num - 1, len(file.content.splitlines()))
+                            # Mark this file as having matches for LRU tracking
+                            files_with_matches.add(file.file_name)
+                            context = self._get_context_lines(formatted_lines, match_line_num=line_num, context_lines=context_lines or 0)
 
                             # Format the match result
                             match_header = f"\n=== {file.file_name}:{line_num} ==="
-                            match_content = "\n".join(context_lines)
+                            match_content = "\n".join(context)
                             results.append(f"{match_header}\n{match_content}")
 
                             file_matches += 1
@@ -312,6 +364,10 @@ class LettaFileToolExecutor(ToolExecutor):
         # Execute with timeout
         await asyncio.wait_for(_search_files(), timeout=self.GREP_TIMEOUT_SECONDS)
 
+        # Mark access for files that had matches
+        if files_with_matches:
+            await self.files_agents_manager.mark_access_bulk(agent_id=agent_state.id, file_names=list(files_with_matches), actor=self.actor)
+
         # Format final results
         if not results or total_matches == 0:
             summary = f"No matches found for pattern: '{pattern}'"
@@ -337,6 +393,7 @@ class LettaFileToolExecutor(ToolExecutor):
 
         return "\n".join(formatted_results)
 
+    @trace_method
     async def search_files(self, agent_state: AgentState, query: str, limit: int = 10) -> str:
         """
         Search for text within attached files using semantic search and return passages with their source filenames.
@@ -360,7 +417,13 @@ class LettaFileToolExecutor(ToolExecutor):
         self.logger.info(f"Semantic search started for agent {agent_state.id} with query '{query}' (limit: {limit})")
 
         # Get semantic search results
-        passages = await self.agent_manager.list_source_passages_async(actor=self.actor, agent_id=agent_state.id, query_text=query)
+        passages = await self.agent_manager.list_source_passages_async(
+            actor=self.actor,
+            agent_id=agent_state.id,
+            query_text=query,
+            embed_query=True,
+            embedding_config=agent_state.embedding_config,
+        )
 
         if not passages:
             return f"No semantic matches found for query: '{query}'"
@@ -401,6 +464,12 @@ class LettaFileToolExecutor(ToolExecutor):
                 passage_content = "\n".join(formatted_lines)
                 results.append(f"{passage_header}\n{passage_content}")
 
+        # Mark access for files that had matches
+        if files_with_passages:
+            matched_file_names = [name for name in files_with_passages.keys() if name != "Unknown File"]
+            if matched_file_names:
+                await self.files_agents_manager.mark_access_bulk(agent_id=agent_state.id, file_names=matched_file_names, actor=self.actor)
+
         # Create summary header
         file_count = len(files_with_passages)
         summary = f"Found {total_passages} semantic matches in {file_count} file{'s' if file_count != 1 else ''} for query: '{query}'"