Vertex AI GeminiでGitリポジトリを丸ごと解析！コンテキスト キャッシュを利用したマルチモーダル活用

DX開発事業部の西田です。

今年の春のGenerative AI Summit Tokyo ’24で巨大なコードベースそのままGeminiのマルチモーダルコンテキストで受け入れて、コードベース全体を考慮した回答を得るスライドがありました。

実際にこれをやろうと思ったらどのようにプログラムで実現すればよいのか、調査する必要が生じたので調べたところこちらのズバリな記事を見つけたので試してみました！

Practical Guide: Using Gemini Context Caching with Large Codebases

Google has once again pushed the boundaries of AI innovation with the introduction of Gemini context caching for Vertex AI and Google AI…

Medium

以降のコードの実行環境はGoogle Colabを使っています。
Vertex AIが利用できるプロジェクトへの権限をもったGoogleアカウントでお試しください。

Gitリポジトリに格納されたファイルとパスを単一のテキストファイルに書き出す

Geminiに渡すためにPythonでリポジトリをクローンしてフォルダ構造とファイルの中身をすべて単一のテキストファイルにおさめます。
ここでは私がメンテナーをしているOSSのツールのリポジトリを使って試すことにしました。

PythonからGitを操作するためにGitPythonのインストール

!pip install GitPython=="3.1.43"

リポジトリをクローンして格納されているファイル情報を取得する関数。

import git
import os

def list_and_read_repo_files(repo_url, branch="main"):
    """
    Clones a Git repository, lists all files (excluding .git folder), and reads their contents.

    Args:
        repo_url (str): URL of the Git repository.
        branch (str, optional): Branch to checkout. Defaults to "main".

    Returns:
        dict: A dictionary where keys are file paths and values are their contents.
    """
    try:
        # Temporary directory for the clone
        repo_dir = "temp_repo"

        # Clone the repository
        print(f"Cloning repository from {repo_url}...")
        git.Repo.clone_from(repo_url, repo_dir, branch=branch)
        print("Cloning complete!")

        file_contents = {}
        for root, _, files in os.walk(repo_dir):
            for file in files:
                ##print(file)
                # Exclude .git folder and its contents
                if ".git" not in root:
                    ##print(".git not in root")
                    file_path = os.path.join(root, file)
                    try:
                      with open(file_path, "r", encoding="utf-8") as f:
                        file_contents[file_path] = f.read()
                    except Exception as e:  # Catch any unexpected errors
                         print(f"An unexpected error occurred: {e}")

        return file_contents

    except git.exc.GitCommandError as e:
        print(f"Error cloning repository: {e}")
    except UnicodeDecodeError as e:
        print(f"Error reading file: {e}")
    except Exception as e:  # Catch any unexpected errors
        print(f"An unexpected error occurred: {e}")
    finally:
        # Clean up the temporary repository directory
        if os.path.exists(repo_dir):
            git.rmtree(repo_dir)

リポジトリ情報を渡して関数を実行し、ファイル情報を取得。

repo_url = "https://github.com/danishi/dynamodb-csv.git"  # Replace with actual URL
branch = "master"  # Replace if different

file_data = list_and_read_repo_files(repo_url, branch)

これを単一のテキストファイルに保存します。

if file_data:
   output_file = "fullcode.text"
   with open(output_file, "w", encoding="utf-8") as outfile:

       for file_path, content in file_data.items():
           outfile.write(f"<file path={file_path}>")
           outfile.write(f"{content}")
           outfile.write("</file>")
       outfile.close()
       #print(f"File: {file_path}\nContent:\n{content}\n---")

テキストファイルの中身はこのようにfileタグで区切られて、パスとファイルの中身がわかるように記述されます。

<file path=temp_repo/app/__init__.py></file>
<file path=temp_repo/app/main.py>import os
import boto3
from boto3.session import Session
import configparser

ファイルサイズはだいたい92KBでした、ちょっとサンプルとしては小さかったかもしれません。

size_in_bytes = os.path.getsize(output_file)

print(f"File Size: {size_in_bytes/1024} KB")

