letta-nightly 0.8.16.dev20250721104533__tar.gz → 0.8.17.dev20250721222456__tar.gz

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/PKG-INFO

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

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/__init__.py

@@ -5,7 +5,7 @@ try:
     __version__ = version("letta")
 except PackageNotFoundError:
     # Fallback for development installations
-    __version__ = "0.8.16"
+    __version__ = "0.8.17"
 
 if os.environ.get("LETTA_VERSION"):
     __version__ = os.environ["LETTA_VERSION"]

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/schemas/letta_message.py

@@ -63,13 +63,13 @@ class LettaMessage(BaseModel):
             dt = dt.replace(tzinfo=timezone.utc)
         return dt.isoformat(timespec="seconds")
 
-    @field_serializer("is_err", when_used="unless-none")
-    def serialize_is_err(self, value: bool | None, _info):
+    @field_serializer("is_err", mode="wrap")
+    def serialize_is_err(self, value: bool | None, handler, _info):
         """
         Only serialize is_err field when it's True (for debugging purposes).
         When is_err is None or False, this field will be excluded from the JSON output.
         """
-        return value if value is True else None
+        return handler(value) if value is True else None
 
 
 class SystemMessage(LettaMessage):

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/schemas/message.py

@@ -269,6 +269,7 @@ class Message(BaseMessage):
                             otid=otid,
                             sender_id=self.sender_id,
                             step_id=self.step_id,
+                            is_err=self.is_err,
                         )
                     )
                 # Otherwise, we may have a list of multiple types
@@ -287,6 +288,7 @@ class Message(BaseMessage):
                                     otid=otid,
                                     sender_id=self.sender_id,
                                     step_id=self.step_id,
+                                    is_err=self.is_err,
                                 )
                             )
                         elif isinstance(content_part, ReasoningContent):
@@ -301,6 +303,7 @@ class Message(BaseMessage):
                                     name=self.name,
                                     otid=otid,
                                     step_id=self.step_id,
+                                    is_err=self.is_err,
                                 )
                             )
                         elif isinstance(content_part, RedactedReasoningContent):
@@ -315,6 +318,7 @@ class Message(BaseMessage):
                                     otid=otid,
                                     sender_id=self.sender_id,
                                     step_id=self.step_id,
+                                    is_err=self.is_err,
                                 )
                             )
                         elif isinstance(content_part, OmittedReasoningContent):
@@ -328,6 +332,7 @@ class Message(BaseMessage):
                                     name=self.name,
                                     otid=otid,
                                     step_id=self.step_id,
+                                    is_err=self.is_err,
                                 )
                             )
                         else:
@@ -355,6 +360,7 @@ class Message(BaseMessage):
                                 otid=otid,
                                 sender_id=self.sender_id,
                                 step_id=self.step_id,
+                                is_err=self.is_err,
                             )
                         )
                     else:
@@ -371,6 +377,7 @@ class Message(BaseMessage):
                                 otid=otid,
                                 sender_id=self.sender_id,
                                 step_id=self.step_id,
+                                is_err=self.is_err,
                             )
                         )
         elif self.role == MessageRole.tool:
@@ -416,6 +423,7 @@ class Message(BaseMessage):
                     otid=Message.generate_otid_from_id(self.id, len(messages)),
                     sender_id=self.sender_id,
                     step_id=self.step_id,
+                    is_err=self.is_err,
                 )
             )
         elif self.role == MessageRole.user:
@@ -437,6 +445,7 @@ class Message(BaseMessage):
                     otid=self.otid,
                     sender_id=self.sender_id,
                     step_id=self.step_id,
+                    is_err=self.is_err,
                 )
             )
         elif self.role == MessageRole.system:

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/server/db.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 import threading
 import time
@@ -116,6 +117,13 @@ class DatabaseRegistry:
         self.config = LettaConfig.load()
         self.logger = get_logger(__name__)
 
+        if settings.db_max_concurrent_sessions:
+            self._db_semaphore = asyncio.Semaphore(settings.db_max_concurrent_sessions)
+            self.logger.info(f"Initialized database throttling with max {settings.db_max_concurrent_sessions} concurrent sessions")
+        else:
+            self.logger.info("Database throttling is disabled")
+            self._db_semaphore = None
+
     def initialize_sync(self, force: bool = False) -> None:
         """Initialize the synchronous database engine if not already initialized."""
         with self._lock:
