生成AI (Cohere)＋LangChain＋Vector Database (PostgreSQL)でRAGを試してみた

生成AIを企業が使う場合、社内データを使った回答を得るにはファインチューニング、もしくは Retrieval-Augmented Generation (RAG、検索拡張生成) を行う必要があります。
企業では毎日データが更新される中で、ファインチューニングを頻繁に行うのはコスト高で現実的ではありません。
そこでRAGを使った方法が注目されています。
ということで、今回は以下の組み合わせでRAGを試してみました。

• 生成AI: Cohere Command
• Vector Database: PostgreSQL (pgvector)
• 生成AIとVector Databaseの連携: LangChain

将来的にはCohereのサービスがOCIで、Vector DatabaseがOracle Database 23cで提供される予定なので、GAとなった際はこれらを使って試したいと思います。
現時点ではリリースされていないため、一旦前述の組み合わせで試してみました。
また動作確認は Windows Subsystem for Linux 上に立てた Oracle Linux 9.3 で行いました。

なお以降のPythonスクリプトで指定するCohereのAPI Keyですが、事前にCohereのアカウントを作成し取得してください。

Windows Subsystem for Linux 上の Oracle Linux 9.3 でsystemdを使う設定

WSLで試す場合は本手順を実行し、systemdを使えるようにしておきます。

$ vim /etc/wsl.conf
## 以下のように修正
## [boot]
## systemd=true

修正したらWSLを再起動します。

wsl --shutdown

必要なPythonパッケージのインストール

Pythonのコードを実行するOSユーザで、今回の検証向けにPython仮想環境を作り、アクティベートします。

$ python -m venv rag
$ source ./rag/bin/activate

Pythonパッケージのインストール作業を進めていきます。

LangChain

$ pip install langchain

Cohereのサービスを使うためのパッケージ

$ pip install cohere

pgvectorを使うためのパッケージ

$ pip install pgvector

PostgreSQLへ接続するためのパッケージ

$ pip install psycopg2-binary

PostgreSQLのインストール

今回はPostgreSQL 13で検証しました。
まずはPostgreSQL用のリポジトリを追加します。

$ sudo dnf install -y https://download.postgresql.org/pub/repos/yum/reporpms/EL-9-x86_64/pgdg-redhat-repo-latest.noarch.rpm

PostgreSQL本体、および開発に必要なパッケージをインストール

$ sudo dnf install postgresql-server postgresql-libs postgresql postgresql-contrib

PostgreSQLを初期化

$ sudo /usr/bin/postgresql-setup --initdb

pgvectorのインストール

まずはPostgreSQLのリポジトリからpgvectorをインストールします。

$ sudo dnf install pgvector_13

pgvectorのextention関連ファイルをPostgreSQLのextention配置先にコピーします。

$ cp -p /usr/pgsql-13/share/extension/vector* /usr/share/pgsql/extension/
$ cp -p /usr/pgsql-13/lib/vector.so /usr/lib64/pgsql/

PostgreSQLを起動します。

$ sudo systemctl start postgresql
$ sudo systemctl status postgresql

PostgreSQLにログインし、エクステンションを有効化します。

$ sudo su - postgres
$ psql
CREATE EXTENSION vector;

PostgreSQLへのホスト名指定でのログインを有効化

PostgreSQLのpostgresユーザにパスワードを設定

$ psql
postgres=# alter role postgres with password 'postgres';
postgres=# \q

PostgreSQLへのログイン設定を修正

$ vim /var/lib/pgsql/data/pg_hba.conf
## TYPEがhostの行を以下のように修正
## host    all             all             127.0.0.1/32            md5

PostgreSQLを再起動

$ sudo systemctl stop postgresql
$ sudo systemctl start postgresql
$ sudo systemctl status postgresql

PostgreSQLにサンプルデータを追加

表を作成します。
今回は架空のおもちゃメーカーの製品情報をサンプルデータとして追加します。
この製品情報を使ってRAGの動作確認をしたいと思います。

CREATE TABLE products (
    product_id SERIAL PRIMARY KEY,
    product_name VARCHAR(255),
    manufacturer_name VARCHAR(255),
    product_overview TEXT,
    price FLOAT,
    release_date DATE
);

架空の製品情報を挿入します。

INSERT INTO products (product_name, manufacturer_name, product_overview, price, release_date)
    VALUES ('Tiny Tunes Piano', 
        'Playful Minds Co.', 
        'A colorful and musical toy piano that lets toddlers play and learn songs. The piano has 12 keys, each with a different animal sound and a corresponding sticker. The piano also has a mode switch that allows toddlers to choose between animal sounds, musical notes, or songs. The piano comes with a songbook that teaches toddlers how to play 10 popular nursery rhymes.', 
       29.99, 
       '2023-12-01');

