AI: add ai_step to replace extract_content on rerun

browser_use/agent/prompts.py

@@ -493,3 +493,50 @@ def get_rerun_summary_message(prompt: str, screenshot_b64: str | None = None) ->
 	else:
 		# Without screenshot: use simple string content
 		return UserMessage(content=prompt)
+
+
+def get_ai_step_system_prompt() -> str:
+	"""
+	Get system prompt for AI step action used during rerun.
+
+	Returns:
+		System prompt string for AI step
+	"""
+	return """
+You are an expert at extracting data from webpages.
+
+<input>
+You will be given:
+1. A query describing what to extract
+2. The markdown of the webpage (filtered to remove noise)
+3. Optionally, a screenshot of the current page state
+</input>
+
+<instructions>
+- Extract information from the webpage that is relevant to the query
+- ONLY use the information available in the webpage - do not make up information
+- If the information is not available, mention that clearly
+- If the query asks for all items, list all of them
+</instructions>
+
+<output>
+- Present ALL relevant information in a concise way
+- Do not use conversational format - directly output the relevant information
+- If information is unavailable, state that clearly
+</output>
+""".strip()
+
+
+def get_ai_step_user_prompt(query: str, stats_summary: str, content: str) -> str:
+	"""
+	Build user prompt for AI step action.
+
+	Args:
+		query: What to extract or analyze
+		stats_summary: Content statistics summary
+		content: Page markdown content
+
+	Returns:
+		Formatted prompt string
+	"""
+	return f'<query>\n{query}\n</query>\n\n<content_stats>\n{stats_summary}\n</content_stats>\n\n<webpage_content>\n{content}\n</webpage_content>'

browser_use/agent/service.py

@@ -2261,6 +2261,111 @@ class Agent(Generic[Context, AgentStructuredOutput]):
 				long_term_memory=f'Rerun completed: {success_count} steps succeeded, {error_count} errors',
 			)
 
