Anda dapat menggunakan PostgreSQL dan pgvector sebagai penerapan retriever. Gunakan
contoh berikut sebagai titik awal dan memodifikasinya agar berfungsi dengan database Anda
skema.
Kami menggunakan database/sql untuk terhubung ke server Postgres, tetapi Anda dapat menggunakan library klien lain pilihan Anda.
Go
func defineRetriever(db *sql.DB, embedder *ai.Embedder) *ai.Retriever {
f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) {
eres, err := embedder.Embed(ctx, &ai.EmbedRequest{Documents: []*ai.Document{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)
}
Berikut cara menggunakan retriever dalam flow:
Go
retriever := defineRetriever(db, embedder)
type input struct {
Question string
Show string
}
genkit.DefineFlow("askQuestion", func(ctx context.Context, in input) (string, error) {
res, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{
Document: &ai.Document{Content: []*ai.Part{ai.NewTextPart(in.Question)}},
Options: in.Show,
})
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
})