Geminiの200万トークンのコンテキスト・ウィンドウなら6万行以上のコードを扱えるとのことなので、もっと巨大なリポジトリでも問題なく試せると思います。

ファイルをコンテキスト キャッシュに格納しGeminiに問い合わせを実行する

これでファイルの準備が整ったのでコードに対して問い合わせを行っていきます。
参考記事と同じようにVertex AIのコンテキスト キャッシュを使うことでコードを書き出したファイルをVertex AI側にキャッシュさせ、コスト効率よく問い合わせに利用していきます。

Vertex AIを呼び出すためのライブラリをインストール。

!pip install --upgrade google-cloud-aiplatform=="1.59.0"

ColabでGoogleの認証を通します。

from google.colab import auth
auth.authenticate_user()

Vertex AIの初期設定をします。プロジェクトIDは任意のものに打ち替えてください。
システムプロンプトを与えてハルシネーションを抑制しています。

import vertexai
from vertexai.preview import caching

project_id = "<USE YOUR PROJECT ID HERE>"

vertexai.init(project=project_id, location="asia-northeast1")

system_instruction = """
あなたは優秀なソフトウェアエンジニアです。提供されたソース内の事実に常に従い、新しい事実を作り上げることは決してありません。
では、このプロジェクトのコードベースを見て、次の質問に答えてください。
"""

先ほど作ったファイルをコンテキストとして渡すために変数に入れます。

with open("fullcode.text", "r", encoding="utf-8") as f:
     fullcode_as_string = f.read()

contents = [
    fullcode_as_string
]

コンテキスト キャッシュを作成します。
ファイルのサイズとキャッシュの永続時間に応じて課金が発生するので調節します。

モデルはGemini 1.5 FlashとGemini 1.5 Proが現在対応しています。

import datetime

cached_content = caching.CachedContent.create(
    model_name="gemini-1.5-pro-001",
    system_instruction=system_instruction,
    contents=contents,
    ttl=datetime.timedelta(minutes=10),
)

作成したコンテキスト キャッシュからモデルのインスタンスを作成。
これでGeminiに質問が行えます。

from vertexai.preview.generative_models import GenerativeModel

cache_id=cached_content.name
cached_content = caching.CachedContent(cached_content_name=cache_id)

model = GenerativeModel.from_cached_content(cached_content=cached_content)

早速聞いてみましょう。

from IPython.display import Markdown

response = model.generate_content("このプロジェクトは何をしていますか？")
Markdown(response.text)

このプロジェクトは、コマンドライン上で DynamoDB に CSV ファイルをインポート/エクスポートできるユーティリティである「DynamoDB CSV utility」です。

## 主な機能

- CSV ファイルのインポート：DynamoDB テーブルに CSV ファイルをインポートできます。バッチ書き込みにより高速処理を実現しています。
- CSV ファイルのエクスポート：DynamoDB テーブルのデータを CSV ファイルにエクスポートできます。インデックスを利用したエクスポートや、クエリ条件を指定したエクスポートにも対応しています。
- テーブルの切り捨て：DynamoDB テーブルのデータをすべて削除する機能です。実験的なインポートで不要なデータを消去する際に便利です。
- テーブルの移動：あるテーブルのすべてのアイテムを別のテーブルに移動する機能です。コピーのような動作をします。

## 使用方法

1. インストール
   - pip install dynamodb-csv
2. 設定ファイルの作成
   - config.ini ファイルを作成し、AWS のアクセスキー、シークレットキー、リージョンなどを設定します。
3. CSV ファイルと CSV 仕様ファイルの作成
   - インポートする CSV ファイルと、そのフォーマットを定義する CSV 仕様ファイルを作成します。
4. DynamoDB テーブルの作成
   - 仕様に合った DynamoDB テーブルを作成します。
5. CSV インポート、エクスポート、切り捨て、移動の実行
   - コマンドラインから dynamodb-csv コマンドを実行し、オプションを指定して各機能を実行します。

## 詳細

詳しくは、プロジェクトの README.md ファイルを参照してください。

## まとめ

