🛏️

RAG環境を簡単プロビジョニング!DBスキーマ質問システムを構築する

2025/01/24に公開

RAG環境を簡単プロビジョニング!DBスキーマ質問システムを構築する

ごあいさつと今年の抱負

あけましておめでとうございます!開発の小倉です!
あけおめ技術ブログということで、アドベントカレンダーに初めて参加させていただくことになりました。
テーマは決めず、各自興味がある分野について好きに書いて良いということで、この度は心のどこかでやりたいと考えていたことに着手するモチベーションが湧く機会を与えていただき感謝です!
他のメンバーもそれぞれの独自の視点で記事のテーマが決められていてとても面白いですね。
そんなメディカルフォースを、今年もどうぞよろしくお願いいたします!

私の今年の抱負は「省力化」で、自分と他のメンバーの仕事が楽になるようなアイデアを考えて、行動していきたいと思います。

やりたいこと

システム開発において、仕様の把握や設計を行うためにDBスキーマを理解する必要が生じる場面は多々あります。
システムの規模が大きくなるほど、キャッチアップにかかるコストも比例して増大します。
この負担を軽減するアプローチとして、AIを活用するのは有効な手段の一つです。

そこで今回は、下記の3点に挑戦します。

  • LLMのデータソースとして、DBスキーマ定義を簡単にアウトプットできるようにする。
  • DBスキーマ定義をVector DBに取り込み、LLMを活用して、スキーマのキャッチアップや設計に関するアドバイスを得られるRAG環境を構築する。
  • RAG環境の構築を迅速に行えるようにする。

大まかにこれからやること

  1. データソースとして、PostgreSQL環境からDBスキーマ定義をエクスポートする。
  2. AWS CDKを使用してRAG環境をプロビジョニングする。
  3. RAG環境でデータソースをインポートし、質問を投げて動作を検証する。

環境構成

環境構成図

事前に準備しておくもの

長くなってしまうので説明は割愛しますが、以下の環境および権限をあらかじめ用意する必要があります。

  • PostgreSQL環境
  • awscli実行環境
  • CDK実行環境
    • node 18.17.1
  • AWSのプロビジョニングやCDKのブートストラップ実行に必要なIAM権限

手順

1. DBからスキーマ定義を取得する

テーブルの作成

サンプルとして、学校をドメインとしたスキーマを定義しました。

ER図

