From cd5b1b64cfc14b2c5c9ebfd8f1a9bfa846592fec Mon Sep 17 00:00:00 2001 From: Samuel Bushi Date: Thu, 4 Jul 2024 19:25:54 +0000 Subject: [PATCH 1/3] Add pgvector for go --- docs/templates/pgvector.md | 257 +++++++++++++++++++++++++++---------- 1 file changed, 192 insertions(+), 65 deletions(-) diff --git a/docs/templates/pgvector.md b/docs/templates/pgvector.md index be8cf3b0cb..91d1c64c0f 100644 --- a/docs/templates/pgvector.md +++ b/docs/templates/pgvector.md @@ -4,71 +4,198 @@ You can use PostgreSQL and `pgvector` as your retriever implementation. Use the following example as a starting point and modify it to work with your database schema. -```js -import { embed } from '@genkit-ai/ai/embedder'; -import { Document, defineRetriever, retrieve } from '@genkit-ai/ai/retriever'; -import { defineFlow } from '@genkit-ai/flow'; -import { textEmbeddingGecko } from '@genkit-ai/vertexai'; -import { toSql } from 'pgvector'; -import postgres from 'postgres'; -import { z } from 'zod'; - -const sql = postgres({ ssl: false, database: 'recaps' }); - -const QueryOptions = z.object({ - show: z.string(), - k: z.number().optional(), -}); - -const sqlRetriever = defineRetriever( - { - name: 'pgvector-myTable', - configSchema: QueryOptions, - }, - async (input, options) => { - const embedding = await embed({ - embedder: textEmbeddingGecko, - content: input, - }); - const results = await sql` - SELECT episode_id, season_number, chunk as content - FROM embeddings - WHERE show_id = ${options.show} - ORDER BY embedding <#> ${toSql(embedding)} LIMIT ${options.k ?? 3} - `; - return { - documents: results.map((row) => { - const { content, ...metadata } = row; - return Document.fromText(content, metadata); - }), - }; - } -); -``` +
+ +
+

Node.js (Typescript)

+

```js + import { embed } from '@genkit-ai/ai/embedder'; + import { Document, defineRetriever, retrieve } from '@genkit-ai/ai/retriever'; + import { defineFlow } from '@genkit-ai/flow'; + import { textEmbeddingGecko } from '@genkit-ai/vertexai'; + import { toSql } from 'pgvector'; + import postgres from 'postgres'; + import { z } from 'zod'; + + const sql = postgres({ ssl: false, database: 'recaps' }); + + const QueryOptions = z.object({ + show: z.string(), + k: z.number().optional(), + }); + + const sqlRetriever = defineRetriever( + { + name: 'pgvector-myTable', + configSchema: QueryOptions, + }, + async (input, options) => { + const embedding = await embed({ + embedder: textEmbeddingGecko, + content: input, + }); + const results = await sql + SELECT episode_id, season_number, chunk as content + FROM embeddings + WHERE show_id = ${options.show} + ORDER BY embedding <#> ${toSql(embedding)} LIMIT ${options.k ?? 3} + ; + return { + documents: results.map((row) => { + const { content, ...metadata } = row; + return Document.fromText(content, metadata); + }), + }; + } + ); + ``` +

+
+
+

Go

+

+ ```go + package main + + import ( + "context" + "fmt" + "os" + + ) + + import ( + "context" + "fmt" + "log" + "os" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/plugins/googleai" + "github.com/jackc/pgx/v5" + "github.com/pgvector/pgvector-go" + ) + + func retrieve(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { + // urlExample := "postgres://username:password@localhost:5432/database_name" + conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL")) + if err != nil { + return nil, fmt.Errorf( "Unable to connect to database: %v\n", err) + } + defer conn.Close(context.Background()) + + // Use the embedder to convert the document we want to + // retrieve into a vector. + ereq := &ai.EmbedRequest{ + Document: req.Document, + } + embedder := googleai.Embedder("embedding-001") + vals, err := embedder.Embed(ctx, ereq) + if err != nil { + return nil, fmt.Errorf("embedding failed: %v", err) + } + + k := 3 + if options, _ := req.Options.(*RetrieverOptions); options != nil { + k = options.K + } + + rows, err := conn.Query(context.Background(), "SELECT chunk as content FROM embeddings WHERE show_id = $1 ORDER BY embedding <#> $2 LIMIT $3", "game_of_thrones", pgvector.NewVector(vals), k) + if err != nil { + return nil, fmt.Errorf("Query failed: %v\n", err) + } + + docs := make([]*ai.Document, 0, 3) + + for rows.Next() { + var content string + rows.Scan(&content) + docs = append(docs, ai.DocumentFromText(content, nil)) + } + + resp := &ai.RetrieverResponse{ + Documents: docs, + } + return resp, nil + } + + + func main() { + // ... Configure genkit + + provider := "pgvector" + name := "tvShows" + retriever := ai.DefineRetriever(provider, name, retrieve) + } + ``` +

