You can use PostgreSQL and pgvector as your retriever implementation. Use the
following examples as a starting point and modify it to work with your database
schema.
We use database/sql to connect to the Postgres server, but you may use another client library of your choice.
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)
}
And here's how to use the retriever in a flow:
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
})