テーブル作成用SQL
CREATE TABLE subjects (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE subjects IS '教科テーブル';
COMMENT ON COLUMN subjects.id IS '教科ID';
COMMENT ON COLUMN subjects.name IS '教科名';
COMMENT ON COLUMN subjects.created_at IS '作成日時';
COMMENT ON COLUMN subjects.updated_at IS '更新日時';

CREATE TABLE teachers (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    subject_id INT REFERENCES subjects(id),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE teachers IS '先生テーブル';
COMMENT ON COLUMN teachers.id IS '先生ID';
COMMENT ON COLUMN teachers.name IS '先生の名前';
COMMENT ON COLUMN teachers.subject_id IS '担当する教科のID';
COMMENT ON COLUMN teachers.created_at IS '作成日時';
COMMENT ON COLUMN teachers.updated_at IS '更新日時';

CREATE TABLE classrooms (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    capacity INT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE classrooms IS '教室テーブル';
COMMENT ON COLUMN classrooms.id IS '教室ID';
COMMENT ON COLUMN classrooms.name IS '教室名';
COMMENT ON COLUMN classrooms.capacity IS '教室の収容人数';
COMMENT ON COLUMN classrooms.created_at IS '作成日時';
COMMENT ON COLUMN classrooms.updated_at IS '更新日時';

CREATE TABLE classes (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    homeroom_teacher_id INT REFERENCES teachers(id),
    classroom_id INT REFERENCES classrooms(id),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE classes IS 'クラステーブル';
COMMENT ON COLUMN classes.id IS 'クラスID';
COMMENT ON COLUMN classes.name IS 'クラス名';
COMMENT ON COLUMN classes.homeroom_teacher_id IS '担任の先生ID';
COMMENT ON COLUMN classes.classroom_id IS '教室ID';
COMMENT ON COLUMN classes.created_at IS '作成日時';
COMMENT ON COLUMN classes.updated_at IS '更新日時';

CREATE TABLE students (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    class_id INT REFERENCES classes(id),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE students IS '生徒テーブル';
COMMENT ON COLUMN students.id IS '生徒ID';
COMMENT ON COLUMN students.name IS '生徒の名前';
COMMENT ON COLUMN students.class_id IS '所属クラスのID';
COMMENT ON COLUMN students.created_at IS '作成日時';
COMMENT ON COLUMN students.updated_at IS '更新日時';

CREATE TABLE schedules (
    id SERIAL PRIMARY KEY,
    classroom_id INT REFERENCES classrooms(id),
    subject_id INT REFERENCES subjects(id),
    teacher_id INT REFERENCES teachers(id),
    start_time TIMESTAMP NOT NULL,
    end_time TIMESTAMP NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
COMMENT ON TABLE schedules IS '時間割テーブル';
COMMENT ON COLUMN schedules.id IS '時間割ID';
COMMENT ON COLUMN schedules.classroom_id IS '教室ID';
COMMENT ON COLUMN schedules.subject_id IS '教科ID';
COMMENT ON COLUMN schedules.teacher_id IS '担当する先生ID';
COMMENT ON COLUMN schedules.start_time IS '授業開始時刻';
COMMENT ON COLUMN schedules.end_time IS '授業終了時刻';
COMMENT ON COLUMN schedules.created_at IS '作成日時';
COMMENT ON COLUMN schedules.updated_at IS '更新日時';

2. VectorDBへの取り込み用にデータを出力する

VectorDBに取り込むためのデータとして、information_schema からDBスキーマの情報を取得します。
カラムやリレーションの情報を保持し、ベクトル化しやすい形式にするため、JSON形式でデータを出力します。
出力したJSONはデータソースとして利用するので、テキストファイルとして保存しておきます。

出力されるJSONには以下のフィールドが含まれます:

  • table_schema: テーブルのスキーマ名
  • table_name: テーブル名
  • column_name: カラム名
  • data_type: カラムのデータ型(例: integer, varchar など)
  • character_maximum_length: カラムが文字列型の場合の最大文字数(NULLの場合もあり)
  • column_comment: カラムに付与されたコメント(NULLの場合もあり)
  • target_schema: 外部キーの参照先スキーマ名(NULLの場合もあり)
  • target_table: 外部キーの参照先テーブル名(NULLの場合もあり)
  • target_column: 外部キーの参照先カラム名(NULLの場合もあり)

なお、以下の出力用クエリは、サンプルデータに限らず他のデータでも実行可能です。

DBスキーマ情報出力用SQL
select json_agg(result)
from (
    WITH columns_data AS (
        SELECT 
            t.table_schema,
            t.table_name,
            json_agg(
                json_build_object(
                    'column_name', c.column_name,
                    'data_type', c.data_type,
                    'character_maximum_length', c.character_maximum_length,
                    'column_comment', col_description(format('%s.%s', t.table_schema, t.table_name)::regclass::oid, c.ordinal_position)
                )
            ) AS columns
        FROM 
            information_schema.tables t
        JOIN 
            information_schema.columns c
            ON t.table_schema = c.table_schema
            AND t.table_name = c.table_name
        WHERE 
            t.table_type = 'BASE TABLE'
            AND t.table_schema NOT IN ('pg_catalog', 'information_schema')
        GROUP BY 
            t.table_schema, t.table_name
    ),
    relations_data AS (
        SELECT 
            tc.table_schema AS source_schema,
            tc.table_name AS source_table,
            json_agg(
                json_build_object(
                    'target_schema', ccu.table_schema,
                    'target_table', ccu.table_name,
                    'target_column', ccu.column_name
                )
            ) AS relations
        FROM 
            information_schema.table_constraints AS tc
        JOIN 
            information_schema.key_column_usage AS kcu
            ON tc.constraint_name = kcu.constraint_name
            AND tc.table_schema = kcu.table_schema
        JOIN 
            information_schema.constraint_column_usage AS ccu
            ON ccu.constraint_name = tc.constraint_name
            AND ccu.table_schema = tc.table_schema
        WHERE 
            tc.constraint_type = 'FOREIGN KEY'
        GROUP BY 
            tc.table_schema, tc.table_name
    )
    SELECT 
        c.table_schema,
        c.table_name,
        c.columns,
        COALESCE(r.relations, '[]'::json) AS relations
    FROM 
        columns_data c
    LEFT JOIN 
        relations_data r
        ON c.table_schema = r.source_schema
        AND c.table_name = r.source_table
    ORDER BY 
        c.table_schema, c.table_name
) as result;

3. Pineconeのセットアップ

Pineconeを使用して、RAG環境で利用するデータベースを準備します。
Dimensions の値は、利用する埋め込みモデルに適合させる必要があります。
今回は実験のため、AWSが提供する titan-embed-text-v2:0 を埋め込みモデルとして使用します。このモデルにおけるベクトルの次元数は1024です。
詳細は公式ドキュメントをご参照ください。

Pineconeの設定

データベース作成後、エンドポイントのURLとAPIキーを取得しておきます。

4. Secrets ManagerでAPIキーを管理する

PineconeのAPIキーをAWS Secrets Managerで管理します。

以下のコマンドを実行して、Secrets ManagerにAPIキーを登録します。

$ aws secretsmanager create-secret --region ap-northeast-1 --name PineconeSecret --secret-string "{\"apiKey\":\"PINECONE_API_KEY\"}" [--profile {profile_name}]

実行後、キーのARNが出力されるので、取得しておきます。

5. 使用可能なモデルのARNを確認する

前述の通り、今回は埋め込みモデルとして titan-embed-text-v2:0 を利用します。
以下のコマンドを実行してARNを確認し、取得しておきます。

$ aws bedrock list-foundation-models [--profile {profile_name}] | grep modelArn | grep titan-embed-text-v2:0

6. 環境変数をセットする

CDKを実行する前に、.env ファイルに以下の環境変数を設定します:

ACCOUNT_ID={AWSアカウントID}
REGION=ap-northeast-1
EMBEDDING_MODEL_ARN=arn:aws:bedrock:ap-northeast-1::foundation-model/amazon.titan-embed-text-v2:0
PINECONE_ENDPOINT={PineconeエンドポイントURL}
PINECONE_SECRET_ARN={Secrets ManagerのPinecone APIキーのARN}

7. RAG環境を構築する

環境はCDKを利用してプロビジョニングします。
ソースコードのリポジトリは以下となります。

https://github.com/ogreverse/cdk/tree/main/simple_rag_with_bedrock

CDK Deploy

以下のコマンドを使用して、CDKでスタックをデプロイします:

$ npx cdk deploy [--profile {profile名}]

デプロイが完了すると、必要なスタックが作成されます。

8. モデルアクセスをリクエストする

Amazon Bedrockでモデルアクセスをリクエストします。

Amazon Bedrock > Bedrock configurations > モデルアクセス

  • 埋め込みモデルとして Titan Text Embeddings V2
  • テキスト生成モデル: Claude 3 Haiku

これらのモデルにアクセス権限を付与してもらいます。

9. データソースを同期する

S3バケットにデータソースをアップロードする

JSONで出力したDBスキーマ定義のテキストファイルを、作成されたS3バケットにアップロードします。

ナレッジベースで同期する。

  1. Amazon Bedrock > ナレッジベースから該当のナレッジベースを選択します。
  2. データソースから該当するデータソースを選択し、「同期」を実行します。

Knowledge Baseでの同期

10. RAG環境を利用する

ここまでで環境構築は完了となり、Bedrockのコンソール上からチャット形式でメッセージを投げることができます。

11. RAG環境を閉じる

環境が不要になったら下記のコマンドで作成したスタックを削除します。

$ npx cdk destroy [--profile {profile名}]

いくつか質問を投げてテストしてみる

ここからは、テーブル構造の確認や設計を行う際に想定される疑問を実際に投げてみて、どのような回答結果が得られるかを記録します。
※ 回答結果は見やすさを考慮して適宜改行を加えています。
※ 回答結果に対して「結果へのコメント」を追記してます。

Q1. テーブルの一覧と概要をください。

テーブルの一覧と概要は以下の通りです:

  • students: 生徒情報を管理するテーブル。生徒の名前、所属クラス、作成日時、更新日時などの情報が含まれています。[1]
  • subjects: 教科情報を管理するテーブル。教科名、作成日時、更新日時などの情報が含まれています。[2]
  • classrooms: 教室情報を管理するテーブル。教室名、収容人数、作成日時、更新日時などの情報が含まれています。[3]
  • classes: クラス情報を管理するテーブル。クラス名、所属教室、担任の先生、作成日時、更新日時などの情報が含まれています。[4]
  • schedules: 時間割情報を管理するテーブル。授業開始時刻、授業終了時刻、担当教科、担当教師、教室などの情報が含まれています。[5]
  • teachers: 教師情報を管理するテーブル。教師の名前、担当教科、作成日時、更新日時などの情報が含まれています。

(結果へのコメント)
内容的にはテーブルが網羅されており、問題なさそうです。

Q2. 生徒のカラムについて教えてください。

A.生徒のカラムには以下のものがあります:

  • id: 生徒ID
  • name: 生徒の名前
  • class_id: 所属クラスのID
  • created_at: 作成日時

(結果へのコメント)
updated_at のカラムが省略されてしまい、完璧な回答ではありません。

Q3. 先生の住所と電話番号を保持したいのですが、どのテーブルにどういったカラムを追加すべきでしょうか?

先生の住所と電話番号を保持するには、「teachers」テーブルに以下のカラムを追加するのが適切です:

  • address: 先生の住所を保持するためのカラム
  • phone_number: 先生の電話番号を保持するためのカラム

(結果へのコメント)
質問が良くないかもしれないですが、型についての情報も欲しいところです。

Q4. 生徒が所属する部活を管理したいのですが、どのようなテーブルを用意すれば良いでしょうか。カラムの提案もしてください。

生徒の部活情報を管理するためのテーブルを以下のように提案します:

  • table_name: clubs
  • columns:
      - id (integer, primary key)
      - name (character varying, 100 characters, 部活名)
      - description (text, 部活の説明)
      - created_at (timestamp without time zone, 作成日時)
      - updated_at (timestamp without time zone, 更新日時)
    このテーブルには部活の基本情報を保持します。

また、生徒と部活の関係を管理するためのテーブルを以下のように提案します:

  • table_name: student_clubs
  • columns:
      - id (integer, primary key)
      - student_id (integer, 生徒ID)
      - club_id (integer, 部活ID)
      - joined_at (timestamp without time zone, 部活に入部した日時)
      - left_at (timestamp without time zone, 部活を退部した日時)
      - created_at (timestamp without time zone, 作成日時)
      - updated_at (timestamp without time zone, 更新日時)
    このテーブルにより、生徒が所属する部活の履歴を管理することができます。

(結果へのコメント)
概ね良い結果になっているかと思います。

Q5. 授業のログを残しておきたいです。テーブル設計をしていただけますか?

授業のログを記録するためのテーブル設計は以下のようになります:
schedules テーブル
  - id: 時間割ID
  - classroom_id: 教室ID
  - subject_id: 教科ID
  - teacher_id: 担当する先生ID
  - start_time: 授業開始時刻
  - end_time: 授業終了時刻
  - created_at: 作成日時
  - updated_at: 更新日時
このテーブルでは、授業の時間割情報を管理することができます。教室、教科、担当教師などの情報を関連付けて保持しています。

(結果へのコメント)
既存のschedulesテーブルと変わらない内容が出てきてしまいました。

DBスキーマは正しく読み取れているようですが、質問によっては回答内容が不適切なものが見受けられます。
この点については、データソースの内容、使用するモデル、およびプロンプトの調整によって改善の余地があると考えられます。

テキスト生成AIのモデルを変更して同じ質問を投げてみる

せっかくなので、テキスト生成AIのモデルを Claude 3 Haiku から Claude 3.5 Sonnet に変更して、同じ質問を投げて回答がどの程度違うかも検証してみました。

Q1. テーブルの一覧と概要をください。

以下は、データベース内のテーブルの一覧とその概要です:

  1. classes: クラス情報を管理するテーブル。クラス名、担任の先生ID、教室IDなどの情報が含まれています。
  2. students: 生徒の情報を管理するテーブル。生徒の名前、所属クラスのIDなどの情報が含まれています。
  3. subjects: 教科に関する情報を管理するテーブル。教科名などの情報が含まれています。
  4. teachers: 教師の情報を管理するテーブル。教師の名前、担当する教科のIDなどの情報が含まれています。
  5. classrooms: 教室の情報を管理するテーブル。教室名、収容人数などの情報が含まれています。
  6. schedules: 時間割情報を管理するテーブル。授業の開始時刻、終了時刻などの情報が含まれています。
    これらのテーブルは互いに関連付けられており、学校の運営に必要な様々な情報を効率的に管理できるように設計されています。

(結果へのコメント)
Haikuの結果と比較すると、作成日時と更新日時が省略されており、本質的なカラムが優先されて紹介されるようになりました。

Q2. 生徒のカラムについて教えてください。

生徒テーブル(students)には以下のカラムがあります:

  1. id: 生徒ID(整数型)
  2. name: 生徒の名前(文字列型、最大100文字)
  3. class_id: 所属クラスのID(整数型)
  4. updated_at: 更新日時(タイムスタンプ型)
  5. created_at: 作成日時(タイムスタンプ型)
    これらのカラムは生徒の基本情報を管理するために使用されます。
    生徒IDは一意の識別子として、名前は生徒の個人情報として、class_idは生徒が所属するクラスとの関連付けに使用されます。
    また、更新日時と作成日時は記録の管理に役立ちます。

(結果へのコメント)
正しくカラム一覧が返されました。
応答内容の情報量も少し増えています。

Q3. 先生の住所と電話番号を保持したいのですが、どのテーブルにどういったカラムを追加すべきでしょうか?

先生の住所と電話番号を保持するには、"teachers" テーブルに新しいカラムを追加するのが適切です。
具体的には以下のカラムを追加することをお勧めします:

  1. address: 先生の住所を保存するための文字列カラム
  2. phone_number: 先生の電話番号を保存するための文字列カラム
    これらのカラムを "teachers" テーブルに追加することで、各先生の連絡先情報を効率的に管理することができます。

[1] カラムの追加にあたっては、以下の点に注意してください:

  1. addressカラムは、十分な長さの文字列型(例:VARCHAR(255))で設定し、詳細な住所情報を格納できるようにします。
  2. phone_numberカラムも文字列型(例:VARCHAR(20))で設定し、国際電話番号形式も含めて対応できるようにします。
  3. 必要に応じて、これらの新しいカラムにNULL制約やデフォルト値を設定することを検討してください。
  4. データベースの正規化原則に従い、必要であれば住所情報を別テーブルに分割することも考慮に入れてください。
    これらの変更により、先生の連絡先情報を適切に管理し、必要に応じて容易にアクセスできるようになります。

(結果へのコメント)
型の情報が追加されました。注意点についても言及されており、設計の際に有益な情報に感じます。

Q4. 生徒が所属する部活を管理したいのですが、どのようなテーブルを用意すれば良いでしょうか。カラムの提案もしてください。

申し訳ありませんが、提供された検索結果には生徒の部活動を管理するためのテーブルやカラムに関する具体的な情報が含まれていません。
しかし、生徒の部活動を管理するためのテーブルとカラムを提案させていただきます。
以下のようなテーブルとカラムを用意することをお勧めします:

  1. clubs(部活動)テーブル:
      - id: 整数型(主キー)
      - name: 文字列型(部活動名)
      - description: 文字列型(部活動の説明)
      - created_at: タイムスタンプ型(作成日時)
      - updated_at: タイムスタンプ型(更新日時)
  2. student_clubs(生徒と部活動の関連)テーブル:
      - id: 整数型(主キー)
      - student_id: 整数型(外部キー、studentsテーブルのidを参照)
      - club_id: 整数型(外部キー、clubsテーブルのidを参照)
      - joined_at: 日付型(入部日)
      - role: 文字列型(部員の役割、例:部長、副部長、一般部員など)
      - created_at: タイムスタンプ型(作成日時)
      - updated_at: タイムスタンプ型(更新日時)
    これらのテーブルを追加することで、生徒と部活動の関係を管理できます。
    student_clubsテーブルは多対多の関係を表現し、一人の生徒が複数の部活動に所属したり、一つの部活動に複数の生徒が所属したりすることができます。

(結果へのコメント)
student_clubsのカラムについて、Haikuとの違いが見受けられます。
退部日について考慮されていないのはマイナス点、リレーションについての情報が追加されたのはプラス点かと思います。

Q5. 授業のログを残しておきたいです。テーブル設計をしていただけますか?

授業のログを残すためのテーブル設計は、既存のデータベース構造を活用して行うことができます。
現在の設計には「schedules」というテーブルがあり、これを使用して授業のログを記録することができます。
「schedules」テーブルには以下のカラムが含まれています:

  1. id: 時間割ID(主キー)
  2. subject_id: 教科ID
  3. teacher_id: 担当する先生ID
  4. classroom_id: 教室ID
  5. start_time: 授業開始時刻
  6. end_time: 授業終了時刻
  7. created_at: 作成日時
  8. updated_at: 更新日時
    このテーブルは既に授業に関する基本的な情報を記録しているため、授業のログを残すのに適しています。
    各授業の開始時刻、終了時刻、担当教師、教科、使用教室などの情報が記録されます。
    ただし、より詳細なログを残したい場合は、以下のようなカラムを追加することを検討できます:
  9. attendance_count: 出席者数
  10. notes: 授業に関する備考や特記事項
    これらの追加カラムにより、より詳細な授業ログを記録することができます。

(結果へのコメント)
既存のテーブルにログとして記録するべきカラムが提案されており、「ログとしてどういったカラムがあるべきか検討が必要」というヒントを与えてくれているのも良いかと思います。

Haikuと比べて、全体的に内容がリッチになっています。
そのおかげで質問内容やテーブル設計の改良のヒントが掴みやすい印象を受けました。

まとめ

というわけで、RAG環境構築と、RAGを利用したDB設計の省力化に挑戦することができました。
DBスキーマをPineconeやAmazon Bedrockを活用してベクトルデータとして管理する方法、モデルアクセスの設定やデータソースの同期方法についても説明し、実際に質問を投げることでテーブル構造の確認や設計支援がどのように行えるかを検証しました。

現時点では、回答内容に改善の余地が見られる部分が見られる結果となりましたが、DBスキーマの読み取り自体は問題なく行えてそうでした。
このあたりはプロンプトの調整でも改善はできるように感じます。

今後は、今回のシステムを実務に載せられるよう環境を整えたり、データソースの出力や同期の自動化をしたり、UIの構築、またDBだけではなくNotion連携などもできたら良いかと考えてます。
DBスキーマ以外の情報もデータソースとして活用できるので、様々な目的に対して今回利用した環境を応用できそうです。
今後も、ニーズがあれば自社の業務に適した環境を構築し、さらなる省力化を目指していければと思います。

参考資料

Discussion