[備忘録]LLMGraphTransformerでドキュメントノードに任意のメタデータを追加する

## 背景

RAGに使う目的でNeo4jのグラフをPythonで作成する際に、LangChainのLLMGraphTransformer(以下LLMGTと表記します)を利用していました。LLMGTは通常のノードの他にもドキュメント本体のノードも作成することができるのですが、そこに参照先や更新日等の情報を付加して保存したいと思ったのがきっかけです。

LLMGTを使って元ドキュメントのノードに追加情報を付与する方法を調べてきましたので、備忘録として残します。

## LLMGT

LLMGTは自身の指定したAIを用いて、ドキュメントからノードとノード同士の関係性を抽出してくれます。使用するLLMを引数として渡すことで初期化できます。

## LLMGTを使ったグラフ作成例

まずはドキュメントのノードにメタデータを追加しない場合のコードについて紹介します。このコードは以下の情報を環境変数に設定していることが前提となります。

| 環境変数名              | 説明                                                    |
|-------------------------|---------------------------------------------------------|
| AZURE_OPENAI_ENDPOINT   | 自身のAzure OpenAIのエンドポイント                       |
| AZURE_OPENAI_API_KEY    | 自身のAzure OpenAIのAPIキー                              |
| CHAT_API_VERSION        | デプロイしているチャットモデルのAPIバージョン             |
| CHAT_MODEL_NAME         | デプロイしているチャットモデルの名前                      |
| NEO4J_URI               | Neo4jのURI(`bolt://localhost:7687`など)                  |
| NEO4J_USERNAME          | Neo4jで設定した名前                                      |
| NEO4J_PASSWORD          | Neo4jで設定したパスワード                                |

```python
import os  
from langchain_community.graphs import Neo4jGraph
from langchain_core.documents import Document
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_openai import AzureChatOpenAI,AzureOpenAIEmbeddings
from dotenv import load_dotenv
load_dotenv()

# チャットモデルをAPIのバージョンとデプロイされたモデル名で初期化する
llm = AzureChatOpenAI(
  openai_api_version=os.getenv("CHAT_API_VERSION"),
  azure_deployment=os.getenv("CHAT_MODEL_NAME"),
  temperature=0.7,
)

# LLMGTの初期化
llm_transformer = LLMGraphTransformer(llm=llm)

# Neo4jドライバーの初期化
driver = Neo4jGraph()

# サンプルドキュメント
doc = "私は社会人です。私は朝07:00に家を出て職場に行きます。"

# ドキュメントをDocumentクラスにする
doc_list = [Document(page_content=doc)]

# ドキュメントのリストをグラフデータに変換
graph_documents = llm_transformer.convert_to_graph_documents(doc_list)

driver.add_graph_documents(
  graph_documents,
  baseEntityLabel=True,
  include_source=True
)
```

### コードの説明

取り込むドキュメントは`langchain_core.documents`の`Document`に変換しています。`driver.add_graph_documents`部分では`include_source`を`True`にすることでドキュメントそのもののノードを作成できます。また、`baseEntityLabel`を`True`にすることで、インデックス作成とクエリのパフォーマンスが向上します。

実行すると以下のようなグラフがAIにより生成されます。ドキュメントそのもののノード(id: `d961be36f93eaa82ad378114c5ca5019`)も作成されていることがわかります。

![neo4j_doc](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3657689/65211772-124c-4c03-c168-d6eb95726e02.png)

## ドキュメントのノードにメタ情報をつける

ドキュメントのノードにメタ情報を追加する方法について説明します。
メタ情報は`Document`型へ変換する際に`metadata`という引数を渡すことで付与することが可能です。以下はドキュメントのタイトルと更新日を追加する例です。

```python
doc_list = [Document(
  page_content=doc,
  metadata={
    "title": "日常の行動",
    "updated_at": "2024-11-28"
    }
)]
```

先ほどのコードを置き換えて実行すると以下のように、ドキュメントそのもののノードにメタデータが追加されます。

![neo4j_doc_metadata](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3657689/082294f5-4363-67fb-200b-409b7d72fef7.png)

### コード全体

メタ情報を追加したコードは以下のようになります。

```python
import os  
from langchain_community.graphs import Neo4jGraph
from langchain_core.documents import Document
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_openai import AzureChatOpenAI
from dotenv import load_dotenv
load_dotenv()

# チャットモデルの初期化
llm = AzureChatOpenAI(
  openai_api_version=os.getenv("CHAT_API_VERSION"),
  azure_deployment=os.getenv("CHAT_MODEL_NAME"),
  temperature=0.7,
)

# LLMGTの初期化
llm_transformer = LLMGraphTransformer(llm=llm)

# Neo4jドライバーの初期化
driver = Neo4jGraph()

# サンプルドキュメント
doc = "私は社会人です。私は朝07:00に家を出て職場に行きます。"

# ドキュメントをDocumentクラスにする
# metadataにタイトルと更新日を追加
doc_list = [Document(
  page_content=doc,
  metadata={
    "title": "日常の行動",
    "updated_at": "2024-11-28"
    }
)]

# ドキュメントのリストをグラフデータに変換
graph_documents = llm_transformer.convert_to_graph_documents(doc_list)

driver.add_graph_documents(
  graph_documents,
  baseEntityLabel=True,
  include_source=True
)

```

## 終わりに

LLMGTはAIを用いてグラフを構築できるので、取り込むドキュメントのグラフの定義が動的で難しい際に便利だなと思いました。ドキュメントそのものに任意のメタ情報を埋め込むことで、アクセス方法を増やすことができる点も便利です。