OpenAI Agents SDKを使ってナレッジベースのベクトル検索をする

はじめに

先日、リリースされた「Agents SDK」を使って、簡単な独自のナレッジベースのベクトル検索を実施しました。
今回は、その実装方法と実行例を紹介します。

ベクトル検索とは

ベクトル検索は、テキストや画像などのデータを数値ベクトル（埋め込み）に変換することによって、意味や文脈の類似性を捉えることができます。
例えば、「自動車の価格」という検索クエリに対して、「車の値段」や「自動車の費用」など、異なる表現でも意味が近い情報を見つけることができます。
OpenAI APIでは、テキストをベクトル化する機能が提供されており、今回はOpenAI APIを使ってURLからテキストをドキュメント化し、そのドキュメントに対して「Agent SDK」を使って自然言語で質問を投げかけます。

実装

ベクトルストアの作成

import sys
from openai import OpenAI

client = OpenAI()

knowledge_base_name = sys.argv[1]

vector_store = client.vector_stores.create(
    name=knowledge_base_name
)

print(f'Vector Store ID: {vector_store.id}')

ベクトル検索を実施するためのナレッジベースとなるベクトルストアを作成します。
このベクトルストアIDを取得し、後続のスクリプトで使用します。

ファイルアップロード

import sys
import requests
import time
from io import BytesIO
from openai import OpenAI

client = OpenAI()

vector_store_id = sys.argv[1]
target = sys.argv[2]

def create_file(client, file_path):
    if file_path.startswith("http://") or file_path.startswith("https://"):
        response = requests.get(file_path)
        file_content = BytesIO(response.content)
        file_name = file_path.split("/")[-1]
        file_tuple = (file_name, file_content)
        result = client.files.create(
            file=file_tuple,
            purpose="assistants"
        )
    else:
        with open(file_path, "rb") as file_content:
            result = client.files.create(
                file=file_content,
                purpose="assistants"
            )
    return result.id

file_id = create_file(client, target)

result = client.vector_stores.files.create(
    vector_store_id=vector_store_id,
    file_id=file_id
)

while True:
    # 最新のステータスを取得
    result = client.vector_stores.files.retrieve(
        vector_store_id=vector_store_id,
        file_id=file_id
    )
    print(f"現在の状態: {result.status}")
    
    # 処理が完了したかをチェック
    if result.status == "completed":
        print("処理が完了しました！")
        break
    elif result.status == "failed" or result.status == "canceled":
        print("エラーが発生しました")
        break
    
    time.sleep(2)
    
print(f'File追加: {result.id}')

このスクリプトは、ローカルファイルまたはURLから指定されたファイルを読み込みファイルをアップロードしてベクトルストアに追加します。
create_file関数は公式ドキュメントのサンプルコードを参考にしています。
ここまでは従来のOpenAI APIを使用しています。

ファイル削除

import sys
from openai import OpenAI

client = OpenAI()

vector_store_id = sys.argv[1]
file_id = sys.argv[2]

result = client.vector_stores.files.delete(
    vector_store_id=vector_store_id,
    file_id=file_id
)

print(f'ファイルの削除: {result.id}')

ベクトルストアに追加されたファイルを削除します。
指定したファイルが削除されると、検索ではそのファイルに関する情報は取得できなくなります。

ベクトル検索

import sys
from openai import OpenAI
from agents import Agent, Runner, FileSearchTool

client = OpenAI()

vector_store_id = sys.argv[1]
message = sys.argv[2]

agent = Agent(
    name="Assistant",
    tools=[
        FileSearchTool(
            max_num_results=3,
            vector_store_ids=[vector_store_id],
        ),
    ],
)

result = Runner.run_sync(agent, message)
print(result.final_output)

OpenAI Agents SDKを使用してベクトルストア内のドキュメントから情報を検索します。

toolsにFileSearchToolを指定することによりベクトル検索を実施できます。

実行例

上記のスクリプトを使って、ベクトルストアにファイルを登録してその内容に対して質問応答を実施してみます。
OPENAI_API_KEYを環境変数に設定して以下のライブラリをインストールしてください。

pip install openai-agents

ベクトルストアの作成

python create_vector_store.py keiba
Vector Store ID: vs_67d93869a6c08191a741a604ab87300a

keibaという名前のベクトルストアが作成されました。
このベクトルストアが検索の対象となるファイルをナレッジとして保持します。

ナレッジの追加

python upload_file.py vs_67d93869a6c08191a741a604ab87300a https://jra.jp/company/about/vision/pdf/keikaku2025.pdf
現在の状態: in_progress
現在の状態: completed
処理が完了しました！
File追加: file-JUd7RXHRWXqdNsaErpoFUc

