GraphRAGでポケモン図鑑の構築に挑戦【LlamaIndex × Neptune 】

## はじめに

この記事は、[@minorun365](https://qiita.com/minorun365)さんの「[GraphRAGをAWS上で動かしてみる](https://qiita.com/minorun365/items/a62dec701bef17cf7562)」を参考に、実際にポケモンWikiの内容でGraphRAGを試してみた実験記録です。

## 注意
筆者はポケモンについて詳しくないです。
本記事は技術検証目的であり、商用利用を意図するものではありません。
- 本記事は技術検証目的の実験記録です
- ポケモンWikiのコンテンツは [出典元URL] より引用
- ポケモンは株式会社ポケモンの登録商標です



### できること
- NeptuneのグラフDBにデータ登録
- プログラムから自然言語でグラフDBを検索

## 環境準備
みのるんさんの記事を参考に、「Neptune Databaseを作成」の章まで実施してください。
クラスター作成まで完了したらIAMロールを修正します。

### SageMakerのIAMロールを修正
ノートブックの名前を選択
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/22960d14-c7c1-e7bc-0c97-415583472212.png)

アクション->SageMakaerで詳細を表示する
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/5a4860c0-8bfb-fce7-4d6a-caa8ecf7bc1d.png)

アクセス許可と暗号化->IAMロールARNを選択
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/724aea4d-0fe7-3524-db33-8e5ca9f4894d.png)

 許可ポリシーに```AmazonBedrockFullAccess```を追加
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/cfd5c0e4-649f-6500-b237-ea17bc5f91d1.png)



## 実装

### 1. データの準備
#### 元データを準備
今回はポケモンwikiからピカチュウのページをPDFで保存します。
ページの左側の「印刷用バージョン」からPDF化できます。
[ピカチュウのページ](https://wiki.xn--rckteqa2e.com/wiki/%E3%83%94%E3%82%AB%E3%83%81%E3%83%A5%E3%82%A6)

#### ノートブックにアップロード
ノートブックを起動したら、直下に```Data```フォルダを作成
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/44eae4d0-84e6-e9f6-a719-7ea48c59001c.png)

先ほど保存したPDFどドラックアンドドロップ！
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/4f939e59-0104-30e1-b02b-5bc28fe120fd.png)


### 2. サンプルプログラムの実行
みのるんさんの記事をベースに少しいじっています。

#### ノーブックを作成
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/a454043d-5b03-b366-e316-c37047228fa7.png)

#### 必要なパッケージのインストール
````Python
%pip install -U boto3 llama-index-llms-bedrock llama-index-graph-stores-neptune llama-index-embeddings-bedrock llama-index-readers-file
````

#### インポート
```Python
from llama_index.llms.bedrock import Bedrock
from llama_index.embeddings.bedrock import BedrockEmbedding
from llama_index.core import StorageContext, SimpleDirectoryReader, KnowledgeGraphIndex, Settings
from llama_index.graph_stores.neptune import NeptuneDatabaseGraphStore
from IPython.display import Markdown, display
```

#### 回答生成用LLMと、埋め込み用モデルを設定
```Python
llm = Bedrock(
    model="anthropic.claude-3-sonnet-20240229-v1:0",
    region="us-east-1"  # Neptuneと同じリージョンを指定
)
embed_model = BedrockEmbedding(model="amazon.titan-embed-text-v2:0")
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512
```

#### ドキュメントとグラフDBを設定
```Python
documents = SimpleDirectoryReader("./Data").load_data()
graph_store = NeptuneDatabaseGraphStore(
    host="<書き込みエンドポイント名>",
    port=8182
)
```

### 重要！ グラフDBへ登録
次のようにシンプルな実装も可能な様なのですが、リクエスト制限に引っかかってしまったので、複雑な実装を行います。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/9ee28cdb-da38-c408-7389-fbbc12f950f1.png)
```plaintext:エラー内容
Retrying llama_index.llms.bedrock.utils.completion_with_retry.<locals>._completion_with_retry in 4.0 seconds as it raised ThrottlingException: An error occurred (ThrottlingException) when calling the InvokeModel operation: Too many requests, please wait before trying again. You have sent too many requests.  Wait before trying again..
```