DynamoDB CSV utility は、コマンドライン上で簡単に DynamoDB に CSV ファイルをインポート/エクスポートできる便利なツールです。CSV ファイルのフォーマットを定義する CSV 仕様ファイルを使用することで、柔軟なデータの取り扱いが可能です。

いい感じです。ツールを作った本人なので説明があってるのは私が保証します。

もうちょっと踏み込んだ処理のところも説明してもらいましょう。

response = model.generate_content("エクスポート処理のロジックを詳細に解説して")
Markdown(response.text)

`app/dynamodb/csv_export.py` ファイルの `csv_export` 関数は、DynamoDB テーブルを CSV ファイルにエクスポートする処理を行います。処理の詳細は以下の通りです。

1. **CSV 仕様ファイルの読み込み**
    - `csv_export` 関数に渡された CSV ファイル名(`file`) に `.spec` を付けたファイルを読み込み、CSV 仕様ファイルとして扱います。
    - CSV 仕様ファイルは `configparser` を用いてパースし、各カラムのデータ型などの情報を取得します。
    - 例外が発生した場合は、"CSV specification file can't read:{e}" というエラーメッセージと、終了コード 1 を返します。

2. **デリミタオプションの取得**
    - CSV 仕様ファイルに `DELIMITER_OPTION` セクションが存在する場合、`DelimiterCharacter` オプションの値をデリミタとして使用します。
    - `DELIMITER_OPTION` セクションが存在しない場合は、デフォルトのデリミタとしてスペース(" ")を使用します。

3. **CSV ファイルへの書き込み**
    - `file` で指定されたパスに、UTF-8 エンコーディングで CSV ファイルを開きます。
    - "please wait {name} exporting {file}" というメッセージを表示し、エクスポート処理の開始を知らせます。
    - 以下のいずれかの方法で DynamoDB テーブルからデータを取得します。
        - **QUERY_OPTION が指定されている場合:**
            - CSV 仕様ファイルに `QUERY_OPTION` セクションが存在する場合、その内容に従って DynamoDB テーブルに対してクエリを実行します。
            - `PKAttribute`, `PKAttributeValue`, `PKAttributeType` オプションが指定されている場合は、パーティションキーに基づいてクエリを実行します。
            - `SKAttribute`, `SKAttributeValues`, `SKAttributeType`, `SKAttributeExpression` オプションが指定されている場合は、ソートキーに基づいてクエリを実行します。
            - `KeyConditionExpression` を用いて、パーティションキーとソートキーの条件を指定します。
            - 例外が発生した場合は、"query option error:{e}" というエラーメッセージと、終了コード 1 を返します。
        - **QUERY_OPTION が指定されていない場合:**
            - `table.scan` メソッドを使用して、DynamoDB テーブルをスキャンし、全てのアイテムを取得します。
            - `LastEvaluatedKey` が存在する場合は、次のページのデータを取得するために `ExclusiveStartKey` に設定します。
    - クライアントエラーが発生した場合は、"aws client error:{e}" というエラーメッセージと、終了コード 1 を返します。
    - テーブルが見つからない場合は、"table not found:{e}" というエラーメッセージと、終了コード 1 を返します。

4. **CSV ヘッダーの書き込み**
    - CSV ヘッダーラベルを CSV 仕様ファイルの `CSV_SPEC` セクションから取得します。
    - `csv.DictWriter` を使用して、CSV ヘッダーラベルをファイルに書き込みます。

5. **アイテムの変換と書き込み**
    - 取得したアイテムをループ処理し、CSV 仕様ファイルに従って各アイテムの値を変換します。
        - `convert_item` 関数を呼び出して、各アイテムの値を CSV 仕様ファイルで指定されたデータ型に変換します。
        - 変換できない属性は削除されます。
    - 変換後のアイテムを `csv.DictWriter` を使用して CSV ファイルに書き込みます。

6. **終了処理**
    - "{name} csv exported {count} items" というメッセージと、終了コード 0 を返します。
    - I/O エラーが発生した場合は、"I/O error:{e}" というエラーメッセージを表示します。
    - その他の例外が発生した場合は、例外メッセージと、終了コード 1 を返します。

