قالب pgvector retriever

می توانید از PostgreSQL و pgvector به عنوان پیاده سازی بازیابی خود استفاده کنید. از مثال های زیر به عنوان نقطه شروع استفاده کنید و آن را برای کار با طرح پایگاه داده خود تغییر دهید.

ما از پایگاه داده/sql برای اتصال به سرور Postgres استفاده می کنیم، اما شما می توانید از کتابخانه مشتری دیگری که انتخاب می کنید استفاده کنید.

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