画像や図表も検索可能に：Multimodal RAGで解決する非構造化データの壁

pip install llama-index-core llama-index-readers-file llama-index-llms-openai llama-index-multi-modal-llms-openai llama-index-embeddings-openai python-dotenv

import logging
import os
import sys
from typing import List, Optional

from dotenv import load_dotenv
from llama_index.core import (
    Settings,
    SimpleDirectoryReader,
    VectorStoreIndex,
    StorageContext,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import BaseNode, ImageNode
from llama_index.multi_modal_llms.openai import OpenAIMultiModal
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# ロギングの設定
logging.basicConfig(
    stream=sys.stdout,
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# 環境変数の読み込み
load_dotenv()

class MultimodalRAGPipeline:
    def __init__(
        self,
        input_dir: str,
        model_name: str = "gpt-4o",
        embed_model_name: str = "text-embedding-3-small",
        persist_dir: str = "./storage"
    ):
        """
        Multimodal RAGパイプラインの初期化
        
        Args:
            input_dir (str): PDFが格納されているディレクトリパス
            model_name (str): 使用するLLMモデル名
            embed_model_name (str): 使用する埋め込みモデル名
            persist_dir (str): インデックスを保存するディレクトリ
        """
        self.input_dir = input_dir
        self.persist_dir = persist_dir
        
        # APIキーのチェック
        if not os.getenv("OPENAI_API_KEY"):
            logger.error("OPENAI_API_KEYが設定されていません。")
            raise ValueError("Missing OpenAI API Key")

        # LLMとEmbeddingモデルの設定
        try:
            self.llm = OpenAI(model=model_name)
            self.embed_model = OpenAIEmbedding(model=embed_model_name)
            
            # マルチモーダルLLM（画像理解用）の設定
            self.multi_modal_llm = OpenAIMultiModal(model=model_name)
            
            Settings.llm = self.llm
            Settings.embed_model = self.embed_model
            
            logger.info(f"モデルの初期化完了: LLM={model_name}, Embedding={embed_model_name}")
            
        except Exception as e:
            logger.error(f"モデル初期化中にエラーが発生しました: {e}")
            raise

    def load_documents(self) -> List[BaseNode]:
        """
        ドキュメントを読み込み、画像とテキストを抽出する
        
        Returns:
            List[BaseNode]: 抽出されたノードのリスト
        """
        logger.info(f"ディレクトリ '{self.input_dir}' からドキュメントを読み込み中...")
        
        try:
            # 画像を自動抽出するリーダー
            reader = SimpleDirectoryReader(
                self.input_dir,
                required_exts=[".pdf", ".jpg", ".png"],
                recursive=True,
                # 画像を抽出してImageNodeとして扱う設定
                file_metadata=lambda x: {"file_name": x}
            )
            documents = reader.load_data()
            logger.info(f"ドキュメント読み込み成功: {len(documents)} 件のドキュメント")
            
            # ノードパーサーの設定（テキスト用）
            text_parser = SentenceSplitter(
                chunk_size=1024,
                chunk_overlap=20
            )
            
            text_nodes = []
            image_nodes = []
            
            for doc in documents:
                if doc.image_embeds is not None or isinstance(doc, ImageNode):
                    image_nodes.append(doc)
                else:
                    # テキストノードの分割処理
                    text_nodes.extend(text_parser.get_nodes_from_documents([doc]))
            
            logger.info(f"ノード分割完了: テキスト={len(text_nodes)}, 画像={len(image_nodes)}")
            
            return text_nodes + image_nodes

        except FileNotFoundError:
            logger.error(f"ディレクトリが見つかりません: {self.input_dir}")
            raise
        except Exception as e:
            logger.error(f"ドキュメント読み込み中に予期せぬエラー: {e}")
            raise

    def create_image_summaries(self, image_nodes: List[ImageNode]) -> List[BaseNode]:
        """
        画像ノードに対して、Vision LLMを用いて要約テキストを生成する
        
        Args:
            image_nodes (List[ImageNode]): 画像ノードのリスト
            
        Returns:
            List[BaseNode]: 要約テキストを含むノードリスト
        """
        if not image_nodes:
            logger.info("要約対象の画像ノードはありません。")
            return []

        logger.info(f"{len(image_nodes)} 枚の画像に対して要約生成を開始します...")
        processed_nodes = []

        for img_node in image_nodes:
            try:
                # 画像パスの取得
                image_path = img_node.metadata.get("file_path")
                if not image_path or not os.path.exists(image_path):
                    logger.warning(f"画像ファイルが見つかりません: {image_path}, スキップします。")
                    continue

                # プロンプトの作成
                prompt = """
                この画像を詳細に説明してください。特に、グラフの場合は数値の推移や傾向、
                表の場合はキーとなるデータポイントを抽出してテキスト化してください。
                説明は日本語で行い、検索しやすいような具体的なキーワードを含めてください。
                """
                
                # Vision LLMによる画像の理解と要約生成
                response = self.multi_modal_llm.complete(
                    prompt=prompt,
                    image_documents=[img_node]
                )
                
                summary_text = response.text
                logger.info(f"画像要約生成成功 ({os.path.basename(image_path)}): {summary_text[:50]}...")
                
                # 要約テキストを新しいノードとして作成し、元の画像への参照を保持
                summary_node = BaseNode(
                    text=summary_text,
                    metadata={
                        "type": "image_summary",
                        "source_image_path": image_path,
                        "file_name": img_node.metadata.get("file_name", "unknown")
                    }
                )
                processed_nodes.append(summary_node)
                
            except Exception as e:
                logger.error(f"画像要約生成中にエラーが発生しました ({img_node.metadata.get('file_name')}): {e}")
                # エラーが発生しても処理を止めずに次へ進む
                continue
                
        return processed_nodes

    def build_index(self, nodes: List[BaseNode]):
        """
        ベクトルインデックスを構築し、保存する
        
        Args:
            nodes (List[BaseNode]): インデックス化するノードリスト
        """
        logger.info("ベクトルインデックスの構築を開始します...")
        
        try:
            # ストレージコンテキストの作成
            storage_context = StorageContext.from_defaults()
            
            # インデックスの作成
            index = VectorStoreIndex(nodes, storage_context=storage_context)
            
            # インデックスの永続化
            index.storage_context.persist(persist_dir=self.persist_dir)
            logger.info(f"インデックスの保存完了: {self.persist_dir}")
            
            return index
            
        except Exception as e:
            logger.error(f"インデックス構築中にエラーが発生しました: {e}")
            raise

    def run_pipeline(self):
        """パイプライン全体を実行する"""
        try:
            # 1. ドキュメント読み込み
            nodes = self.load_documents()
            
            # 画像ノードとテキストノードを分離
            image_nodes = [n for n in nodes if isinstance(n, ImageNode)]
            text_nodes = [n for n in nodes if not isinstance(n, ImageNode)]
            
            # 2. 画像の要約生成
            summary_nodes = self.create_image_summaries(image_nodes)
            
            # 3. ノードの統合（元のテキスト + 画像要約）
            all_nodes = text_nodes + summary_nodes
            
            # 4. インデックス構築
            index = self.build_index(all_nodes)
            
            logger.info("パイプライン完了。クエリエンジンの準備ができました。")
            return index.as_query_engine(similarity_top_k=5)
            
        except Exception as e:
            logger.critical(f"パイプライン実行中に致命的なエラーが発生しました: {e}")
            sys.exit(1)

if __name__ == "__main__":
    # 使用例
    INPUT_DIR = "./data"  # PDFを含むディレクトリ
    STORAGE_DIR = "./storage_multimodal"
    
    try:
        pipeline = MultimodalRAGPipeline(
            input_dir=INPUT_DIR,
            persist_dir=STORAGE_DIR
        )
        query_engine = pipeline.run_pipeline()
        
        # テストクエリ
        response = query_engine.query("財務諸表において、売上高の変動が最も激しい四半期はいつですか？その根拠となるグラフも教えてください。")
        print("\n=== 回答 ===")
        print(response)
        
    except Exception as e:
        logger.error(f"実行に失敗しました: {e}")