mẫu pgvector retriever

Bạn có thể sử dụng PostgreSQL và pgvector làm cách triển khai trình truy xuất. Sử dụng các ví dụ sau làm điểm xuất phát và sửa đổi để hoạt động với cơ sở dữ liệu của bạn giản đồ.

Chúng tôi sử dụng database/sql để kết nối với máy chủ Postgres nhưng bạn có thể sử dụng thư viện ứng dụng khác tùy chọn.

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)
}

Và dưới đây là cách sử dụng chú chó trong một dòng dữ liệu:

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
})