Pdf Datasource

class PdfDocument(BaseDocument):
    """Document representation for PDF file content.

    Extends BaseDocument to handle PDF-specific document processing including
    metadata filtering for embeddings and LLM contexts.

    Attributes:
        included_embed_metadata_keys: Metadata keys to include in embeddings
        included_llm_metadata_keys: Metadata keys to include in LLM context
        text: Document content in text format
        metadata: Extracted PDF metadata
        attachments: Dictionary of document attachments
        excluded_embed_metadata_keys: Metadata keys to exclude from embeddings
        excluded_llm_metadata_keys: Metadata keys to exclude from LLM context
    """

    included_embed_metadata_keys: List[str] = [
        "Title",
        "CreationDate",
        "ModDate",
        "creation_date",
        "client_name",
        "offer_name",
        "project_lead",
    ]
    included_llm_metadata_keys: List[str] = [
        "Title",
        "CreationDate",
        "ModDate",
        "creation_date",
        "client_name",
        "offer_name",
        "project_lead",
    ]

    def __init__(self, text: str, metadata: dict, attachments: dict = None):
        """Initialize PDF document.

        Args:
            text: Extracted text content
            metadata: PDF metadata dictionary
            attachments: Optional dictionary of attachments
        """
        super().__init__()
        self.text = text
        self.metadata = metadata
        self.attachments = attachments or {}
        self.excluded_embed_metadata_keys = self._set_excluded_metadata_keys(
            self.metadata, self.included_embed_metadata_keys
        )
        self.excluded_llm_metadata_keys = self._set_excluded_metadata_keys(
            self.metadata, self.included_llm_metadata_keys
        )

    @staticmethod
    def _set_excluded_metadata_keys(
        metadata: dict, included_keys: List[str]
    ) -> List[str]:
        """Determine metadata keys to exclude from processing.

        Args:
            metadata: Complete metadata dictionary
            included_keys: Keys to include in processing

        Returns:
            List[str]: Keys that should be excluded

        Note:
            Returns any key from metadata that isn't in included_keys
        """
        return [key for key in metadata.keys() if key not in included_keys]

class PdfDatasourceManager(DatasourceManager):
    """Manager for PDF content extraction and processing.

    Handles document retrieval, cleaning, splitting and embedding updates
    for PDF documents. Implements the base DatasourceManager
    interface for PDF-specific processing.
    """

    pass

class DefaultPDFParser:
    def parse(self, file_path: str) -> List[PdfDocument]:
        with open(file_path, "rb") as f:
            reader = pypdf.PdfReader(f)
            text = "\n\n".join(
                page.extract_text() or "" for page in reader.pages
            )
            metadata = self._extract_metadata(
                reader=reader, file_path=file_path
            )
            return [PdfDocument(text=text, metadata=metadata)]

    def _extract_metadata(
        self, reader: pypdf.PdfReader, file_path: str
    ) -> dict:
        """Extract and process PDF metadata.

        Args:
            reader: PDF reader instance

        Returns:
            dict: Processed metadata dictionary

        Note:
            Converts date strings to ISO format where possible
        """
        pdf_metadata = reader.metadata
        metadata = {
            "datasource": "pdf",
            "url": file_path,
            "title": os.path.basename(file_path),
        }
        if pdf_metadata is not None:
            for key, value in pdf_metadata.items():
                clean_key = key.strip("/")
                if clean_key in ["CreationDate", "ModDate"]:
                    date_str = value.strip("D:")
                    try:
                        parsed_date = datetime.strptime(
                            date_str[:14], "%Y%m%d%H%M%S"
                        )
                        metadata[clean_key] = parsed_date.isoformat()
                    except ValueError:
                        metadata[clean_key] = value
                else:
                    metadata[clean_key] = value
        return metadata

