ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/loaders/pdf.py

@@ -0,0 +1,373 @@
+from collections.abc import Callable
+from typing import List, Optional, Union
+import re
+from pathlib import Path, PurePath
+import fitz
+import pymupdf4llm
+from ..stores.models import Document
+from .abstract import AbstractLoader
+
+class PDFLoader(AbstractLoader):
+    """
+    Advanced PDF Loader using PyMuPDF (fitz).
+    - Skips image-only pages.
+    - Combines title-only pages with next content page.
+    - Preserves tables as text for chatbot/RAG KB usage.
+    - Returns a Parrot Document per logical page.
+    - Supports chapter-based splitting for markdown output.
+    """
+
+    extensions: List[str] = {'.pdf'}
+
+    def __init__(
+        self,
+        source: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
+        *,
+        tokenizer: Union[str, Callable] = None,
+        text_splitter: Union[str, Callable] = None,
+        source_type: str = 'file',
+        as_markdown: bool = False,
+        use_chapters: bool = False,
+        use_pages: bool = False,
+        **kwargs
+    ):
+        super().__init__(
+            source,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            **kwargs
+        )
+        self.doctype = 'pdf'
+        self._source_type = source_type
+        self.as_markdown = as_markdown
+        self.use_chapters = use_chapters
+        self.use_pages = use_pages
+
+    def is_title_only(self, text: str, min_len: int = 5, max_len: int = 50) -> bool:
+        """Check if text looks like a title (short, single line, large font)."""
+        lines = [l for l in text.strip().split('\n') if l.strip()]
+        if len(lines) == 1 and min_len <= len(lines[0]) <= max_len:
+            return True
+        return False
+
+    def is_image_only(self, page: fitz.Page) -> bool:
+        """Return True if the page only contains images (no visible text)."""
+        text = page.get_text("text").strip()
+        if text:
+            return False
+        # Has no text, check if images exist
+        img_list = page.get_images(full=True)
+        return len(img_list) > 0
+
+    def is_table_like(self, text: str) -> bool:
+        """Naive check: Table if lines have multiple columns (lots of |, tab, or spaces)."""
+        lines = [l for l in text.split('\n') if l.strip()]
+        if not lines:
+            return False
+        count_table_lines = sum(1 for l in lines if ('|' in l or '\t' in l or (len(l.split()) > 3)))
+        return (count_table_lines > len(lines) // 2) and len(lines) > 2
+
+    def extract_table(self, page: fitz.Page) -> Optional[str]:
+        """Attempt to extract table structure, return as markdown if detected, else None."""
+        # PyMuPDF can't extract structured tables, so fallback to plain text with basic cleanup
+        text = page.get_text("text")
+        lines = [l.strip() for l in text.split('\n') if l.strip()]
+        # Try to join lines with | if possible
+        if not lines:
+            return None
+        # Heuristic: If tab separated or lots of spaces, format as a markdown table
+        table_lines = []
+        for l in lines:
+            if '\t' in l:
+                cells = [c.strip() for c in l.split('\t')]
+                table_lines.append("| " + " | ".join(cells) + " |")
+            elif '|' in l:
+                table_lines.append(l)
+            else:
+                # Split by multiple spaces
+                cells = [c.strip() for c in l.split("  ") if c.strip()]
+                if len(cells) > 2:
+                    table_lines.append("| " + " | ".join(cells) + " |")
+                else:
+                    table_lines.append(l)
+        if table_lines:
+            # Add markdown header if more than 2 columns
+            if len(table_lines) > 1 and table_lines[0].count('|') == table_lines[1].count('|'):
+                ncols = table_lines[0].count('|') - 1
+                if ncols > 1:
+                    header_sep = "| " + " | ".join(['---'] * ncols) + " |"
+                    table_lines.insert(1, header_sep)
+            return "\n".join(table_lines)
+        return None
+
+    def extract_chapters_from_markdown(self, md_text: str) -> List[dict]:
+        """
+        Extract chapters from markdown text based on headers.
+        Returns list of dicts with 'title' and 'content' keys.
+        """
+        chapters = []
+
+        # Split by horizontal rules and headers
+        # Look for patterns like: -----\n**TITLE**\n or # Title
+
+        # First, let's handle the horizontal rule + bold title pattern
+        sections = re.split(r'\n-----+\n', md_text)
+
+        for i, section in enumerate(sections):
+            section = section.strip()
+            if not section:
+                continue
+
+            # Look for bold titles at the beginning of sections
+            title_match = re.match(r'^\*\*([^*]+)\*\*', section)
+            if title_match:
+                title = title_match.group(1).strip()
+                # Get content after the title
+                content = re.sub(r'^\*\*[^*]+\*\*\s*', '', section, count=1).strip()
+            else:
+                # Look for markdown headers (# ## ###)
+                header_match = re.match(r'^(#{1,6})\s*(.+?)$', section, re.MULTILINE)
+                if header_match:
+                    title = header_match.group(2).strip()
+                    # Get content after the header
+                    content = re.sub(r'^#{1,6}\s*.+?$', '', section, count=1, flags=re.MULTILINE).strip()
+                else:
+                    # No clear title found, use section number or first line
+                    lines = section.split('\n')
+                    if lines:
+                        title = f"Section {i+1}" if not lines[0].strip() else lines[0][:50] + "..."
+                        content = section
+                    else:
+                        continue
+
+            # Skip if content is too short (less than 10 characters)
+            if len(content.strip()) < 10:
+                self.logger.info(f"Skipping chapter '{title}' - content too short")
+                continue
+
+            chapters.append({
+                'title': title,
+                'content': content,
+                'chapter_number': len(chapters) + 1
+            })
+
+        return chapters
+
+    def extract_pages_from_markdown(self, md_text: str) -> List[dict]:
+        """
+        Extract pages from markdown text based on page separators.
+        Returns list of dicts with 'title' and 'content' keys.
+        """
+        pages = []
+
+        # Split by page indicators (common patterns)
+        page_patterns = [
+            r'\n-----+\n',  # Horizontal rules
+            r'Slide \d+',   # Slide indicators
+            r'Page \d+',    # Page indicators
+        ]
+
+        # Try to split by the most common pattern first
+        sections = re.split(r'\n-----+\n', md_text)
+
+        for i, section in enumerate(sections):
+            section = section.strip()
+            if not section or len(section) < 10:
+                continue
+
+            # Extract title from the beginning of the page
+            lines = section.split('\n')
+            title = None
+            content_start = 0
+
+            # Look for bold title or header at the beginning
+            for j, line in enumerate(lines[:3]):  # Check first 3 lines
+                line = line.strip()
+                if re.match(r'^\*\*([^*]+)\*\*$', line):
+                    title = re.match(r'^\*\*([^*]+)\*\*$', line).group(1)
+                    content_start = j + 1
+                    break
+                elif re.match(r'^#{1,6}\s*(.+?)$', line):
+                    title = re.match(r'^#{1,6}\s*(.+?)$', line).group(1)
+                    content_start = j + 1
+                    break
+
+            if not title:
+                title = f"Page {i+1}"
+
+            # Get content after title
+            content = '\n'.join(lines[content_start:]).strip()
+
+            if len(content) < 10:
+                continue
+
+            pages.append({
+                'title': title,
+                'content': content,
+                'page_number': i + 1
+            })
+
+        return pages
+
+    async def _load(self, path: PurePath, **kwargs) -> List[Document]:
+        self.logger.info(f"Loading PDF file: {path}")
+        docs = []
+        all_text = []   # ← For summary collection
+        doc = fitz.open(str(path))
+        if self.as_markdown:
+            md_text = pymupdf4llm.to_markdown(path)
+            if self.use_chapters:
+                # Split by chapters
+                chapters = self.extract_chapters_from_markdown(md_text)
+                self.logger.info(f"Found {len(chapters)} chapters")
+                for chapter in chapters:
+                    document_meta = {
+                        "filename": path.name,
+                        "file_path": str(path),
+                        "chapter_title": chapter['title'],
+                        "chapter_number": chapter['chapter_number'],
+                        "content_type": "chapter"
+                    }
+                    meta = self.create_metadata(
+                        path=path,
+                        doctype="pdf",
+                        source_type="pdf_chapter",
+                        doc_metadata=document_meta,
+                    )
+                    # Combine title and content
+                    full_content = f"# {chapter['title']}\n\n{chapter['content']}"
+                    docs.append(
+                        self.create_document(
+                            content=full_content,
+                            path=path,
+                            metadata=meta
+                        )
+                    )
+            elif self.use_pages:
+                # Split by pages
+                pages = self.extract_pages_from_markdown(md_text)
+                self.logger.info(f"Found {len(pages)} pages")
+
+                for page in pages:
+                    document_meta = {
+                        "filename": path.name,
+                        "file_path": str(path),
+                        "page_title": page['title'],
+                        "page_number": page['page_number'],
+                        "content_type": "page"
+                    }
+
+                    meta = self.create_metadata(
+                        path=path,
+                        doctype="pdf",
+                        source_type="pdf_page",
+                        doc_metadata=document_meta,
+                    )
+
+                    # Combine title and content
+                    full_content = f"## {page['title']}\n\n{page['content']}"
+
+                    docs.append(
+                        self.create_document(
+                            content=full_content,
+                            path=path,
+                            metadata=meta
+                        )
+                    )
+            else:
+                # Return whole markdown as single document
+                document_meta = {
+                    "filename": path.name,
+                    "file_path": str(path),
+                    "content_type": "full_document"
+                }
+                meta = self.create_metadata(
+                    path=path,
+                    doctype="pdf",
+                    source_type="pdf_markdown",
+                    doc_metadata=document_meta,
+                )
+                docs.append(
+                    self.create_document(
+                        content=md_text,
+                        path=path,
+                        metadata=meta
+                    )
+                )
+        else:
+            # Use the default text extraction page-based
+            pending_title = None
+            for i, page in enumerate(doc):
+                page_text = page.get_text("text").strip()
+                if self.is_image_only(page):
+                    self.logger.info(f"Page {i+1}: image-only, skipping.")
+                    continue
+
+                # Title-only page: store to prepend to next content
+                if self.is_title_only(page_text):
+                    self.logger.info(f"Page {i+1}: title-only, saving for next page.")
+                    pending_title = page_text
+                    continue
+
+                # Table page: try to preserve structure
+                if self.is_table_like(page_text):
+                    table_md = self.extract_table(page)
+                    if table_md:
+                        content = (pending_title + '\n\n' if pending_title else '') + table_md
+                        pending_title = None
+                    else:
+                        content = (pending_title + '\n\n' if pending_title else '') + page_text
+                        pending_title = None
+                else:
+                    content = (pending_title + '\n\n' if pending_title else '') + page_text
+                    pending_title = None
+
+                document_meta = {
+                    "filename": path.name,
+                    "file_path": str(path),
+                    "page_number": i + 1,
+                    # "title": doc.metadata.get("title", ""),
+                    # "creationDate": doc.metadata.get("creationDate", ""),
+                    # "author": doc.metadata.get("author", ""),
+                }
+                meta = self.create_metadata(
+                    path=path,
+                    doctype="pdf",
+                    source_type="pdf",
+                    doc_metadata=document_meta,
+                )
+                if len(content) < 10:
+                    self.logger.warning(
+                        f"Page {i+1} content too short, skipping."
+                    )
+                    continue
+                docs.append(
+                    self.create_document(
+                        content=content,
+                        path=path,
+                        metadata=meta
+                    )
+                )
+                all_text.append(content)
+            doc.close()
+            # --- Summarization step ---
+            full_text = "\n\n".join(all_text)
+            summary = await self.summary_from_text(full_text)
+            if summary:
+                summary_meta = self.create_metadata(
+                    path=path,
+                    doctype=self.doctype,
+                    source_type=self._source_type,
+                    doc_metadata={
+                        "summary_for_pages": len(docs),
+                    }
+                )
+                docs.append(
+                    self.create_document(
+                        content=f"SUMMARY:\n\n{summary}",
+                        path=path,
+                        metadata=summary_meta
+                    )
+                )
+        return docs

parrot/loaders/pdfmark.py

@@ -0,0 +1,320 @@
+from typing import Any, Union, List
+import logging
+from collections.abc import Callable
+from pathlib import PurePath
+import fitz
+from ..stores.models import Document
+from .basepdf import BasePDF
+# Option 1: Use MarkItDown (Microsoft's universal document converter)
+try:
+    from markitdown import MarkItDown
+    MARKITDOWN_AVAILABLE = True
+except ImportError:
+    MARKITDOWN_AVAILABLE = False
+
+# Option 2: Use pymupdf4llm (updated PyMuPDF library)
+try:
+    import pymupdf4llm
+    PYMUPDF4LLM_AVAILABLE = True
+except ImportError:
+    PYMUPDF4LLM_AVAILABLE = False
+
+
+logger = logging.getLogger('pdfminer').setLevel(logging.INFO)
+
+class PDFMarkdownLoader(BasePDF):
+    """
+    Loader for PDF files converted content to markdown.
+
+    This loader supports multiple backends for PDF to markdown conversion:
+    1. MarkItDown (Microsoft's universal document converter)
+    2. pymupdf4llm (PyMuPDF's markdown converter)
+    3. Fallback manual conversion using PyMuPDF
+    """
+
+    extensions: List[str] = {'.pdf'}
+
+    def __init__(
+        self,
+        source: Union[str, PurePath, List[PurePath]],
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        markdown_backend: str = "auto",  # "markitdown", "pymupdf4llm", "manual", "auto"
+        chunk_size: int = 1024,
+        chunk_overlap: int = 10,
+        preserve_tables: bool = True,
+        extract_images: bool = False,
+        **kwargs
+    ):
+        super().__init__(
+            source=source,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            **kwargs
+        )
+        self._language = language
+        self.markdown_backend = self._select_backend(markdown_backend)
+        self.preserve_tables = preserve_tables
+        self.extract_images = extract_images
+
+        # Initialize markdown splitter
+        self._splitter = self._get_markdown_splitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap
+        )
+
+        # Initialize conversion backend
+        self._setup_conversion_backend()
+
+    def _select_backend(self, preferred: str) -> str:
+        """Select the best available backend for PDF to markdown conversion."""
+        if preferred == "auto":
+            if MARKITDOWN_AVAILABLE:
+                return "markitdown"
+            elif PYMUPDF4LLM_AVAILABLE:
+                return "pymupdf4llm"
+            else:
+                return "manual"
+        elif preferred == "markitdown" and MARKITDOWN_AVAILABLE:
+            return "markitdown"
+        elif preferred == "pymupdf4llm" and PYMUPDF4LLM_AVAILABLE:
+            return "pymupdf4llm"
+        elif preferred == "manual":
+            return "manual"
+        else:
+            # Fallback to available backend
+            self.logger.warning(f"Preferred backend '{preferred}' not available, using fallback")
+            return self._select_backend("auto")
+
+    def _setup_conversion_backend(self):
+        """Initialize the selected conversion backend."""
+        if self.markdown_backend == "markitdown":
+            self.md_converter = MarkItDown()
+            self.logger.info("Using MarkItDown backend for PDF to markdown conversion")
+        elif self.markdown_backend == "pymupdf4llm":
+            self.logger.info("Using pymupdf4llm backend for PDF to markdown conversion")
+        else:
+            self.logger.info("Using manual PyMuPDF backend for PDF to markdown conversion")
+
+    def _convert_to_markdown_markitdown(self, path: Union[str, PurePath]) -> str:
+        """Convert PDF to markdown using MarkItDown."""
+        try:
+            result = self.md_converter.convert(str(path))
+            return result.text_content if result else ""
+        except Exception as e:
+            self.logger.error(f"MarkItDown conversion failed: {e}")
+            return self._convert_to_markdown_manual(path)
+
+    def _convert_to_markdown_pymupdf4llm(self, path: Union[str, PurePath]) -> str:
+        """Convert PDF to markdown using pymupdf4llm."""
+        try:
+            return pymupdf4llm.to_markdown(str(path))
+        except Exception as e:
+            self.logger.error(f"pymupdf4llm conversion failed: {e}")
+            return self._convert_to_markdown_manual(path)
+
+    def _convert_to_markdown_manual(self, path: Union[str, PurePath]) -> str:
+        """Fallback manual conversion using PyMuPDF with basic markdown formatting."""
+        try:
+            doc = fitz.open(str(path))
+            markdown_text = []
+
+            for _, page_num in enumerate(doc):
+                page = doc[page_num]
+
+                # Extract text blocks with formatting
+                blocks = page.get_text("dict")["blocks"]
+
+                for block in blocks:
+                    if "lines" in block:
+                        block_text = []
+                        for line in block["lines"]:
+                            line_text = ""
+                            for span in line["spans"]:
+                                text = span["text"]
+                                font_size = span.get("size", 12)
+                                flags = span.get("flags", 0)
+
+                                # Basic formatting based on font properties
+                                if font_size > 16:
+                                    text = f"# {text}"
+                                elif font_size > 14:
+                                    text = f"## {text}"
+                                elif font_size > 12:
+                                    text = f"### {text}"
+
+                                # Bold text
+                                if flags & 2**4:  # Bold flag
+                                    text = f"**{text}**"
+
+                                # Italic text
+                                if flags & 2**6:  # Italic flag
+                                    text = f"*{text}*"
+
+                                line_text += text
+
+                            if line_text.strip():
+                                block_text.append(line_text)
+
+                        if block_text:
+                            markdown_text.append("\n".join(block_text))
+
+                # Extract tables if requested
+                if self.preserve_tables:
+                    tables = page.find_tables()
+                    for table in tables:
+                        try:
+                            table_data = table.extract()
+                            if table_data:
+                                markdown_table = self._format_table_as_markdown(table_data)
+                                if markdown_table:
+                                    markdown_text.append(markdown_table)
+                        except Exception as e:
+                            self.logger.debug(f"Failed to extract table: {e}")
+
+            doc.close()
+            return "\n\n".join(markdown_text)
+
+        except Exception as e:
+            self.logger.error(f"Manual PDF conversion failed: {e}")
+            return ""
+
+    def _format_table_as_markdown(self, table_data: List[List[str]]) -> str:
+        """Convert table data to markdown format."""
+        if not table_data or len(table_data) < 1:
+            return ""
+
+        markdown_rows = []
+
+        # Header row
+        header_row = " | ".join(str(cell) if cell else "" for cell in table_data[0])
+        markdown_rows.append(f"| {header_row} |")
+
+        # Separator row
+        separator = " | ".join("---" for _ in table_data[0])
+        markdown_rows.append(f"| {separator} |")
+
+        # Data rows
+        for row in table_data[1:]:
+            data_row = " | ".join(str(cell) if cell else "" for cell in row)
+            markdown_rows.append(f"| {data_row} |")
+
+        return "\n".join(markdown_rows)
+
+    async def _load(self, path: Union[str, PurePath, List[PurePath]], **kwargs) -> List[Document]:
+        """
+        Load a PDF file and convert to markdown format.
+
+        Args:
+            path (Union[str, PurePath, List[PurePath]]): The path to the PDF file.
+
+        Returns:
+            List[Document]: A list of AI-Parrot Documents.
+        """
+        self.logger.info(f"Loading PDF file: {path}")
+        docs = []
+
+        # Convert to markdown using selected backend
+        if self.markdown_backend == "markitdown":
+            md_text = self._convert_to_markdown_markitdown(path)
+        elif self.markdown_backend == "pymupdf4llm":
+            md_text = self._convert_to_markdown_pymupdf4llm(path)
+        else:
+            md_text = self._convert_to_markdown_manual(path)
+
+        if not md_text.strip():
+            self.logger.warning(f"No markdown content extracted from {path}")
+            return docs
+
+        # Extract PDF metadata
+        try:
+            pdf = fitz.open(str(path))
+            pdf_metadata = pdf.metadata  # pylint: disable=E1101  # noqa: E1101
+            pdf.close()
+        except Exception as e:
+            self.logger.warning(
+                f"Could not extract PDF metadata: {e}"
+            )
+            pdf_metadata = {}
+
+        # Generate summary if enabled
+        try:
+            summary = await self.summary_from_text(md_text)
+        except Exception as e:
+            self.logger.warning(
+                f"Summary generation failed: {e}"
+            )
+            summary = ''
+
+        # Create base metadata
+        base_metadata = {
+            "url": '',
+            "filename": path.name if hasattr(path, 'name') else str(path).rsplit('/', maxsplit=1)[-1],  # noqa
+            "source": str(path.name if hasattr(path, 'name') else path),
+            "type": 'pdf',
+            "data": {},
+            "category": self.category,
+            "source_type": self._source_type,
+            "conversion_backend": self.markdown_backend,
+            "document_meta": {
+                "title": pdf_metadata.get("title", ""),
+                "creationDate": pdf_metadata.get("creationDate", ""),
+                "author": pdf_metadata.get("author", ""),
+            }
+        }
+
+        # Add summary document if available
+        if summary:
+            summary_metadata = {
+                **base_metadata,
+                "content_type": "summary"
+            }
+            docs.append(
+                Document(
+                    page_content=summary,
+                    metadata=summary_metadata
+                )
+            )
+
+        # Split markdown content into chunks
+        try:
+            chunks = self._splitter.split_text(md_text)
+            self.logger.info(f"Split document into {len(chunks)} chunks")
+        except Exception as e:
+            self.logger.error(
+                f"Failed to split text: {e}"
+            )
+            # Fallback: use the entire text as one chunk
+            chunks = [md_text]
+
+        # Create documents for each chunk
+        for chunk_index, chunk in enumerate(chunks):
+            chunk_metadata = {
+                **base_metadata,
+                "content_type": "chunk",
+                "chunk_index": chunk_index,
+                "total_chunks": len(chunks)
+            }
+
+            docs.append(
+                Document(
+                    page_content=chunk,
+                    metadata=chunk_metadata
+                )
+            )
+
+        return docs
+
+    def get_supported_backends(self) -> List[str]:
+        """Get list of available conversion backends."""
+        backends = ["manual"]  # Always available
+
+        if MARKITDOWN_AVAILABLE:
+            backends.append("markitdown")
+        if PYMUPDF4LLM_AVAILABLE:
+            backends.append("pymupdf4llm")
+
+        return backends