🐘

SemanticKernelとPostgreSQL

に公開

PostgreSQLのVectorStore

AIのRAGを実現するためにVectorStoreを使います。
自分の投稿ではVectorStoreには InMemoryQdrant を使用していましたが、個人的にローカルにもRDBが欲しいなぁ、ということもあり。
すぐに使えるベクター ストア コネクタ (プレビュー) こちらを確認すると、 Postgres Vector Store がありましたので、使い方を見てみました。

サーバーをDockerにして扱えるようにしてみる

接続先をDockerで作ったサーバーにしてSemantic Kernelから使用できるように確認してみました。

docker compose

pgvectorを組み込んだイメージを使います。
初回起動時に初期化を行わせるために、 volumes/pgsql/init にファイルを置いて、読み込ませるようにします。

docker-compose.yaml
services:
  pgsql:
    image: pgvector/pgvector:pg17-trixie
    restart: always
    environment:
      POSTGRES_USER: ${PGUSER:-postgres}
      POSTGRES_PASSWORD: ${PGPASSWORD:-password}
      POSTGRES_DB: ${PGDATABASE:-inquery}
    ports:
      - 5432:5432
    volumes:
      - ./volumes/pgsql/data:/var/lib/postgresql/data
      - ./volumes/pgsql/init:/docker-entrypoint-initdb.d

docker初期設定ファイル

初回起動時に初期化を行うファイル群です。
testというアカウントを作って、そのアカウントでvectorを使うように設定します。
vectorを使えるようにするためにはSUPERUSERになっておく必要があるのでそのようになっています。

volumes/pgsql/init/init01.sql
CREATE ROLE test WITH PASSWORD 'testpass';
ALTER ROLE test LOGIN;
ALTER ROLE test SUPERUSER;
CREATE DATABASE testdb WITH OWNER test;

次は、作成したユーザーでvectorを使えるようにします。

volumes/pgsql/init/init02.sql
\connect testdb

SET ROLE test;
CREATE EXTENSION IF NOT EXISTS vector;

最後に、データベースに対して使用するユーザーでvectorを使えるように設定します。

volumes/pgsql/init/init03.sql
ALTER ROLE test NOSUPERUSER;

ユーザーのSUPERUSERを外します。

使えるようにしたサーバーをSemantic Kernelで使う

一通り動くようにしたC#コードは以下。

Program.cs
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.VectorData;
using Microsoft.SemanticKernel;
using Npgsql;
using System.ClientModel;
using System.Threading;
using static System.Net.Mime.MediaTypeNames;

var builder = Kernel.CreateBuilder();
#pragma warning disable SKEXP0010 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。
builder.Services
    .AddSingleton(sp =>
    {
        var b = new NpgsqlDataSourceBuilder("Host=localhost;Port=5432;Database=testdb;Username=test;Password=testpass");
        b.UseVector();
        return b.Build();
    })
    .AddPostgresVectorStore()
    .AddPostgresCollection<int, QAItemModel>("my_collection")
    .AddOpenAIEmbeddingGenerator(
        "cl-nagoya/ruri-v3-310m",
        new OpenAI.OpenAIClient(new ApiKeyCredential("-"), new OpenAI.OpenAIClientOptions()
        {
            Endpoint = new Uri("http://localhost:7997"),
        }))
    ;
#pragma warning restore SKEXP0010 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。
var kernel = builder.Build();

var collection = kernel.Services.GetRequiredService<VectorStoreCollection<int, QAItemModel>>();
await collection.EnsureCollectionExistsAsync();

var embedding = kernel.Services.GetRequiredService<IEmbeddingGenerator<string, Embedding<float>>>();
var items = new List<QAItemModel>
{
    new QAItemModel
    {
        Id = 1,
        Question = "大きな森の小道を小さな猫が2匹並んで歩いている",
    },
    new QAItemModel
    {
        Id = 2,
        Question = "広い湖のほとりには黄色い小さな花が沢山咲いている",
    }
};
foreach (var item in items)
{
    item.QuestionEmbedding = (await embedding.GenerateAsync(item.Question)).Vector;
}
await collection.UpsertAsync(items);

var searchText = "猫に関する文章は";
var searchEmb = await embedding.GenerateAsync($"検索クエリ: {searchText}");
var results = collection.SearchAsync(searchEmb, 1);
await foreach(var result in results)
{
    Console.WriteLine($"Question: {result.Record.Question}, Score: {result.Score}");
}

return;

public class QAItemModel
{
    [VectorStoreKey]
    public int Id { get; set; }

    [VectorStoreData(IsIndexed = true)]
    public string Question { get; set; } = string.Empty;

    [VectorStoreVector(Dimensions: 768)]
    public ReadOnlyMemory<float> QuestionEmbedding { get; set; }
}

全部載せました。

  • Kernelを設定して
  • Collection作って
  • データ設定してDBに入れて
  • 検索させてみる

ということになっています。
EntityFrameworkでいう、コードファーストですね。
コードからDBのテーブルが作られます。

まとめ

PostgreSQLのVector拡張したものをSemantic Kernelから扱ってみました。
あとはEntityFrameworkと絡めるのもよいですね。

Discussion