diff --git a/backend/utils/elements.py b/backend/utils/elements.py index 42fc8ca..9b3fcef 100644 --- a/backend/utils/elements.py +++ b/backend/utils/elements.py @@ -10,6 +10,8 @@ from langchain_core.documents import Document from pydantic import BaseModel, PrivateAttr, validator +from .image import local_image_to_base64 + class Element(BaseModel): """Abstract base class representing an element with a type, format, and metadata. @@ -311,6 +313,73 @@ class TableImage(Table, Image): format: Literal["image"] = "image" +def _create_text_element(doc: Document, element_class: type[Text]) -> Element: + """Create a text element from a Langchain Document object. + + Args: + doc (Document): Langchain Document object. + element_class (Type[Text]): Text element class to create (Text, TableText). + + Returns: + Element: Text element created from the Document object. + """ + match doc.metadata["source"]: + case "content": + element = element_class( + type=doc.metadata["type"], + format=doc.metadata["format"], + text=doc.page_content, + metadata=doc.metadata, + ) + case "summary": + element = element_class( + type=doc.metadata["type"], + format=doc.metadata["format"], + text="No content available", + metadata=doc.metadata, + ) + element.set_summary(doc.page_content) + case other: + raise ValueError(f"Unsupported element source: {other}") + return element + + +NO_IMAGE = local_image_to_base64("img/no_image.png") + + +def _create_image_element(doc: Document, element_class: type[Image]) -> Element: + """Create an image element from a Langchain Document object. + + Args: + doc (Document): Langchain Document object. + element_class (Type[Image]): Image element class to create (Image, TableImage). + + Returns: + Element: Image element created from the Document object. + """ + match doc.metadata["source"]: + case "content": + element = element_class( + type=doc.metadata["type"], + format=doc.metadata["format"], + base64=doc.page_content, + mime_type=doc.metadata["mime_type"], + metadata=doc.metadata, + ) + case "summary": + element = element_class( + type=doc.metadata["type"], + format=doc.metadata["format"], + base64=NO_IMAGE, + mime_type="image/png", + metadata=doc.metadata, + ) + element.set_summary(doc.page_content) + case other: + raise ValueError(f"Unsupported element source: {other}") + return element + + def convert_documents_to_elements(docs: list[Document]) -> list: """Convert a list of Langchain Document objects to a list of Element objects. @@ -327,37 +396,19 @@ def convert_documents_to_elements(docs: list[Document]) -> list: for doc in docs: match doc.metadata["type"]: case "text": - text_format = doc.metadata["format"] - element = Text( - text=doc.page_content, format=text_format, metadata=doc.metadata - ) + element = _create_text_element(doc, Text) case "image": - element = Image( - base64=doc.page_content, - mime_type=doc.metadata["mime_type"], - metadata=doc.metadata, - ) - + element = _create_image_element(doc, Image) case "table": - table_format = doc.metadata["format"] - match table_format: + match doc.metadata["format"]: case "text" | "html" | "markdown": - element = TableText( - text=doc.page_content, - format=table_format, - metadata=doc.metadata, - ) - + element = _create_text_element(doc, TableText) case "image": - element = TableImage( - base64=doc.page_content, - mime_type=doc.metadata["mime_type"], - metadata=doc.metadata, - ) + element = _create_image_element(doc, TableImage) case other: raise ValueError(f"Unsupported table format: {other}") case other: - raise ValueError(f"Unsupported document type: {other['type']}") + raise ValueError(f"Unsupported document type: {other}") elements.append(element) return elements diff --git a/backend/utils/image.py b/backend/utils/image.py index adc43fc..9d03b38 100644 --- a/backend/utils/image.py +++ b/backend/utils/image.py @@ -2,6 +2,7 @@ import base64 import io +from pathlib import Path from PIL import Image @@ -32,3 +33,16 @@ def resize_base64_image( # Encode the resized image to Base64 return base64.b64encode(buffered.getvalue()).decode("utf-8") + + +def local_image_to_base64(image_path: str) -> str: + """Convert a local image to a Base64 string. + + Args: + image_path (str): Path to the image. + + Returns: + str: Base64 string. + """ + with Path(image_path).open("rb") as image_file: + return base64.b64encode(image_file.read()).decode("utf-8") diff --git a/img/no_image.png b/img/no_image.png new file mode 100644 index 0000000..431b25e Binary files /dev/null and b/img/no_image.png differ