INSERT INTO products (product_name, manufacturer_name, product_overview, price, release_date)
    VALUES ('Tiny Blocks Castle', 
        'Playful Minds Co.', 
        'A set of 50 wooden blocks in various shapes and colors that toddlers can use to build their own castle. The blocks are easy to stack and connect with magnets. The set also includes a drawbridge, a flag, and two knight figures. The blocks are made of natural and eco-friendly materials and are safe for toddlers to play with.', 
       39.99, 
       '2024-01-15');

INSERT INTO products (product_name, manufacturer_name, product_overview, price, release_date)
    VALUES ('Tiny Pals Farm', 
        'Playful Minds Co.', 
        'A fun and interactive toy farm that introduces toddlers to farm animals and their sounds. The farm has a barn, a silo, a tractor, and a fence. The farm also has 10 animal figures, each with a button that plays its sound when pressed. The farm has a sensor that recognizes the animal figures and plays a song or a fact about them when placed on the farm.', 
       49.99, 
       '2024-02-01');

ベクトルデータの生成

Cohereのembedding機能を使ってベクトルデータを生成し、結果をpgvectorを使いPostgreSQLに格納します。

import os
import psycopg2
from langchain.vectorstores.pgvector import PGVector
from langchain.embeddings import CohereEmbeddings
from langchain.docstore.document import Document

## Cohereのembedding機能を使うためにAPI KEYを用意
COHERE_API_KEY = os.environ['COHERE_API_KEY']

## Cohereのembedding機能を使うための設定
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY)

## PostgreSQLへの接続文字列
CONNECTION_STRING = "postgresql+psycopg2://postgres:postgres@localhost:5432/postgres"

## ベクトルデータを格納するコレクション名を指定
COLLECTION_NAME = "products_v"

## 操作するコレクションを設定
store = PGVector(
    collection_name=COLLECTION_NAME,
    connection_string=CONNECTION_STRING,
    embedding_function=embeddings,
)

## データベースへの接続情報を設定
conn = psycopg2.connect(
    host="localhost",
    database="postgres",
    user="postgres",
    password="postgres"
)

## カーソルを作成
cur = conn.cursor()

## ベクトル化する対象の製品データを取得
cur.execute("SELECT product_name, manufacturer_name, product_overview, price, release_date FROM products;")

## 結果を取得
rows = cur.fetchall()

for row in rows:
    ## 各行の製品情報を1つのドキュメントにまとめる
    doc = "product name: " + row[0] + ", manufacturer name: " + row[1] + ", product overview: " + row[2] + ", price: " + str(row[3]) + "$ , release date: " + row[4].strftime('%Y/%m/%d')
    
    ## ドキュメントのベクトル化、およびベクトルデータのコレクションへの保存
    store.add_documents([Document(page_content=doc)])

## カーソルと接続をクローズ
conn.commit()
cur.close()
conn.close()

ちなみにこの時点でPostgreSQLにログインし、作成済みテーブルを確認すると、LangChainによって以下テーブルが作成されているのを確認できます。

• langchain_pg_collection
• langchain_pg_embedding

postgres=# \dt
                  List of relations
 Schema |          Name           | Type  |  Owner
--------+-------------------------+-------+----------
 public | langchain_pg_collection | table | postgres
 public | langchain_pg_embedding  | table | postgres
 public | products                | table | postgres

postgres=# \d langchain_pg_collection
             Table "public.langchain_pg_collection"
  Column   |       Type        | Collation | Nullable | Default
-----------+-------------------+-----------+----------+---------
 name      | character varying |           |          |
 cmetadata | json              |           |          |
 uuid      | uuid              |           | not null |
Indexes:
    "langchain_pg_collection_pkey" PRIMARY KEY, btree (uuid)
Referenced by:
    TABLE "langchain_pg_embedding" CONSTRAINT "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE

postgres=# \d langchain_pg_embedding
               Table "public.langchain_pg_embedding"
    Column     |       Type        | Collation | Nullable | Default
---------------+-------------------+-----------+----------+---------
 collection_id | uuid              |           |          |
 embedding     | vector            |           |          |
 document      | character varying |           |          |
 cmetadata     | json              |           |          |
 custom_id     | character varying |           |          |
 uuid          | uuid              |           | not null |