+	async def _execute_ai_step(
+		self,
+		query: str,
+		include_screenshot: bool = False,
+		extract_links: bool = False,
+		ai_step_llm: BaseChatModel | None = None,
+	) -> ActionResult:
+		"""
+		Execute an AI step during rerun to re-evaluate extract actions.
+		Analyzes full page DOM/markdown + optional screenshot.
+
+		Args:
+			query: What to analyze or extract from the current page
+			include_screenshot: Whether to include screenshot in analysis
+			extract_links: Whether to include links in markdown extraction
+			ai_step_llm: Optional LLM to use. If not provided, uses agent's LLM
+
+		Returns:
+			ActionResult with extracted content
+		"""
+		from browser_use.agent.prompts import get_ai_step_system_prompt, get_ai_step_user_prompt, get_rerun_summary_message
+		from browser_use.llm.messages import SystemMessage, UserMessage
+		from browser_use.utils import sanitize_surrogates
+
+		# Use provided LLM or agent's LLM
+		llm = ai_step_llm or self.llm
+		self.logger.debug(f'Using LLM for AI step: {llm.model}')
+
+		# Extract clean markdown
+		try:
+			from browser_use.dom.markdown_extractor import extract_clean_markdown
+
+			content, content_stats = await extract_clean_markdown(
+				browser_session=self.browser_session, extract_links=extract_links
+			)
+		except Exception as e:
+			return ActionResult(error=f'Could not extract clean markdown: {type(e).__name__}: {e}')
+
+		# Get screenshot if requested
+		screenshot_b64 = None
+		if include_screenshot:
+			try:
+				screenshot = await self.browser_session.take_screenshot(full_page=False)
+				if screenshot:
+					import base64
+
+					screenshot_b64 = base64.b64encode(screenshot).decode('utf-8')
+			except Exception as e:
+				self.logger.warning(f'Failed to capture screenshot for ai_step: {e}')
+
+		# Build prompt with content stats
+		original_html_length = content_stats['original_html_chars']
+		initial_markdown_length = content_stats['initial_markdown_chars']
+		final_filtered_length = content_stats['final_filtered_chars']
+		chars_filtered = content_stats['filtered_chars_removed']
+
+		stats_summary = f"""Content processed: {original_html_length:,} HTML chars → {initial_markdown_length:,} initial markdown → {final_filtered_length:,} filtered markdown"""
+		if chars_filtered > 0:
+			stats_summary += f' (filtered {chars_filtered:,} chars of noise)'
+
+		# Sanitize content
+		content = sanitize_surrogates(content)
+		query = sanitize_surrogates(query)
+
+		# Get prompts from prompts.py
+		system_prompt = get_ai_step_system_prompt()
+		prompt_text = get_ai_step_user_prompt(query, stats_summary, content)
+
+		# Build user message with optional screenshot
+		if screenshot_b64:
+			user_message = get_rerun_summary_message(prompt_text, screenshot_b64)
+		else:
+			user_message = UserMessage(content=prompt_text)
+
+		try:
+			import asyncio
+
+			response = await asyncio.wait_for(llm.ainvoke([SystemMessage(content=system_prompt), user_message]), timeout=120.0)
+
+			current_url = await self.browser_session.get_current_page_url()
+			extracted_content = (
+				f'<url>\n{current_url}\n</url>\n<query>\n{query}\n</query>\n<result>\n{response.completion}\n</result>'
+			)
+
+			# Simple memory handling
+			MAX_MEMORY_LENGTH = 1000
+			if len(extracted_content) < MAX_MEMORY_LENGTH:
+				memory = extracted_content
+				include_extracted_content_only_once = False
+			else:
+				file_name = await self.file_system.save_extracted_content(extracted_content)
+				memory = f'Query: {query}\nContent in {file_name} and once in <read_state>.'
+				include_extracted_content_only_once = True
+
+			self.logger.info(f'🤖 AI Step: {memory}')
+			return ActionResult(
+				extracted_content=extracted_content,
+				include_extracted_content_only_once=include_extracted_content_only_once,
+				long_term_memory=memory,
+			)
+		except Exception as e:
+			self.logger.warning(f'Failed to execute AI step: {e.__class__.__name__}: {e}')
+			self.logger.debug('Full error traceback:', exc_info=True)
+			return ActionResult(error=f'AI step failed: {e}')
+
 	async def rerun_history(
 		self,
 		history: AgentHistoryList,
@@ -2268,6 +2373,7 @@ class Agent(Generic[Context, AgentStructuredOutput]):
 		skip_failures: bool = True,
 		delay_between_actions: float = 2.0,
 		summary_llm: BaseChatModel | None = None,
+		ai_step_llm: BaseChatModel | None = None,
 	) -> list[ActionResult]:
 		"""
 		Rerun a saved history of actions with error handling and retry logic.
@@ -2278,6 +2384,7 @@ class Agent(Generic[Context, AgentStructuredOutput]):
 		                skip_failures: Whether to skip failed actions or stop execution
 		                delay_between_actions: Delay between actions in seconds
 		                summary_llm: Optional LLM to use for generating the final summary. If not provided, uses the agent's LLM
+		                ai_step_llm: Optional LLM to use for AI steps (extract actions). If not provided, uses the agent's LLM
 
 		Returns:
 		                List of action results (including AI summary as the final result)
@@ -2326,7 +2433,7 @@ class Agent(Generic[Context, AgentStructuredOutput]):
 			retry_count = 0
 			while retry_count < max_retries:
 				try:
-					result = await self._execute_history_step(history_item, step_delay)
+					result = await self._execute_history_step(history_item, step_delay, ai_step_llm)
 					results.extend(result)
 					break
 
@@ -2399,28 +2506,66 @@ class Agent(Generic[Context, AgentStructuredOutput]):
 			self.logger.debug('📝 Saved initial actions to history as step 0')
 			self.logger.debug('Initial actions completed')
 
-	async def _execute_history_step(self, history_item: AgentHistory, delay: float) -> list[ActionResult]:
-		"""Execute a single step from history with element validation"""
+	async def _execute_history_step(
+		self, history_item: AgentHistory, delay: float, ai_step_llm: BaseChatModel | None = None
+	) -> list[ActionResult]:
+		"""Execute a single step from history with element validation.
+
+		For extract actions, uses AI to re-evaluate the content since page content may have changed.
+		"""
 		assert self.browser_session is not None, 'BrowserSession is not set up'
 
 		await asyncio.sleep(delay)
 		state = await self.browser_session.get_browser_state_summary(include_screenshot=False)
 		if not state or not history_item.model_output:
 			raise ValueError('Invalid state or model output')
-		updated_actions = []
+
+		results = []
+		pending_actions = []
+
 		for i, action in enumerate(history_item.model_output.action):
-			updated_action = await self._update_action_indices(
-				history_item.state.interacted_element[i],
-				action,
-				state,
-			)
-			updated_actions.append(updated_action)
+			# Check if this is an extract action - use AI step instead
+			action_data = action.model_dump(exclude_unset=True)
+			action_name = next(iter(action_data.keys()), None)
 
-			if updated_action is None:
-				raise ValueError(f'Could not find matching element {i} in current page')
+			if action_name == 'extract':
+				# Execute any pending actions first to maintain correct order
+				# (e.g., if step is [click, extract], click must happen before extract)
+				if pending_actions:
+					batch_results = await self.multi_act(pending_actions)
+					results.extend(batch_results)
+					pending_actions = []
 
-		result = await self.multi_act(updated_actions)
-		return result
+				# Now execute AI step for extract action
+				extract_params = action_data['extract']
+				query = extract_params.get('query', '')
+				extract_links = extract_params.get('extract_links', False)
+
+				self.logger.info(f'🤖 Using AI step for extract action: {query[:50]}...')
+				ai_result = await self._execute_ai_step(
+					query=query,
+					include_screenshot=False,  # Match original extract behavior
+					extract_links=extract_links,
+					ai_step_llm=ai_step_llm,
+				)
+				results.append(ai_result)
+			else:
+				# For non-extract actions, update indices and collect for batch execution
+				updated_action = await self._update_action_indices(
+					history_item.state.interacted_element[i],
+					action,
+					state,
+				)
+				if updated_action is None:
+					raise ValueError(f'Could not find matching element {i} in current page')
+				pending_actions.append(updated_action)
+
+		# Execute any remaining pending actions
+		if pending_actions:
+			batch_results = await self.multi_act(pending_actions)
+			results.extend(batch_results)
+
+		return results
 
 	async def _update_action_indices(
 		self,

examples/features/rerun_history.py

@@ -18,17 +18,29 @@ Note: Initial actions (like opening URLs from tasks) are now automatically
 saved to history and will be replayed during rerun, so you don't need to
 worry about manually specifying URLs when rerunning.
 
-AI Summary:
-The rerun will automatically generate an AI summary at the end that analyzes
-the final screenshot and execution statistics.:
+AI Features During Rerun:
 
+1. AI Step for Extract Actions:
+   When an 'extract' action is replayed, the rerun automatically uses AI to
+   re-analyze the current page content (since it may have changed with new data).
+   This ensures the extracted content reflects the current state, not cached results.
+
+2. AI Summary:
+   At the end of the rerun, an AI summary analyzes the final screenshot and
+   execution statistics to determine success/failure.
+
+Custom LLM Usage:
 	# Option 1: Use agent's LLM (default)
 	results = await agent.load_and_rerun(history_file)
 
-	# Option 2: Use a specific LLM for summary generation
+	# Option 2: Use custom LLMs for AI steps and summary
 	from browser_use.llm import ChatOpenAI
-	summary_llm = ChatOpenAI(model='gpt-4.1-mini')
-	results = await agent.load_and_rerun(history_file, summary_llm=summary_llm)
+	custom_llm = ChatOpenAI(model='gpt-4.1-mini')
+	results = await agent.load_and_rerun(
+		history_file,
+		ai_step_llm=custom_llm,    # For extract action re-evaluation
+		summary_llm=custom_llm,     # For final summary
+	)
 
 The AI summary will be the last item in results and will have:
 	- extracted_content: The summary text
@@ -46,13 +58,16 @@ from browser_use.llm import ChatBrowserUse
 async def main():
 	# Example task to demonstrate history saving and rerunning
 	history_file = Path('agent_history.json')
-	task = 'Go to https://browser-use.github.io/stress-tests/challenges/reference-number-form.html and fill the form with example data and submit.'
+	task = 'Go to https://browser-use.github.io/stress-tests/challenges/reference-number-form.html and fill the form with example data and submit and extract the refernence number.'
 	llm = ChatBrowserUse()
 
-	# Optional: Use a custom LLM for AI summary generation
-	# Uncomment to use a custom LLM for summaries:
+	# Optional: Use custom LLMs for AI features during rerun
+	# Uncomment to use a custom LLM:
 	# from browser_use.llm import ChatOpenAI
-	# summary_llm = ChatOpenAI(model='gpt-4.1-mini')
+	# custom_llm = ChatOpenAI(model='gpt-4.1-mini')
+	# ai_step_llm = custom_llm   # For re-evaluating extract actions
+	# summary_llm = custom_llm   # For final summary
+	ai_step_llm = None  # Set to None to use agent's LLM (default)
 	summary_llm = None  # Set to None to use agent's LLM (default)
 
 	# Step 1: Run the agent and save history
@@ -100,17 +115,22 @@ async def main():
 				old_value = variables[var_name].original_value
 				print(f'  • {var_name}: "{old_value}" → "{new_value}"')
 
-			# Rerun with substituted values and optional custom summary LLM
-			substitute_agent = Agent(task='', llm=llm)
-			results = await substitute_agent.load_and_rerun(history_file, variables=new_values, summary_llm=summary_llm)
+		# Rerun with substituted values and optional custom LLMs
+		substitute_agent = Agent(task='', llm=llm)
+		results = await substitute_agent.load_and_rerun(
+			history_file,
+			variables=new_values,
+			ai_step_llm=ai_step_llm,  # For extract action re-evaluation
+			summary_llm=summary_llm,  # For final summary
+		)
 
-			# Display AI-generated summary (last result)
-			if results and results[-1].is_done:
-				summary = results[-1]
-				print('\n📊 AI Summary:')
-				print(f'  Summary: {summary.extracted_content}')
-				print(f'  Success: {summary.success}')
-			print('✓ History rerun with substituted values complete')
+		# Display AI-generated summary (last result)
+		if results and results[-1].is_done:
+			summary = results[-1]
+			print('\n📊 AI Summary:')
+			print(f'  Summary: {summary.extracted_content}')
+			print(f'  Success: {summary.success}')
+		print('✓ History rerun with substituted values complete')
 	else:
 		print('\n⚠️  No variables detected, skipping substitution rerun')
 

tests/ci/test_ai_step.py

@@ -0,0 +1,120 @@
+"""Tests for AI step private method used during rerun"""
+
+from unittest.mock import AsyncMock
+
+from browser_use.agent.service import Agent
+from browser_use.agent.views import ActionResult
+from tests.ci.conftest import create_mock_llm
+
+
+async def test_execute_ai_step_basic():
+	"""Test that _execute_ai_step extracts content with AI"""
+
+	# Create mock LLM that returns text response
+	async def custom_ainvoke(*args, **kwargs):
+		from browser_use.llm.views import ChatInvokeCompletion
+
+		return ChatInvokeCompletion(completion='Extracted: Test content from page', usage=None)
+
+	mock_llm = AsyncMock()
+	mock_llm.ainvoke.side_effect = custom_ainvoke
+	mock_llm.model = 'mock-model'
+
+	llm = create_mock_llm(actions=None)
+	agent = Agent(task='Test task', llm=llm)
+	await agent.browser_session.start()
+
+	try:
+		# Execute _execute_ai_step with mock LLM
+		result = await agent._execute_ai_step(
+			query='Extract the main heading',
+			include_screenshot=False,
+			extract_links=False,
+			ai_step_llm=mock_llm,
+		)
+
+		# Verify result
+		assert isinstance(result, ActionResult)
+		assert result.extracted_content is not None
+		assert 'Extracted: Test content from page' in result.extracted_content
+		assert result.long_term_memory is not None
+
+	finally:
+		await agent.close()
+
+
+async def test_execute_ai_step_with_screenshot():
+	"""Test that _execute_ai_step includes screenshot when requested"""
+
+	# Create mock LLM
+	async def custom_ainvoke(*args, **kwargs):
+		from browser_use.llm.views import ChatInvokeCompletion
+
+		# Verify that we received a message with image content
+		messages = args[0] if args else []
+		assert len(messages) >= 1, 'Should have at least one message'
+
+		# Check if any message has image content
+		has_image = False
+		for msg in messages:
+			if hasattr(msg, 'content') and isinstance(msg.content, list):
+				for part in msg.content:
+					if hasattr(part, 'type') and part.type == 'image_url':
+						has_image = True
+						break
+
+		assert has_image, 'Should include screenshot in message'
+		return ChatInvokeCompletion(completion='Extracted content with screenshot analysis', usage=None)
+
+	mock_llm = AsyncMock()
+	mock_llm.ainvoke.side_effect = custom_ainvoke
+	mock_llm.model = 'mock-model'
+
+	llm = create_mock_llm(actions=None)
+	agent = Agent(task='Test task', llm=llm)
+	await agent.browser_session.start()
+
+	try:
+		# Execute _execute_ai_step with screenshot
+		result = await agent._execute_ai_step(
+			query='Analyze this page',
+			include_screenshot=True,
+			extract_links=False,
+			ai_step_llm=mock_llm,
+		)
+
+		# Verify result
+		assert isinstance(result, ActionResult)
+		assert result.extracted_content is not None
+		assert 'Extracted content with screenshot analysis' in result.extracted_content
+
+	finally:
+		await agent.close()
+
+
+async def test_execute_ai_step_error_handling():
+	"""Test that _execute_ai_step handles errors gracefully"""
+	# Create mock LLM that raises an error
+	mock_llm = AsyncMock()
+	mock_llm.ainvoke.side_effect = Exception('LLM service unavailable')
+	mock_llm.model = 'mock-model'
+
+	llm = create_mock_llm(actions=None)
+	agent = Agent(task='Test task', llm=llm)
+	await agent.browser_session.start()
+
+	try:
+		# Execute _execute_ai_step - should return ActionResult with error
+		result = await agent._execute_ai_step(
+			query='Extract data',
+			include_screenshot=False,
+			ai_step_llm=mock_llm,
+		)
+
+		# Verify error is in result (not raised)
+		assert isinstance(result, ActionResult)
+		assert result.error is not None
+		assert 'AI step failed' in result.error
+
+	finally:
+		await agent.close()