🧸

RAGのお試し。ベクトルDBを使わずにベクトル検索を実装してみた。

2025/03/05に公開

ベクトル検索

近年、生成系AIの盛り上がりとともにベクトル検索が注目されています。

テキストを埋め込みベクトルとして扱うことで、自然言語ベースの高精度な検索やQAを実現できる点が大きな魅力かと思います。

本格的に運用するなら Pinecone や Weaviate といった専用のベクトルデータベース（Vector DB）を導入するケースが多いかと思いますが、今回は、「ベクトルDBなし」かつ「AWSの標準サービス＋OpenAI API」 で、シンプル＆低コストにベクトル検索を試してみました。

処理フロー概要

1. 外部知識データの登録

ユーザーがテキストファイルを S3 にアップロード
S3 イベント通知を受けて Lambda が起動
Lambda はアップロードされたテキストから OpenAI の埋め込み API を使ってベクトルを生成
生成したベクトルを S3 に保存し、DynamoDB にメタデータを登録

2. クエリ応答

ユーザーのクエリを受け取り、まずはクエリを埋め込みベクトル化
DynamoDB から対象ドキュメントを特定し、そのベクトルを S3 から取り出してコサイン類似度などで照合
最も類似度が高いドキュメントの内容を使い、ChatGPT API（などの LLM）を呼び出して回答を生成

外部知識（ドキュメント）をベクトル化して登録の実装

Lambda 処理（ドキュメントアップロード時 → ベクトル生成・保存）

S3 バケットへのファイルアップロードをトリガーに、Lambda が起動
OpenAI の Embedding API で生成した埋め込みベクトルを S3 に保存し、DynamoDB へメタデータを登録

import json
import boto3
import openai
import uuid
import os

# AWSクライアント
s3 = boto3.client('s3')
dynamodb = boto3.client('dynamodb')

# OpenAI APIキー
openai.api_key = os.getenv('API_KEY')

# パラメータ
S3_BUCKET_NAME = os.getenv('DOC_INPUT_BUCKET')         # テキストファイルのアップロード先
VECTOR_S3_BUCKET_NAME = os.getenv('VECTOR_BUCKET')     # ベクトルを保存するバケット
DYNAMODB_TABLE_NAME = os.getenv('DYNAMODB_TABLE_NAME') # DynamoDBテーブル名

def lambda_handler(event, context):
    try:
        # イベント情報からアップロードされたファイルを特定
        record = event['Records'][0]
        bucket = record['s3']['bucket']['name']
        key = record['s3']['object']['key']

        # S3からファイル（テキスト）を取得
        response = s3.get_object(Bucket=bucket, Key=key)
        text_data = response['Body'].read().decode('utf-8')

        # OpenAI APIで埋め込みベクトルを生成
        embedding_response = openai.embeddings.create(
            input=text_data,
            model="text-embedding-ada-002"
        )

        embedding_vector = embedding_response.data[0].embedding
        
        # 埋め込みベクトルをS3に保存（JSON形式）
        vector_key = f"vectors/{uuid.uuid4()}.json"
        s3.put_object(
            Bucket=VECTOR_S3_BUCKET_NAME,
            Key=vector_key,
            Body=json.dumps({"embedding": embedding_vector}),
            ContentType="application/json"
        )

        # ランダムなDocumentIDを生成
        document_id = str(uuid.uuid4())

        # DynamoDBにメタデータを保存
        dynamodb.put_item(
            TableName=DYNAMODB_TABLE_NAME,
            Item={
                "DocumentID": {"S": document_id},
                "SourceS3Key": {"S": key},
                "VectorS3Key": {"S": vector_key}
            }
        )

        return {
            "statusCode": 200,
            "body": f"Embedding for {document_id} saved successfully."
        }

    except Exception as e:
        print(f"Error processing file: {e}")
        return {
            "statusCode": 500,
            "body": "Error processing the file."
        }

