# coding: utf-8 # Copyright (c) 2025 inclusionAI. import json import os import base64 import tempfile import subprocess from typing import Any, Dict, Tuple from urllib.parse import urlparse from pydantic import BaseModel from aworld.config import ToolConfig from examples.tools.tool_action import DocumentExecuteAction from aworld.core.common import Observation, ActionModel, ActionResult from aworld.core.tool.base import ToolFactory, Tool from aworld.logs.util import logger from examples.tools.document.utils import encode_image_from_file, encode_image_from_url from aworld.utils import import_package, import_packages from aworld.tools.utils import build_observation class InputDocument(BaseModel): document_path: str | None = None @ToolFactory.register(name="document_analysis", desc="document analysis", supported_action=DocumentExecuteAction, conf_file_name=f'document_analysis_tool.yaml') class DocumentTool(Tool): def __init__(self, conf: ToolConfig, **kwargs) -> None: """Init document tool.""" import_package('cv2', install_name='opencv-python') import_packages(['xmltodict', 'pandas', 'docx2markdown', 'PyPDF2', 'numpy']) super(DocumentTool, self).__init__(conf, **kwargs) self.cur_observation = None self.content = None self.keyframes = [] self.init() self.step_finished = True def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[ Observation, dict[str, Any]]: super().reset(seed=seed, options=options) self.close() self.step_finished = True return build_observation(observer=self.name(), ability=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.name), {} def init(self) -> None: self.initialized = True def close(self) -> None: pass def finished(self) -> bool: return self.step_finished def do_step(self, actions: list[ActionModel], **kwargs) -> Tuple[Observation, float, bool, bool, Dict[str, Any]]: self.step_finished = False reward = 0. fail_error = "" observation = build_observation(observer=self.name(), ability=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.name) info = {} try: if not actions: raise ValueError("actions is empty") action = actions[0] document_path = action.params.get("document_path", "") if not document_path: raise ValueError("document path invalid") output, keyframes, error = self.document_analysis(document_path) observation.content = output observation.action_result.append( ActionResult(is_done=True, success=False if error else True, content=f"{output}", error=f"{error}", keep=False)) info['key_frame'] = f"{keyframes}" reward = 1. except Exception as e: fail_error = str(e) finally: self.step_finished = True info["exception"] = fail_error info.update(kwargs) return (observation, reward, kwargs.get("terminated", False), kwargs.get("truncated", False), info) def document_analysis(self, document_path): import xmltodict error = None # Initialize content to empty list to avoid None return self.content = [] try: if any(document_path.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]): parsed_url = urlparse(document_path) is_url = all([parsed_url.scheme, parsed_url.netloc]) if not is_url: base64_image = encode_image_from_file(document_path) else: base64_image = encode_image_from_url(document_path) self.content = f"data:image/jpeg;base64,{base64_image}" if any(document_path.endswith(ext) for ext in ["xls", "xlsx"]): try: try: import pandas as pd except ImportError: error = "pandas library not found. Please install pandas: pip install pandas" return self.content, self.keyframes, error excel_data = {} with pd.ExcelFile(document_path) as xls: sheet_names = xls.sheet_names for sheet_name in sheet_names: df = pd.read_excel(xls, sheet_name=sheet_name) sheet_data = df.to_dict(orient='records') excel_data[sheet_name] = sheet_data self.content = json.dumps(excel_data, ensure_ascii=False) logger.info(f"Successfully processed Excel file: {document_path}") logger.info(f"Found {len(sheet_names)} sheets: {', '.join(sheet_names)}") except Exception as excel_error: error = str(excel_error) if any(document_path.endswith(ext) for ext in ["json", "jsonl", "jsonld"]): with open(document_path, "r", encoding="utf-8") as f: self.content = json.load(f) f.close() if any(document_path.endswith(ext) for ext in ["xml"]): data = None with open(document_path, "r", encoding="utf-8") as f: data = f.read() f.close() try: self.content = xmltodict.parse(data) logger.info(f"The extracted xml data is: {self.content}") except Exception as e: logger.info(f"The raw xml data is: {data}") error = str(e) self.content = data if any(document_path.endswith(ext) for ext in ["doc", "docx"]): from docx2markdown._docx_to_markdown import docx_to_markdown file_name = os.path.basename(document_path) md_file_path = f"{file_name}.md" docx_to_markdown(document_path, md_file_path) with open(md_file_path, "r") as f: self.content = f.read() f.close() if any(document_path.endswith(ext) for ext in ["pdf"]): # try using pypdf to extract text from pdf try: from PyPDF2 import PdfReader # Open file in binary mode for PdfReader f = open(document_path, "rb") reader = PdfReader(f) extracted_text = "" for page in reader.pages: extracted_text += page.extract_text() self.content = extracted_text f.close() except Exception as pdf_error: error = str(pdf_error) # audio if any(document_path.endswith(ext.lower()) for ext in [".mp3", ".wav", ".wave"]): try: # audio-> base64 with open(document_path, "rb") as audio_file: audio_bytes = audio_file.read() audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') # ext ext = os.path.splitext(document_path)[1].lower() mime_type = "audio/mpeg" if ext == ".mp3" else "audio/wav" # data URI self.content = f"data:{mime_type};base64,{audio_base64}" except Exception as audio_error: error = str(audio_error) logger.error(f"Error processing audio file: {error}") # video if any(document_path.endswith(ext.lower()) for ext in [".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv"]): try: try: import cv2 import numpy as np except ImportError: error = "Required libraries not found. Please install opencv-python: pip install opencv-python" return None, None, error # create temp dir temp_dir = tempfile.mkdtemp() # 1.get audio -> base64 audio_path = os.path.join(temp_dir, "extracted_audio.mp3") # get audio by ffmpeg try: subprocess.run([ "ffmpeg", "-i", document_path, "-q:a", "0", "-map", "a", audio_path, "-y" ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # audio->base64 with open(audio_path, "rb") as audio_file: audio_bytes = audio_file.read() audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') audio_data_uri = f"data:audio/mpeg;base64,{audio_base64}" except (subprocess.SubprocessError, FileNotFoundError) as e: logger.warning(f"Failed to extract audio: {str(e)}") audio_data_uri = None # 2. get keyframes cap = cv2.VideoCapture(document_path) if not cap.isOpened(): raise ValueError(f"Could not open video file: {document_path}") # get video message fps = cap.get(cv2.CAP_PROP_FPS) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) duration = frame_count / fps if fps > 0 else 0 # keyframes policy- per duration/10s，max 10 keyframes_count = min(10, int(frame_count)) frames_interval = max(1, int(frame_count / keyframes_count)) self.keyframes = [] frame_index = 0 while True: ret, frame = cap.read() if not ret: break # per frames_interval save if frame_index % frames_interval == 0: # save JPEG -> base64 _, buffer = cv2.imencode(".jpg", frame) img_base64 = base64.b64encode(buffer).decode('utf-8') time_position = frame_index / fps if fps > 0 else 0 self.keyframes.append(f"data:image/jpeg;base64,{img_base64}") if len(self.keyframes) >= keyframes_count: break frame_index += 1 cap.release() self.content = audio_data_uri logger.info(f"Successfully processed video file: {document_path}") logger.info(f"Extracted {len(self.keyframes)} keyframes and audio track") # clean tmp files try: os.remove(audio_path) os.rmdir(temp_dir) except Exception as cleanup_error: logger.warning(f"Error cleaning up temp files: {str(cleanup_error)}") except Exception as video_error: error = str(video_error) logger.error(f"Error processing video file: {error}") if any(document_path.endswith(ext) for ext in ["pptx"]): try: # Initialize content list and empty keyframes self.content = [] self.keyframes = [] # Check if file exists if not os.path.exists(document_path): error = f"File does not exist: {document_path}" return self.content, self.keyframes, error # Check if file is readable if not os.access(document_path, os.R_OK): error = f"File is not readable: {document_path}" return self.content, self.keyframes, error # Check file size try: file_size = os.path.getsize(document_path) if file_size == 0: error = "File is empty" return self.content, self.keyframes, error except Exception as size_error: logger.warning(f"Cannot get file size: {str(size_error)}") try: # Import required libraries from pptx import Presentation from PIL import Image, ImageDraw, ImageFont import io except ImportError as import_error: error = f"Missing required libraries: {str(import_error)}. Please install: pip install python-pptx Pillow" return self.content, self.keyframes, error # Create temporary directory for images try: temp_dir = tempfile.mkdtemp() except Exception as temp_dir_error: error = f"Failed to create temporary directory: {str(temp_dir_error)}" return self.content, self.keyframes, error # Open presentation try: presentation = Presentation(document_path) # Get total slides count total_slides = len(presentation.slides) if total_slides == 0: error = "PPTX file does not contain any slides" return self.content, self.keyframes, error # Process each slide for i, slide in enumerate(presentation.slides): # Generate temporary file path for current slide img_path = os.path.join(temp_dir, f"slide_{i + 1}.jpg") # Get slide dimensions try: slide_width = presentation.slide_width slide_height = presentation.slide_height # PPTX dimensions are in EMU (English Metric Unit) # 1 inch = 914400 EMU, 1 cm = 360000 EMU # Convert to pixels (assuming 96 DPI) slide_width_px = int(slide_width / 914400 * 96 * 10) slide_height_px = int(slide_height / 914400 * 96 * 10) # Ensure dimensions are reasonable positive integers slide_width_px = max(1, min(slide_width_px, 4000)) # Limit max width to 4000px slide_height_px = max(1, min(slide_height_px, 3000)) # Limit max height to 3000px except Exception as size_error: # Use default dimensions slide_width_px = 960 # Default width 960px slide_height_px = 720 # Default height 720px # Create blank image try: # Log operation start # Create blank image try: slide_img = Image.new('RGB', (slide_width_px, slide_height_px), 'white') draw = ImageDraw.Draw(slide_img) except Exception as img_create_error: logger.error( f"Slide {i + 1} blank image creation failed: {str(img_create_error) or 'Unknown error'}") raise # Draw slide number try: font = ImageFont.load_default() draw.text((20, 20), f"Slide {i + 1}/{total_slides}", fill="black", font=font) except Exception as font_error: logger.warning(f"Failed to draw slide number: {str(font_error) or 'Unknown error'}") # Record shape count try: shape_count = len(slide.shapes) except Exception as shape_count_error: logger.warning( f"Failed to get slide {i + 1} shape count: {str(shape_count_error) or 'Unknown error'}") shape_count = 0 # Try to render shapes on image shape_success_count = 0 shape_fail_count = 0 try: for j, shape in enumerate(slide.shapes): try: shape_type = type(shape).__name__ # Process images if hasattr(shape, 'image') and shape.image: try: # Extract image from shape image_stream = io.BytesIO(shape.image.blob) img = Image.open(image_stream) # Calculate position left = shape.left top = shape.top # Paste image onto slide slide_img.paste(img, (left, top)) shape_success_count += 1 except Exception as img_error: logger.warning( f"Failed to process image {j + 1} in slide {i + 1}: {str(img_error) or 'Unknown error'}") if not str(img_error): import traceback logger.warning( f"Image processing stack: {traceback.format_exc()}") shape_fail_count += 1 # Process text elif hasattr(shape, 'text') and shape.text: try: text = shape.text[:30] + "..." if len( shape.text) > 30 else shape.text # Simple text rendering text_left = shape.left text_top = shape.top draw.text((text_left, text_top), shape.text, fill="black", font=font) shape_success_count += 1 except Exception as text_error: logger.warning( f"Failed to process text {j + 1} in slide {i + 1}: {str(text_error) or 'Unknown error'}") if not str(text_error): import traceback logger.warning( f"Text processing stack: {traceback.format_exc()}") shape_fail_count += 1 else: logger.info( f"Shape {j + 1} in slide {i + 1} is neither image nor text, skipping") except Exception as shape_error: if not str(shape_error): import traceback logger.warning(f"Shape processing stack: {traceback.format_exc()}") shape_fail_count += 1 except Exception as shapes_iteration_error: logger.error( f"Failed while iterating through shapes in slide {i + 1}: {str(shapes_iteration_error) or 'Unknown error'}") if not str(shapes_iteration_error): import traceback logger.error(f"Shape iteration stack: {traceback.format_exc()}") # Save slide image try: slide_img.save(img_path, 'JPEG') # Check if image was saved successfully if not os.path.exists(img_path): raise ValueError(f"Saved image file does not exist: {img_path}") file_size = os.path.getsize(img_path) if file_size == 0: raise ValueError( f"Saved image file is empty: {img_path}, size: {file_size} bytes") # Convert to base64 try: base64_image = encode_image_from_file(img_path) self.content.append(f"data:image/jpeg;base64,{base64_image}") except Exception as base64_error: error_msg = str(base64_error) or "Unknown base64 conversion error" if not str(base64_error): import traceback logger.error(f"Base64 conversion stack: {traceback.format_exc()}") raise ValueError(f"Base64 conversion error: {error_msg}") except Exception as save_error: error_msg = str(save_error) or "Unknown save error" logger.error(f"Failed to save slide {i + 1} as image: {error_msg}") if not str(save_error): import traceback logger.error(f"Image save stack: {traceback.format_exc()}") raise ValueError(f"Image save error: {error_msg}") except Exception as slide_render_error: error_msg = str(slide_render_error) or "Unknown rendering error" logger.error(f"Failed to render slide {i + 1}: {error_msg}") if not str(slide_render_error): import traceback logger.error(f"Slide rendering stack: {traceback.format_exc()}") # Continue processing next slide, don't interrupt the entire process continue except Exception as pptx_error: error = f"Failed to process PPTX file: {str(pptx_error)}" import traceback # Clean up temporary files try: for file in os.listdir(temp_dir): try: file_path = os.path.join(temp_dir, file) os.remove(file_path) except Exception as file_error: logger.warning(f"Failed to delete temporary file: {str(file_error)}") os.rmdir(temp_dir) except Exception as cleanup_error: logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}") if len(self.content) > 0: logger.info(f"Extracted {len(self.content)} slides") else: error = error or "Could not extract any slides from PPTX file" logger.error(error) except Exception as outer_error: error = f"Error occurred during PPTX file processing: {str(outer_error)}" import traceback return self.content, self.keyframes, error finally: pass return self.content, self.keyframes, error