Multimedika/Bot_Development
0
1import os2import shutil3import asyncio4from urllib.parse import quote5from dotenv import load_dotenv6from io import BufferedIOBase7from typing import List, Optional, Union8from pathlib import Path9 10from botocore.exceptions import ClientError11from botocore.config import Config12from boto3.session import Session13 14from pydantic import PrivateAttr15from llama_index.core.async_utils import run_jobs16from llama_parse import LlamaParse17from llama_parse.utils import (18 nest_asyncio_err,19 nest_asyncio_msg,20)21 22from llama_index.core.schema import Document23 24 25load_dotenv()26 27FileInput = Union[str, bytes, BufferedIOBase]28 29 30class S3ImageSaver:31 def __init__(self, bucket_name, access_key=None, secret_key=None, region_name=None):32 self.bucket_name = bucket_name33 self.region_name = region_name34 self.session = Session(35 aws_access_key_id=access_key,36 aws_secret_access_key=secret_key,37 region_name=self.region_name,38 )39 self.s3_client = self.session.client(40 "s3", config=Config(signature_version="s3v4", region_name=self.region_name)41 )42 43 def save_image(self, image_path, title):44 """Saves an image to the S3 bucket."""45 try:46 print("---Saving Images---")47 title_encoded = quote(title)48 s3_key = f"images/{title}/{os.path.basename(image_path)}"49 with open(image_path, "rb") as file:50 self.s3_client.upload_fileobj(file, self.bucket_name, s3_key)51 52 s3_url = f"https://{self.bucket_name}.s3.{self.region_name}.amazonaws.com/images/{title_encoded}/{os.path.basename(image_path)}"53 print(f"Image saved to S3 bucket: {s3_url}")54 return s3_url55 except ClientError as e:56 print(f"Error saving image to S3: {e}")57 return None58 59 60class LlamaParseWithS3(LlamaParse):61 _s3_image_saver: S3ImageSaver = PrivateAttr()62 63 def __init__(self, *args, s3_image_saver=None, **kwargs):64 super().__init__(*args, **kwargs)65 self._s3_image_saver = s3_image_saver or S3ImageSaver(66 bucket_name=os.getenv("S3_BUCKET_NAME"),67 access_key=os.getenv("AWS_ACCESS_KEY_ID"),68 secret_key=os.getenv("AWS_SECRET_ACCESS_KEY"),69 region_name="us-west-2",70 )71 72 async def aget_images(73 self, json_result: List[dict], download_path: str74 ) -> List[dict]:75 """Download images from the parsed result."""76 headers = {"Authorization": f"Bearer {self.api_key}"}77 78 # make the download path79 if not os.path.exists(download_path):80 os.makedirs(download_path)81 82 try:83 images = []84 for result in json_result:85 job_id = result["job_id"]86 for page in result["pages"]:87 if self.verbose:88 print(f"> Image for page {page['page']}: {page['images']}")89 for image in page["images"]:90 image_name = image["name"]91 92 # get the full path93 image_path = os.path.join(download_path, f"{image_name}")94 95 # get a valid image path96 if not image_path.endswith(".png"):97 if not image_path.endswith(".jpg"):98 image_path += ".png"99 100 image["path"] = image_path101 image["job_id"] = job_id102 103 image["original_file_path"] = result.get("file_path", None)104 105 image["page_number"] = page["page"]106 with open(image_path, "wb") as f:107 image_url = f"{self.base_url}/api/parsing/job/{job_id}/result/image/{image_name}"108 async with self.client_context() as client:109 res = await client.get(110 image_url, headers=headers, timeout=self.max_timeout111 )112 res.raise_for_status()113 f.write(res.content)114 images.append(image)115 return images116 except Exception as e:117 print("Error while downloading images from the parsed result:", e)118 if self.ignore_errors:119 return []120 else:121 raise e122 123 async def aget_images_s3(self, json_result: List[dict], title) -> List[dict]:124 125 images = await self.aget_images(126 json_result, download_path="tmp/"127 ) # Download to temporary location128 129 # Process each image and upload to S3130 for image in images:131 image_path = image["path"]132 try:133 s3_url = self._s3_image_saver.save_image(image_path, title)134 if s3_url:135 image["image_link"] = s3_url136 except Exception as e:137 print(f"Error saving image to S3: {image_path} - {e}")138 139 # After processing all images, delete the tmp folder140 folder_path = "tmp/"141 try:142 shutil.rmtree(folder_path) # Deletes the folder and all its contents143 print(f"Folder {folder_path} and all its contents were deleted successfully.")144 except Exception as e:145 print(f"Error deleting folder {folder_path}: {e}")146 147 return images148 149 def get_images(self, json_result: List[dict], title) -> List[dict]:150 """Download images from the parsed result and save them to S3."""151 try:152 return asyncio.run(self.aget_images_s3(json_result, title))153 except RuntimeError as e:154 if nest_asyncio_err in str(e):155 raise RuntimeError(nest_asyncio_msg)156 else:157 raise e158 159 def get_single_job_id(json_result):160 # Check if the list is not empty and has at least one result161 if json_result:162 json_id = json_result[0].get("job_id")163 return json_id164 return None # Return None if the list is empty165 166 # The code if we know the job id167 168 # async def _aget_json(169 # self, job_id, file_path: FileInput, extra_info: Optional[dict] = None170 # ) -> List[dict]:171 # """Load data from the input path."""172 # try:173 # if self.verbose:174 # print("Started parsing the file under job_id %s" % job_id)175 # result = await self._get_job_result(job_id, "json")176 # result["job_id"] = job_id177 178 # if not isinstance(file_path, (bytes, BufferedIOBase)):179 # result["file_path"] = str(file_path)180 181 # return [result]182 # except Exception as e:183 # file_repr = file_path if isinstance(file_path, str) else "<bytes/buffer>"184 # print(f"Error while parsing the file '{file_repr}':", e)185 # if self.ignore_errors:186 # return []187 # else:188 # raise e189 190 async def aget_json(191 self,192 file_path: Union[List[FileInput], FileInput],193 extra_info: Optional[dict] = None,194 ) -> List[dict]:195 """Load data from the input path."""196 if isinstance(file_path, (str, Path, bytes, BufferedIOBase)):197 return await self._aget_json(file_path, extra_info=extra_info)198 199 # The code when job id known200 # return await self._aget_json(201 # job_id="cda0870a-b896-4140-84ea-1565e1aa1565",202 # file_path=file_path,203 # extra_info=extra_info,204 # )205 elif isinstance(file_path, list):206 jobs = [self._aget_json(f, extra_info=extra_info) for f in file_path]207 try:208 results = await run_jobs(209 jobs,210 workers=self.num_workers,211 desc="Parsing files",212 show_progress=self.show_progress,213 )214 215 # return flattened results216 return [item for sublist in results for item in sublist]217 except RuntimeError as e:218 if nest_asyncio_err in str(e):219 raise RuntimeError(nest_asyncio_msg)220 else:221 raise e222 else:223 raise ValueError(224 "The input file_path must be a string or a list of strings."225 )226 227 async def _aload_data(228 self,229 job_id,230 extra_info: Optional[dict] = None,231 verbose: bool = False,232 ) -> List[Document]:233 """Load data from the input path."""234 try:235 result = await self._get_job_result(236 job_id, self.result_type.value, verbose=verbose237 )238 239 docs = [240 Document(241 text=result[self.result_type.value],242 metadata=extra_info or {},243 )244 ]245 if self.split_by_page:246 return self._get_sub_docs(docs)247 else:248 return docs249 250 except Exception as e:251 print(f"Error while parsing the file :", e)252 if self.ignore_errors:253 return []254 else:255 raise e256 257 async def aload_data(258 self,259 job_id,260 extra_info: Optional[dict] = None,261 ) -> List[Document]:262 """Load data from the input path."""263 try:264 return await self._aload_data(265 job_id, extra_info=extra_info, verbose=self.verbose266 )267 268 except RuntimeError as e:269 if nest_asyncio_err in str(e):270 raise RuntimeError(nest_asyncio_msg)271 else:272 raise e273 274 def load_data(275 self,276 job_id,277 extra_info: Optional[dict] = None,278 ) -> List[Document]:279 """Load data from the input path."""280 try:281 return asyncio.run(self.aload_data(job_id, extra_info))282 except RuntimeError as e:283 if nest_asyncio_err in str(e):284 raise RuntimeError(nest_asyncio_msg)285 else:286 raise e287