Multimedika/Bot_Development

1import os2import nest_asyncio3from io import BytesIO4from typing import List5 6from dotenv import load_dotenv7from fastapi import UploadFile8 9from llama_index.core.schema import Document10from script.get_metadata import Metadata11from core.prompt import PARSER_INSTRUCTION12from service.llamaparse import S3ImageSaver13from service.llamaparse import LlamaParseWithS314from utils.error_handlers import handle_error, handle_exception15from fastapi.responses import JSONResponse16 17load_dotenv()18 19 20nest_asyncio.apply()21 22 23def get_documents(json_list: List[dict]):24    text_documents = []25    try:26        for idx, page in enumerate(json_list):27            text_document = Document(text=page["md"], metadata={"page": page["page"]})28            text_documents.append(text_document)29        return text_documents30    except Exception as e:31        return handle_error(32            e, "Error processing file in get_documents", status_code=40033        )34 35 36def parse_journal(title, content: bytes, file_name: str, lang: str = "en"):37    """Parse the journal using LlamaParse."""38    try:39        # Initialize the parser40        s3_image_saver = S3ImageSaver(41            bucket_name=os.getenv("S3_BUCKET_NAME"),42            access_key=os.getenv("AWS_ACCESS_KEY_ID"),43            secret_key=os.getenv("AWS_SECRET_ACCESS_KEY"),44            region_name="us-west-2",45        )46        print("s3 image saver",s3_image_saver)47 48        s3_parser = LlamaParseWithS3(49            api_key=os.getenv(50                "LLAMA_PARSE_API_KEY"51            ),  # can also be set in your env as LLAMA_CLOUD_API_KEY52            parsing_instruction=PARSER_INSTRUCTION,53            result_type="markdown",  # "markdown" and "text" are available54            verbose=True,55            language=lang,  # Optionally you can define a language, default=en56            s3_image_saver=s3_image_saver,57        )58        59        md_json_objs = s3_parser.get_json_result(60            content, extra_info={"file_name": file_name}61        )62        63        json_list = md_json_objs[0]["pages"]64 65        image_dicts = s3_parser.get_images(md_json_objs, title)66        67        if isinstance(image_dicts, JSONResponse):68            image_urls=image_dicts  # Return the error response directly69        else:70            image_urls = [71                {"page_number": img["page_number"], "image_link": img["image_link"]}72                for img in image_dicts73                if img["image_link"] is not None74            ]75 76        return json_list, image_urls77 78    except Exception as e:79        return handle_error(80            e, "Error processing file in parse_journal", status_code=40081        )82 83 84async def upload_file(reference, file: UploadFile, lang: str = "en"):85    try:86        # Read the binary content of the uploaded file once87        content = await file.read()88        89        # Store the file content in a BytesIO stream for reuse later90        file_stream = BytesIO(content)91        92        # Parse the journal93        title = reference["title"]94 95        json_list, image_urls = parse_journal(title, content, file.filename, lang)96        parsed_documents = get_parsed_documents(json_list, image_urls)97        98        if isinstance(image_urls, JSONResponse):99            return image_urls  # Return the error response directly100 101        metadata_gen = Metadata(reference)102        documents_with_metadata = metadata_gen.apply_metadata(parsed_documents)103 104        print("Banyak documents : \n", len(documents_with_metadata))105 106        # Return both parsed documents and metadata107        return documents_with_metadata, file_stream108 109    except Exception as e:110        print("error ", e)111        return handle_exception(e)112 113def get_parsed_documents(json_dicts=None, image_links=None):114    try:115        """Split docs into nodes, by separator."""116        parsed_documents = []117 118        # Preprocess metadata119        md_texts = [d["md"] for d in json_dicts] if json_dicts is not None else None120        121        # Create a dictionary to store lists of image links for each page number122        image_link_dict = {}123        if image_links:124            for item in image_links:125                page_number = item["page_number"]126                image_link = item["image_link"]127                if page_number in image_link_dict:128                    image_link_dict[page_number].append(image_link)129                else:130                    image_link_dict[page_number] = [image_link]  131        132        md_texts = [d["md"] for d in json_dicts]133 134        for idx, md_text in enumerate(md_texts):135            page_number = idx + 1136            chunk_metadata = {"page_number": page_number}137 138            # Set the image link if it exists; otherwise, set it to None139            chunk_metadata["image_links"] = image_link_dict.get(page_number, [])140 141            # Add parsed text and create the Document object142            parsed_document = Document(143                text=md_text,144                metadata=chunk_metadata,145            )146 147            parsed_documents.append(parsed_document)148 149        return parsed_documents150    except Exception as e:151        return handle_error(152            e, "Error processing documents in get_text_documents", status_code=400153        )154