はじめに

コードを書いている時のいちばんの悩みどころは変数や関数の命名ですよね。いい名前が浮かばないとモヤモヤしますし、気づいたら時間が溶けています。

LangChain について調べていて、手はじめになにか簡単なものは作れないかと考えていたところ、変数や関数の命名をやらせるというアイディアが降ってきたので、さっそく実践してみました。

LangChain とは

私が説明するまでもないですが、LangChain とは LLM を使ったアプリの作成を簡素化するように設計されたフレームワークです。Python が主流で最も機能が多いですが、TypeScript 版もあります。Go での実装プロジェクトもありますが、機能はまだ少ないようです。

いくつかのモジュールで構成されています。これを見るだけでも私は少しワクワクしました。

モジュール 役割 (かなり端折っています)
Model I/O 言語モデルの呼び出し
Retrieval RAG 関連の機能
Memory 会話履歴の保持、読み込み
Chains モジュールや機能の組み合わせ
Agents 外部データの取り扱い
Callbacks イベント処理

つくってみる

環境

python: 3.11.4
langchain: 0.0.350

素案

Python があまり得意ではないので、ChatGPT に以下のコンセプトを渡して素案を作ってもらいました。

- アプリ名: nsug (意味: name suggest)
- コンセプト:
  - テキストに書かれたコードから変数名、関数名、メソッド名を探し出し、より文脈に適した命名を提案するアプリ
  - Python, LangChain を使って CLI ツールとして実装する
  - 言語モデルは OpenAI を使う
- パラメータ:
  - 第1引数: コードが貼り付けられたテキストファイルのパス (default: code.txt)
  - 第2引数: 第1引数で与えられたコードの中で、提案してほしい現在の変数 (必須、default: なし)
  - 第3引数: いくつ提案してほしいか (default: 5)
- 振る舞い:
  - 環境変数 OPENAI_API_KEY が設定されていない場合は呼び出し元に処理を戻す
  - 第1引数で与えられたファイルを読み込む
    - 読み込みエラーは呼び出し元に処理を戻す
    - ファイルなしエラーは呼び出し元に処理を戻す
    - 空ファイルは呼び出し元に処理を戻す
  - 第2引数で与えられた単語をコード内から探す
    - 複数ヒットしてもよい
    - 見つからない場合エラーとし呼び出し元に処理を戻す
    - 変数名,関数名,メソッド名として検出するためのロジックは以下
      - 単語の左は改行、スペース、タブのいずれかである必要がある
      - 単語の右は "(","[", ":", " ", "=" のいずれかである必要がある
  - コードの文脈を読み取り、第2引数で与えられた名前をさらによくするための提案をする
  - 単語を第3引数で与えられた数値の数だけ、改行区切りで標準出力に出力する

しかし、ChatGPT は LangChain についてはあまり詳しくなく幻覚を多く含んだ回答だったので、LangChain 関連のコードは調査しながら書くことになりました。アウトラインといくつかのヘルパー関数はそのまま使えました。

コード

実際のコードです。

# nsug.py

import os
import sys
import re
import argparse
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.output_parsers import CommaSeparatedListOutputParser


def check_openai_api_key():
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise EnvironmentError("Environment variable not set: OPENAI_API_KEY")
    return api_key


def read_code_from_file(file_path):
    try:
        with open(file_path, "r") as file:
            code = file.read().strip()
        if not code:
            raise ValueError("File is empty")
        return code
    except FileNotFoundError:
        raise FileNotFoundError("File not found")
    except IOError as e:
        raise IOError("Error reading file: " + str(e))


def find_name_in_code(code, variable):
    pattern = r"(?<=[\n\s\t])" + re.escape(variable) + r"(?=[\(\[: =])"
    if not re.search(pattern, code):
        raise ValueError(f"Name '{variable}' not found in code")
    return True


def run(code, name, max):
    output_parser = CommaSeparatedListOutputParser()
    format_instructions = output_parser.get_format_instructions()
    llm = OpenAI(temperature=0)
    prompt = PromptTemplate(
        template="Please analyze the following code and suggest better naming in {name} as it appears in the code, taking into account the context. Please provide {max} suggestions. {format_instructions}\n\n####\n\n{code}\n\n####",
        input_variables=["name", "max", "code"],
        partial_variables={"format_instructions": format_instructions},
    )
    output = llm(prompt.format(code=code, name=name, max=max))
    items = output_parser.parse(output)
    for item in items:
        print(item)


