pgvector リトリーバー テンプレート

リトリーバー実装として PostgreSQL と pgvector を使用できます。次の例を出発点として使用し、お使いのデータベース スキーマで機能するように変更してください。

Postgres サーバーに接続するために database/sql を使用しますが、別の任意のクライアント ライブラリを使用することもできます。

func defineRetriever(db *sql.DB, embedder ai.Embedder) ai.Retriever {
	f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) {
		eres, err := ai.Embed(ctx, embedder, ai.WithEmbedDocs(req.Document))
		if err != nil {
			return nil, err
		}
		rows, err := db.QueryContext(ctx, `
			SELECT episode_id, season_number, chunk as content
			FROM embeddings
			WHERE show_id = $1
		  	ORDER BY embedding <#> $2
		  	LIMIT 2`,
			req.Options, pgv.NewVector(eres.Embeddings[0].Embedding))
		if err != nil {
			return nil, err
		}
		defer rows.Close()

		res := &ai.RetrieverResponse{}
		for rows.Next() {
			var eid, sn int
			var content string
			if err := rows.Scan(&eid, &sn, &content); err != nil {
				return nil, err
			}
			meta := map[string]any{
				"episode_id":    eid,
				"season_number": sn,
			}
			doc := &ai.Document{
				Content:  []*ai.Part{ai.NewTextPart(content)},
				Metadata: meta,
			}
			res.Documents = append(res.Documents, doc)
		}
		if err := rows.Err(); err != nil {
			return nil, err
		}
		return res, nil
	}
	return ai.DefineRetriever(provider, "shows", f)
}

フロー内でこのリトリーバーを使用する方法は次のとおりです。

retriever := defineRetriever(db, embedder)

type input struct {
	Question string
	Show     string
}

genkit.DefineFlow("askQuestion", func(ctx context.Context, in input) (string, error) {
	res, err := ai.Retrieve(ctx, retriever,
		ai.WithRetrieverOpts(in.Show),
		ai.WithRetrieverText(in.Question))
	if err != nil {
		return "", err
	}
	for _, doc := range res.Documents {
		fmt.Printf("%+v %q\n", doc.Metadata, doc.Content[0].Text)
	}
	// Use documents in RAG prompts.
	return "", nil
})