Indexes:
    "langchain_pg_embedding_pkey" PRIMARY KEY, btree (uuid)
Foreign-key constraints:
    "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE

langchain_pg_collection には作成したコレクションの情報が格納されています。

postgres=# select name, uuid from langchain_pg_collection;
    name    |                 uuid
------------+--------------------------------------
 products_v | bf2790ed-521c-44ae-879b-82ffdb8edc3b

langchain_pg_embedding には langchain_pg_collection に格納されているコレクションのベクトルデータ、および元テキストが保存されています。
ベクトルデータ自体は langchain_pg_embedding 表の embedding 列に格納されており、4096個の数値で構成されています。

素の Cohere Command による回答の確認

以下のスクリプトでまずは素の Cohere Command の回答内容を確認します。

import os
from langchain.chains import LLMChain
from langchain.llms import Cohere
from langchain.prompts import PromptTemplate

## Cohereのembedding機能を使うためにAPI KEYを用意
COHERE_API_KEY = os.environ['COHERE_API_KEY']

## 質問文のテンプレートを指定
template = """
Please answer the following question.
question: {question}
answer:
"""

prompt = PromptTemplate(template=template, input_variables=["question"])

## 生成AI(LLM)の設定
llm = Cohere(model="command-nightly", cohere_api_key=COHERE_API_KEY)

llm_chain = LLMChain(prompt=prompt, llm=llm)

## 質問を実行
question = "What is the toy named Tiny Pals Farm?"

answer = llm_chain.run(question)

## 回答を標準出力
print(answer)

回答内容は以下のようになりました。
PostgreSQL内の製品情報は特に参照していないため、適当に生成された回答になっているようです。

 Tiny Pals Farm is a toy farm set that was produced by Fisher-Price during the 1990s. It was a popular toy among children and is now considered a collectible among nostalgic adults who enjoyed it as children. The set typically includes a variety of farm animals, a farmer figure, a barn, and other accessories such as fences, hay bales, and vehicles.

The Tiny Pals Farm was part of the wider Tiny Pals range of toys by Fisher Price, which included various other playsets such as a castle, a school, and a zoo. The toys in this range were designed to be compact and portable, making them ideal for children to carry around and play with wherever they went.

The farm set, in particular, was well-loved for its charming and colorful design, as well as its interactive features. For instance, the barn usually had a swinging gate, a loft with a movable ladder, and doors that opened and closed. The animals could be placed in their designated areas within the barn or roam free around the farmyard.

While the Tiny Pals Farm toy is no longer in production, it continues to hold a special place in the hearts of many who played with it as children and retains a certain collector's value among those

RAGの動作確認

以下がRAGを試した際のPythonスクリプトです。

import os
from langchain.vectorstores.pgvector import PGVector
from langchain.embeddings import CohereEmbeddings
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatCohere
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.agents import Tool
from langchain.chains.qa_with_sources.map_reduce_prompt import QUESTION_PROMPT
from langchain.prompts import PromptTemplate

## Cohereのembedding機能を使うためにAPI KEYを用意
COHERE_API_KEY = os.environ['COHERE_API_KEY']

## Cohereのembedding機能を使うための設定
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY)

## PostgreSQLへの接続文字列
CONNECTION_STRING = "postgresql+psycopg2://postgres:postgres@localhost:5432/postgres"

## ベクトルデータを格納するコレクション名を指定
COLLECTION_NAME = "products_v"

## 操作するコレクションを設定
store = PGVector(
    collection_name=COLLECTION_NAME,
    connection_string=CONNECTION_STRING,
    embedding_function=embeddings,
)

## 生成AI(LLM)の設定
llm =ChatCohere(model="command-nightly")

## LLMとVector Databaseの連携設定
qa = RetrievalQA.from_chain_type(llm, chain_type="map_reduce", retriever=store.as_retriever(), return_source_documents=True)

tools = [
    Tool(
      name = "pgvector_search",
      func=qa,
      description="vector database for products info"
  )
]

agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, handle_parsing_errors=True)

## 質問文のテンプレートを指定
template = """
Please answer the following question.
question: {question}
answer:
"""

prompt = PromptTemplate(
    input_variables=["question"],
    template=template,
)

## 質問を実行
query = "What is the toy named Tiny Pals Farm?"
question = prompt.format(question=query)
agent.run(question)

以下が出力された回答です。

Tiny Pals Farm is a line of children's toys inspired by farm animals. It includes figurines and play sets for imaginative play.

一応Vector Database内の情報を参照した回答になっているようです。
以上がRAGを実行する一連の流れでした。