diff --git a/docs-go/Makefile b/docs-go/Makefile index f8fb71f937..78b8282441 100644 --- a/docs-go/Makefile +++ b/docs-go/Makefile @@ -1,6 +1,7 @@ WEAVE=$(HOME)/go/bin/weave -all: $(WEAVE) get-started-go.md flows.md models.md prompts.md dotprompt.md pgvector.md cloud-run.md +all: $(WEAVE) get-started-go.md flows.md models.md prompts.md dotprompt.md pgvector.md \ + cloud-run.md rag.md $(WEAVE): ../go/internal/cmd/weave/*.go go -C ../go install ./internal/cmd/weave diff --git a/docs-go/rag.md b/docs-go/rag.md new file mode 100644 index 0000000000..beb0aa65c7 --- /dev/null +++ b/docs-go/rag.md @@ -0,0 +1,391 @@ + + +# Retrieval-augmented generation (RAG) + +Firebase Genkit provides abstractions that help you build retrieval-augmented generation +(RAG) flows, as well as plugins that provide integrations with related tools. + +## What is RAG? + +Retrieval-augmented generation is a technique used to incorporate external +sources of information into an LLM’s responses. It's important to be able to do +so because, while LLMs are typically trained on a broad body of +material, practical use of LLMs often requires specific domain knowledge (for +example, you might want to use an LLM to answer customers' questions about your +company’s products). + +One solution is to fine-tune the model using more specific data. However, this +can be expensive both in terms of compute cost and in terms of the effort needed +to prepare adequate training data. + +In contrast, RAG works by incorporating external data sources into a prompt at +the time it's passed to the model. For example, you could imagine the prompt, +"What is Bart's relationship to Lisa?" might be expanded ("augmented") by +prepending some relevant information, resulting in the prompt, "Homer and +Marge's children are named Bart, Lisa, and Maggie. What is Bart's relationship +to Lisa?" + +This approach has several advantages: + +- It can be more cost effective because you don't have to retrain the model. +- You can continuously update your data source and the LLM can immediately make + use of the updated information. +- You now have the potential to cite references in your LLM's responses. + +On the other hand, using RAG naturally means longer prompts, and some LLM API +services charge for each input token you send. Ultimately, you must evaluate the +cost tradeoffs for your applications. + +RAG is a very broad area and there are many different techniques used to achieve +the best quality RAG. The core Genkit framework offers two main abstractions to +help you do RAG: + +- Indexers: add documents to an "index". +- Embedders: transforms documents into a vector representation +- Retrievers: retrieve documents from an "index", given a query. + +These definitions are broad on purpose because Genkit is un-opinionated about +what an "index" is or how exactly documents are retrieved from it. Genkit only +provides a `Document` format and everything else is defined by the retriever or +indexer implementation provider. + +### Indexers + +The index is responsible for keeping track of your documents in such a way that +you can quickly retrieve relevant documents given a specific query. This is most +often accomplished using a vector database, which indexes your documents using +multidimensional vectors called embeddings. A text embedding (opaquely) +represents the concepts expressed by a passage of text; these are generated +using special-purpose ML models. By indexing text using its embedding, a vector +database is able to cluster conceptually related text and retrieve documents +related to a novel string of text (the query). + +Before you can retrieve documents for the purpose of generation, you need to +ingest them into your document index. A typical ingestion flow does the +following: + +1. Split up large documents into smaller documents so that only relevant + portions are used to augment your prompts – "chunking". This is necessary + because many LLMs have a limited context window, making it impractical to + include entire documents with a prompt. + + Genkit doesn't provide built-in chunking libraries; however, there are open + source libraries available that are compatible with Genkit. + +1. Generate embeddings for each chunk. Depending on the database you're using, + you might explicitly do this with an embedding generation model, or you + might use the embedding generator provided by the database. + +1. Add the text chunk and its index to the database. + +You might run your ingestion flow infrequently or only once if you are working +with a stable source of data. On the other hand, if you are working with data +that frequently changes, you might continuously run the ingestion flow (for +example, in a Cloud Firestore trigger, whenever a document is updated). + +### Embedders + +An embedder is a function that takes content (text, images, audio, etc.) and creates a numeric vector that encodes the semantic meaning of the original content. As mentioned above, embedders are leveraged as part of the process of indexing, however, they can also be used independently to create embeddings without an index. + +### Retrievers + +A retriever is a concept that encapsulates logic related to any kind of document +retrieval. The most popular retrieval cases typically include retrieval from +vector stores, however, in Genkit a retriever can be any function that returns data. + +To create a retriever, you can use one of the provided implementations or +create your own. + +## Supported indexers, retrievers, and embedders + +Genkit provides indexer and retriever support through its plugin system. The +following plugins are officially supported: + +- [Pinecone](plugins/pinecone.md) cloud vector database + +In addition, Genkit supports the following vector stores through predefined +code templates, which you can customize for your database configuration and +schema: + +- PostgreSQL with [`pgvector`](templates/pgvector.md) + +Embedding model support is provided through the following plugins: + +| Plugin | Models | +| ------------------------- | -------------------- | +| [Google Generative AI][1] | Gecko text embedding | +| [Google Vertex AI][2] | Gecko text embedding | + +[1]: plugins/google-genai.md +[2]: plugins/vertex-ai.md + +## Defining a RAG Flow + +The following examples show how you could ingest a collection of restaurant menu PDF documents +into a vector database and retrieve them for use in a flow that determines what food items are available. + +### Install dependencies + +In this example, we will use the `textsplitter` library from `langchaingo` and +the `ledongthuc/pdf` PDF parsing Library: + +```posix-terminal +go get github.com/tmc/langchaingo/textsplitter + +go get github.com/ledongthuc/pdf +``` + +### Define an Indexer + +The following example shows how to create an indexer to ingest a collection of PDF documents +and store them in a local vector database. + +It uses the local file-based vector similarity retriever +that Genkit provides out-of-the box for simple testing and prototyping (_do not +use in production_) + +#### Create the indexer + +```go +// Import Genkit's file-based vector retriever, (Don't use in production.) +import "github.com/firebase/genkit/go/plugins/localvec" + +// Vertex AI provides the text-embedding-004 embedder model. +import "github.com/firebase/genkit/go/plugins/vertexai" +``` + +```go +ctx := context.Background() + +err := vertexai.Init(ctx, &vertexai.Config{}) +if err != nil { + log.Fatal(err) +} +err = localvec.Init() +if err != nil { + log.Fatal(err) +} + +menuPDFIndexer, _, err := localvec.DefineIndexerAndRetriever( + "menuQA", + localvec.Config{ + Embedder: vertexai.Embedder("text-embedding-004"), + }, +) +if err != nil { + log.Fatal(err) +} +``` + +#### Create chunking config + +This example uses the `textsplitter` library which provides a simple text +splitter to break up documents into segments that can be vectorized. + +The following definition configures the chunking function to return document +segments of 200 characters, with an overlap between chunks of 20 characters. + +```go +splitter := textsplitter.NewRecursiveCharacter( + textsplitter.WithChunkSize(200), + textsplitter.WithChunkOverlap(20), +) +``` + +More chunking options for this library can be found in the +[`langchaingo` documentation](https://pkg.go.dev/github.com/tmc/langchaingo/textsplitter#Option). + +#### Define your indexer flow + +```go +genkit.DefineFlow( + "indexMenu", + func(ctx context.Context, path string) (any, error) { + // Extract plain text from the PDF. Wrap the logic in Run so it + // appears as a step in your traces. + pdfText, err := genkit.Run(ctx, "extract", func() (string, error) { + return readPDF(path) + }) + if err != nil { + return nil, err + } + + // Split the text into chunks. Wrap the logic in Run so it + // appears as a step in your traces. + docs, err := genkit.Run(ctx, "chunk", func() ([]*ai.Document, error) { + chunks, err := splitter.SplitText(pdfText) + if err != nil { + return nil, err + } + + var docs []*ai.Document + for _, chunk := range chunks { + docs = append(docs, ai.DocumentFromText(chunk, nil)) + } + return docs, nil + }) + if err != nil { + return nil, err + } + + // Add chunks to the index. + err = menuPDFIndexer.Index(ctx, &ai.IndexerRequest{Documents: docs}) + return nil, err + }, +) +``` + +```go +// Helper function to extract plain text from a PDF. Excerpted from +// https://github.com/ledongthuc/pdf +func readPDF(path string) (string, error) { + f, r, err := pdf.Open(path) + if f != nil { + defer f.Close() + } + if err != nil { + return "", err + } + + reader, err := r.GetPlainText() + if err != nil { + return "", err + } + + bytes, err := io.ReadAll(reader) + if err != nil { + return "", err + } + return string(bytes), nil +} +``` + +#### Run the indexer flow + +```posix-terminal +genkit flow:run indexMenu "'menu.pdf'" +``` + +After running the `indexMenu` flow, the vector database will be seeded with +documents and ready to be used in Genkit flows with retrieval steps. + +### Define a flow with retrieval + +The following example shows how you might use a retriever in a RAG flow. Like +the indexer example, this example uses Genkit's file-based vector retriever, +which you should not use in production. + +```go + ctx := context.Background() + + err := vertexai.Init(ctx, &vertexai.Config{}) + if err != nil { + log.Fatal(err) + } + err = localvec.Init() + if err != nil { + log.Fatal(err) + } + + model := vertexai.Model("gemini-1.5-pro") + + _, menuPdfRetriever, err := localvec.DefineIndexerAndRetriever( + "menuQA", + localvec.Config{ + Embedder: vertexai.Embedder("text-embedding-004"), + }, + ) + if err != nil { + log.Fatal(err) + } + + genkit.DefineFlow( + "menuQA", + func(ctx context.Context, question string) (string, error) { + // Retrieve text relevant to the user's question. + docs, err := menuPdfRetriever.Retrieve(ctx, &ai.RetrieverRequest{ + Document: ai.DocumentFromText(question, nil), + }) + if err != nil { + return "", err + } + + // Construct a system message containing the menu excerpts you just + // retrieved. + menuInfo := ai.NewSystemTextMessage("Here's the menu context:") + for _, doc := range docs.Documents { + menuInfo.Content = append(menuInfo.Content, doc.Content...) + } + + // Call Generate, including the menu information in your prompt. + resp, err := model.Generate(ctx, &ai.GenerateRequest{ + Messages: []*ai.Message{ + ai.NewSystemTextMessage(` +You are acting as a helpful AI assistant that can answer questions about the +food available on the menu at Genkit Grub Pub. +Use only the context provided to answer the question. If you don't know, do not +make up an answer. Do not add or change items on the menu.`), + menuInfo, + ai.NewUserTextMessage(question), + }, + }, nil) + if err != nil { + return "", err + } + + return resp.Text() + }) +``` + +## Write your own indexers and retrievers + +It's also possible to create your own retriever. This is useful if your +documents are managed in a document store that is not supported in Genkit (eg: +MySQL, Google Drive, etc.). The Genkit SDK provides flexible methods that let +you provide custom code for fetching documents. + +You can also define custom retrievers that build on top of existing retrievers +in Genkit and apply advanced RAG techniques (such as reranking or prompt +extension) on top. + +For example, suppose you have a custom re-ranking function you want to use. The +following example defines a custom retriever that applies your function to the +menu retriever defined earlier: + +```go +type CustomMenuRetrieverOptions struct { + K int + PreRerankK int +} +advancedMenuRetriever := ai.DefineRetriever( + "custom", + "advancedMenuRetriever", + func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { + // Handle options passed using our custom type. + opts, _ := req.Options.(CustomMenuRetrieverOptions) + // Set fields to default values when either the field was undefined + // or when req.Options is not a CustomMenuRetrieverOptions. + if opts.K == 0 { + opts.K = 3 + } + if opts.PreRerankK == 0 { + opts.PreRerankK = 10 + } + + // Call the retriever as in the simple case. + response, err := menuPDFRetriever.Retrieve(ctx, &ai.RetrieverRequest{ + Document: req.Document, + Options: localvec.RetrieverOptions{K: opts.PreRerankK}, + }) + if err != nil { + return nil, err + } + + // Re-rank the returned documents using your custom function. + rerankedDocs := rerank(response.Documents) + response.Documents = rerankedDocs[:opts.K] + + return response, nil + }, +) +``` diff --git a/docs-go/rag.src b/docs-go/rag.src new file mode 100644 index 0000000000..3f5a963aa7 --- /dev/null +++ b/docs-go/rag.src @@ -0,0 +1,208 @@ +# Retrieval-augmented generation (RAG) + +Firebase Genkit provides abstractions that help you build retrieval-augmented generation +(RAG) flows, as well as plugins that provide integrations with related tools. + +## What is RAG? + +Retrieval-augmented generation is a technique used to incorporate external +sources of information into an LLM’s responses. It's important to be able to do +so because, while LLMs are typically trained on a broad body of +material, practical use of LLMs often requires specific domain knowledge (for +example, you might want to use an LLM to answer customers' questions about your +company’s products). + +One solution is to fine-tune the model using more specific data. However, this +can be expensive both in terms of compute cost and in terms of the effort needed +to prepare adequate training data. + +In contrast, RAG works by incorporating external data sources into a prompt at +the time it's passed to the model. For example, you could imagine the prompt, +"What is Bart's relationship to Lisa?" might be expanded ("augmented") by +prepending some relevant information, resulting in the prompt, "Homer and +Marge's children are named Bart, Lisa, and Maggie. What is Bart's relationship +to Lisa?" + +This approach has several advantages: + +- It can be more cost effective because you don't have to retrain the model. +- You can continuously update your data source and the LLM can immediately make + use of the updated information. +- You now have the potential to cite references in your LLM's responses. + +On the other hand, using RAG naturally means longer prompts, and some LLM API +services charge for each input token you send. Ultimately, you must evaluate the +cost tradeoffs for your applications. + +RAG is a very broad area and there are many different techniques used to achieve +the best quality RAG. The core Genkit framework offers two main abstractions to +help you do RAG: + +- Indexers: add documents to an "index". +- Embedders: transforms documents into a vector representation +- Retrievers: retrieve documents from an "index", given a query. + +These definitions are broad on purpose because Genkit is un-opinionated about +what an "index" is or how exactly documents are retrieved from it. Genkit only +provides a `Document` format and everything else is defined by the retriever or +indexer implementation provider. + +### Indexers + +The index is responsible for keeping track of your documents in such a way that +you can quickly retrieve relevant documents given a specific query. This is most +often accomplished using a vector database, which indexes your documents using +multidimensional vectors called embeddings. A text embedding (opaquely) +represents the concepts expressed by a passage of text; these are generated +using special-purpose ML models. By indexing text using its embedding, a vector +database is able to cluster conceptually related text and retrieve documents +related to a novel string of text (the query). + +Before you can retrieve documents for the purpose of generation, you need to +ingest them into your document index. A typical ingestion flow does the +following: + +1. Split up large documents into smaller documents so that only relevant + portions are used to augment your prompts – "chunking". This is necessary + because many LLMs have a limited context window, making it impractical to + include entire documents with a prompt. + + Genkit doesn't provide built-in chunking libraries; however, there are open + source libraries available that are compatible with Genkit. + +1. Generate embeddings for each chunk. Depending on the database you're using, + you might explicitly do this with an embedding generation model, or you + might use the embedding generator provided by the database. + +1. Add the text chunk and its index to the database. + +You might run your ingestion flow infrequently or only once if you are working +with a stable source of data. On the other hand, if you are working with data +that frequently changes, you might continuously run the ingestion flow (for +example, in a Cloud Firestore trigger, whenever a document is updated). + +### Embedders + +An embedder is a function that takes content (text, images, audio, etc.) and creates a numeric vector that encodes the semantic meaning of the original content. As mentioned above, embedders are leveraged as part of the process of indexing, however, they can also be used independently to create embeddings without an index. + +### Retrievers + +A retriever is a concept that encapsulates logic related to any kind of document +retrieval. The most popular retrieval cases typically include retrieval from +vector stores, however, in Genkit a retriever can be any function that returns data. + +To create a retriever, you can use one of the provided implementations or +create your own. + +## Supported indexers, retrievers, and embedders + +Genkit provides indexer and retriever support through its plugin system. The +following plugins are officially supported: + +- [Pinecone](plugins/pinecone.md) cloud vector database + +In addition, Genkit supports the following vector stores through predefined +code templates, which you can customize for your database configuration and +schema: + +- PostgreSQL with [`pgvector`](templates/pgvector.md) + +Embedding model support is provided through the following plugins: + +| Plugin | Models | +| ------------------------- | -------------------- | +| [Google Generative AI][1] | Gecko text embedding | +| [Google Vertex AI][2] | Gecko text embedding | + +[1]: plugins/google-genai.md +[2]: plugins/vertex-ai.md + +## Defining a RAG Flow + +The following examples show how you could ingest a collection of restaurant menu PDF documents +into a vector database and retrieve them for use in a flow that determines what food items are available. + +### Install dependencies + +In this example, we will use the `textsplitter` library from `langchaingo` and +the `ledongthuc/pdf` PDF parsing Library: + +```posix-terminal +go get github.com/tmc/langchaingo/textsplitter + +go get github.com/ledongthuc/pdf +``` + +### Define an Indexer + +The following example shows how to create an indexer to ingest a collection of PDF documents +and store them in a local vector database. + +It uses the local file-based vector similarity retriever +that Genkit provides out-of-the box for simple testing and prototyping (_do not +use in production_) + +#### Create the indexer + +```go +// Import Genkit's file-based vector retriever, (Don't use in production.) +import "github.com/firebase/genkit/go/plugins/localvec" + +// Vertex AI provides the text-embedding-004 embedder model. +import "github.com/firebase/genkit/go/plugins/vertexai" +``` + +%include ../go/internal/doc-snippets/rag/main.go vec + +#### Create chunking config + +This example uses the `textsplitter` library which provides a simple text +splitter to break up documents into segments that can be vectorized. + +The following definition configures the chunking function to return document +segments of 200 characters, with an overlap between chunks of 20 characters. + +%include ../go/internal/doc-snippets/rag/main.go splitcfg + +More chunking options for this library can be found in the +[`langchaingo` documentation](https://pkg.go.dev/github.com/tmc/langchaingo/textsplitter#Option). + +#### Define your indexer flow + +%include ../go/internal/doc-snippets/rag/main.go indexflow + +%include ../go/internal/doc-snippets/rag/main.go readpdf + +#### Run the indexer flow + +```posix-terminal +genkit flow:run indexMenu "'menu.pdf'" +``` + +After running the `indexMenu` flow, the vector database will be seeded with +documents and ready to be used in Genkit flows with retrieval steps. + +### Define a flow with retrieval + +The following example shows how you might use a retriever in a RAG flow. Like +the indexer example, this example uses Genkit's file-based vector retriever, +which you should not use in production. + +%include ../go/internal/doc-snippets/rag/main.go retrieve + +## Write your own indexers and retrievers + +It's also possible to create your own retriever. This is useful if your +documents are managed in a document store that is not supported in Genkit (eg: +MySQL, Google Drive, etc.). The Genkit SDK provides flexible methods that let +you provide custom code for fetching documents. + +You can also define custom retrievers that build on top of existing retrievers +in Genkit and apply advanced RAG techniques (such as reranking or prompt +extension) on top. + +For example, suppose you have a custom re-ranking function you want to use. The +following example defines a custom retriever that applies your function to the +menu retriever defined earlier: + +%include ../go/internal/doc-snippets/rag/main.go customret diff --git a/go/internal/doc-snippets/rag/main.go b/go/internal/doc-snippets/rag/main.go new file mode 100644 index 0000000000..7b4a155f77 --- /dev/null +++ b/go/internal/doc-snippets/rag/main.go @@ -0,0 +1,247 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package rag + +import ( + "context" + "io" + "log" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" + "github.com/firebase/genkit/go/plugins/localvec" + "github.com/firebase/genkit/go/plugins/vertexai" + + // "github.com/ledongthuc/pdf" + // "github.com/tmc/langchaingo/textsplitter" + "github.com/firebase/genkit/go/internal/doc-snippets/rag/pdf" + "github.com/firebase/genkit/go/internal/doc-snippets/rag/textsplitter" +) + +func main() { + //!+vec + ctx := context.Background() + + err := vertexai.Init(ctx, &vertexai.Config{}) + if err != nil { + log.Fatal(err) + } + err = localvec.Init() + if err != nil { + log.Fatal(err) + } + + menuPDFIndexer, _, err := localvec.DefineIndexerAndRetriever( + "menuQA", + localvec.Config{ + Embedder: vertexai.Embedder("text-embedding-004"), + }, + ) + if err != nil { + log.Fatal(err) + } + //!-vec + //!+splitcfg + splitter := textsplitter.NewRecursiveCharacter( + textsplitter.WithChunkSize(200), + textsplitter.WithChunkOverlap(20), + ) + //!-splitcfg + //!+indexflow + genkit.DefineFlow( + "indexMenu", + func(ctx context.Context, path string) (any, error) { + // Extract plain text from the PDF. Wrap the logic in Run so it + // appears as a step in your traces. + pdfText, err := genkit.Run(ctx, "extract", func() (string, error) { + return readPDF(path) + }) + if err != nil { + return nil, err + } + + // Split the text into chunks. Wrap the logic in Run so it + // appears as a step in your traces. + docs, err := genkit.Run(ctx, "chunk", func() ([]*ai.Document, error) { + chunks, err := splitter.SplitText(pdfText) + if err != nil { + return nil, err + } + + var docs []*ai.Document + for _, chunk := range chunks { + docs = append(docs, ai.DocumentFromText(chunk, nil)) + } + return docs, nil + }) + if err != nil { + return nil, err + } + + // Add chunks to the index. + err = menuPDFIndexer.Index(ctx, &ai.IndexerRequest{Documents: docs}) + return nil, err + }, + ) + //!-indexflow + + err = genkit.Init(ctx, nil) + if err != nil { + log.Fatal(err) + } +} + +// !+readpdf +// Helper function to extract plain text from a PDF. Excerpted from +// https://github.com/ledongthuc/pdf +func readPDF(path string) (string, error) { + f, r, err := pdf.Open(path) + if f != nil { + defer f.Close() + } + if err != nil { + return "", err + } + + reader, err := r.GetPlainText() + if err != nil { + return "", err + } + + bytes, err := io.ReadAll(reader) + if err != nil { + return "", err + } + return string(bytes), nil +} + +//!-readpdf + +func menuQA() { + //!+retrieve + ctx := context.Background() + + err := vertexai.Init(ctx, &vertexai.Config{}) + if err != nil { + log.Fatal(err) + } + err = localvec.Init() + if err != nil { + log.Fatal(err) + } + + model := vertexai.Model("gemini-1.5-pro") + + _, menuPdfRetriever, err := localvec.DefineIndexerAndRetriever( + "menuQA", + localvec.Config{ + Embedder: vertexai.Embedder("text-embedding-004"), + }, + ) + if err != nil { + log.Fatal(err) + } + + genkit.DefineFlow( + "menuQA", + func(ctx context.Context, question string) (string, error) { + // Retrieve text relevant to the user's question. + docs, err := menuPdfRetriever.Retrieve(ctx, &ai.RetrieverRequest{ + Document: ai.DocumentFromText(question, nil), + }) + if err != nil { + return "", err + } + + // Construct a system message containing the menu excerpts you just + // retrieved. + menuInfo := ai.NewSystemTextMessage("Here's the menu context:") + for _, doc := range docs.Documents { + menuInfo.Content = append(menuInfo.Content, doc.Content...) + } + + // Call Generate, including the menu information in your prompt. + resp, err := model.Generate(ctx, &ai.GenerateRequest{ + Messages: []*ai.Message{ + ai.NewSystemTextMessage(` +You are acting as a helpful AI assistant that can answer questions about the +food available on the menu at Genkit Grub Pub. +Use only the context provided to answer the question. If you don't know, do not +make up an answer. Do not add or change items on the menu.`), + menuInfo, + ai.NewUserTextMessage(question), + }, + }, nil) + if err != nil { + return "", err + } + + return resp.Text() + }) + //!-retrieve +} + +func customret() { + _, menuPDFRetriever, _ := localvec.DefineIndexerAndRetriever( + "menuQA", + localvec.Config{ + Embedder: vertexai.Embedder("text-embedding-004"), + }, + ) + + //!+customret + type CustomMenuRetrieverOptions struct { + K int + PreRerankK int + } + advancedMenuRetriever := ai.DefineRetriever( + "custom", + "advancedMenuRetriever", + func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { + // Handle options passed using our custom type. + opts, _ := req.Options.(CustomMenuRetrieverOptions) + // Set fields to default values when either the field was undefined + // or when req.Options is not a CustomMenuRetrieverOptions. + if opts.K == 0 { + opts.K = 3 + } + if opts.PreRerankK == 0 { + opts.PreRerankK = 10 + } + + // Call the retriever as in the simple case. + response, err := menuPDFRetriever.Retrieve(ctx, &ai.RetrieverRequest{ + Document: req.Document, + Options: localvec.RetrieverOptions{K: opts.PreRerankK}, + }) + if err != nil { + return nil, err + } + + // Re-rank the returned documents using your custom function. + rerankedDocs := rerank(response.Documents) + response.Documents = rerankedDocs[:opts.K] + + return response, nil + }, + ) + //!-customret + + _ = advancedMenuRetriever +} + +func rerank(document []*ai.Document) []*ai.Document { + panic("unimplemented") +} diff --git a/go/internal/doc-snippets/rag/pdf/pdf.go b/go/internal/doc-snippets/rag/pdf/pdf.go new file mode 100644 index 0000000000..3cb2c48d8f --- /dev/null +++ b/go/internal/doc-snippets/rag/pdf/pdf.go @@ -0,0 +1,30 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pdf + +import ( + "io" + "os" +) + +type Reader struct {} + +func (Reader) GetPlainText() (io.Reader, error) { + panic("Stub.") +} + +func Open(file string) (*os.File, *Reader, error) { + panic("Stub.") +} \ No newline at end of file diff --git a/go/internal/doc-snippets/rag/textsplitter/textsplitter.go b/go/internal/doc-snippets/rag/textsplitter/textsplitter.go new file mode 100644 index 0000000000..c827cd2e88 --- /dev/null +++ b/go/internal/doc-snippets/rag/textsplitter/textsplitter.go @@ -0,0 +1,33 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package textsplitter + +type Splitter struct{} + +func (Splitter) SplitText(t string) ([]string, error) { + panic("Stub.") +} + +func NewRecursiveCharacter(args ...any) Splitter { + panic("Stub.") +} + +func WithChunkSize(n int) any { + panic("Stub.") +} + +func WithChunkOverlap(n int) any { + panic("Stub.") +}