modello retriever pgvector

Puoi utilizzare PostgreSQL e pgvector come implementazione del retriever. Utilizza i seguenti esempi come punto di partenza e modificali in base allo schema del tuo database.

Utilizziamo database/sql per connetterci al server Postgres, ma puoi utilizzare un'altra libreria client a tua scelta.

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

Ecco come utilizzare il retriever in un flusso:

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