データカタログのコメントをLLMに考えさせてみる

# 要約

LLMを使って
これを
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/01e19540-86db-ad18-f501-c2890f9fbef3.png)



こうします
![2024-03-28-15-15-26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/40a799c3-d17d-e3f8-5fd3-ce7c2e35101a.png)





# はじめに
今回はBedrockの利用ユースケース案の簡単な検証です。

冒頭示したように、Glue Data CatalogにはGlueでETLジョブを走らせるためのデータスキーマが登録されています。各カラムには名前とデータタイプのほかにコメントを格納できます。コメント自体は機能に何ら影響は与えません。メモ書きです。

あまり意識することはないかもしれませんが、AthenaなどS3をソースにテーブル作成する際には、バックグラウンドで自動的にData Catalog テーブルが作成されます。今回対象としているテーブルもその一つで、先日以下の記事でAthenaを利用する際に作られたテーブルです。このほかにもELBのログ分析のためAthenaを使ったことがある人もそのテーブルが作られているはずです。

https://qiita.com/kanuazut/items/82d3d761e4bb8c6e3e03


このテーブル自体は何てことないテーブルですし、この前作られたばかりなので各カラムが何を示しているか私にはすぐわかります。記事を読んでくれた人にも多分分かります。しかし、数か月後の自分がこのテーブルを見た時、input, outputが何を意味しているか、すぐに思いつかないかもしれません。初見の他人がみればなおさらです。

そこでLLMにデータカラムが何のデータなのか、コメント付けをさせてみようというのが今回の記事です。データカタログの運用・管理の自動化、人間の負荷軽減に使えるのではないかなと思ったので検証してみます。






# 手順

今回のざっくりとした構成は以下の通りとなります。
ローカルのboto3からGlue Data CatalogとBedrockを利用して全ての操作を行います。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/ed3bbe0e-3581-39da-f051-29c2dbfc155f.png)



##  ライブラリインストール
今回は以下のライブラリを使いました。実行環境にインストールしておきます。
```bash
pip install boto3 litellm
```

## スキーマの取得
まずはDataCatalogからテーブルのスキーマを取得します。

ローカルのPython環境で認証を通します。

```python
import os
os.environ["AWS_ACCESS_KEY_ID"] = ""  # Access key
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
os.environ["AWS_REGION_NAME"] = "us-east-1" # Region
```
<br>

boto3 clientを使ってテーブル一覧を取得してます。今回は`bedrock`データベース内のテーブルを対象にしています。
```python
import boto3
client = boto3.client('glue', region_name = 'us-east-1')

db_name = "bedrock"
response = client.get_tables(
    DatabaseName=db_name,
)
```
<br>

以下はレスポンスの抜粋です。
TableList.StorageDescriptor.Columnsがテーブルスキーマの本体です。
```json: responseデータの抜粋
{
    "TableList": [
        {
            "Name": "bedrock_logs",
            "DatabaseName": "bedrock",
            "CreateTime": "",
            "UpdateTime": "",
            "Retention": 0,
            "StorageDescriptor": {
                "Columns": [
                    {
                        "Name": "identity",
                        "Type": "string",
                        "Comment": ""
                    },
                    {
                        "Name": "region",
                        "Type": "string",
                        "Comment": ""
                    },
                    {
                        "Name": "operation",
                        "Type": "string",
                        "Comment": ""
                    },
                    {
                        "Name": "modelid",
                        "Type": "string",
                        "Comment": ""
                    },
                    {
                        "Name": "requestid",
                        "Type": "string",
                        "Comment": ""
                    },
                    {
                        "Name": "input",
                        "Type": "struct<inputTokenCount:int>",
                        "Comment": ""
                    },
                    {
                        "Name": "output",
                        "Type": "struct<outputTokenCount:int>",
                        "Comment": ""
                    },
                    {
                        "Name": "timestamp",
                        "Type": "string",
                        "Comment": ""
                    }
                ],
                "Location": "s3://bedrock-test-20230000/AWSLogs/000000000000/BedrockModelInvocationLogs",
                "InputFormat": "com.amazon.ionhiveserde.formats.IonInputFormat",
                "OutputFormat": "com.amazon.ionhiveserde.formats.IonOutputFormat",
                "Compressed": False,
                "NumberOfBuckets": 0,
                "SortColumns": [],
                "StoredAsSubDirectories": False
            },
            "CreatedBy": "arn:aws:iam::000000000000:user/username",
            "IsRegisteredWithLakeFormation": False,
            "CatalogId": "000000000000",
            "VersionId": "7"
        }
    ]
}
```
<br>



今回は`bedrock_logs`というテーブルを対象とするので、テーブルリストから単体のテーブル部分を抽出しておきます。
```python
tbl_name = "bedrock_logs"
for tb in response["TableList"]:
    if tb["Name"] == tbl_name:
        table = tb
```


## Bedrockでカラムコメント作成

litellmを使ってBedrockのCalude 3 Sonnetにカラムコメントを作成させます。特にlitellmである必要はないのでboto3がいい人はboto3.client("bedrock")を使ってください。

{info}で与えているのは取得したテーブルの情報全量(スキーマ含む)、{schema}で与えているのはテーブルスキーマ部分のみです。雑に与えています。
```python
import litellm
prompt = """
あなたは優秀なデータエンジニアです。
与えられるテーブル定義を見て、各ColumnのCommentフィールドに日本語で説明文を入れることが仕事です。
<schema>タグに与えられた値を加工してCommentフィールドを埋めたJSONを出力してください。
またテーブルの補足情報を<info></info>に与えるので参考情報としてください。

<info>
{info}
</info>

<schema>
{schema}
</schema>
""".format(schema=table["StorageDescriptor"]["Columns"], info=table)
print("入力: \n", prompt)

resp = litellm.completion(
    model = "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
    stream = False,
    max_tokens=4096
)
print("出力: \n", resp.choices[0].message.content)
```

