agentskills.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881
  1. """
  2. agentskills.py
  3. This module provides various file manipulation skills for the OpenDevin agent.
  4. Functions:
  5. - open_file(path: str, line_number: int | None = 1, context_lines: int = 100): Opens a file and optionally moves to a specific line.
  6. - goto_line(line_number): Moves the window to show the specified line number.
  7. - scroll_down(): Moves the window down by the number of lines specified in WINDOW.
  8. - scroll_up(): Moves the window up by the number of lines specified in WINDOW.
  9. - create_file(filename): Creates and opens a new file with the given name.
  10. - search_dir(search_term, dir_path='./'): Searches for a term in all files in the specified directory.
  11. - search_file(search_term, file_path=None): Searches for a term in the specified file or the currently open file.
  12. - find_file(file_name, dir_path='./'): Finds all files with the given name in the specified directory.
  13. - edit_file(file_name, start, end, content): Replaces lines in a file with the given content.
  14. - append_file(file_name, content): Appends given content to a file.
  15. """
  16. import base64
  17. import functools
  18. import os
  19. import shutil
  20. import subprocess
  21. import tempfile
  22. from inspect import signature
  23. from typing import Optional
  24. import docx
  25. import PyPDF2
  26. from openai import OpenAI
  27. from pptx import Presentation
  28. from pylatexenc.latex2text import LatexNodes2Text
  29. CURRENT_FILE: str | None = None
  30. CURRENT_LINE = 1
  31. WINDOW = 100
  32. ENABLE_AUTO_LINT = os.getenv('ENABLE_AUTO_LINT', 'false').lower() == 'true'
  33. # This is also used in unit tests!
  34. MSG_FILE_UPDATED = '[File updated. Please review the changes and make sure they are correct (correct indentation, no duplicate lines, etc). Edit the file again if necessary.]'
  35. # OPENAI
  36. OPENAI_API_KEY = os.getenv(
  37. 'OPENAI_API_KEY', os.getenv('SANDBOX_ENV_OPENAI_API_KEY', '')
  38. )
  39. OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL', 'https://api.openai.com/v1')
  40. OPENAI_MODEL = os.getenv('OPENAI_MODEL', 'gpt-4o-2024-05-13')
  41. MAX_TOKEN = os.getenv('MAX_TOKEN', 500)
  42. OPENAI_PROXY = f'{OPENAI_BASE_URL}/chat/completions'
  43. client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
  44. # Define the decorator using the functionality of UpdatePwd
  45. def update_pwd_decorator(func):
  46. @functools.wraps(func)
  47. def wrapper(*args, **kwargs):
  48. old_pwd = os.getcwd()
  49. jupyter_pwd = os.environ.get('JUPYTER_PWD', None)
  50. if jupyter_pwd:
  51. os.chdir(jupyter_pwd)
  52. try:
  53. return func(*args, **kwargs)
  54. finally:
  55. os.chdir(old_pwd)
  56. return wrapper
  57. def _is_valid_filename(file_name) -> bool:
  58. if not file_name or not isinstance(file_name, str) or not file_name.strip():
  59. return False
  60. invalid_chars = '<>:"/\\|?*'
  61. if os.name == 'nt': # Windows
  62. invalid_chars = '<>:"/\\|?*'
  63. elif os.name == 'posix': # Unix-like systems
  64. invalid_chars = '\0'
  65. for char in invalid_chars:
  66. if char in file_name:
  67. return False
  68. return True
  69. def _is_valid_path(path) -> bool:
  70. if not path or not isinstance(path, str):
  71. return False
  72. try:
  73. return os.path.exists(os.path.normpath(path))
  74. except PermissionError:
  75. return False
  76. def _create_paths(file_name) -> bool:
  77. try:
  78. dirname = os.path.dirname(file_name)
  79. if dirname:
  80. os.makedirs(dirname, exist_ok=True)
  81. return True
  82. except PermissionError:
  83. return False
  84. def _check_current_file(file_path: str | None = None) -> bool:
  85. global CURRENT_FILE
  86. if not file_path:
  87. file_path = CURRENT_FILE
  88. if not file_path or not os.path.isfile(file_path):
  89. raise ValueError('No file open. Use the open_file function first.')
  90. return True
  91. def _clamp(value, min_value, max_value):
  92. return max(min_value, min(value, max_value))
  93. def _lint_file(file_path: str) -> tuple[Optional[str], Optional[int]]:
  94. """
  95. Lint the file at the given path and return a tuple with a boolean indicating if there are errors,
  96. and the line number of the first error, if any.
  97. Returns:
  98. tuple[str, Optional[int]]: (lint_error, first_error_line_number)
  99. """
  100. if file_path.endswith('.py'):
  101. # Define the flake8 command with selected error codes
  102. command = [
  103. 'flake8',
  104. '--isolated',
  105. '--select=F821,F822,F831,E112,E113,E999,E902',
  106. file_path,
  107. ]
  108. # Run the command using subprocess and redirect stderr to stdout
  109. result = subprocess.run(
  110. command,
  111. stdout=subprocess.PIPE,
  112. stderr=subprocess.STDOUT,
  113. )
  114. if result.returncode == 0:
  115. # Linting successful. No issues found.
  116. return None, None
  117. # Extract the line number from the first error message
  118. error_message = result.stdout.decode().strip()
  119. lint_error = 'ERRORS:\n' + error_message
  120. first_error_line = None
  121. for line in error_message.splitlines(True):
  122. if line.strip():
  123. # The format of the error message is: <filename>:<line>:<column>: <error code> <error message>
  124. parts = line.split(':')
  125. if len(parts) >= 2:
  126. try:
  127. first_error_line = int(parts[1])
  128. break
  129. except ValueError:
  130. # Not a valid line number, continue to the next line
  131. continue
  132. return lint_error, first_error_line
  133. # Not a python file, skip linting
  134. return None, None
  135. def _print_window(file_path, targeted_line, WINDOW, return_str=False):
  136. global CURRENT_LINE
  137. _check_current_file(file_path)
  138. with open(file_path) as file:
  139. content = file.read()
  140. # Ensure the content ends with a newline character
  141. if not content.endswith('\n'):
  142. content += '\n'
  143. lines = content.splitlines(True) # Keep all line ending characters
  144. total_lines = len(lines)
  145. # cover edge cases
  146. CURRENT_LINE = _clamp(targeted_line, 1, total_lines)
  147. half_window = max(1, WINDOW // 2)
  148. # Ensure at least one line above and below the targeted line
  149. start = max(1, CURRENT_LINE - half_window)
  150. end = min(total_lines, CURRENT_LINE + half_window)
  151. # Adjust start and end to ensure at least one line above and below
  152. if start == 1:
  153. end = min(total_lines, start + WINDOW - 1)
  154. if end == total_lines:
  155. start = max(1, end - WINDOW + 1)
  156. output = ''
  157. # only display this when there's at least one line above
  158. if start > 1:
  159. output += f'({start - 1} more lines above)\n'
  160. for i in range(start, end + 1):
  161. _new_line = f'{i}|{lines[i-1]}'
  162. if not _new_line.endswith('\n'):
  163. _new_line += '\n'
  164. output += _new_line
  165. if end < total_lines:
  166. output += f'({total_lines - end} more lines below)\n'
  167. output = output.rstrip()
  168. if return_str:
  169. return output
  170. else:
  171. print(output)
  172. def _cur_file_header(CURRENT_FILE, total_lines) -> str:
  173. if not CURRENT_FILE:
  174. return ''
  175. return f'[File: {os.path.abspath(CURRENT_FILE)} ({total_lines} lines total)]\n'
  176. @update_pwd_decorator
  177. def open_file(
  178. path: str, line_number: int | None = 1, context_lines: int | None = 100
  179. ) -> None:
  180. """
  181. Opens the file at the given path in the editor. If line_number is provided, the window will be moved to include that line.
  182. It only shows the first 100 lines by default! Max `context_lines` supported is 2000, use `scroll up/down`
  183. to view the file if you want to see more.
  184. Args:
  185. path: str: The path to the file to open, preferredly absolute path.
  186. line_number: int | None = 1: The line number to move to. Defaults to 1.
  187. context_lines: int | None = 100: Only shows this number of lines in the context window (usually from line 1), with line_number as the center (if possible). Defaults to 100.
  188. """
  189. global CURRENT_FILE, CURRENT_LINE, WINDOW
  190. if not os.path.isfile(path):
  191. raise FileNotFoundError(f'File {path} not found')
  192. CURRENT_FILE = os.path.abspath(path)
  193. with open(CURRENT_FILE) as file:
  194. total_lines = max(1, sum(1 for _ in file))
  195. if not isinstance(line_number, int) or line_number < 1 or line_number > total_lines:
  196. raise ValueError(f'Line number must be between 1 and {total_lines}')
  197. CURRENT_LINE = line_number
  198. # Override WINDOW with context_lines
  199. if context_lines is None or context_lines < 1:
  200. context_lines = 100
  201. WINDOW = _clamp(context_lines, 1, 2000)
  202. output = _cur_file_header(CURRENT_FILE, total_lines)
  203. output += _print_window(CURRENT_FILE, CURRENT_LINE, WINDOW, return_str=True)
  204. print(output)
  205. @update_pwd_decorator
  206. def goto_line(line_number: int) -> None:
  207. """
  208. Moves the window to show the specified line number.
  209. Args:
  210. line_number: int: The line number to move to.
  211. """
  212. global CURRENT_FILE, CURRENT_LINE, WINDOW
  213. _check_current_file()
  214. with open(str(CURRENT_FILE)) as file:
  215. total_lines = max(1, sum(1 for _ in file))
  216. if not isinstance(line_number, int) or line_number < 1 or line_number > total_lines:
  217. raise ValueError(f'Line number must be between 1 and {total_lines}')
  218. CURRENT_LINE = _clamp(line_number, 1, total_lines)
  219. output = _cur_file_header(CURRENT_FILE, total_lines)
  220. output += _print_window(CURRENT_FILE, CURRENT_LINE, WINDOW, return_str=True)
  221. print(output)
  222. @update_pwd_decorator
  223. def scroll_down() -> None:
  224. """Moves the window down by 100 lines.
  225. Args:
  226. None
  227. """
  228. global CURRENT_FILE, CURRENT_LINE, WINDOW
  229. _check_current_file()
  230. with open(str(CURRENT_FILE)) as file:
  231. total_lines = max(1, sum(1 for _ in file))
  232. CURRENT_LINE = _clamp(CURRENT_LINE + WINDOW, 1, total_lines)
  233. output = _cur_file_header(CURRENT_FILE, total_lines)
  234. output += _print_window(CURRENT_FILE, CURRENT_LINE, WINDOW, return_str=True)
  235. print(output)
  236. @update_pwd_decorator
  237. def scroll_up() -> None:
  238. """Moves the window up by 100 lines.
  239. Args:
  240. None
  241. """
  242. global CURRENT_FILE, CURRENT_LINE, WINDOW
  243. _check_current_file()
  244. with open(str(CURRENT_FILE)) as file:
  245. total_lines = max(1, sum(1 for _ in file))
  246. CURRENT_LINE = _clamp(CURRENT_LINE - WINDOW, 1, total_lines)
  247. output = _cur_file_header(CURRENT_FILE, total_lines)
  248. output += _print_window(CURRENT_FILE, CURRENT_LINE, WINDOW, return_str=True)
  249. print(output)
  250. @update_pwd_decorator
  251. def create_file(filename: str) -> None:
  252. """Creates and opens a new file with the given name.
  253. Args:
  254. filename: str: The name of the file to create.
  255. """
  256. if os.path.exists(filename):
  257. raise FileExistsError(f"File '{filename}' already exists.")
  258. with open(filename, 'w') as file:
  259. file.write('\n')
  260. open_file(filename)
  261. print(f'[File {filename} created.]')
  262. def _edit_or_append_file(
  263. file_name: str,
  264. start: int | None = None,
  265. end: int | None = None,
  266. content: str = '',
  267. is_append: bool = False,
  268. ) -> None:
  269. """Internal method to handle common logic for edit_/append_file methods.
  270. Args:
  271. file_name: str: The name of the file to edit or append to.
  272. start: int | None = None: The start line number for editing. Ignored if is_append is True.
  273. end: int | None = None: The end line number for editing. Ignored if is_append is True.
  274. content: str: The content to replace the lines with or to append.
  275. is_append: bool = False: Whether to append content to the file instead of editing.
  276. """
  277. global CURRENT_FILE, CURRENT_LINE, WINDOW
  278. ERROR_MSG = f'[Error editing file {file_name}. Please confirm the file is correct.]'
  279. ERROR_MSG_SUFFIX = (
  280. 'Your changes have NOT been applied. Please fix your edit command and try again.\n'
  281. 'You either need to 1) Open the correct file and try again or 2) Specify the correct start/end line arguments.\n'
  282. 'DO NOT re-run the same failed edit command. Running it again will lead to the same error.'
  283. )
  284. if not _is_valid_filename(file_name):
  285. raise FileNotFoundError('Invalid file name.')
  286. if not _is_valid_path(file_name):
  287. raise FileNotFoundError('Invalid path or file name.')
  288. if not _create_paths(file_name):
  289. raise PermissionError('Could not access or create directories.')
  290. if not os.path.isfile(file_name):
  291. raise FileNotFoundError(f'File {file_name} not found.')
  292. # Use a temporary file to write changes
  293. content = str(content or '')
  294. temp_file_path = ''
  295. src_abs_path = os.path.abspath(file_name)
  296. first_error_line = None
  297. try:
  298. # Create a temporary file
  299. with tempfile.NamedTemporaryFile('w', delete=False) as temp_file:
  300. temp_file_path = temp_file.name
  301. # Read the original file and check if empty and for a trailing newline
  302. with open(file_name) as original_file:
  303. lines = original_file.readlines()
  304. if is_append:
  305. if lines and not (len(lines) == 1 and lines[0].strip() == ''):
  306. if not lines[-1].endswith('\n'):
  307. lines[-1] += '\n'
  308. content_lines = content.splitlines(keepends=True)
  309. new_lines = lines + content_lines
  310. content = ''.join(new_lines)
  311. else:
  312. # Handle cases where start or end are None
  313. if start is None:
  314. start = 1 # Default to the beginning
  315. if end is None:
  316. end = len(lines) # Default to the end
  317. # Check arguments
  318. if not (1 <= start <= len(lines)):
  319. print(
  320. f'{ERROR_MSG}\n'
  321. f'Invalid start line number: {start}. Line numbers must be between 1 and {len(lines)} (inclusive).\n'
  322. f'{ERROR_MSG_SUFFIX}'
  323. )
  324. return
  325. if not (1 <= end <= len(lines)):
  326. print(
  327. f'{ERROR_MSG}\n'
  328. f'Invalid end line number: {end}. Line numbers must be between 1 and {len(lines)} (inclusive).\n'
  329. f'{ERROR_MSG_SUFFIX}'
  330. )
  331. return
  332. if start > end:
  333. print(
  334. f'{ERROR_MSG}\n'
  335. f'Invalid line range: {start}-{end}. Start must be less than or equal to end.\n'
  336. f'{ERROR_MSG_SUFFIX}'
  337. )
  338. return
  339. if not content.endswith('\n'):
  340. content += '\n'
  341. content_lines = content.splitlines(True)
  342. new_lines = lines[: start - 1] + content_lines + lines[end:]
  343. content = ''.join(new_lines)
  344. if not content.endswith('\n'):
  345. content += '\n'
  346. # Write the new content to the temporary file
  347. temp_file.write(content)
  348. # Replace the original file with the temporary file atomically
  349. shutil.move(temp_file_path, src_abs_path)
  350. # Handle linting
  351. if ENABLE_AUTO_LINT:
  352. # BACKUP the original file
  353. original_file_backup_path = os.path.join(
  354. os.path.dirname(file_name),
  355. f'.backup.{os.path.basename(file_name)}',
  356. )
  357. with open(original_file_backup_path, 'w') as f:
  358. f.writelines(lines)
  359. lint_error, first_error_line = _lint_file(file_name)
  360. if lint_error is not None:
  361. if first_error_line is not None:
  362. CURRENT_LINE = int(first_error_line)
  363. print(
  364. '[Your proposed edit has introduced new syntax error(s). Please understand the errors and retry your edit command.]'
  365. )
  366. print(lint_error)
  367. print('[This is how your edit would have looked if applied]')
  368. print('-------------------------------------------------')
  369. _print_window(file_name, CURRENT_LINE, 10)
  370. print('-------------------------------------------------\n')
  371. print('[This is the original code before your edit]')
  372. print('-------------------------------------------------')
  373. _print_window(original_file_backup_path, CURRENT_LINE, 10)
  374. print('-------------------------------------------------')
  375. print(
  376. 'Your changes have NOT been applied. Please fix your edit command and try again.\n'
  377. 'You either need to 1) Specify the correct start/end line arguments or 2) Correct your edit code.\n'
  378. 'DO NOT re-run the same failed edit command. Running it again will lead to the same error.'
  379. )
  380. # recover the original file
  381. with open(original_file_backup_path) as fin, open(
  382. file_name, 'w'
  383. ) as fout:
  384. fout.write(fin.read())
  385. os.remove(original_file_backup_path)
  386. return
  387. except FileNotFoundError as e:
  388. print(f'File not found: {e}')
  389. except IOError as e:
  390. print(f'An error occurred while handling the file: {e}')
  391. except ValueError as e:
  392. print(f'Invalid input: {e}')
  393. except Exception as e:
  394. # Clean up the temporary file if an error occurs
  395. if temp_file_path and os.path.exists(temp_file_path):
  396. os.remove(temp_file_path)
  397. print(f'An unexpected error occurred: {e}')
  398. raise e
  399. # Update the file information and print the updated content
  400. with open(file_name, 'r', encoding='utf-8') as file:
  401. n_total_lines = max(1, len(file.readlines()))
  402. if first_error_line is not None and int(first_error_line) > 0:
  403. CURRENT_LINE = first_error_line
  404. else:
  405. if is_append:
  406. CURRENT_LINE = max(1, len(lines)) # end of original file
  407. else:
  408. CURRENT_LINE = start or n_total_lines or 1
  409. print(
  410. f'[File: {os.path.abspath(file_name)} ({n_total_lines} lines total after edit)]'
  411. )
  412. CURRENT_FILE = file_name
  413. _print_window(CURRENT_FILE, CURRENT_LINE, WINDOW)
  414. print(MSG_FILE_UPDATED)
  415. @update_pwd_decorator
  416. def edit_file(file_name: str, start: int, end: int, content: str) -> None:
  417. """Edit a file.
  418. Replaces in given file `file_name` the lines `start` through `end` (inclusive) with the given text `content`.
  419. If a line must be inserted, an already existing line must be passed in `content` with new content accordingly!
  420. Args:
  421. file_name: str: The name of the file to edit.
  422. start: int: The start line number. Must satisfy start >= 1.
  423. end: int: The end line number. Must satisfy start <= end <= number of lines in the file.
  424. content: str: The content to replace the lines with.
  425. """
  426. _edit_or_append_file(
  427. file_name, start=start, end=end, content=content, is_append=False
  428. )
  429. @update_pwd_decorator
  430. def append_file(file_name: str, content: str) -> None:
  431. """Append content to the given file.
  432. It appends text `content` to the end of the specified file.
  433. Args:
  434. file_name: str: The name of the file to append to.
  435. content: str: The content to append to the file.
  436. """
  437. _edit_or_append_file(file_name, start=1, end=None, content=content, is_append=True)
  438. @update_pwd_decorator
  439. def search_dir(search_term: str, dir_path: str = './') -> None:
  440. """Searches for search_term in all files in dir. If dir is not provided, searches in the current directory.
  441. Args:
  442. search_term: str: The term to search for.
  443. dir_path: Optional[str]: The path to the directory to search.
  444. """
  445. if not os.path.isdir(dir_path):
  446. raise FileNotFoundError(f'Directory {dir_path} not found')
  447. matches = []
  448. for root, _, files in os.walk(dir_path):
  449. for file in files:
  450. if file.startswith('.'):
  451. continue
  452. file_path = os.path.join(root, file)
  453. with open(file_path, 'r', errors='ignore') as f:
  454. for line_num, line in enumerate(f, 1):
  455. if search_term in line:
  456. matches.append((file_path, line_num, line.strip()))
  457. if not matches:
  458. print(f'No matches found for "{search_term}" in {dir_path}')
  459. return
  460. num_matches = len(matches)
  461. num_files = len(set(match[0] for match in matches))
  462. if num_files > 100:
  463. print(
  464. f'More than {num_files} files matched for "{search_term}" in {dir_path}. Please narrow your search.'
  465. )
  466. return
  467. print(f'[Found {num_matches} matches for "{search_term}" in {dir_path}]')
  468. for file_path, line_num, line in matches:
  469. print(f'{file_path} (Line {line_num}): {line}')
  470. print(f'[End of matches for "{search_term}" in {dir_path}]')
  471. @update_pwd_decorator
  472. def search_file(search_term: str, file_path: Optional[str] = None) -> None:
  473. """Searches for search_term in file. If file is not provided, searches in the current open file.
  474. Args:
  475. search_term: str: The term to search for.
  476. file_path: Optional[str]: The path to the file to search.
  477. """
  478. global CURRENT_FILE
  479. if file_path is None:
  480. file_path = CURRENT_FILE
  481. if file_path is None:
  482. raise FileNotFoundError(
  483. 'No file specified or open. Use the open_file function first.'
  484. )
  485. if not os.path.isfile(file_path):
  486. raise FileNotFoundError(f'File {file_path} not found')
  487. matches = []
  488. with open(file_path) as file:
  489. for i, line in enumerate(file, 1):
  490. if search_term in line:
  491. matches.append((i, line.strip()))
  492. if matches:
  493. print(f'[Found {len(matches)} matches for "{search_term}" in {file_path}]')
  494. for match in matches:
  495. print(f'Line {match[0]}: {match[1]}')
  496. print(f'[End of matches for "{search_term}" in {file_path}]')
  497. else:
  498. print(f'[No matches found for "{search_term}" in {file_path}]')
  499. @update_pwd_decorator
  500. def find_file(file_name: str, dir_path: str = './') -> None:
  501. """Finds all files with the given name in the specified directory.
  502. Args:
  503. file_name: str: The name of the file to find.
  504. dir_path: Optional[str]: The path to the directory to search.
  505. """
  506. if not os.path.isdir(dir_path):
  507. raise FileNotFoundError(f'Directory {dir_path} not found')
  508. matches = []
  509. for root, _, files in os.walk(dir_path):
  510. for file in files:
  511. if file_name in file:
  512. matches.append(os.path.join(root, file))
  513. if matches:
  514. print(f'[Found {len(matches)} matches for "{file_name}" in {dir_path}]')
  515. for match in matches:
  516. print(f'{match}')
  517. print(f'[End of matches for "{file_name}" in {dir_path}]')
  518. else:
  519. print(f'[No matches found for "{file_name}" in {dir_path}]')
  520. @update_pwd_decorator
  521. def parse_pdf(file_path: str) -> None:
  522. """Parses the content of a PDF file and prints it.
  523. Args:
  524. file_path: str: The path to the file to open.
  525. """
  526. print(f'[Reading PDF file from {file_path}]')
  527. content = PyPDF2.PdfReader(file_path)
  528. text = ''
  529. for page_idx in range(len(content.pages)):
  530. text += (
  531. f'@@ Page {page_idx + 1} @@\n'
  532. + content.pages[page_idx].extract_text()
  533. + '\n\n'
  534. )
  535. print(text.strip())
  536. @update_pwd_decorator
  537. def parse_docx(file_path: str) -> None:
  538. """
  539. Parses the content of a DOCX file and prints it.
  540. Args:
  541. file_path: str: The path to the file to open.
  542. """
  543. print(f'[Reading DOCX file from {file_path}]')
  544. content = docx.Document(file_path)
  545. text = ''
  546. for i, para in enumerate(content.paragraphs):
  547. text += f'@@ Page {i + 1} @@\n' + para.text + '\n\n'
  548. print(text)
  549. @update_pwd_decorator
  550. def parse_latex(file_path: str) -> None:
  551. """
  552. Parses the content of a LaTex file and prints it.
  553. Args:
  554. file_path: str: The path to the file to open.
  555. """
  556. print(f'[Reading LaTex file from {file_path}]')
  557. with open(file_path) as f:
  558. data = f.read()
  559. text = LatexNodes2Text().latex_to_text(data)
  560. print(text.strip())
  561. def _base64_img(file_path: str) -> str:
  562. with open(file_path, 'rb') as image_file:
  563. encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
  564. return encoded_image
  565. def _base64_video(file_path: str, frame_interval: int = 10) -> list[str]:
  566. import cv2
  567. video = cv2.VideoCapture(file_path)
  568. base64_frames = []
  569. frame_count = 0
  570. while video.isOpened():
  571. success, frame = video.read()
  572. if not success:
  573. break
  574. if frame_count % frame_interval == 0:
  575. _, buffer = cv2.imencode('.jpg', frame)
  576. base64_frames.append(base64.b64encode(buffer).decode('utf-8'))
  577. frame_count += 1
  578. video.release()
  579. return base64_frames
  580. def _prepare_image_messages(task: str, base64_image: str):
  581. return [
  582. {
  583. 'role': 'user',
  584. 'content': [
  585. {'type': 'text', 'text': task},
  586. {
  587. 'type': 'image_url',
  588. 'image_url': {'url': f'data:image/jpeg;base64,{base64_image}'},
  589. },
  590. ],
  591. }
  592. ]
  593. @update_pwd_decorator
  594. def parse_audio(file_path: str, model: str = 'whisper-1') -> None:
  595. """
  596. Parses the content of an audio file and prints it.
  597. Args:
  598. file_path: str: The path to the audio file to transcribe.
  599. model: Optional[str]: The audio model to use for transcription. Defaults to 'whisper-1'.
  600. """
  601. print(f'[Transcribing audio file from {file_path}]')
  602. try:
  603. # TODO: record the COST of the API call
  604. with open(file_path, 'rb') as audio_file:
  605. transcript = client.audio.translations.create(model=model, file=audio_file)
  606. print(transcript.text)
  607. except Exception as e:
  608. print(f'Error transcribing audio file: {e}')
  609. @update_pwd_decorator
  610. def parse_image(
  611. file_path: str, task: str = 'Describe this image as detail as possible.'
  612. ) -> None:
  613. """
  614. Parses the content of an image file and prints the description.
  615. Args:
  616. file_path: str: The path to the file to open.
  617. task: Optional[str]: The task description for the API call. Defaults to 'Describe this image as detail as possible.'.
  618. """
  619. print(f'[Reading image file from {file_path}]')
  620. # TODO: record the COST of the API call
  621. try:
  622. base64_image = _base64_img(file_path)
  623. response = client.chat.completions.create(
  624. model=OPENAI_MODEL,
  625. messages=_prepare_image_messages(task, base64_image),
  626. max_tokens=MAX_TOKEN,
  627. )
  628. content = response.choices[0].message.content
  629. print(content)
  630. except Exception as error:
  631. print(f'Error with the request: {error}')
  632. @update_pwd_decorator
  633. def parse_video(
  634. file_path: str,
  635. task: str = 'Describe this image as detail as possible.',
  636. frame_interval: int = 30,
  637. ) -> None:
  638. """
  639. Parses the content of an image file and prints the description.
  640. Args:
  641. file_path: str: The path to the video file to open.
  642. task: Optional[str]: The task description for the API call. Defaults to 'Describe this image as detail as possible.'.
  643. frame_interval: Optional[int]: The interval between frames to analyze. Defaults to 30.
  644. """
  645. print(
  646. f'[Processing video file from {file_path} with frame interval {frame_interval}]'
  647. )
  648. task = task or 'This is one frame from a video, please summarize this frame.'
  649. base64_frames = _base64_video(file_path)
  650. selected_frames = base64_frames[::frame_interval]
  651. if len(selected_frames) > 30:
  652. new_interval = len(base64_frames) // 30
  653. selected_frames = base64_frames[::new_interval]
  654. print(f'Totally {len(selected_frames)} would be analyze...\n')
  655. idx = 0
  656. for base64_frame in selected_frames:
  657. idx += 1
  658. print(f'Process the {file_path}, current No. {idx * frame_interval} frame...')
  659. # TODO: record the COST of the API call
  660. try:
  661. response = client.chat.completions.create(
  662. model=OPENAI_MODEL,
  663. messages=_prepare_image_messages(task, base64_frame),
  664. max_tokens=MAX_TOKEN,
  665. )
  666. content = response.choices[0].message.content
  667. current_frame_content = f"Frame {idx}'s content: {content}\n"
  668. print(current_frame_content)
  669. except Exception as error:
  670. print(f'Error with the request: {error}')
  671. @update_pwd_decorator
  672. def parse_pptx(file_path: str) -> None:
  673. """
  674. Parses the content of a pptx file and prints it.
  675. Args:
  676. file_path: str: The path to the file to open.
  677. """
  678. print(f'[Reading PowerPoint file from {file_path}]')
  679. try:
  680. pres = Presentation(str(file_path))
  681. text = []
  682. for slide_idx, slide in enumerate(pres.slides):
  683. text.append(f'@@ Slide {slide_idx + 1} @@')
  684. for shape in slide.shapes:
  685. if hasattr(shape, 'text'):
  686. text.append(shape.text)
  687. print('\n'.join(text))
  688. except Exception as e:
  689. print(f'Error reading PowerPoint file: {e}')
  690. __all__ = [
  691. # file operation
  692. 'open_file',
  693. 'goto_line',
  694. 'scroll_down',
  695. 'scroll_up',
  696. 'create_file',
  697. 'append_file',
  698. 'edit_file',
  699. 'search_dir',
  700. 'search_file',
  701. 'find_file',
  702. # readers
  703. 'parse_pdf',
  704. 'parse_docx',
  705. 'parse_latex',
  706. 'parse_pptx',
  707. ]
  708. if OPENAI_API_KEY and OPENAI_BASE_URL:
  709. __all__ += ['parse_audio', 'parse_video', 'parse_image']
  710. DOCUMENTATION = ''
  711. for func_name in __all__:
  712. func = globals()[func_name]
  713. cur_doc = func.__doc__
  714. # remove indentation from docstring and extra empty lines
  715. cur_doc = '\n'.join(filter(None, map(lambda x: x.strip(), cur_doc.split('\n'))))
  716. # now add a consistent 4 indentation
  717. cur_doc = '\n'.join(map(lambda x: ' ' * 4 + x, cur_doc.split('\n')))
  718. fn_signature = f'{func.__name__}' + str(signature(func))
  719. DOCUMENTATION += f'{fn_signature}:\n{cur_doc}\n\n'