datapizza-ai-parsers-azure 0.0.2__py3-none-any.whl

datapizza/modules/parsers/azure/__init__.py

@@ -0,0 +1,3 @@
+from .azure_parser import AzureParser
+
+__all__ = ["AzureParser"]

datapizza/modules/parsers/azure/azure_parser.py

@@ -0,0 +1,391 @@
+import json
+from typing import Any
+
+import aiofiles
+from datapizza.core.modules.parser import Parser
+from datapizza.core.utils import extract_media
+from datapizza.type import Media, MediaNode, Node, NodeType
+
+from azure.ai.documentintelligence import DocumentIntelligenceClient
+from azure.ai.documentintelligence.aio import (
+    DocumentIntelligenceClient as AsyncDocumentIntelligenceClient,
+)
+from azure.ai.documentintelligence.models import AnalyzeDocumentRequest, AnalyzeResult
+from azure.core.credentials import AzureKeyCredential
+
+
+class AzureParser(Parser):
+    """
+    Parser that creates a hierarchical tree structure from Azure AI Document Intelligence response.
+    The hierarchy goes from document -> pages -> paragraphs/tables -> lines/cells -> words.
+
+    params:
+        api_key: str
+        endpoint: str
+        result_type: str = "markdown", "text"
+    """
+
+    def __init__(self, api_key: str, endpoint: str, result_type: str = "text"):
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.result_type = result_type
+        self.parser = None  # self._create_parser()
+        self.a_parser = None  # self._create_a_parser()
+
+    def _create_parser(self):
+        document_intelligence_client = DocumentIntelligenceClient(
+            endpoint=self.endpoint, credential=AzureKeyCredential(self.api_key)
+        )
+        return document_intelligence_client
+
+    def _create_a_parser(self):
+        parser = AsyncDocumentIntelligenceClient(
+            endpoint=self.endpoint, credential=AzureKeyCredential(self.api_key)
+        )
+        return parser
+
+    def _get_parser(self):
+        if not self.parser:
+            self.parser = self._create_parser()
+        return self.parser
+
+    def _get_a_parser(self):
+        if not self.a_parser:
+            self.a_parser = self._create_a_parser()
+        return self.a_parser
+
+    def _parse_file(self, file_path: str) -> Node:
+        """Parse an Azure Document Intelligence JSON file into a Node structure."""
+        with open(file_path) as file:
+            json_data = json.load(file)
+
+        return self._parse_json(json_data, file_path=file_path)
+
+    def _get_missing_paragraphs(self, json_data: dict) -> list[str]:
+        """Get missing paragraphs from the Azure Document Intelligence JSON data."""
+
+        sections = json_data.get("sections", [])
+        figures = json_data.get("figures", [])
+        tables = json_data.get("tables", [])
+
+        all_paragraphs = [
+            "/paragraphs/" + str(x) for x in range(len(json_data.get("paragraphs", [])))
+        ]
+
+        elements = []
+
+        def _process_section(section):
+            for element in section.get("elements", []):
+                if "paragraph" in element:
+                    elements.append(element)
+                elif "section" in element:
+                    section_idx = element.split("/")[2]
+                    next_section = sections[int(section_idx)]
+                    _process_section(next_section)
+
+        for section in sections:
+            _process_section(section)
+
+        def _process_figure(figure):
+            for element in figure.get("elements", []):
+                if "paragraph" in element:
+                    elements.append(element)
+                elif "section" in element:
+                    section_idx = element.split("/")[2]
+                    next_section = sections[int(section_idx)]
+                    _process_section(next_section)
+
+        for figure in figures:
+            _process_figure(figure)
+
+        def _process_table(table):
+            for element in table.get("elements", []):
+                if "paragraph" in element:
+                    elements.append(element)
+                elif "section" in element:
+                    section_idx = element.split("/")[2]
+                    next_section = sections[int(section_idx)]
+                    _process_section(next_section)
+
+        for table in tables:
+            _process_table(table)
+
+        missing = [x for x in all_paragraphs if x not in elements]
+
+        return missing
+
+    def _insert_missing_paragraphs(self, json_data: dict) -> dict:
+        """Insert missing paragraphs into the Azure Document Intelligence JSON data."""
+
+        missing = self._get_missing_paragraphs(json_data)
+
+        def _insert_paragraph_recursive(section, p_idx, p):
+            for i, element in enumerate(section.get("elements", [])):
+                if "paragraph" in element:
+                    if int(element.split("/")[2]) > int(p_idx):
+                        section["elements"].insert(i, p)
+                        return True
+                elif "section" in element:
+                    section_idx = element.split("/")[2]
+                    next_section = json_data["sections"][int(section_idx)]
+                    if _insert_paragraph_recursive(next_section, p_idx, p):
+                        return True
+            return False
+
+        for p in missing:
+            idx = int(p.split("/")[2])
+
+            for section in json_data.get("sections", []):
+                if _insert_paragraph_recursive(section, idx, p):
+                    break
+        return json_data
+
+    def _parse_json(self, json_data: dict, file_path: str) -> Node:
+        """
+        Parse Azure Document Intelligence JSON into a hierarchical Node structure.
+
+        Args:
+            json_data: The Azure Document Intelligence JSON response
+
+        Returns:
+            A Node representing the document with hierarchical structure
+        """
+        # Create root document node
+
+        json_data = self._insert_missing_paragraphs(json_data)
+
+        document_node = Node(
+            children=[],
+            metadata=self._extract_document_metadata(json_data),
+            node_type=NodeType.DOCUMENT,
+        )
+
+        # Process each page in the document
+        analyze_result = json_data  # .get('analyzeResult', {})
+        sections = analyze_result.get("sections", [])
+
+        document_node.children = self._process_children_elements(
+            sections[0], analyze_result, file_path=file_path
+        )
+
+        return document_node
+
+    def _process_children_elements(
+        self,
+        parent_object: dict[str, Any],
+        analyze_result: dict[str, Any],
+        file_path: str,
+    ) -> list[Node]:
+        """Process children elements of a section."""
+        children_nodes = []
+        elements = parent_object.get("elements", [])
+        for _element_idx, element in enumerate(elements):
+            if "paragraph" in element:
+                paragrap_index = element.split("/")[2]
+
+                paragraph = analyze_result.get("paragraphs", [])[int(paragrap_index)]
+                paragraph_node = self._create_paragraph_node(paragraph)
+                paragraph_node.children = self._process_children_elements(
+                    paragraph, analyze_result, file_path=file_path
+                )
+                children_nodes.append(paragraph_node)
+
+            elif "table" in element:
+                table_index = element.split("/")[2]
+                table = analyze_result.get("tables", [])[int(table_index)]
+                table_node = self._create_media_node(
+                    media=table,
+                    node_type=NodeType.TABLE,
+                    content_result=analyze_result.get("content", ""),
+                    file_path=file_path,
+                )
+                table_node.children = self._process_children_elements(
+                    table, analyze_result, file_path=file_path
+                )
+                children_nodes.append(table_node)
+
+            elif "figures" in element:
+                image_index = element.split("/")[2]
+                image = analyze_result.get("figures", [])[int(image_index)]
+                image_node = self._create_media_node(
+                    media=image,
+                    node_type=NodeType.FIGURE,
+                    content_result=analyze_result.get("content", ""),
+                    file_path=file_path,
+                )
+                image_node.children = self._process_children_elements(
+                    image, analyze_result, file_path=file_path
+                )
+                children_nodes.append(image_node)
+
+            elif "section" in element:
+                section_index = element.split("/")[2]
+                section = analyze_result.get("sections", [])[int(section_index)]
+                section_node = Node(children=[], node_type=NodeType.SECTION)
+                section_node.children = self._process_children_elements(
+                    section, analyze_result, file_path=file_path
+                )
+                children_nodes.append(section_node)
+
+        return children_nodes
+
+    def _transform_cells_to_markdown(
+        self, table_data: dict[str, Any], content_result: str
+    ) -> str:
+        """Transforms table cells from Azure response to a markdown table string."""
+        cells = table_data.get("cells", [])
+        if not cells:
+            return ""
+
+        offset = table_data.get("spans", [{}])[0].get("offset")
+        length = table_data.get("spans", [{}])[0].get("length")
+        if offset is None or length is None:
+            return ""
+
+        markdown_table = content_result[offset : offset + length]
+
+        return markdown_table
+
+    def _create_media_node(
+        self,
+        media: dict[str, Any],
+        node_type: NodeType,
+        content_result: str,
+        file_path: str,
+    ) -> Node:
+        """Create a node for an media with its child elements."""
+        # Get bounding regions
+        bounding_regions = media.get("boundingRegions", [])
+
+        if file_path and bounding_regions:
+            base64_image = extract_media(
+                coordinates=bounding_regions[0]["polygon"],
+                file_path=file_path,
+                page_number=bounding_regions[0]["pageNumber"],
+            )
+
+            media_obj = Media(
+                media_type="image",
+                source=base64_image,
+                source_type="base64",
+                extension="png",
+            )
+        else:
+            raise ValueError("No bounding regions found for media")
+
+        content = None
+        metadata = {
+            "boundingRegions": bounding_regions,
+        }
+        if node_type == NodeType.TABLE:
+            content = self._transform_cells_to_markdown(media, content_result)
+            metadata["rowCount"] = media.get("rowCount")
+            metadata["columnCount"] = media.get("columnCount")
+
+        # Create MediaNode with bounding regions metadata
+        image_node = MediaNode(
+            media=media_obj,
+            children=[],
+            node_type=node_type,
+            metadata=metadata,
+            content=content,
+        )
+        return image_node
+
+    def _extract_document_metadata(self, json_data: dict[str, Any]) -> dict[str, Any]:
+        """Extract document-level metadata from the Azure response."""
+        metadata = {}
+        analyze_result = json_data.get("analyzeResult", {})
+
+        # Add document-level metadata
+        if "documentResults" in analyze_result:
+            doc_results = analyze_result["documentResults"]
+            if doc_results and len(doc_results) > 0:
+                metadata.update(doc_results[0].get("fields", {}))
+
+        # Add model information if available
+        metadata["modelId"] = analyze_result.get("modelId")
+        metadata["apiVersion"] = analyze_result.get("apiVersion")
+
+        return metadata
+
+    # def _create_table_node(self, table: Dict[str, Any]) -> Node:
+    #     """Create a node for a table with its child lines and words."""
+    #     table_node = Node(
+    #         children=[],
+    #         node_type=NodeType.TABLE,
+    #         content=table.get("content", ""),
+    #         metadata={
+    #             "boundingRegions": table.get("boundingRegions", []),
+    #         },
+    #     )
+    #     return table_node
+
+    def _create_paragraph_node(self, paragraph: dict[str, Any]) -> Node:
+        """Create a node for a paragraph with its child lines and words."""
+        para_node = Node(
+            children=[],
+            node_type=NodeType.PARAGRAPH,
+            content=paragraph.get("content", ""),
+            metadata={
+                "boundingRegions": paragraph.get("boundingRegions", {}),
+            },
+        )
+        return para_node
+
+    def parse_with_azure_ai(self, file_path: str) -> dict:
+        """
+        Parse a Document with Azure AI Document Intelligence into a json dictionary.
+
+        Args:
+            file_path: Path to the document
+
+        Returns:
+            A dictionary with the Azure AI Document Intelligence response
+        """
+
+        with open(file_path, "rb") as file:
+            file_content = file.read()
+
+        parser = self._get_parser()
+        poller = parser.begin_analyze_document(
+            "prebuilt-layout",
+            AnalyzeDocumentRequest(bytes_source=file_content),
+            output_content_format=self.result_type,
+        )
+        result: AnalyzeResult = poller.result()
+        return result.as_dict()
+
+    async def a_parse_with_azure_ai(self, file_path: str) -> dict:
+        async with aiofiles.open(file_path, "rb") as file:
+            file_content = await file.read()
+
+        parser = self._get_a_parser()
+        async with parser:
+            poller = await parser.begin_analyze_document(
+                "prebuilt-layout",
+                AnalyzeDocumentRequest(bytes_source=file_content),
+                output_content_format=self.result_type,
+            )
+            result: AnalyzeResult = await poller.result()
+            return result.as_dict()
+
+    def parse(self, file_path: str) -> Node:
+        """
+        Parse a Document with Azure AI Document Intelligence into a Node structure.
+
+        Args:
+            file_path: Path to the document
+
+        Returns:
+            A Node representing the document with hierarchical structure
+        """
+        result_dict = self.parse_with_azure_ai(file_path)
+        return self._parse_json(result_dict, file_path=file_path)
+
+    def __call__(self, file_path: str) -> Node:
+        return self.parse(file_path)
+
+    async def a_parse(self, file_path: str) -> Node:
+        result_dict = await self.a_parse_with_azure_ai(file_path)
+        return self._parse_json(result_dict, file_path=file_path)

