From 0b61750a1cbd4f86a94fe923f946bde88b7e858c Mon Sep 17 00:00:00 2001
From: magmueller <mamagnus00@gmail.com>
Date: Tue, 4 Feb 2025 22:57:06 -0800
Subject: [PATCH] Ability to read files

---
 .gitignore                               |  3 ++
 browser_use/browser/context.py           |  2 +-
 examples/custom-functions/file_upload.py | 40 ++++++++++++++----------
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/.gitignore b/.gitignore
index 907bdff..dcdf281 100644
--- a/.gitignore
+++ b/.gitignore
@@ -181,3 +181,6 @@ gcp-login.json
 .vscode
 .ruff_cache
 .idea
+*.txt
+*.pdf
+*.csv
\ No newline at end of file
diff --git a/browser_use/browser/context.py b/browser_use/browser/context.py
index 06d3b1e..80cdbbb 100644
--- a/browser_use/browser/context.py
+++ b/browser_use/browser/context.py
@@ -1049,7 +1049,7 @@ class BrowserContext:
 		element_handle = await self.get_locate_element(selector_map[index])
 		return element_handle
 
-	async def get_dom_element_by_index(self, index: int) -> DOMElementNode | None:
+	async def get_dom_element_by_index(self, index: int) -> DOMElementNode:
 		selector_map = await self.get_selector_map()
 		return selector_map[index]
 
diff --git a/examples/custom-functions/file_upload.py b/examples/custom-functions/file_upload.py
index 3772c77..f1efdf6 100644
--- a/examples/custom-functions/file_upload.py
+++ b/examples/custom-functions/file_upload.py
@@ -27,7 +27,7 @@ controller = Controller()
 
 
 @controller.action(
-	'Upload file to element ',
+	'Upload file to interactive element with file path ',
 )
 async def upload_file(index: int, path: str, browser: BrowserContext, available_file_paths: list[str]):
 	if path not in available_file_paths:
@@ -38,47 +38,53 @@ async def upload_file(index: int, path: str, browser: BrowserContext, available_
 
 	dom_el = await browser.get_dom_element_by_index(index)
 
-	if dom_el is None:
-		return ActionResult(error=f'No element found at index {index}')
-
 	file_upload_dom_el = dom_el.get_file_upload_element()
 
 	if file_upload_dom_el is None:
-		logger.info(f'No file upload element found at index {index}')
-		return ActionResult(error=f'No file upload element found at index {index}')
+		msg = f'No file upload element found at index {index}'
+		logger.info(msg)
+		return ActionResult(error=msg)
 
 	file_upload_el = await browser.get_locate_element(file_upload_dom_el)
 
 	if file_upload_el is None:
-		logger.info(f'No file upload element found at index {index}')
-		return ActionResult(error=f'No file upload element found at index {index}')
+		msg = f'No file upload element found at index {index}'
+		logger.info(msg)
+		return ActionResult(error=msg)
 
 	try:
 		await file_upload_el.set_input_files(path)
 		msg = f'Successfully uploaded file to index {index}'
 		logger.info(msg)
-		return ActionResult(extracted_content=msg)
+		return ActionResult(extracted_content=msg, include_in_memory=True)
 	except Exception as e:
-		logger.debug(f'Error in set_input_files: {str(e)}')
-		return ActionResult(error=f'Failed to upload file to index {index}')
+		msg = f'Failed to upload file to index {index}: {str(e)}'
+		logger.info(msg)
+		return ActionResult(error=msg)
 
 
-@controller.action('Close file dialog')
-async def close_file_dialog(browser: BrowserContext):
-	page = await browser.get_current_page()
-	await page.keyboard.press('Escape')
+@controller.action('Read the file content of a file given a path')
+async def read_file(path: str, available_file_paths: list[str]):
+	if path not in available_file_paths:
+		return ActionResult(error=f'File path {path} is not available')
+
+	with open(path, 'r') as f:
+		content = f.read()
+	msg = f'File content: {content}'
+	logger.info(msg)
+	return ActionResult(extracted_content=msg, include_in_memory=True)
 
 
 def create_file(file_type: str = 'txt'):
 	with open(f'tmp.{file_type}', 'w') as f:
 		f.write('test')
 	file_path = Path.cwd() / f'tmp.{file_type}'
-	print(f'Created file: {file_path}')
+	logger.info(f'Created file: {file_path}')
 	return str(file_path)
 
 
 async def main():
-	task = f'Go to https://kzmpmkh2zfk1ojnpxfn1.lite.vusercontent.net/ and upload one file to each upload field'
+	task = f'Go to https://kzmpmkh2zfk1ojnpxfn1.lite.vusercontent.net/ and - read the file content and upload them to fields'
 
 	available_file_paths = [create_file('txt'), create_file('pdf'), create_file('csv')]