`convert_item` 関数は、DynamoDB アイテムの値を CSV 仕様ファイルで指定されたデータ型に変換する処理を行います。変換処理は、CSV 仕様ファイルの `CSV_SPEC` セクションで指定されたデータ型に基づいて行われます。

- `S` (String): `str` 関数を使用して文字列に変換します。
- `I` (Integer): `int` 関数を使用して整数に変換します。
- `D` (Decimal): `float` 関数を使用して浮動小数点数に変換します。
- `B` (Boolean): 値が空の場合は空文字列を返します。
- `J` (Json): `json.dumps` 関数を使用して JSON 文字列に変換します。
- `SL` (StringList) または `SS` (StringSet): `delimiter` で指定されたデリミタで結合した文字列を返します。
- `DL` (DecimalList) または `DS` (DecimalSet): 各要素を文字列に変換し、`delimiter` で指定されたデリミタで結合した文字列を返します。

`decimal_encode` 関数は、Decimal オブジェクトを float 型に変換します。これは、`json.dumps` 関数が Decimal オブジェクトを直接処理できないためです。

これもバッチリあっています！

最後にリファクタの提案もしてもらおうと思います。

response = model.generate_content("このプロジェクトでリファクタリングできそうな部分を教えて")
Markdown(response.text)

このプロジェクトのコードベースを基に、リファクタリングできる可能性のある部分をいくつかご紹介します。

* **エラー処理の標準化**: `app.dynamodb` の各モジュールでは、エラー処理方法が異なります。一貫したアプローチ、例えばカスタム例外クラスを使用することで、エラー処理の一貫性と可読性を向上させることができます。
* **関数の一本化**: 各モジュールには、`ClientError` と一般的な例外を処理するほぼ同じ `try-except` ブロックがあります。この共通処理を別関数に抽出し、各モジュールで再利用することでコードの重複を減らすことができます。
* **型ヒントの活用**: 一部の関数では型ヒントが使用されていますが、プロジェクト全体で一貫して適用することで、コードの理解とメンテナンスが容易になります。
* **設定ファイル読み込みの効率化**: `config_read_and_get_table` 関数では、設定ファイルを読み込むたびに `configparser` オブジェクトを生成しています。設定ファイル読み込み部分を初期化時に一度だけ行い、その結果を保持することで、効率性を向上させることができます。
* **関数の分割**: `convert_column` と `convert_item` 関数は、複数のデータ型変換ロジックを含んでいます。これらのロジックを個別の関数に分割し、より理解しやすい構造にすることができます。
* **進捗表示の抽象化**: `tqdm` を使用した進捗表示は各モジュールで個別に行われています。共通の進捗表示関数を作成し、各モジュールで再利用することで、コードの重複を減らすことができます。

これらの改善により、コードの可読性、保守性、堅牢性が向上すると考えられます。

うーん耳が痛い…。これもしっかり的を得たフィードバックでした。

まとめ

試したのは小さなリポジトリですが、Geminiの巨大なコンテキストウィンドウを使うことでもっと大きなリポジトリでもコードベース全体を俯瞰して問い合わせに回答させることができるでしょう。
毎回大きなファイルをマルチモーダルのプロンプトで扱うのはコストが気になりますが、コンテキスト キャッシュを利用することでコストを抑えて何回も違った角度からの問い合わせに対応できます。

他にも以下のようなユースケースにも活用できそうな気がします。

• README.mdを生成してもらう
• Mermaidで処理フローを生成してもらう
• テストコードを生成してもらう
• 改修の影響範囲の調査をしてもらう
• ワークフローに組み込んでPRのレビューをしてもらう

コードに対して生成AIを利用して提案を行ってくれる製品は有名どころでGitHub CopilotやGemini Code Assistがありますが、これらはユーザーごとの月額課金制です。
諸事情でこれらが使えないときや、スポットでこのリポジトリについて調べたいといった場面ではこのやり方のほうがコスト的にも有利に働きそうです！