@@ -364,16 +372,33 @@ class DatabaseRegistry:
     @trace_method
     @asynccontextmanager
     async def async_session(self, name: str = "default") -> AsyncGenerator[AsyncSession, None]:
-        """Async context manager for database sessions."""
-        session_factory = self.get_async_session_factory(name)
-        if not session_factory:
-            raise ValueError(f"No async session factory found for '{name}' or async database is not configured")
+        """Async context manager for database sessions with throttling."""
+        if self._db_semaphore:
+            async with self._db_semaphore:
+                session_factory = self.get_async_session_factory(name)
+                if not session_factory:
+                    raise ValueError(f"No async session factory found for '{name}' or async database is not configured")
+
+                session = session_factory()
+                try:
+                    yield session
+                finally:
+                    await session.close()
+        else:
+            session_factory = self.get_async_session_factory(name)
+            if not session_factory:
+                raise ValueError(f"No async session factory found for '{name}' or async database is not configured")
 
-        session = session_factory()
-        try:
-            yield session
-        finally:
-            await session.close()
+            session = session_factory()
+            try:
+                yield session
+            finally:
+                await session.close()
+
+    @trace_method
+    def session_caller_trace(self, caller_info: str):
+        """Trace sync db caller information for debugging purposes."""
+        pass  # wrapper used for otel tracing only
 
     @trace_method
     def session_caller_trace(self, caller_info: str):

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/server/rest_api/streaming_response.py

@@ -125,14 +125,7 @@ class StreamingResponseWithStatusCode(StreamingResponse):
                         if not isinstance(content, bytes):
                             content = content.encode(self.charset)
                         more_body = False
-                        await send(
-                            {
-                                "type": "http.response.body",
-                                "body": content,
-                                "more_body": more_body,
-                            }
-                        )
-                        raise Exception(f"An exception occurred mid-stream with status code {status_code}", detail={"content": content})
+                        raise Exception(f"An exception occurred mid-stream with status code {status_code} with content {content}")
                 else:
                     content = chunk
 

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/server/rest_api/utils.py

