LLMのRAGとは、Retrieval-Augmented Generation（検索拡張生成）の略称で、大規模言語モデル（LLM）の応答精度を向上させるための技術です🤖

RAGは、LLMが学習データにない最新情報や特定の知識を必要とする質問に答える際に特に役立ちます。

この記事では、Ubuntu22.04 + PostgreSQL(pgvector)でベクトルDBを簡単に構築する手順と、そのベクトルDBを活用して、Pythonで最も簡単なRAGを試したいと思います。

RAG処理の全体の流れ

外部知識をベクトル化してベクトルDBに保存
文書や FAQ、ナレッジ記事を embedding（数値ベクトル）化
ユーザー入力をベクトル化する
クエリも embedding に変換
類似検索（ベクトルDB）で関連文書を取得
pgvectorで近い文書を検索
コサイン類似度や L2 距離で上位 N 件を取得
取得した文書を Transformer に入力

目次

1. PostgreSQL + pgvectorでベクトルDBの構築
2. PythonでベクトルDBを利用してRAGを実行

1.PostgreSQL + pgvectorでベクトルDBの構築

1-1.pgvector をインストール

$ make --version
GNU Make 4.3
$ gcc --version
gcc (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
$ sudo systemctl status postgresql
$ psql --version
psql (PostgreSQL) 14.15 (Ubuntu 14.15-0ubuntu0.22.04.1)
$ sudo apt install postgresql-server-dev-14

$ git clone --branch v0.7.4 https://github.com/pgvector/pgvector.git
$ cd pgvector
$ make
$ sudo make install

1-2.PostgreSQL にログインして拡張を有効化

$ sudo -u postgres psql

postgres=# CREATE EXTENSION vector;
CREATE EXTENSION

1-3.ベクトル型のテーブル作成

3次元ベクトルのテーブルを作成する。※OpenAIだと1536次元

CREATE TABLE items (
    id SERIAL PRIMARY KEY,
    name TEXT,
    embedding VECTOR(3)
);

1-4.データ挿入

postgres=# INSERT INTO items (name, embedding) VALUES
('猫', '[0.1, 0.2, 0.3]'),
('犬', '[0.2, 0.1, 0.4]'),
('車', '[0.9, 0.8, 0.7]');

このベクトルの数値は、

1つの次元が「動物 vs. 乗り物」軸
別の次元が「猫っぽさ vs. 犬っぽさ」軸
もう1つの次元が「かわいい度合い」
などのように特徴軸を表します

1-5.類似検索

postgres=# SELECT id, name, embedding
FROM items
ORDER BY embedding <=> '[0.1, 0.2, 0.25]'   -- クエリベクトル
LIMIT 2;

id | name | embedding
----+------+---------------
1 | 猫 | [0.1,0.2,0.3]
2 | 犬 | [0.2,0.1,0.4]
(2 rows)

猫が最も近いと出てきてますね。いい感じです。

2.PythonでベクトルDBを利用してRAGを実行

2-1.psql側の準備

OpenAIの1536次元ベクトルで作成します。

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding VECTOR(1536) 
);

2-2.Pythonで実行

import os
import psycopg2
from openai import OpenAI

DB_CONFIG = {
    "dbname": os.getenv("DB_NAME"),
    "user": os.getenv("DB_USER"),
    "password": os.getenv("DB_PASSWORD"),
    "host": os.getenv("DB_HOST", "localhost"),
    "port": os.getenv("DB_PORT", 5432)
}

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def embed_text(text):
    resp = client.embeddings.create(
        model="text-embedding-ada-002",
        input=text
    )
    return resp.data[0].embedding

def insert_document(content):
    """文書をベクトル化して DB に保存"""
    vector = embed_text(content)
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO documents (content, embedding) VALUES (%s, %s)",
        (content, vector)
    )
    conn.commit()
    cur.close()
    conn.close()

def search_documents(query, top_k=3):
    """クエリ embedding でベクトルDB検索"""
    q_vector = embed_text(query)
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    cur.execute(f"""
        SELECT content
        FROM documents
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (q_vector, top_k))
    results = [r[0] for r in cur.fetchall()]
    cur.close()
    conn.close()
    return results

def generate_answer(query):
    related_docs = search_documents(query)
    context = "\n".join(related_docs)
    prompt = f"以下の文脈を参考にして質問に答えてください。\n\n文脈:\n{context}\n\n質問: {query}\n回答:"

    completion = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": prompt}
        ]
    )
    return completion.choices[0].message.content

if __name__ == "__main__":
    # 文書の挿入例
    insert_document("Pythonは高水準のプログラミング言語です。")
    insert_document("PostgreSQLはオープンソースのリレーショナルデータベースです。")
    insert_document("OpenAIはAI研究と展開を行う企業です。")

    # 質問応答の例
    question = "Pythonとは何ですか？"
    answer = generate_answer(question)
    print("質問:", question)
    print("回答:", answer)

実行結果
$ python3 rag.py
質問: Pythonとは何ですか？
回答: Pythonは高水準のプログラミング言語です。

おおおお、いい感じ！

まとめ

今回は、psqlにpgvectorでvectorを拡張して、PythonでRAGを実行する手順を紹介しました。

なお、LangChainを使うことで、RAGを簡単に実装することもできます。参考までにLangChainを使ったRAGのサンプルも載せておきます。

import os
from dotenv import load_dotenv
from langchain_openai import OpenAI
from langchain.chains import RetrievalQA
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings

load_dotenv()

api_key = os.getenv("OPENAI_API_KEY")

embeddings = OpenAIEmbeddings(api_key=api_key)

docs = [
    "LangChainは、LLMを活用したアプリ開発を効率化するフレームワークです。",
    "RAGは、外部の知識を検索してLLMに渡すことで精度の高い応答を生成する手法です。",
    "FAISSはベクトル検索用の高速ライブラリで、LangChainでRAGを実装する際によく使われます。"
]

vectorstore = FAISS.from_texts(docs, embeddings)

llm = OpenAI(temperature=0.7, api_key=api_key)

qa = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectorstore.as_retriever()
)

question = "RAGとは何ですか？"
answer = qa.run(question)
print("質問:", question)
print("回答:", answer)