+
+ +
And here's how to use the retriever in a flow: -```js -// Simple flow to use the sqlRetriever -export const askQuestionsOnGoT = defineFlow( - { - name: 'askQuestionsOnGoT', - inputSchema: z.string(), - outputSchema: z.string(), - }, - async (inputQuestion) => { - const docs = await retrieve({ - retriever: sqlRetriever, - query: inputQuestion, - options: { - show: 'Game of Thrones', - }, - }); - console.log(docs); - - // Continue with using retrieved docs - // in RAG prompts. - //... - } -); -``` +
+ +
+

Node.js (Typescript)

+

+ ```js + // Simple flow to use the sqlRetriever + export const askQuestionsOnGoT = defineFlow( + { + name: 'askQuestionsOnGoT', + inputSchema: z.string(), + outputSchema: z.string(), + }, + async (inputQuestion) => { + const docs = await retrieve({ + retriever: sqlRetriever, + query: inputQuestion, + options: { + show: 'Game of Thrones', + }, + }); + console.log(docs); + + // Continue with using retrieved docs + // in RAG prompts. + //... + } + ); + ``` +

+
+
+

Go

+

+ ```go + func main() { + // ... + // Setup genkit and get the pgvector retriever as shown above + + // Use your retriever. + tvShowSuggestionFlow := genkit.DefineFlow( + "tvShowSuggestionFlow", + func(ctx context.Context, inputQuery string) (string, error) { + dRequest := ai.DocumentFromText(inputQuery, nil) + retrieverReq := &ai.RetrieverRequest{ + Document: dRequest, + } + response, err := retriever.Retrieve(ctx, retrieverReq) + if err != nil { + return "", err + } + + // Use retrieved docs in `response` for RAG + // ... + + return suggestion, nil + }, + ) + } + ``` +

+
+ +
From 3dc9a7a0f247e54065ba417fdfc137acd2c1b6c8 Mon Sep 17 00:00:00 2001 From: Samuel Bushi Date: Thu, 4 Jul 2024 19:29:58 +0000 Subject: [PATCH 2/3] [Go] Add snippet for pgvector --- docs/templates/pgvector.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/docs/templates/pgvector.md b/docs/templates/pgvector.md index 91d1c64c0f..280364c746 100644 --- a/docs/templates/pgvector.md +++ b/docs/templates/pgvector.md @@ -1,9 +1,12 @@ # pgvector retriever template You can use PostgreSQL and `pgvector` as your retriever implementation. Use the -following example as a starting point and modify it to work with your database +following examples as a starting point and modify it to work with your database schema. +For the Golang snippet, we use [pgx](https://github.com/jackc/pgx) as the +Postgres client, but you may use another client libaray of your choice. +
@@ -174,7 +177,7 @@ And here's how to use the retriever in a flow: // ... // Setup genkit and get the pgvector retriever as shown above - // Use your retriever. + // Use your retriever. tvShowSuggestionFlow := genkit.DefineFlow( "tvShowSuggestionFlow", func(ctx context.Context, inputQuery string) (string, error) { @@ -189,10 +192,10 @@ And here's how to use the retriever in a flow: // Use retrieved docs in `response` for RAG // ... - + return suggestion, nil }, - ) + ) } ```

From aaee97126d80a1026f368dd2122b8291d5fe2544 Mon Sep 17 00:00:00 2001 From: Samuel Bushi Date: Thu, 4 Jul 2024 19:33:35 +0000 Subject: [PATCH 3/3] format --- docs/templates/pgvector.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/templates/pgvector.md b/docs/templates/pgvector.md index 280364c746..44b68f3002 100644 --- a/docs/templates/pgvector.md +++ b/docs/templates/pgvector.md @@ -4,7 +4,7 @@ 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. -For the Golang snippet, we use [pgx](https://github.com/jackc/pgx) as the +For the Golang snippet, we use [pgx](https://github.com/jackc/pgx) as the Postgres client, but you may use another client libaray of your choice.
@@ -133,6 +133,7 @@ Postgres client, but you may use another client libaray of your choice. ```

+
@@ -200,5 +201,6 @@ And here's how to use the retriever in a flow: ```

+