URLからJRAの経営計画2025年のPDFファイルがアップロードされ、ベクトルストアに追加されました。
実際にいくつかの質問を投げかけて、PDF内の情報を使用した検索が機能するか確認してみましょう。

質問1: JRAアニバーサリーの開催日

python search_vector.py vs_67d93869a6c08191a741a604ab87300a "JRAアニバーサリーはいつ実施されますか？"
JRAアニバーサリーは9月15日（祝日・月曜日）に実施されます.

エージェントは、PDFから該当する情報を正確に抽出し回答しました。

質問2: JRAの経営方針

python search_vector.py vs_67d93869a6c08191a741a604ab87300a "JRAの経営方針について教えてください"
JRAの経営方針についての情報は以下の通りです：

1. **経営目標**:
   - 将来にわたる事業運営の安定化と経営基盤の強化。
   - 競馬の施行を通じた社会貢献を果たすことを目指しています。
   - 「開催日数 288日（36開催）」の競馬の着実な実施と「競馬の魅力の向上」を基本目標としています。

2. **施策**:
   - 国際的なスポーツエンターテインメントとして競馬をより魅力的にするための施策を実施。
   - 競馬番組の充実、質の高い競走の提供、競馬への参加促進及び販売促進などを行ないます。
   - 信頼確保に向けた取り組みの強化。

3. **持続可能な開発目標への取り組み**:
   - SDGsや気候変動、馬の福祉などの中央競馬を取り巻く課題に積極的に取り組みます。
   - デジタルトランスフォーメーション（DX）の推進により、事業運営の効率化・安定化を図ります。

総じて、JRAは中央競馬事業の長期的な発展を目指し、社会および経済的な影響を考慮しながら社会的基盤を強化する方針を掲げています。

PDF内の情報を抽出し、経営方針に関する情報を詳細に回答しました。

質問3: 2025年の日本ダービーの開催日(1回目)

python search_vector.py vs_67d93869a6c08191a741a604ab87300a "2025年の日本ダービーの開催日はいつですか？"
2025年の日本ダービーについての具体的な開催日は、提供された資料の中には記載されていないようです。そのため、今年度のJRAの公式発表か公式ウェブサイトを確認することをお勧めします。

ベクトルストアに含まれていない情報については、エージェントは正直に「情報が見つからない」と回答しました。

ナレッジの追加

python upload_file.py vs_67d93869a6c08191a741a604ab87300a https://www.jra.go.jp/datafile/seiseki/replay/2025/jyusyo.html
現在の状態: in_progress
現在の状態: completed
処理が完了しました！
File追加: file-1ZR3PCpJKHxZm5EQ1Pa6cW

直接JRAの重賞レース一覧URLから先ほどのベクトルストアに対して情報を追加してみます。

質問4: 2025年の日本ダービーの開催日(2回目)

python search_vector.py vs_67d93869a6c08191a741a604ab87300a "2025年の日本ダービーの開催日はいつですか？"
2025年の日本ダービーは6月1日（日曜日）に開催されます。

ベクトルストアに新しい情報が追加されたため、正確な回答をしました。

ファイルの削除

python delete_vector_store_file.py vs_67d93869a6c08191a741a604ab87300a file-1ZR3PCpJKHxZm5EQ1Pa6cW
ファイルの削除: file-1ZR3PCpJKHxZm5EQ1Pa6cW

先ほどの重賞レース一覧のファイルを削除します。
これで、ベクトルストアに追加された情報は削除され、検索結果には含まれなくなります。

質問5: 2025年の日本ダービーの開催日(3回目)

python search_vector.py vs_67d93869a6c08191a741a604ab87300a "2025年の日本ダービーの開催日はいつですか？"
2025年の日本ダービーの開催日はファイル内にはありませんでした。一般には、日本ダービーは毎年5月の最後の日曜日に開催されますので、2025年は5月25日になる可能性が高いです。ただし、正確な日程は公式発表を確認してください。

ベクトルストアからファイルが削除されたため、情報が見つからない旨を回答しました。
実際はベクトルストアとの連携が解除されただけで、情報となるファイル自体は削除されません。

まとめ

OpenAI Agents SDK使って、特定のドキュメントに対するベクトル検索を簡単に実装できました。
今回の実装は基本的なものですが、WebSearchToolなどを組み合わせることにより機能の拡張ができるのではないでしょうか。
社内文書のQAシステムや、特定分野のドキュメントに特化した検索システムなど、さまざまな用途に応用できそうなので、ぜひ試してみてください。

参考資料

• OpenAI Agents SDK
• OpenAI API
• GitHub: OpenAI Agents Python