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