#### (sonnet v2が。。。)試行錯誤したグラフDBへの登録処理
```Python
import time
import logging
from typing import List
from llama_index.core.schema import Document
from llama_index.core import KnowledgeGraphIndex

from llama_index.core import StorageContext

# storage_contextの初期化
storage_context = StorageContext.from_defaults(graph_store=graph_store)

# ログ設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def process_documents_in_batches(
    documents: List[Document],
    storage_context,
    batch_size: int = 2,     # より小さいバッチサイズ
    sleep_time: int = 10,    # より長い待機時間
    max_retries: int = 5     # リトライ回数を増加
) -> KnowledgeGraphIndex:
    """
    ドキュメントを非常に慎重にバッチ処理する
    """
    final_index = None
    processed_count = 0

    try:
        total_batches = len(documents) // batch_size + (1 if len(documents) % batch_size else 0)

        for i in range(0, len(documents), batch_size):
            batch = documents[i:i+batch_size]
            current_batch = i//batch_size + 1

            logger.info(f"Starting batch {current_batch}/{total_batches}")
            logger.info(f"Documents in current batch: {len(batch)}")

            success = False
            retry_count = 0

            while not success and retry_count < max_retries:
                try:
                    if final_index is None:
                        logger.info("Creating new index...")
                        final_index = KnowledgeGraphIndex.from_documents(
                            batch,
                            storage_context=storage_context,
                            max_triplets_per_chunk=2
                        )
                    else:
                        logger.info("Adding to existing index...")
                        final_index.insert_nodes(batch)

                    success = True
                    processed_count += len(batch)
                    logger.info(f"Successfully processed batch. Total documents processed: {processed_count}")

                except Exception as e:
                    retry_count += 1
                    logger.warning(f"Error on attempt {retry_count}: {str(e)}")

                    if retry_count < max_retries:
                        wait_time = sleep_time * retry_count  # 指数的バックオフ
                        logger.info(f"Waiting {wait_time} seconds before retry...")
                        time.sleep(wait_time)
                    else:
                        logger.error(f"Failed to process batch after {max_retries} attempts")
                        raise

            # バッチ間の待機時間
            logger.info(f"Waiting {sleep_time} seconds before next batch...")
            time.sleep(sleep_time)

    except Exception as e:
        logger.error(f"Fatal error: {str(e)}")
        logger.info(f"Processed {processed_count} documents before error")
        raise

    logger.info("Processing completed successfully!")
    return final_index

# 使用例
try:
    index = process_documents_in_batches(
        documents,
        storage_context=storage_context,
        batch_size=2,    # 非常に小さいバッチサイズ
        sleep_time=10,   # 十分な待機時間
        max_retries=5    # 多めのリトライ回数
    )
except Exception as e:
    logger.error(f"Failed to process documents: {str(e)}")
```

### データ登録準備完了！
succesfullyとなればサクセスです！
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/bf74f4a6-34a8-545b-22ff-891b6efb79b0.png)

## 検証
#### まずはピカチュウについて聞いてみる
おおお〜それっぽい！
しかし、ピカチュウについては一般常識として、モデルに含まれている可能性も。。。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/c904ece2-f028-8d40-8a7a-c728a90f2c6e.png)

帽子をかぶるが謎ですが、ちゃんと元データに情報ありました。
ただ、「人間臭い一面」について記載がないので、この子(AI)の解釈みたいですね笑
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/511155f9-2469-0bd3-6f02-1987214dcd56.png)

#### 種族値を聞いてみる
さっき*種族値*という謎ワードがあったので聞いてみる
260または320・・・
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/ff681c21-71fc-82e5-8243-a39c3489e5a0.png)

#### 元データを確認
第一世代は260、第6世代は320となっているので正解です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/9872c40b-737a-f079-6133-bdd7a8f4e2ce.png)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/020c90c6-55d9-3a3f-8468-abda4bd6ae16.png)


#### 世代を指定して種族値を聞いてみる
正解です！
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/e429fb19-b3f9-433c-c1d5-aaec6a8401d2.png)

## データを可視化
それっぽい結果が返ってきたので、グラフDBの中身を可視化してみます。

### ノートブック->アクション->Graph Explorerを選択
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/ef7c8716-d6d4-dc08-0fd9-7c2089b29351.png)

### Entityを選択
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/042a81c9-165b-7e5f-9455-25d7fc6f2d77.png)


### 可視化したい項目を```Send to Explorer```
ピカチュウとライチュウを選択
選択終わったら、Open Graph Explorerを選択
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/f4072b5a-00fa-3160-6fdb-b82787ee2732.png)


### 可視化！
ノードをダブルクリックしてみます。
左側がピカチュウ、右側がライチュウになっていて、「進化する」で紐づいています。
うまくいけてそうですね。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88001/29b2b3fd-8267-e0dd-5f8d-fc029e869074.png)


## まとめ
AWS上でGraphRAGを試してみました。
今回はピカチュウだけだったので、それっぽい結果になったと思います。
他のポケモンなど、複数の情報を同時に食わせた時にどのようなグラフDBが生成されて、どれ位の精度になるか試してみようと思います。