LoginSignup
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@naozo-se

PythonでPostgreSQLベクトルデータベース活用RAGのプログラム

Posted at

はじめに

タイトルの通りです。
この環境のPostgreSQLを利用する python コードです。
LlamaIndex のパッケージを利用します。

確認環境

  • WSL2(Ubuntu22.04.3 LTS)
  • Docker (27.5.1)
  • Docker Compose(v2.32.4)
  • python (3.13.3)
  • PostgreSQL (17.4-1)
  • OpenAI API

フォルダ構成

[projectfolder]
 ├─ main.py
 └─ docs
      ├─ story1.txt
      └─ story2.txt

※docsはベクトルデータにする対象のテキスト

ベクトルデータにする対象のファイルについて

  • AIの知識に左右されないため、適当なストーリーをAI(Gemini)で生成します

  • 以下２つのプロンプト実行結果をstoryのテキストファイルとして利用しています

    実行プロンプト 
    以下の条件で物語を作って欲しいです。

文字数は5000文字前後にしてください。
犬を出してください。名前はカタカナにしてくださ
赤、青、金の鬼を出してください。

以上を満たせば何でもよいです。

    この話の第２部を作成してください。

メインのキャラクターはそのまま利用してください。
鬼の色は緑、黄色、黒を出してください。

文字数は同じく5000文字前後で大丈夫です。

パッケージインストール

pip install llama-index \
llama-index-vector-stores-postgres  \
llama-index-embeddings-openai  \
llama-index-llms-openai  \
python-dotenv  \
openai  \
psycopg2-binary

コード

main.py

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    Settings,
)
from llama_index.vector_stores.postgres import PGVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# ------ 1. 必要な初期設定
# ---- OpenAI APIの設定
OPENAI_API_KEY = "sk-..." # APIキー
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIM = 1536  # text-embedding-3-small の次元数
LLM_MODEL = "gpt-4o-mini"

# ---- Postgresqlの設定
PG_HOST = "host.docker.internal" # 接続先を設定
PG_PORT = 5432
PG_USER = "postgres"
PG_PASSWORD = "postgres"
PG_DATABASE = "postgres"

# ---- LlamaIndex設定
Settings.embed_model = OpenAIEmbedding(
    model=EMBEDDING_MODEL, api_key=OPENAI_API_KEY, dimensions=EMBEDDING_DIM
)
Settings.llm = OpenAI(model=LLM_MODEL, api_key=OPENAI_API_KEY)


# ------ 2. テキストファイルをベクトルデータとして登録
def add_vector_data(folder_path: str, table_name: str) -> None:
    """
    指定されたフォルダ内のドキュメントを読み込み、PostgreSQLの指定テーブルにベクトルデータを追加する関数。
    Args:
        folder_path (str): ドキュメントが格納されているフォルダのパス。
        table_name (str): PostgreSQLのテーブル名。
    """
    # ドキュメントの読み込み
    reader = SimpleDirectoryReader(input_dir=folder_path)
    documents = reader.load_data()

    # PGVectorStore の準備
    vector_store = PGVectorStore.from_params(
        database=PG_DATABASE,
        host=PG_HOST,
        password=PG_PASSWORD,
        port=PG_PORT,
        user=PG_USER,
        table_name=table_name,
        embed_dim=EMBEDDING_DIM,
    )
    # テーブル存在確認と作成は PGVectorStore が内部で行います
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # インデックスの作成とデータの追加
    VectorStoreIndex.from_documents(
        documents, storage_context=storage_context, show_progress=True
    )
    # index = VectorStoreIndex.from_documents(
    #     documents, storage_context=storage_context, show_progress=True
    # )
    # 追加したものをそのまま利用することもできる
    # return index


# ------ 3. インデックス取得
def get_vector_index(table_name: str) -> VectorStoreIndex:
    """
    指定されたPostgreSQLテーブル名からLlamaIndexのインデックスを取得する関数。

    Args:
        table_name (str): 取得元のPostgreSQLテーブル名。

    Returns:
        VectorStoreIndex
    """
    # PGVectorStore の初期化 
    vector_store = PGVectorStore.from_params(
        database=PG_DATABASE,
        host=PG_HOST,
        password=PG_PASSWORD,
        port=PG_PORT,
        user=PG_USER,
        table_name=table_name,
        embed_dim=EMBEDDING_DIM,
    )

    # VectorStoreからインデックスをロード
    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
    return index

# ------ 4. インデックスから、RAGを利用してプロンプトに対する応答を生成する
def query_with_rag(index: VectorStoreIndex, prompt: str) -> str | None:
    """
    取得したインデックスを使用してRAG形式でプロンプトをリクエストする関数。

    Args:
        index (VectorStoreIndex): get_index_from_postgresで取得したインデックスオブジェクト。
        prompt (str): LLMに送信するプロンプト (質問)。

    Returns:
        str
    """
    query_engine = index.as_query_engine(similarity_top_k=3)
    response = query_engine.query(prompt)
    # return str(response)
    return response
    


# ---- 使用例 
if __name__ == "__main__":

    # ベクトルデータに追加するドキュメントファイルのフォルダ
    data_folder = "docs"
    # ベクトルデータのテーブル名
    table_name = "vector_index_test"

    # INDEXの追加
    add_vector_data(data_folder, table_name)
    # 追加時にインデックスを取得することも可能
    # index_instance = add_vector_data(data_folder, table_name)

    # インデックスの取得
    index_instance = get_vector_index(table_name)
    # 質問内容
    question = "出てくる鬼の色を教えてください。" 
    # RAG形式でプロンプトをリクエスト
    answer = query_with_rag(index_instance, question)
    # 内容確認
    print(answer)

実行

$ python main.py 
Parsing nodes: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  7.13it/s]
Generating embeddings: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [00:02<00:00,  8.38it/s]
出てくる鬼の色は赤、青、金、そして黄色です。

構成

  1. 必要な初期設定
  2. テキストファイルをベクトルデータとして登録(関数)
  3. インデックス取得(関数)
  4. インデックスから、RAGを利用してプロンプトに対する応答を生成する(関数)
    となっています

ざっくり注意書き

  • AIで作ってもらったコードを最低限の内容に削って展開しています
  • 基本的に、設定内容は環境変数や定数にしたほうが良いです
  • APIキーはOpenAI APIのものを設定してください
  • 次元数は512～1536 の間の数値を指定できます（text-embedding-3-smallの場合）
  • ベクトルデータ登録時に、戻り値でインデックスのデータを受け取れます（この場合、取得関数は不要です）
  • 事前のチェック処理（ファイルやテーブル）や例外処理がありませんので、お好みで追加してください
  • similarity_top_k は 「回答に使う情報の数」のようなイメージで、 3 から10 程度が一般的のようです(AI調べ)

最後に

何回やっても、黒が出てきません。（テキストファイル内の存在は確認済み）
なんらかAIの配慮なのか…。緑や黄色は出たりでなかったりです。
パラメータのチューニングをしたり、モデルの性能を上げると出てくるかと思います。（未確認）

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?