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

# はじめに
タイトルの通りです。
[この環境](https://qiita.com/naozo-se/items/0730c8ea650eaa0d51c8)の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のテキストファイルとして利用しています
    <details><summary>実行プロンプト</summary>

    ```
    以下の条件で物語を作って欲しいです。
    
    文字数は5000文字前後にしてください。
    犬を出してください。名前はカタカナにしてくださ
    赤、青、金の鬼を出してください。
    
    以上を満たせば何でもよいです。
    ```
    
    ```
    この話の第２部を作成してください。

    メインのキャラクターはそのまま利用してください。
    鬼の色は緑、黄色、黒を出してください。
    
    文字数は同じく5000文字前後で大丈夫です。
    ```
</details>

# パッケージインストール
```shell
pip install llama-index \
llama-index-vector-stores-postgres  \
llama-index-embeddings-openai  \
llama-index-llms-openai  \
python-dotenv  \
openai  \
psycopg2-binary
```

# コード
```python: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)
```


# 実行
```shell
$ 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. 必要な初期設定
1. テキストファイルをベクトルデータとして登録(関数)
1. インデックス取得(関数)
1. インデックスから、RAGを利用してプロンプトに対する応答を生成する(関数)
となっています

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


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