def main():
    parser = argparse.ArgumentParser(description="nsug: a tool for suggesting better names for variables, functions, and methods in code.")
    parser.add_argument("--code", "-c", type=str, default="code.txt", help="path to the code file")
    parser.add_argument("--name", "-n", type=str, required=True, help="name of the variable, function, or method to suggest names for")
    parser.add_argument("--max", "-m", type=int, default=5, help="maximum number of suggestions to return")
    args = parser.parse_args()

    try:
        check_openai_api_key()
        code = read_code_from_file(args.code)
        find_name_in_code(code, args.name)
        run(code, args.name, args.max)
    except Exception as e:
        print(f"error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

以下では、与えられたワードが変数名、関数名、メソッド名としてもっともらしいかを判定します。かなり適当で、いろんな言語をカバーできていないと思います。

  • 単語の左は改行、スペース、タブのいずれか
  • 単語の右は "(", "[", ":", " ", "=" のいずれか
def find_name_in_code(code, variable):
    pattern = r"(?<=[\n\s\t])" + re.escape(variable) + r"(?=[\(\[: =])"
    if not re.search(pattern, code):
        raise ValueError(f"Name '{variable}' not found in code")
    return True

以下では、出力をカンマ区切りのリストで返すためのパーサーを初期化し、そのためのプロンプト文字列を変数に入れています。

output_parser = CommaSeparatedListOutputParser()
format_instructions = output_parser.get_format_instructions()

format_instructions には以下が格納されます。これをプロンプトに埋め込むことで、リストで返すように指示します。

Your response should be a list of comma separated values, eg: `foo, bar, baz`

プロンプトには変数を埋め込むことが可能です。{code} には第 1 引数で渡されたテキストファイルから読み込んだコードを、{name} には第 2 引数で渡されたワードを、{max} では第 3 引数で渡された出力数をそれぞれ埋め込みます。最後にカンマ区切りのリストで返すように指示しています。
作ったプロンプトを llm(prompt.format(code=code, name=name, max=max)) でリクエストし、出力を得ています。

prompt = PromptTemplate(
    template="Please analyze the following code and suggest better naming in {name} as it appears in the code, taking into account the context. This {name} is a variable name, function name, method name, etc. Please provide {max} suggestions. {format_instructions}\n\n####\n\n{code}\n\n####",
    input_variables=["name", "max", "code"],
    partial_variables={"format_instructions": format_instructions},
)
output = llm(prompt.format(code=code, name=name, max=max))

雑ではありますが、これだけで目的のコードが書けました。

使ってみる

では実際に使ってみましょう。code.txt に以下の Go コードを貼り付けます。これはスライスの中から重複を削除する関数です。

func uniq(slice []string) []string {
    seen := make(map[string]struct{})
    res := make([]string, 0, len(slice))
    for _, s := range slice {
        if _, ok := seen[s]; !ok {
            seen[s] = struct{}{}
            res = append(res, s)
        }
    }
    return res
}

この中の変数名 seen が気に入らないとします。この変数は、コンテキストとしてはチェック済みの要素をキーとして保持するための map です。map なので必ずユニークです。

$ python3 nsug.py --name "seen"
uniqueStrings
distinctStrings
uniqueValues
distinctValues
uniqueSet

ちゃんとコンテキストに沿った回答がリストで返ってきていますね。要素数を増やしてみましょう。

$ python3 nsug.py --name "seen" --max 10
uniqueStrings
distinctStrings
uniqueValues
distinctValues
uniqueElements
distinctElements
uniqueItems
distinctItems
uniqueSet
distinctSet

似たようなのが 10 件返ってきました。次に、関数名を考えさせてみましょう。

$ python3 nsug.py --name "uniq" --max 10
unique
distinct
filterDuplicates
removeDuplicates
deduplicate
distinctify
distinctValues
uniqueValues
filterUnique
removeDuplicateValues

ちゃんとコンテキストに沿った名前を提案してきますね。

プロンプトをいじってみる

精度の高いプロンプトを作るために、FewShotPromptMessage を使ってみます。プロンプトを以下のように変更します。

prompt = FewShotPromptTemplate(
    examples=[
        {
            "code": """def foo(str):
print("Hello " + str + "!!")""",
            "name": "foo",
            "max": 5,
            "output": "printGreeting, sayHello, greetWithName, displayHelloMessage, outputGreeting",
        }
    ],
    example_prompt=PromptTemplate(
        template="input1:\n{code}\n\ninput2:\n{name}\n\ninput3:\n{max}\n\noutput:\n{output}\n",
        input_variables=["code", "name", "max", "output"],
    ),
    prefix="Please analyze input1 below and suggest a better naming of input2 as it appears in the code, taking into account the context. Please provide input3 suggestions. "
    + format_instructions,
    suffix="input1:\n{code}\n\ninput2:\n{name}\n\ninput3:\n{max}\n\noutput:",
    input_variables=["code", "name", "max"],
)

こうすることで、以下のプロンプトが生成されます。

Please analyze input1 below and suggest a better naming of input2 as it appears in the code, taking into account the context. Please provide input3 suggestions. Your response should be a list of comma separated values, eg: `foo, bar, baz`

input1:
def foo(str):
    print("Hello " + str + "!!")

input2:
foo

input3:
5

output:
printGreeting, sayHello, greetWithName, displayHelloMessage, outputGreeting


input1:
func uniq(slice []string) []string {
        seen := make(map[string]struct{})
        res := make([]string, 0, len(slice))
        for _, s := range slice {
                if _, ok := seen[s]; !ok {
                        seen[s] = struct{}{}
                        res = append(res, s)
                }
        }
        return res
}

input2:
seen

input3:
5

output:

input1,input2, input3 という入力パラメータを与えると output という結果が返ることを例示してより具体的にコンテキストを理解させ、その流れで回答を求める手法です。
例に使っているのは引数で渡した文字列 (人の名前) に対して Hello と print するだけの関数なので、printGreeting とかそういう出力を期待することを例として教えてやります。
戻り値は以下。少しだけ出力が変わった気がします (temperature=0 で同じ質問には毎回同じ回答をさせるように設定しています)

$ python3 nsug.py --name "seen"
uniqueStrings
distinctStrings
uniqueValues
distinctValues
filterDuplicates

おわりに

プロンプトの準備、LLM の呼び出し、出力の整形と、最低限のことしかしていませんが、簡単に目的の結果を得ることができました。LangChain をフル活用すれば RAG の構築を加速させることができそうですし、IaC と組み合わせればインフラと一緒に管理もできて夢は広がるなと思いました。
昨今では RAG を構築するための敷居が下がってきているのを感じます。RAG を構築する土台であるインフラにも焦点を当ててみたいと思います。

さいごに、OpenAI API は従量課金なので、使いすぎに注意しましょう。