@@ -120,22 +120,22 @@ async def sse_async_generator(
                 yield sse_formatter(usage.model_dump(exclude={"steps_messages"}))
 
             except ContextWindowExceededError as e:
-                log_error_to_sentry(e)
+                capture_sentry_exception(e)
                 logger.error(f"ContextWindowExceededError error: {e}")
                 yield sse_formatter({"error": f"Stream failed: {e}", "code": str(e.code.value) if e.code else None})
 
             except RateLimitExceededError as e:
-                log_error_to_sentry(e)
+                capture_sentry_exception(e)
                 logger.error(f"RateLimitExceededError error: {e}")
                 yield sse_formatter({"error": f"Stream failed: {e}", "code": str(e.code.value) if e.code else None})
 
             except Exception as e:
-                log_error_to_sentry(e)
+                capture_sentry_exception(e)
                 logger.error(f"Caught unexpected Exception: {e}")
                 yield sse_formatter({"error": "Stream failed (internal error occurred)"})
 
     except Exception as e:
-        log_error_to_sentry(e)
+        capture_sentry_exception(e)
         logger.error(f"Caught unexpected Exception: {e}")
         yield sse_formatter({"error": "Stream failed (decoder encountered an error)"})
 

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/services/agent_file_manager.py

@@ -149,7 +149,7 @@ class AgentFileManager:
         return sorted(unique_blocks.values(), key=lambda x: x.label)
 
     async def _extract_unique_sources_and_files_from_agents(
-        self, agent_states: List[AgentState], actor: User
+        self, agent_states: List[AgentState], actor: User, files_agents_cache: dict = None
     ) -> tuple[List[Source], List[FileMetadata]]:
         """Extract unique sources and files from agent states using bulk operations"""
 
@@ -160,6 +160,10 @@ class AgentFileManager:
             files_agents = await self.file_agent_manager.list_files_for_agent(
                 agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False
             )
+            # cache the results for reuse during conversion
+            if files_agents_cache is not None:
+                files_agents_cache[agent_state.id] = files_agents
+
             for file_agent in files_agents:
                 all_source_ids.add(file_agent.source_id)
                 all_file_ids.add(file_agent.file_id)
@@ -168,13 +172,18 @@ class AgentFileManager:
 
         return sources, files
 
-    async def _convert_agent_state_to_schema(self, agent_state: AgentState, actor: User) -> AgentSchema:
+    async def _convert_agent_state_to_schema(self, agent_state: AgentState, actor: User, files_agents_cache: dict = None) -> AgentSchema:
         """Convert AgentState to AgentSchema with ID remapping"""
 
         agent_file_id = self._map_db_to_file_id(agent_state.id, AgentSchema.__id_prefix__)
-        files_agents = await self.file_agent_manager.list_files_for_agent(
-            agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False
-        )
+
+        # use cached file-agent data if available, otherwise fetch
+        if files_agents_cache is not None and agent_state.id in files_agents_cache:
+            files_agents = files_agents_cache[agent_state.id]
+        else:
+            files_agents = await self.file_agent_manager.list_files_for_agent(
+                agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False
+            )
         agent_schema = await AgentSchema.from_agent_state(
             agent_state, message_manager=self.message_manager, files_agents=files_agents, actor=actor
         )
@@ -264,15 +273,21 @@ class AgentFileManager:
                 missing_ids = [agent_id for agent_id in agent_ids if agent_id not in found_ids]
                 raise AgentFileExportError(f"The following agent IDs were not found: {missing_ids}")
 
+            # cache for file-agent relationships to avoid duplicate queries
+            files_agents_cache = {}  # Maps agent_id to list of file_agent relationships
+
             # Extract unique entities across all agents
             tool_set = self._extract_unique_tools(agent_states)
             block_set = self._extract_unique_blocks(agent_states)
 
-            # Extract sources and files from agent states BEFORE conversion
-            source_set, file_set = await self._extract_unique_sources_and_files_from_agents(agent_states, actor)
+            # Extract sources and files from agent states BEFORE conversion (with caching)
+            source_set, file_set = await self._extract_unique_sources_and_files_from_agents(agent_states, actor, files_agents_cache)
 
-            # Convert to schemas with ID remapping
-            agent_schemas = [await self._convert_agent_state_to_schema(agent_state, actor=actor) for agent_state in agent_states]
+            # Convert to schemas with ID remapping (reusing cached file-agent data)
+            agent_schemas = [
+                await self._convert_agent_state_to_schema(agent_state, actor=actor, files_agents_cache=files_agents_cache)
+                for agent_state in agent_states
+            ]
             tool_schemas = [self._convert_tool_to_schema(tool) for tool in tool_set]
             block_schemas = [self._convert_block_to_schema(block) for block in block_set]
             source_schemas = [self._convert_source_to_schema(source) for source in source_set]
@@ -332,32 +347,66 @@ class AgentFileManager:
             # Import in dependency order
             imported_count = 0
             file_to_db_ids = {}  # Maps file IDs to new database IDs
-
-            # 1. Create tools first (no dependencies)
-            for tool_schema in schema.tools:
-                # Convert ToolSchema back to ToolCreate
-                created_tool = await self.tool_manager.create_or_update_tool_async(
-                    pydantic_tool=Tool(**tool_schema.model_dump(exclude={"id"})), actor=actor
-                )
-                file_to_db_ids[tool_schema.id] = created_tool.id
-                imported_count += 1
-
-            # 2. Create blocks (no dependencies)
-            for block_schema in schema.blocks:
-                # Convert BlockSchema back to CreateBlock
-                block = Block(**block_schema.model_dump(exclude={"id"}))
-                created_block = await self.block_manager.create_or_update_block_async(block, actor)
-                file_to_db_ids[block_schema.id] = created_block.id
-                imported_count += 1
-
-            # 3. Create sources (no dependencies)
-            for source_schema in schema.sources:
-                # Convert SourceSchema back to Source
-                source_data = source_schema.model_dump(exclude={"id", "embedding", "embedding_chunk_size"})
-                source = Source(**source_data)
-                created_source = await self.source_manager.create_source(source, actor)
-                file_to_db_ids[source_schema.id] = created_source.id
-                imported_count += 1
+            # in-memory cache for file metadata to avoid repeated db calls
+            file_metadata_cache = {}  # Maps database file ID to FileMetadata
+
+            # 1. Create tools first (no dependencies) - using bulk upsert for efficiency
+            if schema.tools:
+                # convert tool schemas to pydantic tools
+                pydantic_tools = []
+                for tool_schema in schema.tools:
+                    pydantic_tools.append(Tool(**tool_schema.model_dump(exclude={"id"})))
+
+                # bulk upsert all tools at once
+                created_tools = await self.tool_manager.bulk_upsert_tools_async(pydantic_tools, actor)
+
+                # map file ids to database ids
+                # note: tools are matched by name during upsert, so we need to match by name here too
+                created_tools_by_name = {tool.name: tool for tool in created_tools}
+                for tool_schema in schema.tools:
+                    created_tool = created_tools_by_name.get(tool_schema.name)
+                    if created_tool:
+                        file_to_db_ids[tool_schema.id] = created_tool.id
+                        imported_count += 1
+                    else:
+                        logger.warning(f"Tool {tool_schema.name} was not created during bulk upsert")
+
+            # 2. Create blocks (no dependencies) - using batch create for efficiency
+            if schema.blocks:
+                # convert block schemas to pydantic blocks (excluding IDs to create new blocks)
+                pydantic_blocks = []
+                for block_schema in schema.blocks:
+                    pydantic_blocks.append(Block(**block_schema.model_dump(exclude={"id"})))
+
+                # batch create all blocks at once
+                created_blocks = await self.block_manager.batch_create_blocks_async(pydantic_blocks, actor)
+
+                # map file ids to database ids
+                for block_schema, created_block in zip(schema.blocks, created_blocks):
+                    file_to_db_ids[block_schema.id] = created_block.id
+                    imported_count += 1
+
+            # 3. Create sources (no dependencies) - using bulk upsert for efficiency
+            if schema.sources:
+                # convert source schemas to pydantic sources
+                pydantic_sources = []
+                for source_schema in schema.sources:
+                    source_data = source_schema.model_dump(exclude={"id", "embedding", "embedding_chunk_size"})
+                    pydantic_sources.append(Source(**source_data))
+
+                # bulk upsert all sources at once
+                created_sources = await self.source_manager.bulk_upsert_sources_async(pydantic_sources, actor)
+
+                # map file ids to database ids
+                # note: sources are matched by name during upsert, so we need to match by name here too
+                created_sources_by_name = {source.name: source for source in created_sources}
+                for source_schema in schema.sources:
+                    created_source = created_sources_by_name.get(source_schema.name)
+                    if created_source:
+                        file_to_db_ids[source_schema.id] = created_source.id
+                        imported_count += 1
+                    else:
+                        logger.warning(f"Source {source_schema.name} was not created during bulk upsert")
 
             # 4. Create files (depends on sources)
             for file_schema in schema.files:
@@ -383,8 +432,10 @@ class AgentFileManager:
                     file_db_id = file_to_db_ids[file_schema.id]
                     source_db_id = file_to_db_ids[file_schema.source_id]
 
-                    # Get the created file metadata
-                    file_metadata = await self.file_manager.get_file_by_id(file_db_id, actor)
+                    # Get the created file metadata (with caching)
+                    if file_db_id not in file_metadata_cache:
+                        file_metadata_cache[file_db_id] = await self.file_manager.get_file_by_id(file_db_id, actor)
+                    file_metadata = file_metadata_cache[file_db_id]
 
                     # Save the db call of fetching content again
                     file_metadata.content = file_schema.content
@@ -447,7 +498,11 @@ class AgentFileManager:
 
                     for file_agent_schema in agent_schema.files_agents:
                         file_db_id = file_to_db_ids[file_agent_schema.file_id]
-                        file_metadata = await self.file_manager.get_file_by_id(file_db_id, actor)
+
+                        # Use cached file metadata if available
+                        if file_db_id not in file_metadata_cache:
+                            file_metadata_cache[file_db_id] = await self.file_manager.get_file_by_id(file_db_id, actor)
+                        file_metadata = file_metadata_cache[file_db_id]
                         files_for_agent.append(file_metadata)
 
                         if file_agent_schema.visible_content:
@@ -530,9 +585,16 @@ class AgentFileManager:
             entity_ids = [entity.id for entity in entities]
 
             # Check for duplicates within this entity type
-            if len(entity_ids) != len(set(entity_ids)):
-                duplicates = [id for id in entity_ids if entity_ids.count(id) > 1]
-                errors.append(f"Duplicate {entity_type} IDs found: {set(duplicates)}")
+            seen = set()
+            duplicates = set()
+            for entity_id in entity_ids:
+                if entity_id in seen:
+                    duplicates.add(entity_id)
+                else:
+                    seen.add(entity_id)
+
+            if duplicates:
+                errors.append(f"Duplicate {entity_type} IDs found: {duplicates}")
 
             # Check for duplicates across all entity types
             for entity_id in entity_ids:
@@ -545,9 +607,16 @@ class AgentFileManager:
             message_ids = [msg.id for msg in agent.messages]
 
             # Check for duplicates within agent messages
-            if len(message_ids) != len(set(message_ids)):
-                duplicates = [id for id in message_ids if message_ids.count(id) > 1]
-                errors.append(f"Duplicate message IDs in agent {agent.id}: {set(duplicates)}")
+            seen = set()
+            duplicates = set()
+            for message_id in message_ids:
+                if message_id in seen:
+                    duplicates.add(message_id)
+                else:
+                    seen.add(message_id)
+
+            if duplicates:
+                errors.append(f"Duplicate message IDs in agent {agent.id}: {duplicates}")
 
             # Check for duplicates across all entity types
             for message_id in message_ids:

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/services/source_manager.py

@@ -58,6 +58,123 @@ class SourceManager:
                 await source.create_async(session, actor=actor)
                 return source.to_pydantic()
 
+    @enforce_types
+    @trace_method
+    async def bulk_upsert_sources_async(self, pydantic_sources: List[PydanticSource], actor: PydanticUser) -> List[PydanticSource]:
+        """
+        Bulk create or update multiple sources in a single database transaction.
+
+        Uses optimized PostgreSQL bulk upsert when available, falls back to individual
+        upserts for SQLite. This is much more efficient than calling create_source
+        in a loop.
+
+        IMPORTANT BEHAVIOR NOTES:
+        - Sources are matched by (name, organization_id) unique constraint, NOT by ID
+        - If a source with the same name already exists for the organization, it will be updated
+          regardless of any ID provided in the input source
+        - The existing source's ID is preserved during updates
+        - If you provide a source with an explicit ID but a name that matches an existing source,
+          the existing source will be updated and the provided ID will be ignored
+        - This matches the behavior of create_source which also checks by ID first
+
+        PostgreSQL optimization:
+        - Uses native ON CONFLICT (name, organization_id) DO UPDATE for atomic upserts
+        - All sources are processed in a single SQL statement for maximum efficiency
+
+        SQLite fallback:
+        - Falls back to individual create_source calls
+        - Still benefits from batched transaction handling
+
+        Args:
+            pydantic_sources: List of sources to create or update
+            actor: User performing the action
+
+        Returns:
+            List of created/updated sources
+        """
+        if not pydantic_sources:
+            return []
+
+        from letta.settings import settings
+
+        if settings.letta_pg_uri_no_default:
+            # use optimized postgresql bulk upsert
+            async with db_registry.async_session() as session:
+                return await self._bulk_upsert_postgresql(session, pydantic_sources, actor)
+        else:
+            # fallback to individual upserts for sqlite
+            return await self._upsert_sources_individually(pydantic_sources, actor)
+
+    @trace_method
+    async def _bulk_upsert_postgresql(self, session, source_data_list: List[PydanticSource], actor: PydanticUser) -> List[PydanticSource]:
+        """Hyper-optimized PostgreSQL bulk upsert using ON CONFLICT DO UPDATE."""
+        from sqlalchemy import func, select
+        from sqlalchemy.dialects.postgresql import insert
+
+        # prepare data for bulk insert
+        table = SourceModel.__table__
+        valid_columns = {col.name for col in table.columns}
+
+        insert_data = []
+        for source in source_data_list:
+            source_dict = source.model_dump(to_orm=True)
+            # set created/updated by fields
+
+            if actor:
+                source_dict["_created_by_id"] = actor.id
+                source_dict["_last_updated_by_id"] = actor.id
+                source_dict["organization_id"] = actor.organization_id
+
+            # filter to only include columns that exist in the table
+            filtered_dict = {k: v for k, v in source_dict.items() if k in valid_columns}
+            insert_data.append(filtered_dict)
+
+        # use postgresql's native bulk upsert
+        stmt = insert(table).values(insert_data)
+
+        # on conflict, update all columns except id, created_at, and _created_by_id
+        excluded = stmt.excluded
+        update_dict = {}
+        for col in table.columns:
+            if col.name not in ("id", "created_at", "_created_by_id"):
+                if col.name == "updated_at":
+                    update_dict[col.name] = func.now()
+                else:
+                    update_dict[col.name] = excluded[col.name]
+
+        upsert_stmt = stmt.on_conflict_do_update(index_elements=["name", "organization_id"], set_=update_dict)
+
+        await session.execute(upsert_stmt)
+        await session.commit()
+
+        # fetch results
+        source_names = [source.name for source in source_data_list]
+        result_query = select(SourceModel).where(
+            SourceModel.name.in_(source_names), SourceModel.organization_id == actor.organization_id, SourceModel.is_deleted == False
+        )
+        result = await session.execute(result_query)
+        return [source.to_pydantic() for source in result.scalars()]
+
+    @trace_method
+    async def _upsert_sources_individually(self, source_data_list: List[PydanticSource], actor: PydanticUser) -> List[PydanticSource]:
+        """Fallback to individual upserts for SQLite."""
+        sources = []
+        for source in source_data_list:
+            # try to get existing source by name
+            existing_source = await self.get_source_by_name(source.name, actor)
+            if existing_source:
+                # update existing source
+                from letta.schemas.source import SourceUpdate
+
+                update_data = source.model_dump(exclude={"id"}, exclude_none=True)
+                updated_source = await self.update_source(existing_source.id, SourceUpdate(**update_data), actor)
+                sources.append(updated_source)
+            else:
+                # create new source
+                created_source = await self.create_source(source, actor)
+                sources.append(created_source)
+        return sources
+
     @enforce_types
     @trace_method
     async def update_source(self, source_id: str, source_update: SourceUpdate, actor: PydanticUser) -> PydanticSource:

{letta_nightly-0.8.16.dev20250721104533 → letta_nightly-0.8.17.dev20250721222456}/letta/services/tool_executor/builtin_tool_executor.py

@@ -176,7 +176,9 @@ class LettaBuiltinToolExecutor(ToolExecutor):
         app = AsyncFirecrawlApp(api_key=firecrawl_api_key)
 
         # Process all search tasks in parallel
-        search_task_coroutines = [self._process_single_search_task(app, task, limit, return_raw, api_key_source) for task in search_tasks]
+        search_task_coroutines = [
+            self._process_single_search_task(app, task, limit, return_raw, api_key_source, agent_state) for task in search_tasks
+        ]
 
         # Execute all searches concurrently
         search_results = await asyncio.gather(*search_task_coroutines, return_exceptions=True)
@@ -205,7 +207,7 @@ class LettaBuiltinToolExecutor(ToolExecutor):
 
     @trace_method
     async def _process_single_search_task(
-        self, app: "AsyncFirecrawlApp", task: SearchTask, limit: int, return_raw: bool, api_key_source: str
+        self, app: "AsyncFirecrawlApp", task: SearchTask, limit: int, return_raw: bool, api_key_source: str, agent_state: "AgentState"
     ) -> Dict[str, Any]:
         """Process a single search task."""
         from firecrawl import ScrapeOptions
@@ -246,7 +248,9 @@ class LettaBuiltinToolExecutor(ToolExecutor):
                 for result in search_result.get("data"):
                     if result.get("markdown"):
                         # Create async task for OpenAI analysis
-                        analysis_task = self._analyze_document_with_openai(client, result["markdown"], task.query, task.question)
+                        analysis_task = self._analyze_document_with_openai(
+                            client, result["markdown"], task.query, task.question, agent_state
+                        )
                         analysis_tasks.append(analysis_task)
                         results_with_markdown.append(result)
                     else:
@@ -300,7 +304,9 @@ class LettaBuiltinToolExecutor(ToolExecutor):
         return {"query": task.query, "question": task.question, "raw_results": search_result}
 
     @trace_method
-    async def _analyze_document_with_openai(self, client, markdown_content: str, query: str, question: str) -> Optional[DocumentAnalysis]:
+    async def _analyze_document_with_openai(
+        self, client, markdown_content: str, query: str, question: str, agent_state: "AgentState"
+    ) -> Optional[DocumentAnalysis]:
         """Use OpenAI to analyze a document and extract relevant passages using line numbers."""
         original_length = len(markdown_content)
 
@@ -324,7 +330,9 @@ class LettaBuiltinToolExecutor(ToolExecutor):
         # Time the OpenAI request
         start_time = time.time()
 
-        model = os.getenv(WEB_SEARCH_MODEL_ENV_VAR_NAME, WEB_SEARCH_MODEL_ENV_VAR_DEFAULT_VALUE)
+        # Check agent state env vars first, then fall back to os.getenv
+        agent_state_tool_env_vars = agent_state.get_agent_env_vars_as_dict()
+        model = agent_state_tool_env_vars.get(WEB_SEARCH_MODEL_ENV_VAR_NAME) or WEB_SEARCH_MODEL_ENV_VAR_DEFAULT_VALUE
         logger.info(f"Using model {model} for web search result parsing")
         response = await client.beta.chat.completions.parse(
             model=model,