OpenAIのAPIで簡単なRAGをPythonで試す

はじめに

最近、何かとRAG(Retrieval-Augmented Generation)をきくので、
動作を確認するプログラムを作成しました。

想定読者

RAGプログラムの組み方を知りたい方

作業環境

macOS Sonoma 14.5 / Python 3.10.x

作成の流れ

1. Pythonのインストール
2. OpenAIのAPIキー取得
3. プログラミング

1. Pythonのインストール

AI関係は何かと3.10.xが便利なので、関連するサイトを引用します。

macOSの場合

Windowsの場合

2. OpenAIのAPIキー取得

こちらも下記サイトを引用します。

3. プログラミング

作成したサンプルプログラム

sample_program.py

from openai import OpenAI
import numpy as np
import os
from sklearn.metrics.pairwise import cosine_similarity

# OpenAI APIキーの設定
os.environ["OPENAI_API_KEY"] = "xxxxxxxxxxxxxxxxxxxxxxxxx"

# OpenAIクライアントの初期化
client = OpenAI()

def vectorize_text(text):
    response = client.embeddings.create(
        input = text,
        # モデルの指定 (様々なモデルが利用可能)
        model = "text-embedding-3-small"
    )
    return response.data[0].embedding

# 回答データベース
answers = [
    "システム運用事業部",
    "販売管理システム",
    "第一システム部",
    "システム開発事業部"
]

# 回答の埋め込みベクトルを取得
answers_vector = [vectorize_text(answer) for answer in answers]

def rag_sample(question):

    # 質問の埋め込みベクトルを取得
    question_vector = vectorize_text(question)

    # コサイン類似度が最も高い回答を取得
    max_similarity = 0
    most_similar_index = 0
    for index, vector in enumerate(answers_vector):
        similarity = cosine_similarity([question_vector], [vector])[0][0]
        print(f"コサイン類似度: {similarity.round(4)}:{answers[index]}")
        # 取り出したコサイン類似度が最大のものを保存
        if similarity > max_similarity:
            max_similarity = similarity
            most_similar_index = index
    print(f"\n質問: {question}\n回答: {answers[most_similar_index]}\n")

if __name__ == '__main__':
    # 質問文
    question1 = "株式会社ABCの開発を行う事業部は？"
    question2 = "株式会社ABCの運用を行う事業部は？"
    question3 = "売上を管理するシステムは？"
    rag_sample(question1)
    rag_sample(question2)
    rag_sample(question3)

解説

はじめに必要なモジュールをインストールします

pip3 install openai
pip3 install -U scikit-learn

個別に修正や調整を行う箇所は下記になります。

• 前項で取得したAPIキーの設定

# OpenAI APIキーの設定
os.environ["OPENAI_API_KEY"] = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

• 回答の候補を検討します(最終的にはDBからの取得などになるかと思います)

# 回答データベース
answers = [
    "システム運用事業部",
    "販売管理システム",
    "第一システム部",
    "システム開発事業部"
]

• 質問文を考えて、rag_sample関数を呼び出します

    # 質問文
    question1 = "株式会社ABCの開発を行う事業部は？"
    question2 = "株式会社ABCの運用を行う事業部は？"
    question3 = "売上を管理するシステムは？"
    rag_sample(question1)
    rag_sample(question2)
    rag_sample(question3)

プログラムを実行すると、質問文に記載した質問に対し、回答データベースの中から
最もらしい回答(コサイン類似度が最大のもの)を回答します。
例では下記のような結果がターミナルに表示されます。

コサイン類似度: 0.4355:システム運用事業部
コサイン類似度: 0.3906:販売管理システム
コサイン類似度: 0.346:第一システム部
コサイン類似度: 0.5741:システム開発事業部

質問: 株式会社ABCの開発を行う事業部は？
回答: システム開発事業部

コサイン類似度: 0.5202:システム運用事業部
コサイン類似度: 0.4209:販売管理システム
コサイン類似度: 0.3507:第一システム部
コサイン類似度: 0.5153:システム開発事業部

質問: 株式会社ABCの運用を行う事業部は？
回答: システム運用事業部

コサイン類似度: 0.4558:システム運用事業部
コサイン類似度: 0.8353:販売管理システム
コサイン類似度: 0.3792:第一システム部
コサイン類似度: 0.4624:システム開発事業部

質問: 売上を管理するシステムは？
回答: 販売管理システム

まとめ

OpenAIの力をお借りして、専用のチャットボットが、かなり簡単にできそうです。

参考文献

RAG（検索拡張生成）とは？仕組みや生成AIとの関係性をわかりやすく解説

Python 環境構築(Mac編)

Python 3.10 の導入 on Windows

OpenAI APIキー取得方法

@straygizmo さんにリンクいただきました。ありがとうございます！
こちらはローカル環境で試す方向けの記事となります。
Ollamaで簡単なRAGをPythonで試す