出力は以下の通りです。
```json: 出力
 [
  {
    "Name": "identity",
    "Type": "string",
    "Comment": "顧客IDまたはデータ主体を一意に識別する値"
  },
  {
    "Name": "region",
    "Type": "string",
    "Comment": "モデルがデプロイされているAWSリージョン"
  },
  {
    "Name": "operation",
    "Type": "string",
    "Comment": "呼び出されたAmazon Bedrock の操作名"
  },
  {
    "Name": "modelid",
    "Type": "string",
    "Comment": "呼び出されたモデルのID"
  },
  {
    "Name": "requestid",
    "Type": "string",
    "Comment": "呼び出しリクエストの一意のID"
  },
  {
    "Name": "input",
    "Type": "struct<inputTokenCount:int>",
    "Comment": "モデル入力の構造体。inputTokenCountはモデル入力のトークン数。"
  },
  {
    "Name": "output",
    "Type": "struct<outputTokenCount:int>",
    "Comment": "モデル出力の構造体。outputTokenCountはモデル出力のトークン数。"
  },
  {
    "Name": "timestamp",
    "Type": "string",
    "Comment": "モデル呼び出し時のタイムスタンプ"
  }
]
```

いかがでしょうか。個人的にはまあまあいい感じだと思います。恐らく筆者が手動でつけても似たようなコメントになるでしょう。強いて言うなら一つ目の`顧客IDまたはデータ主体を一意に識別する値`は、IAMエンティティとしてくれればもっとよかったですが、おおむね問題なさそうです。

見てわかる通り、`region`が「AWSリージョン」であることや`operation`が「呼び出されたAmazon Bedrock の操作名」であることがLLMには認識できています。これはTableListの中に含まれていたS3のLocationや`bedrock_logs`というテーブル名とBedrockがAWSのサービスであるという事前知識からLLMが推測している内容です。

LLMがより高度なコメント付けを行うには、Column statisticsの情報を一緒に与えたり、S3のLocationパス配下からサンプルデータを一つとってきてファイル内容を渡したりしてみてもいいかもしれません。


## Schemaの更新
LLMが作成したスキーマを更新するだけですが、実行します。
```python
import re
import json

# テキストのコードブロック内だけ抽出する
def extract_code_blocks(text: string):
    code_blocks = re.findall(r'```(?:.*?)\n(.*?)\n```', text, re.DOTALL)
    return code_blocks

# LLMの出力をdictionaryにする
def read_content(d: dict):
    try:
        return json.loads(d)
    except json.decoder.JSONDecodeError as e:
        return json.loads(extract_code_blocks(d)[0])

# 更新
response = client.update_table(
    CatalogId=table["CatalogId"],
    DatabaseName=table["DatabaseName"],
    TableInput={
        "Name": table["Name"],
        "StorageDescriptor": {
            "Columns": read_content(resp.choices[0].message.content)["Columns"]
        }
    }
)
print(response)
```
`read_content`はLLMのテキスト出力を辞書型に直しているだけです。時々、LLMの出力がJSON単体だけでなく「以下にJSONを出力します```json \n 」のように始めてしまうことがあるので、json.loadsできないときはコードブロック内の文字列だけ抽出するようにしています。

スキーマの更新は単純で、スキーマ取得時と同じ辞書型にして`client.update_table`を使うだけです。


これにて冒頭の画像のようにClaude 3の考えたコメントがGlue Data Catalogのテーブルスキーマに反映されました。
![2024-03-28-15-15-26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/0418ce97-080f-2dba-656b-8d755e8f4131.png)


実務上は人がつけたコメントかLLMがつけたコメントか分かるように[Commented by Calude]のような注釈を入れたり、既についているコメントは上書きしないというルールが必要になってくると思います。

# おわりに
今回は検証、というよりユースケース提案みたいな形になりましたが、ちょっと思いついたことAWSのサービスを使って試してみました。

今回は余力がなくてテーブル名、データソース、カラム名、スキーマしかLLMに渡していませんでしたが、生データや統計情報を渡すことで、より正確でより情報量の多いコメントを出力できるのではないかと思います。

また、カタログ作成をトリガーにしてLambde+Bedrockでコメントを自動設定したり、Agents for Amazon Bedrockなどを使って高度化もできそうな気がします。

データカタログに限らず、基盤のマネジメントタスクにLLMを活用していくユースケースは沢山あると思うので、また気になることはやっていきたいです。

## 蛇足ポエム
ポエム色が強くなりますが、LLMの活用が進むにつれ、あらゆる場所で自然文の重要度が増していると思います。

特にAgentやFunction Caliingといった、LLMが次のアクションを推論から導く仕組み(=ReAct)が充実してきているので、データカタログ・データマネジメントの重要性は拍車をかけて上がっているように思います。要はLLM様が見て用途や仕様を判断できる情報をデータ定義にも残しておけよ、ということです。もとより人間のために必要だったものですが、LLMのためにもなるようになってきました。

データカタログの作成や管理は人間がやれればそれが一番確実ですが、データマネジメントは規模に応じて非常に労力のいる仕事だと思うので、こういった仕組みでLLMに力を発揮させる土壌を整えていくことが重要かなと思っています。



# 参考文献


[Glue \- Boto3 1\.34\.71 documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue.html)