PDF読み取りPython

指定のフォルダのPDFファイルを順次読んで、そのメタデータをテーブルに保存する。
そして、その内容をMDファイルに変換し、単元でチャンクして子テーブルとして保存する。(1対多)

getpdf.py

import os
import re
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

import pymupdf
import pymupdf4llm
import sqlcipher3 as sqlcipher
from dotenv import load_dotenv
from langchain_text_splitters import MarkdownHeaderTextSplitter


def clean_markdown(text):
    text = re.sub(r'<sup>(.*?)</sup>', r'[\1]', text)
    return text

def init_db(db_path="pdf_info.db"):
    load_dotenv("/workspaces/envDBPwdDir/.env")
    PASSWORD = os.getenv("DB_PASSWORD")
    conn = sqlcipher.connect(db_path)
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA key = '{PASSWORD}';")
    cursor = conn.cursor()
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS pdf_meta (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            file_name TEXT,
            file_path TEXT UNIQUE,
            title TEXT,
            author TEXT,
            subject TEXT,
            keywords TEXT,
            creator TEXT,
            producer TEXT,
            page_count INTEGER,
            scanned_at TEXT
        )
    """
    )
    conn.commit()

    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS pdf_chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            parent_id INTEGER,
            chunk_index INTEGER,
            content TEXT,
            section_title TEXT,
            FOREIGN KEY (parent_id) REFERENCES pdf_meta(id) ON DELETE CASCADE
        )
    """
    )
    conn.commit()

    return conn


def save_pdfs_to_db(folder_path, db_path="pdf_info.db"):
    target_dir = Path(folder_path)
    if not target_dir.exists():
        print("指定されたフォルダが存在しません。")
        return

    global conn
    conn = init_db(db_path)
    global cursor
    cursor = conn.cursor()
    success_count = 0

    for file_path in target_dir.glob("**/*.pdf"):
        abs_path = str(file_path.resolve())
        cursor.execute(
            "SELECT id FROM pdf_meta WHERE file_path = ?", (abs_path,)
        )
        if cursor.fetchone():
            continue
        try:
            doc = pymupdf.open(file_path)
            meta = doc.metadata  # メタデータを辞書形式で取得
            page_count = len(doc)
            doc.close()  # 読み終わったら閉じる

            title = meta.get("title") or None
            author = meta.get("author") or None
            subject = meta.get("subject") or None
            keywords = meta.get("keywords") or None
            creator = meta.get("creator") or None
            producer = meta.get("producer") or None

            current_time = datetime.now(ZoneInfo("Asia/Tokyo")).strftime("%Y-%m-%d %H:%M:%S")
            cursor.execute(
                """
                INSERT INTO pdf_meta (
                    file_name, file_path, title, author, subject, keywords, creator, producer, page_count, scanned_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
                (
                    file_path.name,
                    abs_path,
                    title,
                    author,
                    subject,
                    keywords,
                    creator,
                    producer,
                    page_count,
                    current_time,
                ),
            )
            success_count += 1
            cursor.execute(f"SELECT id FROM pdf_meta WHERE file_path = '{abs_path}';")
            p_result = cursor.fetchone()
            if p_result:
                pdf_id = p_result[0]
            md_text = pymupdf4llm.to_markdown(abs_path)
            md_text = clean_markdown(md_text)
            insert_chunks_to_db(pdf_id,md_text)

        except Exception as e: # noqa: BLE001
            print(f"エラー(スキップします) {file_path.name}: {e}")

    conn.commit()
    conn.close()
    print(f"\n処理完了: {success_count} 件のPDFを新しくDBに登録しました。")

def insert_chunks_to_db(parent_id, full_md_text, db_path="pdf_info.db"):
    # ① どの見出しレベルで分割するかを設定
    headers_to_split_on = [
        ("#", "Header_1"),
        ("##", "Header_2"),
        ("###", "Header_3"),
    ]
    markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)

    # ② マークダウンテキストを見出しごとにチャンク分割
    chunks = markdown_splitter.split_text(full_md_text)
    # ③ SQLiteへ一括インサートするためのデータリストを作成
    insert_data = []
    for index, chunk in enumerate(chunks):
        # 見出し情報を結合して「第1章 > 概要」のようなタイトル文字列を作る
        meta = chunk.metadata
        titles = [meta[key] for key in ["Header_1", "Header_2", "Header_3"] if key in meta]
        section_title = " > ".join(titles) if titles else "(見出しなし)"

        # 1レコード分のデータ(親ID, チャンク連番, 本文, 見出しタイトル)
        insert_data.append((
            parent_id,
            index,
            chunk.page_content,
            section_title
        ))
    # ④ データベースに高速に一括挿入(executemany)
    if insert_data:
        print(insert_data[1])
        cursor.executemany("""
            INSERT INTO pdf_chunks (parent_id, chunk_index, content, section_title)
            VALUES (?, ?, ?, ?);
        """, insert_data)
        conn.commit()
        print(f"親ID [{parent_id}]: {len(insert_data)} 件のチャンクを明細テーブルに登録しました。")

# --- 実行 ---
if __name__ == "__main__":
    target_folder = "../pdfs"
    save_pdfs_to_db(target_folder)