class NLMPDFParser:
    # Field patterns as class constant
    FIELDS_TO_EXTRACT = [
        {
            "name": "valid_until",
            "search_patterns": r"(?:Gültig bis|Valid until)\s*[:\s]*([\d/]+)",
        },
        {
            "name": "client_name",
            "search_patterns": r"(?:Client|Kunde)\s*[:\s]*([\S ]+?)(?=\s*\b(?:Quote No\.|Quote|Angebotsnummer|Date|Contact|Conta ct|Contents|Project Lead|Projektnummer)\b)",
        },
        {
            "name": "offer_name",
            "search_patterns": r"(?:Angebot|Quote)\s*[:\s]*([\S ]+?)(?=\s*\b(?:Datum|Date|Valid until|Contact|Project Lead|Projektleiter|Projektnummer)\b)",
        },
        {
            "name": "project_lead",
            "search_patterns": r"(?:Project\s*Lead|Projektleiter)\s*[:\s]*([\w\s.]+?)(?=\s*(Contact|Kontakt|Project Number|Quote Number|Valid until|$))",
        },
    ]

    def __init__(self, api_base: str):
        self.reader = LayoutPDFReader(api_base)

    def parse(self, file_path: str) -> List[PdfDocument]:
        """
        Parses the given PDF file and enriches its metadata with additional fields.

        Args:
            file_path (str): Path to the PDF file.

        Returns:
            List[PdfDocument]: List of enriched PdfDocument objects.
        """
        doc = self.reader.read_pdf(file_path)
        metadata = default_file_metadata_func(file_path)
        additional_metadata = self._extract_page_metadata(file_path)
        documents = []

        for chunk in doc.chunks():
            chunk_metadata = metadata.copy()
            chunk_metadata["page_label"] = chunk.page_idx
            enriched_metadata = {**chunk_metadata, **additional_metadata}
            documents.append(
                PdfDocument(
                    text=chunk.to_context_text(),
                    metadata=enriched_metadata,
                )
            )

        return documents

    def _extract_page_metadata(self, file_path: str) -> dict:
        """Extract metadata from first pages of PDF.

        Args:
            file_path: Path to PDF file

        Returns:
            dict: Extracted metadata fields
        """
        reader = pypdf.PdfReader(file_path)
        text = "".join(page.extract_text() or "" for page in reader.pages[:2])
        text = preprocess_text(text)
        return self._extract_fields(text, self.FIELDS_TO_EXTRACT)

    def _extract_fields(self, text: str, fields_to_extract: List[dict]) -> dict:
        extracted_fields = {}

        for field in fields_to_extract:
            match = re.search(field["search_patterns"], text, re.IGNORECASE)
            if match:
                extracted_fields[field["name"]] = match.group(1).strip()

        # Fallback/default values
        extracted_fields.setdefault("valid_until", "01/01/2024")
        extracted_fields.setdefault("client_name", "Unknown Client")
        extracted_fields.setdefault("offer_name", "Generic Offer")
        extracted_fields.setdefault("project_lead", "Not Assigned")

        return extracted_fields

class PdfReader(BaseReader[PdfDocument]):
    """Reader for extracting content from PDF files.

    Implements document extraction from PDF files with support for
    text and metadata extraction.

    Attributes:
        export_limit: Maximum number of documents to process
        base_path: Root directory containing PDF files
    """

    def __init__(self, configuration: PdfDatasourceConfiguration):
        """Initialize PDF reader.

        Args:
            configuration: Settings for PDF processing
        """
        super().__init__()
        self.export_limit = configuration.export_limit
        self.base_path = configuration.base_path

        if configuration.nlm_parser_enabled:
            self.parser = NLMPDFParser(configuration.nlm_parser_api_base)
        else:
            self.parser = DefaultPDFParser()

    def get_all_documents(self) -> List[PdfDocument]:
        documents = []
        pdf_files = [
            f for f in os.listdir(self.base_path) if f.endswith(".pdf")
        ]
        files_to_load = (
            pdf_files
            if self.export_limit is None
            else pdf_files[: self.export_limit]
        )

        for file_name in tqdm(files_to_load, desc="Loading PDFs"):
            file_path = os.path.join(self.base_path, file_name)
            if os.path.isfile(file_path):
                try:
                    parsed_docs = self.parser.parse(file_path)
                    documents.extend(parsed_docs)
                except Exception as e:
                    logging.error(f"Failed to load PDF {file_name}: {str(e)}")

        return documents

    async def get_all_documents_async(self) -> List[PdfDocument]:
        """Load documents asynchronously from configured path.

        Returns:
            List[PdfDocument]: Collection of processed documents

        Note:
            Currently calls synchronous implementation
        """
        return self.get_all_documents()

class PdfCleanerBuilder:
    """Builder for creating PDF content cleaner instances.

    Provides factory method to create Cleaner objects for PDF content.
    """

    @staticmethod
    @inject
    def build() -> Cleaner:
        """Creates a content cleaner for PDFs.

        Returns:
            Cleaner: Configured cleaner instance
        """
        return Cleaner()

class PdfDatasourceManagerBuilder:
    """Builder for creating PDF datasource manager instances.

    Provides factory method to create configured PdfDatasourceManager
    with required components for content processing.
    """

    @staticmethod
    @inject
    def build(
        configuration: PdfDatasourceConfiguration,
        reader: PdfReader,
        cleaner: Cleaner,
        splitter: BoundEmbeddingModelMarkdownSplitter,
    ) -> PdfDatasourceManager:
        """Creates a configured PDF datasource manager.

        Args:
            configuration: Settings for PDF processing
            reader: Component for reading PDF content
            cleaner: Component for cleaning extracted content
            splitter: Component for splitting content into chunks

        Returns:
            PdfDatasourceManager: Configured manager instance
        """
        return PdfDatasourceManager(
            configuration=configuration,
            reader=reader,
            cleaner=cleaner,
            splitter=splitter,
        )

class PdfReaderBuilder:
    """Builder for creating PDF reader instances.

    Provides factory method to create configured PdfReader objects.
    """

    @staticmethod
    @inject
    def build(
        configuration: PdfDatasourceConfiguration,
    ) -> PdfReader:
        """Creates a configured PDF reader.

        Args:
            configuration: Settings for PDF processing

        Returns:
            PdfReader: Configured reader instance
        """
        return PdfReader(
            configuration=configuration,
        )