pgvector 검색기 템플릿

PostgreSQL 및 pgvector를 검색기 구현으로 사용할 수 있습니다. 다음 예시를 시작점으로 사용하고 데이터베이스 스키마와 작동하도록 수정하세요.

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