From 963e4df30a19438ea0f1cc6497d3c83d3bf217a0 Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Wed, 7 May 2025 19:48:00 +0530 Subject: [PATCH 1/9] addition of info and config schema for embedder --- go/ai/embedder.go | 9 +- go/ai/gen.go | 25 ++++++ go/core/action.go | 6 +- go/genkit/genkit.go | 4 +- go/plugins/googlegenai/gemini.go | 32 ++++++- go/plugins/googlegenai/googlegenai.go | 16 ++-- go/plugins/googlegenai/models.go | 124 ++++++++++++++++++++++---- 7 files changed, 185 insertions(+), 31 deletions(-) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index 8666d60904..adb18e5825 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -41,9 +41,16 @@ type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}] func DefineEmbedder( r *registry.Registry, provider, name string, + options *EmbedderOptions, embed func(context.Context, *EmbedRequest) (*EmbedResponse, error), ) Embedder { - return (*embedder)(core.DefineAction(r, provider, name, atype.Embedder, nil, embed)) + metadata := map[string]any{} + metadata["type"] = "embedder" + metadata["info"] = options.Info + if options.ConfigSchema != nil { + metadata["embedder"] = map[string]any{"customOptions": options.ConfigSchema} + } + return (*embedder)(core.DefineAction(r, provider, name, atype.Embedder, metadata, embed)) } // LookupEmbedder looks up an [Embedder] registered by [DefineEmbedder]. diff --git a/go/ai/gen.go b/go/ai/gen.go index 8281084597..62fd0ea870 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -390,3 +390,28 @@ type TraceMetadata struct { Paths []*PathMetadata `json:"paths,omitempty"` Timestamp float64 `json:"timestamp,omitempty"` } + +// EmbedderInfo represents the structure of the embedder information object. +// It mirrors the Zod schema EmbedderInfoSchema. +type EmbedderInfo struct { + // Friendly label for this model (e.g. "Google AI - Gemini Pro") + Label string `json:"label,omitempty"` + + Supports *EmbedderSupports `json:"supports,omitempty"` + + Dimensions int `json:"dimensions,omitempty"` +} + +// EmbedderSupports represents the supported capabilities of the embedder model. +type EmbedderSupports struct { + // Model can input this type of data. + // Expected values could be "text", "image", "video", but the struct + Input []string `json:"input,omitempty"` + + Multilingual bool `json:"multilingual,omitempty"` +} + +type EmbedderOptions struct { + ConfigSchema map[string]any `json:"configSchema,omitempty"` + Info *EmbedderInfo `json:"info,omitempty"` +} diff --git a/go/core/action.go b/go/core/action.go index a5cd85f855..a1d060a16e 100644 --- a/go/core/action.go +++ b/go/core/action.go @@ -73,7 +73,8 @@ type noStream = func(context.Context, struct{}) error // DefineAction creates a new non-streaming Action and registers it. func DefineAction[In, Out any]( r *registry.Registry, - provider, name string, + provider, + name string, atype atype.ActionType, metadata map[string]any, fn Func[In, Out], @@ -116,7 +117,8 @@ func DefineActionWithInputSchema[Out any]( // defineAction creates an action and registers it with the given Registry. func defineAction[In, Out, Stream any]( r *registry.Registry, - provider, name string, + provider, + name string, atype atype.ActionType, metadata map[string]any, inputSchema *jsonschema.Schema, diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 2473b0dd68..2239612202 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -703,8 +703,8 @@ func LookupIndexer(g *Genkit, provider, name string) ai.Indexer { // The `provider` and `name` form the unique identifier. The `embed` function // contains the logic to process an [ai.EmbedRequest] (containing documents or a query) // and return an [ai.EmbedResponse] (containing the corresponding embeddings). -func DefineEmbedder(g *Genkit, provider, name string, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder { - return ai.DefineEmbedder(g.reg, provider, name, embed) +func DefineEmbedder(g *Genkit, provider string, name string, options *ai.EmbedderOptions, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder { + return ai.DefineEmbedder(g.reg, provider, name, options, embed) } // LookupEmbedder retrieves a registered [ai.Embedder] by its provider and name. diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 9320f55ad4..d91e3bc883 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -79,6 +79,29 @@ type EmbedOptions struct { TaskType string `json:"task_type,omitempty"` } +type GeminiEmbeddingConfigSchema struct { + // Override the API key provided at plugin initialization. + APIKey string `json:"apiKey,omitempty"` + + // The `task_type` parameter is defined as the intended downstream application + // to help the model produce better quality embeddings. + // NOTE: Assuming TaskTypeSchema resolves to a string. If it's a different + // complex type, this field would need to be adjusted accordingly (e.g., a struct). + TaskType *string `json:"taskType,omitempty"` + + Title string `json:"title,omitempty"` + + Version string `json:"version,omitempty"` + + // The `outputDimensionality` parameter allows you to specify the dimensionality + // of the embedding output. By default, the model generates embeddings + // with 768 dimensions. Models such as `text-embedding-004`, `text-embedding-005`, + // and `text-multilingual-embedding-002` allow the output dimensionality + // to be adjusted between 1 and 768. + // NOTE: The min(1) and max(768) constraints + OutputDimensionality int `json:"outputDimensionality,omitempty"` +} + // configToMap converts a config struct to a map[string]any. func configToMap(config any) map[string]any { r := jsonschema.Reflector{ @@ -272,13 +295,18 @@ func defineModel(g *genkit.Genkit, client *genai.Client, name string, info ai.Mo // DefineEmbedder defines embeddings for the provided contents and embedder // model -func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string) ai.Embedder { +func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOptions ai.EmbedderOptions) ai.Embedder { provider := googleAIProvider if client.ClientConfig().Backend == genai.BackendVertexAI { provider = vertexAIProvider } - return genkit.DefineEmbedder(g, provider, name, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { + emdOpts := &ai.EmbedderOptions{ + Info: embedOptions.Info, + ConfigSchema: configToMap(&GeminiEmbeddingConfigSchema{}), + } + + return genkit.DefineEmbedder(g, provider, name, emdOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { var content []*genai.Content var embedConfig *genai.EmbedContentConfig diff --git a/go/plugins/googlegenai/googlegenai.go b/go/plugins/googlegenai/googlegenai.go index 6a3a0cd07c..7810ad6ee0 100644 --- a/go/plugins/googlegenai/googlegenai.go +++ b/go/plugins/googlegenai/googlegenai.go @@ -108,8 +108,8 @@ func (ga *GoogleAI) Init(ctx context.Context, g *genkit.Genkit) (err error) { if err != nil { return err } - for _, e := range embedders { - defineEmbedder(g, ga.gclient, e) + for e, eOpts := range embedders { + defineEmbedder(g, ga.gclient, e, eOpts) } return nil @@ -182,8 +182,8 @@ func (v *VertexAI) Init(ctx context.Context, g *genkit.Genkit) (err error) { if err != nil { return err } - for _, e := range embedders { - defineEmbedder(g, v.gclient, e) + for e, eOpts := range embedders { + defineEmbedder(g, v.gclient, e, eOpts) } return nil @@ -248,23 +248,23 @@ func (v *VertexAI) DefineModel(g *genkit.Genkit, name string, info *ai.ModelInfo } // DefineEmbedder defines an embedder with a given name. -func (ga *GoogleAI) DefineEmbedder(g *genkit.Genkit, name string) (ai.Embedder, error) { +func (ga *GoogleAI) DefineEmbedder(g *genkit.Genkit, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { ga.mu.Lock() defer ga.mu.Unlock() if !ga.initted { return nil, errors.New("GoogleAI plugin not initialized") } - return defineEmbedder(g, ga.gclient, name), nil + return defineEmbedder(g, ga.gclient, name, embedOptions), nil } // DefineEmbedder defines an embedder with a given name. -func (v *VertexAI) DefineEmbedder(g *genkit.Genkit, name string) (ai.Embedder, error) { +func (v *VertexAI) DefineEmbedder(g *genkit.Genkit, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { v.mu.Lock() defer v.mu.Unlock() if !v.initted { return nil, errors.New("VertexAI plugin not initialized") } - return defineEmbedder(g, v.gclient, name), nil + return defineEmbedder(g, v.gclient, name, embedOptions), nil } // IsDefinedEmbedder reports whether the named [Embedder] is defined by this plugin. diff --git a/go/plugins/googlegenai/models.go b/go/plugins/googlegenai/models.go index 6999ca58c9..5e2b2e98c5 100644 --- a/go/plugins/googlegenai/models.go +++ b/go/plugins/googlegenai/models.go @@ -24,9 +24,17 @@ const ( gemini25FlashPreview0417 = "gemini-2.5-flash-preview-04-17" - gemini25ProExp0325 = "gemini-2.5-pro-exp-03-25" - gemini25ProPreview0325 = "gemini-2.5-pro-preview-03-25" - gemini25ProPreview0506 = "gemini-2.5-pro-preview-05-06" + gemini25ProExp0325 = "gemini-2.5-pro-exp-03-25" + gemini25ProPreview0325 = "gemini-2.5-pro-preview-03-25" + gemini25ProPreview0506 = "gemini-2.5-pro-preview-05-06" + textembedding004 = "text-embedding-004" + embedding001 = "embedding-001" + textembeddinggecko003 = "textembedding-gecko@003" + textembeddinggecko002 = "textembedding-gecko@002" + textembeddinggecko001 = "textembedding-gecko@001" + textembeddinggeckomultilingual001 = "textembedding-gecko-multilingual@001" + textmultilingualembedding002 = "text-multilingual-embedding-002" + multimodalembedding = "multimodalembedding" ) var ( @@ -163,18 +171,97 @@ var ( } googleAIEmbedders = []string{ - "text-embedding-004", - "embedding-001", + textembedding004, + embedding001, + } + + googleAIEmbedderConfig = map[string]ai.EmbedderOptions{ + textembedding004: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Google Gen AI - Text Embedding 001", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + embedding001: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Google Gen AI - Text Embedding Gecko (Legacy)", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + textembeddinggecko003: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Google Gen AI - Text Embedding Gecko 003", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + textembeddinggecko002: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Vertex AI - Text Embedding Gecko 002", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + textembeddinggecko001: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Vertex AI - Text Embedding Gecko 001", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + textembeddinggeckomultilingual001: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Vertex AI - Text Embedding Gecko Multilingual 001", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + textmultilingualembedding002: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Vertex AI - Text Multilingual Embedding 001", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + multimodalembedding: { + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "Google Gen AI - Text Embedding Gecko (Legacy)", + Supports: &ai.EmbedderSupports{ // Supports object is present + Input: []string{ + "text", + "image", + "video", + }, + }, + }, + }, } vertexAIEmbedders = []string{ - "textembedding-gecko@003", - "textembedding-gecko@002", - "textembedding-gecko@001", - "text-embedding-004", - "textembedding-gecko-multilingual@001", - "text-multilingual-embedding-002", - "multimodalembedding", + textembeddinggecko003, + textembeddinggecko002, + textembeddinggecko001, + textembedding004, + textembeddinggeckomultilingual001, + textmultilingualembedding002, + multimodalembedding, } ) @@ -212,17 +299,22 @@ func listModels(provider string) (map[string]ai.ModelInfo, error) { // listEmbedders returns a list of supported embedders based on the // detected backend -func listEmbedders(backend genai.Backend) ([]string, error) { - embedders := []string{} +func listEmbedders(backend genai.Backend) (map[string]ai.EmbedderOptions, error) { + embeddersNames := []string{} switch backend { case genai.BackendGeminiAPI: - embedders = googleAIEmbedders + embeddersNames = googleAIEmbedders case genai.BackendVertexAI: - embedders = vertexAIEmbedders + embeddersNames = vertexAIEmbedders default: return nil, fmt.Errorf("embedders for backend %s not found", backend) } + embedders := make(map[string]ai.EmbedderOptions, 0) + for _, n := range embeddersNames { + embedders[n] = googleAIEmbedderConfig[n] + } + return embedders, nil } From cd69efb5197b0cb2f9100a1a863d9a15dcee1d4a Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Fri, 30 May 2025 21:46:55 +0530 Subject: [PATCH 2/9] chaneg in custom config for embedder --- go/ai/embedder.go | 17 ++++++++-- go/ai/gen.go | 5 ++- go/core/action.go | 8 ++--- go/plugins/googlegenai/gemini.go | 57 +++++++++----------------------- 4 files changed, 37 insertions(+), 50 deletions(-) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index adb18e5825..8595e6a577 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -22,6 +22,7 @@ import ( "github.com/firebase/genkit/go/core" "github.com/firebase/genkit/go/internal/atype" + "github.com/firebase/genkit/go/internal/base" "github.com/firebase/genkit/go/internal/registry" ) @@ -36,6 +37,12 @@ type Embedder interface { // An embedder is used to convert a document to a multidimensional vector. type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}] +func configToMap(config any) map[string]any { + schema := base.InferJSONSchema(config) + result := base.SchemaAsMap(schema) + return result +} + // DefineEmbedder registers the given embed function as an action, and returns an // [Embedder] that runs it. func DefineEmbedder( @@ -48,9 +55,15 @@ func DefineEmbedder( metadata["type"] = "embedder" metadata["info"] = options.Info if options.ConfigSchema != nil { - metadata["embedder"] = map[string]any{"customOptions": options.ConfigSchema} + metadata["embedder"] = map[string]any{"customOptions": configToMap(options.ConfigSchema)} + } + inputSchema := base.InferJSONSchema(EmbedRequest{}) + if inputSchema.Properties != nil && options.ConfigSchema != nil { + if _, ok := inputSchema.Properties.Get("options"); ok { + inputSchema.Properties.Set("options", base.InferJSONSchema(options.ConfigSchema)) + } } - return (*embedder)(core.DefineAction(r, provider, name, atype.Embedder, metadata, embed)) + return (*embedder)(core.DefineActionWithInputSchema(r, provider, name, atype.Embedder, metadata, inputSchema, embed)) } // LookupEmbedder looks up an [Embedder] registered by [DefineEmbedder]. diff --git a/go/ai/gen.go b/go/ai/gen.go index 2654660699..0ed71f6310 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -399,7 +399,6 @@ type TraceMetadata struct { } // EmbedderInfo represents the structure of the embedder information object. -// It mirrors the Zod schema EmbedderInfoSchema. type EmbedderInfo struct { // Friendly label for this model (e.g. "Google AI - Gemini Pro") Label string `json:"label,omitempty"` @@ -419,6 +418,6 @@ type EmbedderSupports struct { } type EmbedderOptions struct { - ConfigSchema map[string]any `json:"configSchema,omitempty"` - Info *EmbedderInfo `json:"info,omitempty"` + ConfigSchema any `json:"configSchema,omitempty"` + Info *EmbedderInfo `json:"info,omitempty"` } diff --git a/go/core/action.go b/go/core/action.go index a1d060a16e..0e2ccb9c99 100644 --- a/go/core/action.go +++ b/go/core/action.go @@ -100,16 +100,16 @@ func DefineStreamingAction[In, Out, Stream any]( // This differs from DefineAction in that the input schema is // defined dynamically; the static input type is "any". // This is used for prompts. -func DefineActionWithInputSchema[Out any]( +func DefineActionWithInputSchema[In, Out any]( r *registry.Registry, provider, name string, atype atype.ActionType, metadata map[string]any, inputSchema *jsonschema.Schema, - fn Func[any, Out], -) *ActionDef[any, Out, struct{}] { + fn Func[In, Out], +) *ActionDef[In, Out, struct{}] { return defineAction(r, provider, name, atype, metadata, inputSchema, - func(ctx context.Context, in any, _ noStream) (Out, error) { + func(ctx context.Context, in In, _ noStream) (Out, error) { return fn(ctx, in) }) } diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 617d770caa..0aa5f0059d 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -28,13 +28,11 @@ import ( "strconv" "strings" - "github.com/firebase/genkit/go/core" - "github.com/firebase/genkit/go/internal/base" - "github.com/invopop/jsonschema" - "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/core" "github.com/firebase/genkit/go/genkit" "github.com/firebase/genkit/go/internal" + "github.com/firebase/genkit/go/internal/base" "github.com/firebase/genkit/go/plugins/internal/uri" "google.golang.org/genai" ) @@ -84,36 +82,9 @@ type EmbedOptions struct { TaskType string `json:"task_type,omitempty"` } -type GeminiEmbeddingConfigSchema struct { - // Override the API key provided at plugin initialization. - APIKey string `json:"apiKey,omitempty"` - - // The `task_type` parameter is defined as the intended downstream application - // to help the model produce better quality embeddings. - // NOTE: Assuming TaskTypeSchema resolves to a string. If it's a different - // complex type, this field would need to be adjusted accordingly (e.g., a struct). - TaskType *string `json:"taskType,omitempty"` - - Title string `json:"title,omitempty"` - - Version string `json:"version,omitempty"` - - // The `outputDimensionality` parameter allows you to specify the dimensionality - // of the embedding output. By default, the model generates embeddings - // with 768 dimensions. Models such as `text-embedding-004`, `text-embedding-005`, - // and `text-multilingual-embedding-002` allow the output dimensionality - // to be adjusted between 1 and 768. - // NOTE: The min(1) and max(768) constraints - OutputDimensionality int `json:"outputDimensionality,omitempty"` -} - // configToMap converts a config struct to a map[string]any. func configToMap(config any) map[string]any { - r := jsonschema.Reflector{ - DoNotReference: true, // Prevent $ref usage - ExpandedStruct: true, // Include all fields directly - } - schema := r.Reflect(config) + schema := base.InferJSONSchema(config) result := base.SchemaAsMap(schema) return result } @@ -318,21 +289,24 @@ func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOp emdOpts := &ai.EmbedderOptions{ Info: embedOptions.Info, - ConfigSchema: configToMap(&GeminiEmbeddingConfigSchema{}), + ConfigSchema: genai.EmbedContentConfig{}, } return genkit.DefineEmbedder(g, provider, name, emdOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { var content []*genai.Content var embedConfig *genai.EmbedContentConfig - // check if request options matches VertexAI configuration - if opts, _ := req.Options.(*EmbedOptions); opts != nil { - if provider == googleAIProvider { - return nil, fmt.Errorf("wrong options provided for %s provider, got %T", provider, opts) - } - embedConfig = &genai.EmbedContentConfig{ - Title: opts.Title, - TaskType: opts.TaskType, + if req.Options != nil { + if optsMap, ok := req.Options.(map[string]any); ok { + optionsBytes, err := json.Marshal(optsMap) + if err != nil { + fmt.Println("Error marshaling options map to bytes:", err) + } + var tempConfig genai.EmbedContentConfig + err = json.Unmarshal(optionsBytes, &tempConfig) + if err == nil { + embedConfig = &tempConfig + } } } @@ -898,6 +872,7 @@ func toGeminiParts(parts []*ai.Part) ([]*genai.Part, error) { // toGeminiPart converts a [ai.Part] to a [genai.Part]. func toGeminiPart(p *ai.Part) (*genai.Part, error) { + switch { case p.IsText(): return genai.NewPartFromText(p.Text), nil From adde87d8b738a01d2b494d19b78788479f441e27 Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Wed, 4 Jun 2025 21:10:08 +0530 Subject: [PATCH 3/9] config changes --- go/ai/retriever.go | 25 ++++++++++++- go/genkit/genkit.go | 4 +- go/internal/fakeembedder/fakeembedder_test.go | 12 +++++- go/plugins/compat_oai/compat_oai.go | 4 +- go/plugins/googlegenai/gemini.go | 7 ++-- go/plugins/localvec/localvec.go | 13 ++++++- go/plugins/localvec/localvec_test.go | 37 +++++++++++++++++-- go/plugins/ollama/embed.go | 6 ++- go/plugins/pinecone/genkit.go | 11 +++++- go/plugins/pinecone/genkit_test.go | 13 ++++++- go/plugins/weaviate/weaviate_test.go | 13 ++++++- go/samples/rag/main.go | 6 ++- 12 files changed, 129 insertions(+), 22 deletions(-) diff --git a/go/ai/retriever.go b/go/ai/retriever.go index 67ac28e930..36b07b1820 100644 --- a/go/ai/retriever.go +++ b/go/ai/retriever.go @@ -35,12 +35,33 @@ type Retriever interface { Retrieve(ctx context.Context, req *RetrieverRequest) (*RetrieverResponse, error) } +type RetrieverInfo struct { + Label string `json:"label,omitempty"` + Supports *MediaSupports `json:"supports,omitempty"` +} + +// `Supports` is a nested struct equivalent to the nested `z.object` for the `supports` field. +type MediaSupports struct { + Media bool `json:"media,omitempty"` +} + +type RetrieverOptions struct { + ConfigSchema any + Info *RetrieverInfo +} + type retriever core.ActionDef[*RetrieverRequest, *RetrieverResponse, struct{}] // DefineRetriever registers the given retrieve function as an action, and returns a // [Retriever] that runs it. -func DefineRetriever(r *registry.Registry, provider, name string, fn RetrieverFunc) Retriever { - return (*retriever)(core.DefineAction(r, provider, name, core.ActionTypeRetriever, nil, fn)) +func DefineRetriever(r *registry.Registry, provider, name string, options *RetrieverOptions, fn RetrieverFunc) Retriever { + metadata := map[string]any{} + metadata["type"] = "retriever" + metadata["info"] = options.Info + if options.ConfigSchema != nil { + metadata["retriever"] = map[string]any{"customOptions": configToMap(options.ConfigSchema)} + } + return (*retriever)(core.DefineAction(r, provider, name, core.ActionTypeRetriever, metadata, fn)) } // LookupRetriever looks up a [Retriever] registered by [DefineRetriever]. diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index c933cacf88..a27adfd28f 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -651,8 +651,8 @@ func GenerateData[Out any](ctx context.Context, g *Genkit, opts ...ai.GenerateOp // The `provider` and `name` form the unique identifier. The `ret` function // contains the logic to process an [ai.RetrieverRequest] (containing the query) // and return an [ai.RetrieverResponse] (containing the relevant documents). -func DefineRetriever(g *Genkit, provider, name string, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever { - return ai.DefineRetriever(g.reg, provider, name, ret) +func DefineRetriever(g *Genkit, provider, name string, options *ai.RetrieverOptions, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever { + return ai.DefineRetriever(g.reg, provider, name, options, ret) } // LookupRetriever retrieves a registered [ai.Retriever] by its provider and name. diff --git a/go/internal/fakeembedder/fakeembedder_test.go b/go/internal/fakeembedder/fakeembedder_test.go index be57fa06ab..a0099bde12 100644 --- a/go/internal/fakeembedder/fakeembedder_test.go +++ b/go/internal/fakeembedder/fakeembedder_test.go @@ -32,7 +32,17 @@ func TestFakeEmbedder(t *testing.T) { } embed := New() - emb := ai.DefineEmbedder(r, "fake", "embed", embed.Embed) + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 32, + Label: "embed", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + emb := ai.DefineEmbedder(r, "fake", "embed", emdOpts, embed.Embed) d := ai.DocumentFromText("fakeembedder test", nil) vals := []float32{1, 2} diff --git a/go/plugins/compat_oai/compat_oai.go b/go/plugins/compat_oai/compat_oai.go index 47c902fcde..627ac14fee 100644 --- a/go/plugins/compat_oai/compat_oai.go +++ b/go/plugins/compat_oai/compat_oai.go @@ -125,14 +125,14 @@ func (o *OpenAICompatible) DefineModel(g *genkit.Genkit, provider, name string, } // DefineEmbedder defines an embedder with a given name. -func (o *OpenAICompatible) DefineEmbedder(g *genkit.Genkit, provider, name string) (ai.Embedder, error) { +func (o *OpenAICompatible) DefineEmbedder(g *genkit.Genkit, provider, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { o.mu.Lock() defer o.mu.Unlock() if !o.initted { return nil, errors.New("OpenAICompatible.Init not called") } - return genkit.DefineEmbedder(g, provider, name, func(ctx context.Context, input *ai.EmbedRequest) (*ai.EmbedResponse, error) { + return genkit.DefineEmbedder(g, provider, name, &embedOptions, func(ctx context.Context, input *ai.EmbedRequest) (*ai.EmbedResponse, error) { var data openaiGo.EmbeddingNewParamsInputArrayOfStrings for _, doc := range input.Input { for _, p := range doc.Content { diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 3eb820eb36..0abc297e51 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -287,12 +287,11 @@ func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOp provider = vertexAIProvider } - emdOpts := &ai.EmbedderOptions{ - Info: embedOptions.Info, - ConfigSchema: genai.EmbedContentConfig{}, + if embedOptions.ConfigSchema == nil { + embedOptions.ConfigSchema = genai.EmbedContentConfig{} } - return genkit.DefineEmbedder(g, provider, name, emdOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { + return genkit.DefineEmbedder(g, provider, name, &embedOptions, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { var content []*genai.Content var embedConfig *genai.EmbedContentConfig diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go index 7185782c5d..4a3cc4f2e1 100644 --- a/go/plugins/localvec/localvec.go +++ b/go/plugins/localvec/localvec.go @@ -56,8 +56,17 @@ func DefineRetriever(g *genkit.Genkit, name string, cfg Config) (*DocStore, ai.R if err != nil { return nil, nil, err } + retOpts := &ai.RetrieverOptions{ + ConfigSchema: RetrieverOptions{}, + Info: &ai.RetrieverInfo{ + Label: "local vec", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } return ds, - genkit.DefineRetriever(g, provider, name, ds.retrieve), + genkit.DefineRetriever(g, provider, name, retOpts, ds.retrieve), nil } @@ -167,7 +176,7 @@ func (ds *DocStore) retrieve(ctx context.Context, req *ai.RetrieverRequest) (*ai k = min(k, len(scoredDocs)) docs := make([]*ai.Document, 0, k) - for i := range k { + for i := 0; i < k; i++ { docs = append(docs, scoredDocs[i].doc) } diff --git a/go/plugins/localvec/localvec_test.go b/go/plugins/localvec/localvec_test.go index 60bf759a88..fc3d04a209 100644 --- a/go/plugins/localvec/localvec_test.go +++ b/go/plugins/localvec/localvec_test.go @@ -58,7 +58,17 @@ func TestLocalVec(t *testing.T) { embedder.Register(d1, v1) embedder.Register(d2, v2) embedder.Register(d3, v3) - embedAction := genkit.DefineEmbedder(g, "fake", "embedder1", embedder.Embed) + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 32, + Label: "", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + embedAction := genkit.DefineEmbedder(g, "fake", "embedder1", emdOpts, embedder.Embed) ds, err := newDocStore(t.TempDir(), "testLocalVec", embedAction, nil) if err != nil { t.Fatal(err) @@ -120,7 +130,18 @@ func TestPersistentIndexing(t *testing.T) { embedder.Register(d1, v1) embedder.Register(d2, v2) embedder.Register(d3, v3) - embedAction := genkit.DefineEmbedder(g, "fake", "embedder2", embedder.Embed) + + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 32, + Label: "", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + embedAction := genkit.DefineEmbedder(g, "fake", "embedder2", emdOpts, embedder.Embed) tDir := t.TempDir() @@ -196,7 +217,17 @@ func TestInit(t *testing.T) { if err != nil { t.Fatal(err) } - embedder := genkit.DefineEmbedder(g, "fake", "e", fakeembedder.New().Embed) + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + embedder := genkit.DefineEmbedder(g, "fake", "e", emdOpts, fakeembedder.New().Embed) if err := Init(); err != nil { t.Fatal(err) } diff --git a/go/plugins/ollama/embed.go b/go/plugins/ollama/embed.go index 993db8a4dc..096b8d032e 100644 --- a/go/plugins/ollama/embed.go +++ b/go/plugins/ollama/embed.go @@ -134,7 +134,11 @@ func (o *Ollama) DefineEmbedder(g *genkit.Genkit, serverAddress string, model st if !o.initted { panic("ollama.Init not called") } - return genkit.DefineEmbedder(g, provider, serverAddress, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{}, + ConfigSchema: nil, + } + return genkit.DefineEmbedder(g, provider, serverAddress, emdOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { if req.Options == nil { req.Options = &EmbedOptions{Model: model} } diff --git a/go/plugins/pinecone/genkit.go b/go/plugins/pinecone/genkit.go index 6d2f4fef36..0e604323e2 100644 --- a/go/plugins/pinecone/genkit.go +++ b/go/plugins/pinecone/genkit.go @@ -99,7 +99,16 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config) (*Docsto if err != nil { return nil, nil, err } - return ds, genkit.DefineRetriever(g, provider, cfg.IndexID, ds.Retrieve), nil + retOpts := &ai.RetrieverOptions{ + ConfigSchema: RetrieverOptions{}, + Info: &ai.RetrieverInfo{ + Label: "local Vec", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + return ds, genkit.DefineRetriever(g, provider, cfg.IndexID, retOpts, ds.Retrieve), nil } // IsDefinedRetriever reports whether the named [Retriever] is defined by this plugin. diff --git a/go/plugins/pinecone/genkit_test.go b/go/plugins/pinecone/genkit_test.go index 81e152600d..d0914bd3ec 100644 --- a/go/plugins/pinecone/genkit_test.go +++ b/go/plugins/pinecone/genkit_test.go @@ -83,9 +83,20 @@ func TestGenkit(t *testing.T) { if err := (&Pinecone{APIKey: *testAPIKey}).Init(ctx, g); err != nil { t.Fatal(err) } + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + cfg := Config{ IndexID: *testIndex, - Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", embedder.Embed), + Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", emdOpts, embedder.Embed), } ds, retriever, err := DefineRetriever(ctx, g, cfg) diff --git a/go/plugins/weaviate/weaviate_test.go b/go/plugins/weaviate/weaviate_test.go index 24a170e14f..4e587228d8 100644 --- a/go/plugins/weaviate/weaviate_test.go +++ b/go/plugins/weaviate/weaviate_test.go @@ -89,9 +89,20 @@ func TestGenkit(t *testing.T) { t.Fatal(err) } + emdOpts := &ai.EmbedderOptions{ + Info: &ai.EmbedderInfo{ + Dimensions: 768, + Label: "", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + ConfigSchema: nil, + } + classCfg := ClassConfig{ Class: *testClass, - Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", embedder.Embed), + Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", emdOpts, embedder.Embed), } ds, retriever, err := DefineRetriever(ctx, g, classCfg) if err != nil { diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go index d293ed031a..aaa67c11d3 100644 --- a/go/samples/rag/main.go +++ b/go/samples/rag/main.go @@ -47,6 +47,7 @@ import ( "github.com/firebase/genkit/go/plugins/evaluators" "github.com/firebase/genkit/go/plugins/googlegenai" "github.com/firebase/genkit/go/plugins/localvec" + "github.com/firebase/genkit/go/plugins/pinecone" ) const simpleQaPromptTemplate = ` @@ -82,7 +83,7 @@ func main() { }, } g, err := genkit.Init(ctx, - genkit.WithPlugins(&googlegenai.GoogleAI{}, &evaluators.GenkitEval{Metrics: metrics}), + genkit.WithPlugins(&googlegenai.GoogleAI{}, &evaluators.GenkitEval{Metrics: metrics}, &pinecone.Pinecone{}), ) if err != nil { log.Fatal(err) @@ -155,10 +156,11 @@ func main() { genkit.DefineFlow(g, "simpleQaFlow", func(ctx context.Context, input *simpleQaInput) (string, error) { d1 := ai.DocumentFromText("Paris is the capital of France", nil) + d4 := ai.DocumentFromText("India will become a new capital of France", nil) d2 := ai.DocumentFromText("USA is the largest importer of coffee", nil) d3 := ai.DocumentFromText("Water exists in 3 states - solid, liquid and gas", nil) - err := localvec.Index(ctx, []*ai.Document{d1, d2, d3}, docStore) + err := localvec.Index(ctx, []*ai.Document{d1, d2, d3, d4}, docStore) if err != nil { return "", err } From 9776f2ce4dc0f4839398bad4f0932b24023bc5a7 Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Thu, 5 Jun 2025 00:56:52 +0530 Subject: [PATCH 4/9] addition of embedder config for all plugins and addition of retriever config --- go/ai/embedder.go | 24 ++++++++++ go/ai/gen.go | 24 ---------- go/ai/retriever.go | 1 - go/internal/doc-snippets/pinecone.go | 14 +++++- go/internal/doc-snippets/rag/main.go | 48 ++++++++++++++++++++ go/plugins/compat_oai/openai/openai.go | 62 ++++++++++++++++++++++---- go/plugins/firebase/retriever.go | 12 ++++- go/plugins/googlegenai/models.go | 8 ++++ go/plugins/localvec/localvec.go | 14 ++---- go/plugins/localvec/localvec_test.go | 13 +++++- go/plugins/pinecone/genkit.go | 39 +++++++--------- go/plugins/pinecone/genkit_test.go | 20 +++++++-- go/plugins/weaviate/weaviate.go | 4 +- go/plugins/weaviate/weaviate_test.go | 13 +++++- go/samples/menu/main.go | 14 +++++- go/samples/pgvector/main.go | 15 +++++-- go/samples/rag/main.go | 16 +++++-- 17 files changed, 257 insertions(+), 84 deletions(-) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index ce7717c044..f795cd57bf 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -33,6 +33,30 @@ type Embedder interface { Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) } +// EmbedderInfo represents the structure of the embedder information object. +type EmbedderInfo struct { + // Friendly label for this model (e.g. "Google AI - Gemini Pro") + Label string `json:"label,omitempty"` + + Supports *EmbedderSupports `json:"supports,omitempty"` + + Dimensions int `json:"dimensions,omitempty"` +} + +// EmbedderSupports represents the supported capabilities of the embedder model. +type EmbedderSupports struct { + // Model can input this type of data. + // Expected values could be "text", "image", "video", but the struct + Input []string `json:"input,omitempty"` + + Multilingual bool `json:"multilingual,omitempty"` +} + +type EmbedderOptions struct { + ConfigSchema any `json:"configSchema,omitempty"` + Info *EmbedderInfo `json:"info,omitempty"` +} + // An embedder is used to convert a document to a multidimensional vector. type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}] diff --git a/go/ai/gen.go b/go/ai/gen.go index 0ed71f6310..c7c38c2836 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -397,27 +397,3 @@ type TraceMetadata struct { Paths []*PathMetadata `json:"paths,omitempty"` Timestamp float64 `json:"timestamp,omitempty"` } - -// EmbedderInfo represents the structure of the embedder information object. -type EmbedderInfo struct { - // Friendly label for this model (e.g. "Google AI - Gemini Pro") - Label string `json:"label,omitempty"` - - Supports *EmbedderSupports `json:"supports,omitempty"` - - Dimensions int `json:"dimensions,omitempty"` -} - -// EmbedderSupports represents the supported capabilities of the embedder model. -type EmbedderSupports struct { - // Model can input this type of data. - // Expected values could be "text", "image", "video", but the struct - Input []string `json:"input,omitempty"` - - Multilingual bool `json:"multilingual,omitempty"` -} - -type EmbedderOptions struct { - ConfigSchema any `json:"configSchema,omitempty"` - Info *EmbedderInfo `json:"info,omitempty"` -} diff --git a/go/ai/retriever.go b/go/ai/retriever.go index 36b07b1820..b358e1ef6a 100644 --- a/go/ai/retriever.go +++ b/go/ai/retriever.go @@ -40,7 +40,6 @@ type RetrieverInfo struct { Supports *MediaSupports `json:"supports,omitempty"` } -// `Supports` is a nested struct equivalent to the nested `z.object` for the `supports` field. type MediaSupports struct { Media bool `json:"media,omitempty"` } diff --git a/go/internal/doc-snippets/pinecone.go b/go/internal/doc-snippets/pinecone.go index a15aafe467..6f9cc474c3 100644 --- a/go/internal/doc-snippets/pinecone.go +++ b/go/internal/doc-snippets/pinecone.go @@ -48,10 +48,22 @@ func pineconeEx(ctx context.Context) error { var docChunks []*ai.Document // [START defineretriever] + retOpts := &ai.RetrieverOptions{ + ConfigSchema: pinecone.PineconeRetrieverOptions{ + K: 1, + Namespace: "Test", + }, + Info: &ai.RetrieverInfo{ + Label: "Pinecone", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } ds, menuRetriever, err := pinecone.DefineRetriever(ctx, g, pinecone.Config{ IndexID: "menu_data", // Your Pinecone index Embedder: googlegenai.GoogleAIEmbedder(g, "text-embedding-004"), // Embedding model of your choice - }) + }, retOpts) if err != nil { return err } diff --git a/go/internal/doc-snippets/rag/main.go b/go/internal/doc-snippets/rag/main.go index d82f58dff6..3319a351c8 100644 --- a/go/internal/doc-snippets/rag/main.go +++ b/go/internal/doc-snippets/rag/main.go @@ -49,6 +49,17 @@ func main() { if err != nil { log.Fatal(err) } + retOpts := &ai.RetrieverOptions{ + ConfigSchema: localvec.RetrieverOptions{ + K: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "menuQA", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } docStore, _, err := localvec.DefineRetriever( g, @@ -56,6 +67,7 @@ func main() { localvec.Config{ Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"), }, + retOpts, ) if err != nil { log.Fatal(err) @@ -155,12 +167,25 @@ func menuQA() { model := googlegenai.VertexAIModel(g, "gemini-1.5-flash") + retOpts := &ai.RetrieverOptions{ + ConfigSchema: localvec.RetrieverOptions{ + K: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "menuQA", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + _, menuPdfRetriever, err := localvec.DefineRetriever( g, "menuQA", localvec.Config{ Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"), }, + retOpts, ) if err != nil { log.Fatal(err) @@ -207,12 +232,25 @@ func customret() { log.Fatal(err) } + retOpts := &ai.RetrieverOptions{ + ConfigSchema: localvec.RetrieverOptions{ + K: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "menuQA", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + _, menuPDFRetriever, _ := localvec.DefineRetriever( g, "menuQA", localvec.Config{ Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"), }, + retOpts, ) // [START customret] @@ -220,10 +258,20 @@ func customret() { K int PreRerankK int } + genRetOpts := &ai.RetrieverOptions{ + ConfigSchema: CustomMenuRetrieverOptions{}, + Info: &ai.RetrieverInfo{ + Label: "advancedMenuRetriever", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } advancedMenuRetriever := genkit.DefineRetriever( g, "custom", "advancedMenuRetriever", + genRetOpts, func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { // Handle options passed using our custom type. opts, _ := req.Options.(CustomMenuRetrieverOptions) diff --git a/go/plugins/compat_oai/openai/openai.go b/go/plugins/compat_oai/openai/openai.go index 0db9c8636f..71f6971c93 100644 --- a/go/plugins/compat_oai/openai/openai.go +++ b/go/plugins/compat_oai/openai/openai.go @@ -28,6 +28,18 @@ import ( const provider = "openai" +type TextEmbeddingConfig struct { + Dimensions int `json:"dimensions,omitempty"` + EncodingFormat openaiGo.EmbeddingNewParamsEncodingFormat `json:"encodingFormat,omitempty"` +} + +// EmbedderRef represents the main structure for an embedding model's definition. +type EmbedderRef struct { + Name string + ConfigSchema TextEmbeddingConfig // Represents the schema, can be used for default config + Info *ai.EmbedderInfo +} + var ( // Supported models: https://platform.openai.com/docs/models supportedModels = map[string]ai.ModelInfo{ @@ -118,11 +130,40 @@ var ( }, } - // Known embedders: https://platform.openai.com/docs/guides/embeddings - knownEmbedders = []string{ - openaiGo.EmbeddingModelTextEmbedding3Small, - openaiGo.EmbeddingModelTextEmbedding3Large, - openaiGo.EmbeddingModelTextEmbeddingAda002, + supportedEmbeddingModels = map[string]EmbedderRef{ + openaiGo.EmbeddingModelTextEmbeddingAda002: { + Name: "openai/text-embedding-ada-002", + ConfigSchema: TextEmbeddingConfig{}, + Info: &ai.EmbedderInfo{ + Dimensions: 1536, + Label: "Open AI - Text Embedding ADA 002", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + openaiGo.EmbeddingModelTextEmbedding3Large: { + Name: "openai/text-embedding-3-large", + ConfigSchema: TextEmbeddingConfig{}, + Info: &ai.EmbedderInfo{ + Dimensions: 3072, + Label: "Open AI - Text Embedding 3 Large", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, + openaiGo.EmbeddingModelTextEmbedding3Small: { + Name: "openai/text-embedding-3-small", + ConfigSchema: TextEmbeddingConfig{}, // Represents the configurable options + Info: &ai.EmbedderInfo{ + Dimensions: 1536, + Label: "Open AI - Text Embedding 3 Small", + Supports: &ai.EmbedderSupports{ + Input: []string{"text"}, + }, + }, + }, } ) @@ -179,8 +220,8 @@ func (o *OpenAI) Init(ctx context.Context, g *genkit.Genkit) error { } // define default embedders - for _, embedder := range knownEmbedders { - if _, err := o.DefineEmbedder(g, embedder); err != nil { + for _, embedder := range supportedEmbeddingModels { + if _, err := o.DefineEmbedder(g, embedder.Name, embedder.Info, embedder.ConfigSchema); err != nil { return err } } @@ -196,8 +237,11 @@ func (o *OpenAI) DefineModel(g *genkit.Genkit, name string, info ai.ModelInfo) ( return o.openAICompatible.DefineModel(g, provider, name, info) } -func (o *OpenAI) DefineEmbedder(g *genkit.Genkit, name string) (ai.Embedder, error) { - return o.openAICompatible.DefineEmbedder(g, provider, name) +func (o *OpenAI) DefineEmbedder(g *genkit.Genkit, name string, modelInfo *ai.EmbedderInfo, configSchema TextEmbeddingConfig) (ai.Embedder, error) { + return o.openAICompatible.DefineEmbedder(g, provider, name, ai.EmbedderOptions{ + Info: modelInfo, + ConfigSchema: configSchema, + }) } func (o *OpenAI) Embedder(g *genkit.Genkit, name string) ai.Embedder { diff --git a/go/plugins/firebase/retriever.go b/go/plugins/firebase/retriever.go index 5f0e20a302..aab0608d0c 100644 --- a/go/plugins/firebase/retriever.go +++ b/go/plugins/firebase/retriever.go @@ -128,7 +128,17 @@ func defineFirestoreRetriever(g *genkit.Genkit, cfg RetrieverOptions, client *fi return &ai.RetrieverResponse{Documents: documents}, nil } - return genkit.DefineRetriever(g, provider, cfg.Name, retrieve), nil + retOpts := &ai.RetrieverOptions{ + ConfigSchema: cfg, + Info: &ai.RetrieverInfo{ + Label: cfg.Name, + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + + return genkit.DefineRetriever(g, provider, cfg.Name, retOpts, retrieve), nil } // resolveFirestoreCollection resolves the Firestore collection name from the environment if necessary diff --git a/go/plugins/googlegenai/models.go b/go/plugins/googlegenai/models.go index cc173e6bd0..113f633eea 100644 --- a/go/plugins/googlegenai/models.go +++ b/go/plugins/googlegenai/models.go @@ -193,6 +193,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, embedding001: { Info: &ai.EmbedderInfo{ @@ -202,6 +203,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, textembeddinggecko003: { Info: &ai.EmbedderInfo{ @@ -211,6 +213,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, textembeddinggecko002: { Info: &ai.EmbedderInfo{ @@ -220,6 +223,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, textembeddinggecko001: { Info: &ai.EmbedderInfo{ @@ -229,6 +233,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, textembeddinggeckomultilingual001: { Info: &ai.EmbedderInfo{ @@ -238,6 +243,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, textmultilingualembedding002: { Info: &ai.EmbedderInfo{ @@ -247,6 +253,7 @@ var ( Input: []string{"text"}, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, multimodalembedding: { Info: &ai.EmbedderInfo{ @@ -260,6 +267,7 @@ var ( }, }, }, + ConfigSchema: genai.EmbedContentConfig{}, }, } diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go index 4a3cc4f2e1..ba5b811b88 100644 --- a/go/plugins/localvec/localvec.go +++ b/go/plugins/localvec/localvec.go @@ -51,20 +51,12 @@ func Init() error { return nil } // DefineRetriever defines a Retriever and docStore which is also used by the retriever. // The name uniquely identifies the Retriever in the registry. -func DefineRetriever(g *genkit.Genkit, name string, cfg Config) (*DocStore, ai.Retriever, error) { +func DefineRetriever(g *genkit.Genkit, name string, cfg Config, retOpts *ai.RetrieverOptions) (*DocStore, ai.Retriever, error) { ds, err := newDocStore(cfg.Dir, name, cfg.Embedder, cfg.EmbedderOptions) if err != nil { return nil, nil, err } - retOpts := &ai.RetrieverOptions{ - ConfigSchema: RetrieverOptions{}, - Info: &ai.RetrieverInfo{ - Label: "local vec", - Supports: &ai.MediaSupports{ - Media: false, - }, - }, - } + return ds, genkit.DefineRetriever(g, provider, name, retOpts, ds.retrieve), nil @@ -176,7 +168,7 @@ func (ds *DocStore) retrieve(ctx context.Context, req *ai.RetrieverRequest) (*ai k = min(k, len(scoredDocs)) docs := make([]*ai.Document, 0, k) - for i := 0; i < k; i++ { + for i := range k { docs = append(docs, scoredDocs[i].doc) } diff --git a/go/plugins/localvec/localvec_test.go b/go/plugins/localvec/localvec_test.go index fc3d04a209..cec11f20e5 100644 --- a/go/plugins/localvec/localvec_test.go +++ b/go/plugins/localvec/localvec_test.go @@ -232,7 +232,18 @@ func TestInit(t *testing.T) { t.Fatal(err) } const name = "mystore" - _, ret, err := DefineRetriever(g, name, Config{Embedder: embedder}) + retOpts := &ai.RetrieverOptions{ + ConfigSchema: RetrieverOptions{ + K: 2, + }, + Info: &ai.RetrieverInfo{ + Label: name, + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + _, ret, err := DefineRetriever(g, name, Config{Embedder: embedder}, retOpts) if err != nil { t.Fatal(err) } diff --git a/go/plugins/pinecone/genkit.go b/go/plugins/pinecone/genkit.go index 0e604323e2..bcaa6a9658 100644 --- a/go/plugins/pinecone/genkit.go +++ b/go/plugins/pinecone/genkit.go @@ -88,26 +88,29 @@ type Config struct { TextKey string // Metadata key to use to store document text in Pinecone; the default is "_content". } +type SparseVector struct { + Indices []int `json:"indices"` + Values []float32 `json:"values"` +} + +type PineconeRetrieverOptions struct { + K int `json:"k"` + Namespace string `json:"namespace,omitempty"` + Filter map[string]any `json:"filter,omitempty"` + SparseVector *SparseVector `json:"sparseVector,omitempty"` +} + // DefineRetriever defines a Retriever with the given configuration. -func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config) (*Docstore, ai.Retriever, error) { +func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config, retOpts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { p := genkit.LookupPlugin(g, provider).(*Pinecone) if p == nil { - return nil, nil, errors.New("pinecone plugin not found; did you call genkit.Init with the pinecone plugin?") + return nil, nil, errors.New("pinecone plugin not found; did you call genkit.Init with the pinecone plugin") } ds, err := p.newDocstore(ctx, cfg) if err != nil { return nil, nil, err } - retOpts := &ai.RetrieverOptions{ - ConfigSchema: RetrieverOptions{}, - Info: &ai.RetrieverInfo{ - Label: "local Vec", - Supports: &ai.MediaSupports{ - Media: false, - }, - }, - } return ds, genkit.DefineRetriever(g, provider, cfg.IndexID, retOpts, ds.Retrieve), nil } @@ -153,14 +156,6 @@ func Retriever(g *genkit.Genkit, name string) ai.Retriever { return genkit.LookupRetriever(g, provider, name) } -// RetrieverOptions may be passed in the Options field -// of [ai.RetrieverRequest] to pass options to Pinecone. -// The Options field should be either nil or a value of type *RetrieverOptions. -type RetrieverOptions struct { - Namespace string `json:"namespace,omitempty"` // Pinecone namespace to use - Count int `json:"count,omitempty"` // maximum number of values to retrieve -} - // Docstore implements the genkit [ai.DocumentStore] interface. type Docstore struct { Index *index @@ -263,12 +258,12 @@ func (ds *Docstore) Retrieve(ctx context.Context, req *ai.RetrieverRequest) (*ai // TODO: This is plausible when called directly // from Go code, but what will it look like when // run from a resumed flow? - ropt, ok := req.Options.(*RetrieverOptions) + ropt, ok := req.Options.(*PineconeRetrieverOptions) if !ok { - return nil, fmt.Errorf("pinecone.Retrieve options have type %T, want %T", req.Options, &RetrieverOptions{}) + return nil, fmt.Errorf("pinecone.Retrieve options have type %T, want %T", req.Options, &PineconeRetrieverOptions{}) } namespace = ropt.Namespace - count = ropt.Count + count = ropt.K } // Use the embedder to convert the document we want to diff --git a/go/plugins/pinecone/genkit_test.go b/go/plugins/pinecone/genkit_test.go index d0914bd3ec..f5ab71c87e 100644 --- a/go/plugins/pinecone/genkit_test.go +++ b/go/plugins/pinecone/genkit_test.go @@ -39,7 +39,7 @@ func TestGenkit(t *testing.T) { ctx := context.Background() - g, err := genkit.Init(context.Background()) + g, err := genkit.Init(context.Background(), genkit.WithPlugins(&Pinecone{APIKey: *testAPIKey})) if err != nil { t.Fatal(err) } @@ -99,7 +99,19 @@ func TestGenkit(t *testing.T) { Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", emdOpts, embedder.Embed), } - ds, retriever, err := DefineRetriever(ctx, g, cfg) + retOpts := &ai.RetrieverOptions{ + ConfigSchema: PineconeRetrieverOptions{ + K: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "embedder3", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + + ds, retriever, err := DefineRetriever(ctx, g, cfg, retOpts) if err != nil { t.Fatal(err) } @@ -133,8 +145,8 @@ func TestGenkit(t *testing.T) { } }() - retrieverOptions := &RetrieverOptions{ - Count: 2, + retrieverOptions := &PineconeRetrieverOptions{ + K: 2, Namespace: namespace, } retrieverResp, err := ai.Retrieve(ctx, retriever, diff --git a/go/plugins/weaviate/weaviate.go b/go/plugins/weaviate/weaviate.go index 6dc4b413ef..0f1f8b5edd 100644 --- a/go/plugins/weaviate/weaviate.go +++ b/go/plugins/weaviate/weaviate.go @@ -147,7 +147,7 @@ type ClassConfig struct { // DefineRetriever defines [ai.Retriever] // that use the same class. // The name uniquely identifies the Retriever in the registry. -func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig) (*Docstore, ai.Retriever, error) { +func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig, retOpts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { if cfg.Embedder == nil { return nil, nil, errors.New("weaviate: Embedder required") } @@ -164,7 +164,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig) (*D if err != nil { return nil, nil, err } - retriever := genkit.DefineRetriever(g, provider, cfg.Class, ds.Retrieve) + retriever := genkit.DefineRetriever(g, provider, cfg.Class, retOpts, ds.Retrieve) return ds, retriever, nil } diff --git a/go/plugins/weaviate/weaviate_test.go b/go/plugins/weaviate/weaviate_test.go index 4e587228d8..f6fb01283e 100644 --- a/go/plugins/weaviate/weaviate_test.go +++ b/go/plugins/weaviate/weaviate_test.go @@ -104,7 +104,18 @@ func TestGenkit(t *testing.T) { Class: *testClass, Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", emdOpts, embedder.Embed), } - ds, retriever, err := DefineRetriever(ctx, g, classCfg) + retOpts := &ai.RetrieverOptions{ + ConfigSchema: RetrieverOptions{ + Count: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "weaviate", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + ds, retriever, err := DefineRetriever(ctx, g, classCfg, retOpts) if err != nil { t.Fatal(err) } diff --git a/go/samples/menu/main.go b/go/samples/menu/main.go index f8edba1838..1b9202fed7 100644 --- a/go/samples/menu/main.go +++ b/go/samples/menu/main.go @@ -20,6 +20,7 @@ import ( "context" "log" + "github.com/firebase/genkit/go/ai" "github.com/firebase/genkit/go/genkit" "github.com/firebase/genkit/go/plugins/googlegenai" "github.com/firebase/genkit/go/plugins/localvec" @@ -82,9 +83,20 @@ func main() { if err != nil { log.Fatal(err) } + retOpts := &ai.RetrieverOptions{ + ConfigSchema: localvec.RetrieverOptions{ + K: 3, + }, + Info: &ai.RetrieverInfo{ + Label: "go-menu_items", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } docStore, retriever, err := localvec.DefineRetriever(g, "go-menu_items", localvec.Config{ Embedder: embedder, - }) + }, retOpts) if err != nil { log.Fatal(err) } diff --git a/go/samples/pgvector/main.go b/go/samples/pgvector/main.go index 7d87f4fa2b..e8b427b0e5 100644 --- a/go/samples/pgvector/main.go +++ b/go/samples/pgvector/main.go @@ -90,7 +90,16 @@ func run(g *genkit.Genkit) error { } // [START use-retr] - retriever := defineRetriever(g, db, embedder) + retOpts := &ai.RetrieverOptions{ + ConfigSchema: nil, + Info: &ai.RetrieverInfo{ + Label: "pgVector", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + retriever := defineRetriever(g, db, embedder, retOpts) type input struct { Question string @@ -119,7 +128,7 @@ func run(g *genkit.Genkit) error { const provider = "pgvector" // [START retr] -func defineRetriever(g *genkit.Genkit, db *sql.DB, embedder ai.Embedder) ai.Retriever { +func defineRetriever(g *genkit.Genkit, db *sql.DB, embedder ai.Embedder, retOpts *ai.RetrieverOptions) ai.Retriever { f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) { eres, err := ai.Embed(ctx, embedder, ai.WithDocs(req.Query)) if err != nil { @@ -159,7 +168,7 @@ func defineRetriever(g *genkit.Genkit, db *sql.DB, embedder ai.Embedder) ai.Retr } return res, nil } - return genkit.DefineRetriever(g, provider, "shows", f) + return genkit.DefineRetriever(g, provider, "shows", retOpts, f) } // [END retr] diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go index aaa67c11d3..bb156a155a 100644 --- a/go/samples/rag/main.go +++ b/go/samples/rag/main.go @@ -47,7 +47,6 @@ import ( "github.com/firebase/genkit/go/plugins/evaluators" "github.com/firebase/genkit/go/plugins/googlegenai" "github.com/firebase/genkit/go/plugins/localvec" - "github.com/firebase/genkit/go/plugins/pinecone" ) const simpleQaPromptTemplate = ` @@ -83,7 +82,7 @@ func main() { }, } g, err := genkit.Init(ctx, - genkit.WithPlugins(&googlegenai.GoogleAI{}, &evaluators.GenkitEval{Metrics: metrics}, &pinecone.Pinecone{}), + genkit.WithPlugins(&googlegenai.GoogleAI{}, &evaluators.GenkitEval{Metrics: metrics}), ) if err != nil { log.Fatal(err) @@ -97,7 +96,18 @@ func main() { if err := localvec.Init(); err != nil { log.Fatal(err) } - docStore, retriever, err := localvec.DefineRetriever(g, "simpleQa", localvec.Config{Embedder: embedder}) + retOpts := &ai.RetrieverOptions{ + ConfigSchema: localvec.RetrieverOptions{ + K: 1, + }, + Info: &ai.RetrieverInfo{ + Label: "simpleQa", + Supports: &ai.MediaSupports{ + Media: false, + }, + }, + } + docStore, retriever, err := localvec.DefineRetriever(g, "simpleQa", localvec.Config{Embedder: embedder}, retOpts) if err != nil { log.Fatal(err) } From 4016989b4aeda78209fe7c2e06ad5209230f0f6c Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Wed, 11 Jun 2025 21:55:50 +0530 Subject: [PATCH 5/9] change in config schema --- go/internal/doc-snippets/pinecone.go | 5 +---- go/internal/doc-snippets/rag/main.go | 12 +++--------- go/plugins/googlegenai/gemini.go | 16 +++------------- go/plugins/localvec/localvec.go | 4 ++++ go/plugins/localvec/localvec_test.go | 4 +--- go/plugins/pinecone/genkit_test.go | 4 +--- go/plugins/weaviate/weaviate_test.go | 4 +--- go/samples/menu/main.go | 4 +--- go/samples/rag/main.go | 7 +++---- 9 files changed, 18 insertions(+), 42 deletions(-) diff --git a/go/internal/doc-snippets/pinecone.go b/go/internal/doc-snippets/pinecone.go index 6f9cc474c3..f527ee49cc 100644 --- a/go/internal/doc-snippets/pinecone.go +++ b/go/internal/doc-snippets/pinecone.go @@ -49,10 +49,7 @@ func pineconeEx(ctx context.Context) error { // [START defineretriever] retOpts := &ai.RetrieverOptions{ - ConfigSchema: pinecone.PineconeRetrieverOptions{ - K: 1, - Namespace: "Test", - }, + ConfigSchema: pinecone.PineconeRetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "Pinecone", Supports: &ai.MediaSupports{ diff --git a/go/internal/doc-snippets/rag/main.go b/go/internal/doc-snippets/rag/main.go index 3319a351c8..b2f4553fd7 100644 --- a/go/internal/doc-snippets/rag/main.go +++ b/go/internal/doc-snippets/rag/main.go @@ -50,9 +50,7 @@ func main() { log.Fatal(err) } retOpts := &ai.RetrieverOptions{ - ConfigSchema: localvec.RetrieverOptions{ - K: 3, - }, + ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", Supports: &ai.MediaSupports{ @@ -168,9 +166,7 @@ func menuQA() { model := googlegenai.VertexAIModel(g, "gemini-1.5-flash") retOpts := &ai.RetrieverOptions{ - ConfigSchema: localvec.RetrieverOptions{ - K: 3, - }, + ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", Supports: &ai.MediaSupports{ @@ -233,9 +229,7 @@ func customret() { } retOpts := &ai.RetrieverOptions{ - ConfigSchema: localvec.RetrieverOptions{ - K: 3, - }, + ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", Supports: &ai.MediaSupports{ diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 0abc297e51..b2c115cefd 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -295,18 +295,8 @@ func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOp var content []*genai.Content var embedConfig *genai.EmbedContentConfig - if req.Options != nil { - if optsMap, ok := req.Options.(map[string]any); ok { - optionsBytes, err := json.Marshal(optsMap) - if err != nil { - fmt.Println("Error marshaling options map to bytes:", err) - } - var tempConfig genai.EmbedContentConfig - err = json.Unmarshal(optionsBytes, &tempConfig) - if err == nil { - embedConfig = &tempConfig - } - } + if options, _ := req.Options.(*genai.EmbedContentConfig); options != nil { + embedConfig = options } for _, doc := range req.Input { @@ -331,7 +321,7 @@ func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOp }) } -// Generate requests a generate call to the specified model with the provided +// Generate requests generate call to the specified model with the provided // configuration func generate( ctx context.Context, diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go index ba5b811b88..ca42740463 100644 --- a/go/plugins/localvec/localvec.go +++ b/go/plugins/localvec/localvec.go @@ -57,6 +57,10 @@ func DefineRetriever(g *genkit.Genkit, name string, cfg Config, retOpts *ai.Retr return nil, nil, err } + if retOpts != nil && retOpts.ConfigSchema == nil { + retOpts.ConfigSchema = RetrieverOptions{} + } + return ds, genkit.DefineRetriever(g, provider, name, retOpts, ds.retrieve), nil diff --git a/go/plugins/localvec/localvec_test.go b/go/plugins/localvec/localvec_test.go index cec11f20e5..d02ed70a40 100644 --- a/go/plugins/localvec/localvec_test.go +++ b/go/plugins/localvec/localvec_test.go @@ -233,9 +233,7 @@ func TestInit(t *testing.T) { } const name = "mystore" retOpts := &ai.RetrieverOptions{ - ConfigSchema: RetrieverOptions{ - K: 2, - }, + ConfigSchema: RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: name, Supports: &ai.MediaSupports{ diff --git a/go/plugins/pinecone/genkit_test.go b/go/plugins/pinecone/genkit_test.go index f5ab71c87e..069076bf61 100644 --- a/go/plugins/pinecone/genkit_test.go +++ b/go/plugins/pinecone/genkit_test.go @@ -100,9 +100,7 @@ func TestGenkit(t *testing.T) { } retOpts := &ai.RetrieverOptions{ - ConfigSchema: PineconeRetrieverOptions{ - K: 3, - }, + ConfigSchema: PineconeRetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "embedder3", Supports: &ai.MediaSupports{ diff --git a/go/plugins/weaviate/weaviate_test.go b/go/plugins/weaviate/weaviate_test.go index f6fb01283e..ecf8db72b0 100644 --- a/go/plugins/weaviate/weaviate_test.go +++ b/go/plugins/weaviate/weaviate_test.go @@ -105,9 +105,7 @@ func TestGenkit(t *testing.T) { Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", emdOpts, embedder.Embed), } retOpts := &ai.RetrieverOptions{ - ConfigSchema: RetrieverOptions{ - Count: 3, - }, + ConfigSchema: RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "weaviate", Supports: &ai.MediaSupports{ diff --git a/go/samples/menu/main.go b/go/samples/menu/main.go index 1b9202fed7..d988a509d7 100644 --- a/go/samples/menu/main.go +++ b/go/samples/menu/main.go @@ -84,9 +84,7 @@ func main() { log.Fatal(err) } retOpts := &ai.RetrieverOptions{ - ConfigSchema: localvec.RetrieverOptions{ - K: 3, - }, + ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "go-menu_items", Supports: &ai.MediaSupports{ diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go index bb156a155a..8439627d85 100644 --- a/go/samples/rag/main.go +++ b/go/samples/rag/main.go @@ -97,9 +97,7 @@ func main() { log.Fatal(err) } retOpts := &ai.RetrieverOptions{ - ConfigSchema: localvec.RetrieverOptions{ - K: 1, - }, + ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "simpleQa", Supports: &ai.MediaSupports{ @@ -176,7 +174,8 @@ func main() { } dRequest := ai.DocumentFromText(input.Question, nil) - response, err := ai.Retrieve(ctx, retriever, ai.WithDocs(dRequest)) + response, err := ai.Retrieve(ctx, retriever, ai.WithDocs(dRequest), + ai.WithConfig(&localvec.RetrieverOptions{K: 2})) if err != nil { return "", err } From f24d11efb6720fa0bbcab50e1acbad37b8c9c28b Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Mon, 16 Jun 2025 10:20:52 +0530 Subject: [PATCH 6/9] Provide default embedder options in defineEmbedder --- go/plugins/googlegenai/gemini.go | 1 + go/plugins/googlegenai/googlegenai.go | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index e1ea46137b..6a8cc1e548 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -35,6 +35,7 @@ import ( "github.com/firebase/genkit/go/internal" "github.com/firebase/genkit/go/internal/base" "github.com/firebase/genkit/go/plugins/internal/uri" + "github.com/invopop/jsonschema" "google.golang.org/genai" ) diff --git a/go/plugins/googlegenai/googlegenai.go b/go/plugins/googlegenai/googlegenai.go index ee0dc5e7d1..579a80182f 100644 --- a/go/plugins/googlegenai/googlegenai.go +++ b/go/plugins/googlegenai/googlegenai.go @@ -360,7 +360,7 @@ func (ga *GoogleAI) ListActions(ctx context.Context) []core.ActionDesc { func (ga *GoogleAI) ResolveAction(g *genkit.Genkit, atype core.ActionType, name string) error { switch atype { case core.ActionTypeEmbedder: - defineEmbedder(g, ga.gclient, name) + defineEmbedder(g, ga.gclient, name, ai.EmbedderOptions{}) case core.ActionTypeModel: var supports *ai.ModelSupports if strings.Contains(name, "gemini") || strings.Contains(name, "gemma") { @@ -423,7 +423,7 @@ func (v *VertexAI) ListActions(ctx context.Context) []core.ActionDesc { func (v *VertexAI) ResolveAction(g *genkit.Genkit, atype core.ActionType, name string) error { switch atype { case core.ActionTypeEmbedder: - defineEmbedder(g, v.gclient, name) + defineEmbedder(g, v.gclient, name, ai.EmbedderOptions{}) case core.ActionTypeModel: var supports *ai.ModelSupports if strings.Contains(name, "gemini") { From 5510c0a298a38a4701c7d7d7a5819e72b88adba2 Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Fri, 27 Jun 2025 17:02:54 +0530 Subject: [PATCH 7/9] Change in embedder and retriver as per the PR comments --- go/ai/embedder.go | 36 ++++++++++++-------------- go/ai/retriever.go | 26 ++++++++++++------- go/genkit/genkit.go | 8 +++--- go/internal/base/json.go | 6 +++++ go/internal/doc-snippets/pinecone.go | 2 +- go/internal/doc-snippets/rag/main.go | 8 +++--- go/plugins/compat_oai/compat_oai.go | 4 +-- go/plugins/compat_oai/openai/openai.go | 2 +- go/plugins/firebase/retriever.go | 2 +- go/plugins/googlegenai/gemini.go | 8 +++--- go/plugins/googlegenai/googlegenai.go | 16 ++++++------ go/plugins/localvec/localvec.go | 8 +++--- go/plugins/localvec/localvec_test.go | 2 +- go/plugins/ollama/embed.go | 8 ++---- go/plugins/pinecone/genkit.go | 4 +-- go/plugins/pinecone/genkit_test.go | 2 +- go/plugins/weaviate/weaviate.go | 4 +-- go/plugins/weaviate/weaviate_test.go | 2 +- go/samples/menu/main.go | 2 +- go/samples/pgvector/main.go | 2 +- go/samples/rag/main.go | 2 +- 21 files changed, 80 insertions(+), 74 deletions(-) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index f795cd57bf..389f5a103c 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -35,55 +35,51 @@ type Embedder interface { // EmbedderInfo represents the structure of the embedder information object. type EmbedderInfo struct { - // Friendly label for this model (e.g. "Google AI - Gemini Pro") + // Label is a user-friendly name for the embedder model (e.g., "Google AI - Gemini Pro"). Label string `json:"label,omitempty"` - + // Supports defines the capabilities of the embedder, such as input types and multilingual support. Supports *EmbedderSupports `json:"supports,omitempty"` - + // Dimensions specifies the number of dimensions in the embedding vector. Dimensions int `json:"dimensions,omitempty"` } // EmbedderSupports represents the supported capabilities of the embedder model. type EmbedderSupports struct { - // Model can input this type of data. - // Expected values could be "text", "image", "video", but the struct + // Input lists the types of data the model can process (e.g., "text", "image", "video"). Input []string `json:"input,omitempty"` - + // Multilingual indicates whether the model supports multiple languages. Multilingual bool `json:"multilingual,omitempty"` } +// EmbedderOptions represents the configuration options for an embedder. type EmbedderOptions struct { - ConfigSchema any `json:"configSchema,omitempty"` - Info *EmbedderInfo `json:"info,omitempty"` + // ConfigSchema defines the schema for the embedder's configuration options. + ConfigSchema any `json:"configSchema,omitempty"` + // Info contains metadata about the embedder, such as its label and capabilities. + Info *EmbedderInfo `json:"info,omitempty"` } // An embedder is used to convert a document to a multidimensional vector. type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}] -func configToMap(config any) map[string]any { - schema := base.InferJSONSchema(config) - result := base.SchemaAsMap(schema) - return result -} - // DefineEmbedder registers the given embed function as an action, and returns an // [Embedder] that runs it. func DefineEmbedder( r *registry.Registry, provider, name string, - options *EmbedderOptions, + opts *EmbedderOptions, embed func(context.Context, *EmbedRequest) (*EmbedResponse, error), ) Embedder { metadata := map[string]any{} metadata["type"] = "embedder" - metadata["info"] = options.Info - if options.ConfigSchema != nil { - metadata["embedder"] = map[string]any{"customOptions": configToMap(options.ConfigSchema)} + metadata["info"] = opts.Info + if opts.ConfigSchema != nil { + metadata["embedder"] = map[string]any{"customOptions": base.ToSchemaMap(opts.ConfigSchema)} } inputSchema := base.InferJSONSchema(EmbedRequest{}) - if inputSchema.Properties != nil && options.ConfigSchema != nil { + if inputSchema.Properties != nil && opts.ConfigSchema != nil { if _, ok := inputSchema.Properties.Get("options"); ok { - inputSchema.Properties.Set("options", base.InferJSONSchema(options.ConfigSchema)) + inputSchema.Properties.Set("options", base.InferJSONSchema(opts.ConfigSchema)) } } return (*embedder)(core.DefineActionWithInputSchema(r, provider, name, core.ActionTypeEmbedder, metadata, inputSchema, embed)) diff --git a/go/ai/retriever.go b/go/ai/retriever.go index b358e1ef6a..2702463f75 100644 --- a/go/ai/retriever.go +++ b/go/ai/retriever.go @@ -22,6 +22,7 @@ import ( "fmt" "github.com/firebase/genkit/go/core" + "github.com/firebase/genkit/go/internal/base" "github.com/firebase/genkit/go/internal/registry" ) @@ -35,30 +36,37 @@ type Retriever interface { Retrieve(ctx context.Context, req *RetrieverRequest) (*RetrieverResponse, error) } +// RetrieverInfo contains metadata about the retriever, such as its label and capabilities. type RetrieverInfo struct { - Label string `json:"label,omitempty"` - Supports *MediaSupports `json:"supports,omitempty"` + // Label is a user-friendly name for the retriever. + Label string `json:"label,omitempty"` + // Supports defines the capabilities of the retriever, such as media support. + Supports *RetrieverSupports `json:"supports,omitempty"` } -type MediaSupports struct { +// RetrieverSupports defines the supported capabilities of the retriever. +type RetrieverSupports struct { + // Media indicates whether the retriever supports media content. Media bool `json:"media,omitempty"` } +// RetrieverOptions represents the configuration options for a retriever. type RetrieverOptions struct { + // ConfigSchema holds the configuration schema for the retriever. ConfigSchema any - Info *RetrieverInfo + // Info contains metadata about the retriever, such as its label and capabilities. + Info *RetrieverInfo } - type retriever core.ActionDef[*RetrieverRequest, *RetrieverResponse, struct{}] // DefineRetriever registers the given retrieve function as an action, and returns a // [Retriever] that runs it. -func DefineRetriever(r *registry.Registry, provider, name string, options *RetrieverOptions, fn RetrieverFunc) Retriever { +func DefineRetriever(r *registry.Registry, provider, name string, opts *RetrieverOptions, fn RetrieverFunc) Retriever { metadata := map[string]any{} metadata["type"] = "retriever" - metadata["info"] = options.Info - if options.ConfigSchema != nil { - metadata["retriever"] = map[string]any{"customOptions": configToMap(options.ConfigSchema)} + metadata["info"] = opts.Info + if opts.ConfigSchema != nil { + metadata["retriever"] = map[string]any{"customOptions": base.InferJSONSchema(opts.ConfigSchema)} } return (*retriever)(core.DefineAction(r, provider, name, core.ActionTypeRetriever, metadata, fn)) } diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 1823647062..d2a7e88194 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -746,8 +746,8 @@ func GenerateData[Out any](ctx context.Context, g *Genkit, opts ...ai.GenerateOp // The `provider` and `name` form the unique identifier. The `ret` function // contains the logic to process an [ai.RetrieverRequest] (containing the query) // and return an [ai.RetrieverResponse] (containing the relevant documents). -func DefineRetriever(g *Genkit, provider, name string, options *ai.RetrieverOptions, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever { - return ai.DefineRetriever(g.reg, provider, name, options, ret) +func DefineRetriever(g *Genkit, provider, name string, opts *ai.RetrieverOptions, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever { + return ai.DefineRetriever(g.reg, provider, name, opts, ret) } // LookupRetriever retrieves a registered [ai.Retriever] by its provider and name. @@ -764,8 +764,8 @@ func LookupRetriever(g *Genkit, provider, name string) ai.Retriever { // The `provider` and `name` form the unique identifier. The `embed` function // contains the logic to process an [ai.EmbedRequest] (containing documents or a query) // and return an [ai.EmbedResponse] (containing the corresponding embeddings). -func DefineEmbedder(g *Genkit, provider string, name string, options *ai.EmbedderOptions, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder { - return ai.DefineEmbedder(g.reg, provider, name, options, embed) +func DefineEmbedder(g *Genkit, provider string, name string, opts *ai.EmbedderOptions, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder { + return ai.DefineEmbedder(g.reg, provider, name, opts, embed) } // LookupEmbedder retrieves a registered [ai.Embedder] by its provider and name. diff --git a/go/internal/base/json.go b/go/internal/base/json.go index 81e0b12214..cc6c39628f 100644 --- a/go/internal/base/json.go +++ b/go/internal/base/json.go @@ -160,3 +160,9 @@ func GetJsonObjectLines(text string) []string { // Return the slice containing the filtered and trimmed lines. return result } + +func ToSchemaMap(config any) map[string]any { + schema := InferJSONSchema(config) + result := SchemaAsMap(schema) + return result +} diff --git a/go/internal/doc-snippets/pinecone.go b/go/internal/doc-snippets/pinecone.go index f527ee49cc..c810b8bd51 100644 --- a/go/internal/doc-snippets/pinecone.go +++ b/go/internal/doc-snippets/pinecone.go @@ -52,7 +52,7 @@ func pineconeEx(ctx context.Context) error { ConfigSchema: pinecone.PineconeRetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "Pinecone", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/internal/doc-snippets/rag/main.go b/go/internal/doc-snippets/rag/main.go index b2f4553fd7..2454aa8d95 100644 --- a/go/internal/doc-snippets/rag/main.go +++ b/go/internal/doc-snippets/rag/main.go @@ -53,7 +53,7 @@ func main() { ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, @@ -169,7 +169,7 @@ func menuQA() { ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, @@ -232,7 +232,7 @@ func customret() { ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "menuQA", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, @@ -256,7 +256,7 @@ func customret() { ConfigSchema: CustomMenuRetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "advancedMenuRetriever", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/plugins/compat_oai/compat_oai.go b/go/plugins/compat_oai/compat_oai.go index 627ac14fee..ee46c3a42e 100644 --- a/go/plugins/compat_oai/compat_oai.go +++ b/go/plugins/compat_oai/compat_oai.go @@ -125,14 +125,14 @@ func (o *OpenAICompatible) DefineModel(g *genkit.Genkit, provider, name string, } // DefineEmbedder defines an embedder with a given name. -func (o *OpenAICompatible) DefineEmbedder(g *genkit.Genkit, provider, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { +func (o *OpenAICompatible) DefineEmbedder(g *genkit.Genkit, provider, name string, embedOpts *ai.EmbedderOptions) (ai.Embedder, error) { o.mu.Lock() defer o.mu.Unlock() if !o.initted { return nil, errors.New("OpenAICompatible.Init not called") } - return genkit.DefineEmbedder(g, provider, name, &embedOptions, func(ctx context.Context, input *ai.EmbedRequest) (*ai.EmbedResponse, error) { + return genkit.DefineEmbedder(g, provider, name, embedOpts, func(ctx context.Context, input *ai.EmbedRequest) (*ai.EmbedResponse, error) { var data openaiGo.EmbeddingNewParamsInputArrayOfStrings for _, doc := range input.Input { for _, p := range doc.Content { diff --git a/go/plugins/compat_oai/openai/openai.go b/go/plugins/compat_oai/openai/openai.go index 71f6971c93..41d293ccc4 100644 --- a/go/plugins/compat_oai/openai/openai.go +++ b/go/plugins/compat_oai/openai/openai.go @@ -238,7 +238,7 @@ func (o *OpenAI) DefineModel(g *genkit.Genkit, name string, info ai.ModelInfo) ( } func (o *OpenAI) DefineEmbedder(g *genkit.Genkit, name string, modelInfo *ai.EmbedderInfo, configSchema TextEmbeddingConfig) (ai.Embedder, error) { - return o.openAICompatible.DefineEmbedder(g, provider, name, ai.EmbedderOptions{ + return o.openAICompatible.DefineEmbedder(g, provider, name, &ai.EmbedderOptions{ Info: modelInfo, ConfigSchema: configSchema, }) diff --git a/go/plugins/firebase/retriever.go b/go/plugins/firebase/retriever.go index aab0608d0c..42e0fc8c9c 100644 --- a/go/plugins/firebase/retriever.go +++ b/go/plugins/firebase/retriever.go @@ -132,7 +132,7 @@ func defineFirestoreRetriever(g *genkit.Genkit, cfg RetrieverOptions, client *fi ConfigSchema: cfg, Info: &ai.RetrieverInfo{ Label: cfg.Name, - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index c95c904bad..68ea880a4e 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -186,17 +186,17 @@ func defineModel(g *genkit.Genkit, client *genai.Client, name string, info ai.Mo // DefineEmbedder defines embeddings for the provided contents and embedder // model -func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOptions ai.EmbedderOptions) ai.Embedder { +func defineEmbedder(g *genkit.Genkit, client *genai.Client, name string, embedOpts *ai.EmbedderOptions) ai.Embedder { provider := googleAIProvider if client.ClientConfig().Backend == genai.BackendVertexAI { provider = vertexAIProvider } - if embedOptions.ConfigSchema == nil { - embedOptions.ConfigSchema = genai.EmbedContentConfig{} + if embedOpts.ConfigSchema == nil { + embedOpts.ConfigSchema = genai.EmbedContentConfig{} } - return genkit.DefineEmbedder(g, provider, name, &embedOptions, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { + return genkit.DefineEmbedder(g, provider, name, embedOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { var content []*genai.Content var embedConfig *genai.EmbedContentConfig diff --git a/go/plugins/googlegenai/googlegenai.go b/go/plugins/googlegenai/googlegenai.go index 579a80182f..f5c29183c8 100644 --- a/go/plugins/googlegenai/googlegenai.go +++ b/go/plugins/googlegenai/googlegenai.go @@ -112,7 +112,7 @@ func (ga *GoogleAI) Init(ctx context.Context, g *genkit.Genkit) (err error) { return err } for e, eOpts := range embedders { - defineEmbedder(g, ga.gclient, e, eOpts) + defineEmbedder(g, ga.gclient, e, &eOpts) } return nil @@ -186,7 +186,7 @@ func (v *VertexAI) Init(ctx context.Context, g *genkit.Genkit) (err error) { return err } for e, eOpts := range embedders { - defineEmbedder(g, v.gclient, e, eOpts) + defineEmbedder(g, v.gclient, e, &eOpts) } return nil @@ -251,23 +251,23 @@ func (v *VertexAI) DefineModel(g *genkit.Genkit, name string, info *ai.ModelInfo } // DefineEmbedder defines an embedder with a given name. -func (ga *GoogleAI) DefineEmbedder(g *genkit.Genkit, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { +func (ga *GoogleAI) DefineEmbedder(g *genkit.Genkit, name string, embedOpts *ai.EmbedderOptions) (ai.Embedder, error) { ga.mu.Lock() defer ga.mu.Unlock() if !ga.initted { return nil, errors.New("GoogleAI plugin not initialized") } - return defineEmbedder(g, ga.gclient, name, embedOptions), nil + return defineEmbedder(g, ga.gclient, name, embedOpts), nil } // DefineEmbedder defines an embedder with a given name. -func (v *VertexAI) DefineEmbedder(g *genkit.Genkit, name string, embedOptions ai.EmbedderOptions) (ai.Embedder, error) { +func (v *VertexAI) DefineEmbedder(g *genkit.Genkit, name string, embedOpts *ai.EmbedderOptions) (ai.Embedder, error) { v.mu.Lock() defer v.mu.Unlock() if !v.initted { return nil, errors.New("VertexAI plugin not initialized") } - return defineEmbedder(g, v.gclient, name, embedOptions), nil + return defineEmbedder(g, v.gclient, name, embedOpts), nil } // IsDefinedEmbedder reports whether the named [Embedder] is defined by this plugin. @@ -360,7 +360,7 @@ func (ga *GoogleAI) ListActions(ctx context.Context) []core.ActionDesc { func (ga *GoogleAI) ResolveAction(g *genkit.Genkit, atype core.ActionType, name string) error { switch atype { case core.ActionTypeEmbedder: - defineEmbedder(g, ga.gclient, name, ai.EmbedderOptions{}) + defineEmbedder(g, ga.gclient, name, &ai.EmbedderOptions{}) case core.ActionTypeModel: var supports *ai.ModelSupports if strings.Contains(name, "gemini") || strings.Contains(name, "gemma") { @@ -423,7 +423,7 @@ func (v *VertexAI) ListActions(ctx context.Context) []core.ActionDesc { func (v *VertexAI) ResolveAction(g *genkit.Genkit, atype core.ActionType, name string) error { switch atype { case core.ActionTypeEmbedder: - defineEmbedder(g, v.gclient, name, ai.EmbedderOptions{}) + defineEmbedder(g, v.gclient, name, &ai.EmbedderOptions{}) case core.ActionTypeModel: var supports *ai.ModelSupports if strings.Contains(name, "gemini") { diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go index ca42740463..29895c8d56 100644 --- a/go/plugins/localvec/localvec.go +++ b/go/plugins/localvec/localvec.go @@ -51,18 +51,18 @@ func Init() error { return nil } // DefineRetriever defines a Retriever and docStore which is also used by the retriever. // The name uniquely identifies the Retriever in the registry. -func DefineRetriever(g *genkit.Genkit, name string, cfg Config, retOpts *ai.RetrieverOptions) (*DocStore, ai.Retriever, error) { +func DefineRetriever(g *genkit.Genkit, name string, cfg Config, opts *ai.RetrieverOptions) (*DocStore, ai.Retriever, error) { ds, err := newDocStore(cfg.Dir, name, cfg.Embedder, cfg.EmbedderOptions) if err != nil { return nil, nil, err } - if retOpts != nil && retOpts.ConfigSchema == nil { - retOpts.ConfigSchema = RetrieverOptions{} + if opts != nil && opts.ConfigSchema == nil { + opts.ConfigSchema = RetrieverOptions{} } return ds, - genkit.DefineRetriever(g, provider, name, retOpts, ds.retrieve), + genkit.DefineRetriever(g, provider, name, opts, ds.retrieve), nil } diff --git a/go/plugins/localvec/localvec_test.go b/go/plugins/localvec/localvec_test.go index d02ed70a40..98044cc969 100644 --- a/go/plugins/localvec/localvec_test.go +++ b/go/plugins/localvec/localvec_test.go @@ -236,7 +236,7 @@ func TestInit(t *testing.T) { ConfigSchema: RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: name, - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/plugins/ollama/embed.go b/go/plugins/ollama/embed.go index 096b8d032e..5519a79a07 100644 --- a/go/plugins/ollama/embed.go +++ b/go/plugins/ollama/embed.go @@ -128,17 +128,13 @@ func concatenateText(doc *ai.Document) string { } // DefineEmbedder defines an embedder with a given server address. -func (o *Ollama) DefineEmbedder(g *genkit.Genkit, serverAddress string, model string) ai.Embedder { +func (o *Ollama) DefineEmbedder(g *genkit.Genkit, serverAddress string, model string, embedOpts *ai.EmbedderOptions) ai.Embedder { o.mu.Lock() defer o.mu.Unlock() if !o.initted { panic("ollama.Init not called") } - emdOpts := &ai.EmbedderOptions{ - Info: &ai.EmbedderInfo{}, - ConfigSchema: nil, - } - return genkit.DefineEmbedder(g, provider, serverAddress, emdOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { + return genkit.DefineEmbedder(g, provider, serverAddress, embedOpts, func(ctx context.Context, req *ai.EmbedRequest) (*ai.EmbedResponse, error) { if req.Options == nil { req.Options = &EmbedOptions{Model: model} } diff --git a/go/plugins/pinecone/genkit.go b/go/plugins/pinecone/genkit.go index bcaa6a9658..db2fb1e197 100644 --- a/go/plugins/pinecone/genkit.go +++ b/go/plugins/pinecone/genkit.go @@ -101,7 +101,7 @@ type PineconeRetrieverOptions struct { } // DefineRetriever defines a Retriever with the given configuration. -func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config, retOpts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { +func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config, opts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { p := genkit.LookupPlugin(g, provider).(*Pinecone) if p == nil { return nil, nil, errors.New("pinecone plugin not found; did you call genkit.Init with the pinecone plugin") @@ -111,7 +111,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg Config, retOpts if err != nil { return nil, nil, err } - return ds, genkit.DefineRetriever(g, provider, cfg.IndexID, retOpts, ds.Retrieve), nil + return ds, genkit.DefineRetriever(g, provider, cfg.IndexID, opts, ds.Retrieve), nil } // IsDefinedRetriever reports whether the named [Retriever] is defined by this plugin. diff --git a/go/plugins/pinecone/genkit_test.go b/go/plugins/pinecone/genkit_test.go index 069076bf61..0249f85df9 100644 --- a/go/plugins/pinecone/genkit_test.go +++ b/go/plugins/pinecone/genkit_test.go @@ -103,7 +103,7 @@ func TestGenkit(t *testing.T) { ConfigSchema: PineconeRetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "embedder3", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/plugins/weaviate/weaviate.go b/go/plugins/weaviate/weaviate.go index 0f1f8b5edd..0933712a10 100644 --- a/go/plugins/weaviate/weaviate.go +++ b/go/plugins/weaviate/weaviate.go @@ -147,7 +147,7 @@ type ClassConfig struct { // DefineRetriever defines [ai.Retriever] // that use the same class. // The name uniquely identifies the Retriever in the registry. -func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig, retOpts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { +func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig, opts *ai.RetrieverOptions) (*Docstore, ai.Retriever, error) { if cfg.Embedder == nil { return nil, nil, errors.New("weaviate: Embedder required") } @@ -164,7 +164,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, cfg ClassConfig, ret if err != nil { return nil, nil, err } - retriever := genkit.DefineRetriever(g, provider, cfg.Class, retOpts, ds.Retrieve) + retriever := genkit.DefineRetriever(g, provider, cfg.Class, opts, ds.Retrieve) return ds, retriever, nil } diff --git a/go/plugins/weaviate/weaviate_test.go b/go/plugins/weaviate/weaviate_test.go index ecf8db72b0..36ef8847d4 100644 --- a/go/plugins/weaviate/weaviate_test.go +++ b/go/plugins/weaviate/weaviate_test.go @@ -108,7 +108,7 @@ func TestGenkit(t *testing.T) { ConfigSchema: RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "weaviate", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/samples/menu/main.go b/go/samples/menu/main.go index d988a509d7..e5e59ec815 100644 --- a/go/samples/menu/main.go +++ b/go/samples/menu/main.go @@ -87,7 +87,7 @@ func main() { ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "go-menu_items", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/samples/pgvector/main.go b/go/samples/pgvector/main.go index e8b427b0e5..5c4067a04d 100644 --- a/go/samples/pgvector/main.go +++ b/go/samples/pgvector/main.go @@ -94,7 +94,7 @@ func run(g *genkit.Genkit) error { ConfigSchema: nil, Info: &ai.RetrieverInfo{ Label: "pgVector", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go index 8439627d85..4e469f4fca 100644 --- a/go/samples/rag/main.go +++ b/go/samples/rag/main.go @@ -100,7 +100,7 @@ func main() { ConfigSchema: localvec.RetrieverOptions{}, Info: &ai.RetrieverInfo{ Label: "simpleQa", - Supports: &ai.MediaSupports{ + Supports: &ai.RetrieverSupports{ Media: false, }, }, From 9c707cce341c09b5a08b14ecbe20a047160cc98b Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Fri, 27 Jun 2025 17:30:02 +0530 Subject: [PATCH 8/9] add nil config for embedder and retriver for postgressql --- go/plugins/postgresql/genkit.go | 2 +- go/plugins/postgresql/genkit_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/go/plugins/postgresql/genkit.go b/go/plugins/postgresql/genkit.go index ed04c66fe1..4eff9fe40c 100644 --- a/go/plugins/postgresql/genkit.go +++ b/go/plugins/postgresql/genkit.go @@ -89,7 +89,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, p *Postgres, cfg *Co return nil, nil, err } - return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, ds.Retrieve), nil + return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, nil, ds.Retrieve), nil } // Retriever returns the retriever with the given index name. diff --git a/go/plugins/postgresql/genkit_test.go b/go/plugins/postgresql/genkit_test.go index 2fd4b2a686..1d0af49946 100644 --- a/go/plugins/postgresql/genkit_test.go +++ b/go/plugins/postgresql/genkit_test.go @@ -170,7 +170,7 @@ func TestPostgres(t *testing.T) { IDColumn: CustomIdColumn, MetadataJSONColumn: CustomMetadataColumn, IgnoreMetadataColumns: []string{"created_at", "updated_at"}, - Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", embedder.Embed), + Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", nil, embedder.Embed), EmbedderOptions: nil, } From 7c54f9700c232f956131c7063e3952bb4a70143e Mon Sep 17 00:00:00 2001 From: Sahdev Garg Date: Fri, 27 Jun 2025 17:43:12 +0530 Subject: [PATCH 9/9] add nil config for embedder and retriver for alloydb --- go/plugins/alloydb/genkit.go | 2 +- go/plugins/alloydb/genkit_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/go/plugins/alloydb/genkit.go b/go/plugins/alloydb/genkit.go index a1189ce1c4..13f5f0f2fd 100644 --- a/go/plugins/alloydb/genkit.go +++ b/go/plugins/alloydb/genkit.go @@ -88,7 +88,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, p *Postgres, cfg *Co return nil, nil, err } - return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, ds.Retrieve), nil + return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, nil, ds.Retrieve), nil } // Retriever returns the retriever with the given index name. diff --git a/go/plugins/alloydb/genkit_test.go b/go/plugins/alloydb/genkit_test.go index 810fcef4ad..d8d4ab9172 100644 --- a/go/plugins/alloydb/genkit_test.go +++ b/go/plugins/alloydb/genkit_test.go @@ -170,7 +170,7 @@ func TestPostgres(t *testing.T) { IDColumn: CustomIdColumn, MetadataJSONColumn: CustomMetadataColumn, IgnoreMetadataColumns: []string{"created_at", "updated_at"}, - Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", embedder.Embed), + Embedder: genkit.DefineEmbedder(g, "fake", "embedder3", nil, embedder.Embed), EmbedderOptions: nil, }