diff --git a/docs-go/Makefile b/docs-go/Makefile index 2497cc9daf..1e3caecf23 100644 --- a/docs-go/Makefile +++ b/docs-go/Makefile @@ -1,6 +1,6 @@ WEAVE=$(HOME)/go/bin/weave -all: $(WEAVE) get-started-go.md flows.md models.md prompts.md dotprompt.md +all: $(WEAVE) get-started-go.md flows.md models.md prompts.md dotprompt.md pgvector.md $(WEAVE): ../go/internal/cmd/weave/*.go go -C ../go install ./internal/cmd/weave diff --git a/docs-go/pgvector.md b/docs-go/pgvector.md new file mode 100644 index 0000000000..14cfd1cc21 --- /dev/null +++ b/docs-go/pgvector.md @@ -0,0 +1,84 @@ + + +# pgvector retriever template + +You can use PostgreSQL and `pgvector` as your retriever implementation. Use the +following examples as a starting point and modify it to work with your database +schema. + +We use [database/sql](https://pkg.go.dev/database/sql) to connect to the Postgres server, but you may use another client library of your choice. + +- {Go} + + ```go + func defineRetriever(db *sql.DB, embedder *ai.Embedder) *ai.Retriever { + f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { + eres, err := embedder.Embed(ctx, &ai.EmbedRequest{Documents: []*ai.Document{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) + } + ``` + +And here's how to use the retriever in a flow: + +- {Go} + + ```go + retriever := defineRetriever(db, embedder) + + type input struct { + Question string + Show string + } + + genkit.DefineFlow("askQuestion", func(ctx context.Context, in input) (string, error) { + res, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{ + Document: &ai.Document{Content: []*ai.Part{ai.NewTextPart(in.Question)}}, + Options: in.Show, + }) + 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 + }) + ``` diff --git a/docs-go/pgvector.src b/docs-go/pgvector.src new file mode 100644 index 0000000000..b9126ce946 --- /dev/null +++ b/docs-go/pgvector.src @@ -0,0 +1,17 @@ +# pgvector retriever template + +You can use PostgreSQL and `pgvector` as your retriever implementation. Use the +following examples as a starting point and modify it to work with your database +schema. + +We use [database/sql](https://pkg.go.dev/database/sql) to connect to the Postgres server, but you may use another client library of your choice. + +- {Go} + + %include ../go/samples/pgvector/main.go retr + +And here's how to use the retriever in a flow: + +- {Go} + + %include ../go/samples/pgvector/main.go use-retr diff --git a/go/samples/pgvector/main.go b/go/samples/pgvector/main.go index de3f6bc692..023b70cc91 100644 --- a/go/samples/pgvector/main.go +++ b/go/samples/pgvector/main.go @@ -84,6 +84,7 @@ func run() error { } } + // !+use-retr retriever := defineRetriever(db, embedder) type input struct { @@ -105,12 +106,14 @@ func run() error { // Use documents in RAG prompts. return "", nil }) + // !-use-retr return genkit.Init(ctx, nil) } const provider = "pgvector" +// !+retr func defineRetriever(db *sql.DB, embedder *ai.Embedder) *ai.Retriever { f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { eres, err := embedder.Embed(ctx, &ai.EmbedRequest{Documents: []*ai.Document{req.Document}}) @@ -154,6 +157,8 @@ func defineRetriever(db *sql.DB, embedder *ai.Embedder) *ai.Retriever { return ai.DefineRetriever(provider, "shows", f) } +// !-retr + func defineIndexer(db *sql.DB, embedder *ai.Embedder) *ai.Indexer { // The indexer assumes that each Document has a single part, to be embedded, and metadata fields // for the table primary key: show_id, season_number, episode_id.