クエリ → ベクトル検索 → ChatGPT応答生成の実装

Lambda 処理（ユーザーのクエリ → ベクトル検索 → LLM応答）

ユーザーからのクエリを受け取り、OpenAI の Embedding API を使ってクエリをベクトル化
DynamoDB から全ドキュメントのメタデータを取得し、S3 に保存されたベクトルファイルと比較して最も類似度が高いドキュメントを特定
ドキュメントの本文を元に ChatGPT API で回答を生成して返す

import json
import boto3
import openai
import numpy as np
import os

# AWSクライアント
s3 = boto3.client('s3')
dynamodb = boto3.client('dynamodb')

# OpenAI APIキー
openai.api_key = os.getenv('API_KEY')

# パラメータ
S3_BUCKET_NAME = os.getenv('DOC_INPUT_BUCKET')         # ドキュメントを保存しているバケット
VECTOR_S3_BUCKET_NAME = os.getenv('VECTOR_BUCKET')     # ベクトルを保存しているバケット
DYNAMODB_TABLE_NAME = os.getenv('DYNAMODB_TABLE_NAME') # DynamoDBテーブル名

def cosine_similarity(vec1, vec2):
    """コサイン類似度を計算"""
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

def lambda_handler(event, context):
    try:
        # リクエスト本文からクエリを取得
        body = json.loads(event['body'])
        query = body['query']

        # クエリをOpenAI APIでベクトル化
        embedding_response = openai.embeddings.create(
            input=query,
            model="text-embedding-ada-002"
        )
        query_vector = embedding_response.data[0].embedding

        # DynamoDBからメタデータをスキャン（全件取得）
        response = dynamodb.scan(TableName=DYNAMODB_TABLE_NAME)
        items = response['Items']

        # 類似度計算
        results = []
        for item in items:
            document_id = item['DocumentID']['S']
            vector_key = item['VectorS3Key']['S']
            source_s3_key = item['SourceS3Key']['S']

            # S3から保存済みベクトルを取得
            s3_response = s3.get_object(Bucket=VECTOR_S3_BUCKET_NAME, Key=vector_key)
            vector_data = json.loads(s3_response['Body'].read())
            saved_vector = vector_data['embedding']

            # コサイン類似度を算出
            similarity = cosine_similarity(query_vector, saved_vector)

            results.append({
                "DocumentID": document_id,
                "Similarity": similarity,
                "S3Key": source_s3_key
            })

        # 類似度が最も高いドキュメントを選択
        top_result = sorted(results, key=lambda x: x['Similarity'], reverse=True)[0]

        # もっとも類似したドキュメントの本文を取得
        s3_key = top_result['S3Key']
        s3_response = s3.get_object(Bucket=S3_BUCKET_NAME, Key=s3_key)
        doc_content = s3_response['Body'].read().decode('utf-8') 

        # ChatGPT API で回答を生成
        chat_response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {
                    "role": "user", 
                    "content": f"""
                    以下の情報を参考にして、ユーザーの質問に答えてください。
                    情報: {doc_content}
                    質問: {query}
                    """
                }
            ],
            max_tokens=1000
        )

        # ChatGPTからの応答を取得
        answer = chat_response.choices[0].message.content.strip()
        top_source_id = top_result['DocumentID']

        # 結果を返却
        return {
            "statusCode": 200,
            "headers": {"Content-Type": "application/json"},
            "body": json.dumps({
                "query": query,
                "answer": answer,
                "source_document": top_source_id
            })
        }

    except Exception as e:
        print(f"Error processing search: {e}")
        return {
            "statusCode": 500,
            "body": "Error processing the search request."
        }

まとめ

今回は、「ベクトルDBなし」かつ「AWS上の標準サービス＋OpenAI API」でベクトル検索を実装してみました。

本番運用で使用するのは難しいかと思いますが、手順もシンプルでコストもあまりかからないのでお試しにはちょうど良いかなと思います。