はじめに

以前、からあげさんの書かれた記事「Google Colab上で手軽にAdvanced Data Analysis(Code Interpreter)的機能を実現する方法」や、生ビールさんの作られたノートブックを見て、おーこれぞ未来！！！と思いました。

というわけで、今日は、これらをちょっと改造した「Colab上で対話しながら、データ分析を進められるエージェント」を作った話を書きます。

エージェント作るのめちゃ面白いですね～。私は今月末に開催されるPyCon Kyushuで登壇するのですが、トークの後半でこの辺りのお話もできたらと思っています。

使ったのは次のような環境です。

google colaboratory
openai==1.30.1
pandas
numpy

どんな感じで分析できるか？

全体のコードは次のリポジトリにあります。

次のようなサンプル用のデータを対話しながら、分析します。

np.random.seed(42)
data = {
    'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
    'age': np.random.randint(20, 40, 5),
    'income': np.random.normal(50000, 10000, 5),
    'city': ['New York', 'London', 'Paris', 'Tokyo', 'Sydney']
}

# DataFrameを作成
dff = pd.DataFrame(data)

# 'income'列を整数に変換
dff['income'] = dff['income'].astype(int)

次の動画のように対話的な分析が行えます。

コード

コードは次のような感じです。ポイントは変数contextに辞書を持って、execに渡しているところです。この方法はClaudeに教えてもらいました。

model = 'gpt-3.5-turbo'  # 使用するGPTモデルを指定
prompt_history = []  # プロンプトの履歴を保存するリスト
context = {}  # 実行時のコンテキストを保存する辞書

def llm_response(messages, model):
    """
    OpenAIのChatCompletionsAPIを使用して、メッセージに対するLLMの応答を取得する関数。
    
    Args:
        messages (list): メッセージのリスト
        model (str): 使用するGPTモデルの名前
        
    Returns:
        dict: APIレスポンスの辞書
    """
    res = oai_client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0
    )
    return res

def run_agent(user_prompt, filename="", system_call=system_call, model=model):
    """
    ユーザープロンプトに基づいてエージェントを実行する関数。
    
    Args:
        user_prompt (str): ユーザーからの入力プロンプト
        filename (str): 入力データのファイルパス（オプション）
        system_call (str): システムコール（デフォルトはシステムコール関数）
        model (str): 使用するGPTモデルの名前（デフォルトはグローバル変数のモデル）
        
    Returns:
        None
    """
    global context
    
    oai_client = OpenAI(api_key = openai_api_key)
    # プロンプト履歴が空の場合、システムコールをプロンプト履歴に追加
    if prompt_history == []:
        prompt_history.append({"role": "system", "content": system_call})
    
    # ファイル名が指定されている場合、ユーザープロンプトにファイルパスを追加
    if filename != '':
        user_prompt += f""" # 入力データ:ファイルパス {filename} 入力データがあればそれも活用してください。 """
    
    # ユーザープロンプトをプロンプト履歴に追加
    prompt_history.append({"role": "user", "content": user_prompt})
    
    try:
        # LLMの応答を取得
        res = llm_response(prompt_history, model)
    except Exception as e:
        print(e)
    
    # LLMの応答メッセージを取得
    res_message = res.choices[0].message
    
    # LLMの応答をプロンプト履歴に追加
    prompt_history.append({'role': res_message.role, 'content': res_message.content})
    
    # LLMの応答にPythonコードが含まれている場合
    if "```python" in res_message.content:
        # Pythonコードを抽出
        python_codes = res_message.content.split("```python")[1:]
        
        for code in python_codes:
            code = code.split("```")[0]
            try:
                print(res_message.content)
                # Pythonコードを実行
                exec(code, context)
            except Exception as e:
                print(f"Error: {e}")
    else:
        # LLMの応答を表示
        print(res_message.content)

関数の解説です。

このコードは、OpenAIのChatCompletionsAPIを使用してLLMとのやり取りを行うためのものです。主な機能は以下の通りです(by Claude)。

llm_response関数: OpenAIのAPIを使用してLLMの応答を取得する関数。
run_agent関数: ユーザープロンプトに基づいてエージェントを実行する関数。以下の処理を行う。

プロンプト履歴が空の場合、システムコールをプロンプト履歴に追加。
ファイル名が指定されている場合、ユーザープロンプトにファイルパスを追加。
ユーザープロンプトをプロンプト履歴に追加。
LLMの応答を取得し、プロンプト履歴に追加。
LLMの応答にPythonコードが含まれている場合、そのコードを実行。
LLMの応答を表示。

execの解説は公式ドキュメントを参考にして下さい。

システムプロンプトは次のような感じです。これもClaudeに作ってもらっています。

system_call = """"
あなたは、データ分析タスクを支援するためのPythonコードを生成するAIアシスタントです。
ユーザーからの要求に基づいて、与えられたデータを活用し、以下の要件に従ってPythonコードを提供してください。

コード生成の要件:
1. コードはシンプルで読みやすいものにしてください。
2. 変数名や関数名は明確で説明的なものを使用してください。
3. 必要に応じて、コメントを追加して、コードの機能や目的を説明してください。
4. 標準ライブラリを積極的に活用してください。外部ライブラリを使用する場合は、その理由を明確にしてください。
5. ユーザー入力が必要な場合は、入力の形式と例を明示し、入力のバリデーションを行ってください。
6. エラーハンドリングを適切に行い、わかりやすいエラーメッセージを提供してください。
7. コードの実行結果を分かりやすく出力または説明してください。

実行環境について:
- 実行したコードはglobalにcontextを持っています。
- contextには、タスクに関連するデータやオブジェクトが格納されています。
- 必要なデータはcontextから取得し、効果的に活用してください。
- データの形式や内容が不明な場合は、clarificationを求めてください。

コードブロックを作成する際は、以下の形式に従ってください。
```python
# ここにPythonコードを記述```

データ分析のベストプラクティスに従い、効率的で再現可能なコードを心がけてください。
ユーザーとのやり取りを通じて、タスクの要件を明確にし、適切なソリューションを提供してください。

"""

まとめ

というわけで、オレオレ対話データ分析エージェントを作ってみました。ちょっとシステムプロンプトを変えると挙動が変わったりかなり面白いです。

注意としては exec は何でも実行してしまうので、危険な場合があるという点でしょう。

ちなみに、同じコードをGPT-4oでも動かしていますが、なんか話が長いというか親切すぎるというか・・・。この辺り、LLM毎にもうちょいシステムコールを変えたりすべきな点かなと思っています。