datapizza_ai_parsers_azure-0.0.2.dist-info/METADATA

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: datapizza-ai-parsers-azure
+Version: 0.0.2
+Summary: Azure Document Intelligence parser for the datapizza-ai framework
+Author-email: Datapizza <datapizza@datapizza.tech>
+License: MIT
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: <4,>=3.10.0
+Requires-Dist: aiofiles>=24.1.0
+Requires-Dist: azure-ai-documentintelligence<2.0.0,>=1.0.1
+Requires-Dist: datapizza-ai-core==0.0.1
+Requires-Dist: pillow>=11.3.0
+Requires-Dist: pymupdf<2.0.0,>=1.25.4

datapizza_ai_parsers_azure-0.0.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+datapizza/modules/parsers/azure/__init__.py,sha256=78Cmh6Swk3zlqhpXwy8RtGw97Yjw8E-_FARfUzzBgho,65
+datapizza/modules/parsers/azure/azure_parser.py,sha256=ZcdH2VJ0t2xox8sCz9n3X_0jsV4UfjZl2qVj-gZLGr4,14335
+datapizza_ai_parsers_azure-0.0.2.dist-info/METADATA,sha256=BcLLFTpsSkuis9UNRv8_TybXJOtMZAlQQSKeUW2htEY,583
+datapizza_ai_parsers_azure-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+datapizza_ai_parsers_azure-0.0.2.dist-info/RECORD,,

datapizza_ai_parsers_azure-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any