mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-09 07:28:05 -05:00
feat: Add embeddingModel support (#2121)
First part of the implementation to support semantic search in tools. Second part: https://github.com/googleapis/genai-toolbox/pull/2151
This commit is contained in:
29
internal/embeddingmodels/embeddingmodels.go
Normal file
29
internal/embeddingmodels/embeddingmodels.go
Normal file
@@ -0,0 +1,29 @@
|
||||
// Copyright 2026 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 embeddingmodels
|
||||
|
||||
import "context"
|
||||
|
||||
// EmbeddingModelConfig is the interface for configuring embedding models.
|
||||
type EmbeddingModelConfig interface {
|
||||
EmbeddingModelConfigKind() string
|
||||
Initialize(context.Context) (EmbeddingModel, error)
|
||||
}
|
||||
|
||||
type EmbeddingModel interface {
|
||||
EmbeddingModelKind() string
|
||||
ToConfig() EmbeddingModelConfig
|
||||
EmbedParameters(context.Context, []string) ([][]float32, error)
|
||||
}
|
||||
122
internal/embeddingmodels/gemini/gemini.go
Normal file
122
internal/embeddingmodels/gemini/gemini.go
Normal file
@@ -0,0 +1,122 @@
|
||||
// Copyright 2026 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 gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"google.golang.org/genai"
|
||||
)
|
||||
|
||||
const EmbeddingModelKind string = "gemini"
|
||||
|
||||
// validate interface
|
||||
var _ embeddingmodels.EmbeddingModelConfig = Config{}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Model string `yaml:"model" validate:"required"`
|
||||
ApiKey string `yaml:"apiKey"`
|
||||
Dimension int32 `yaml:"dimension"`
|
||||
}
|
||||
|
||||
// Returns the embedding model kind
|
||||
func (cfg Config) EmbeddingModelConfigKind() string {
|
||||
return EmbeddingModelKind
|
||||
}
|
||||
|
||||
// Initialize a Gemini embedding model
|
||||
func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingModel, error) {
|
||||
// Get client configs
|
||||
configs := &genai.ClientConfig{}
|
||||
if cfg.ApiKey != "" {
|
||||
configs.APIKey = cfg.ApiKey
|
||||
}
|
||||
|
||||
// Create new Gemini API client
|
||||
client, err := genai.NewClient(ctx, configs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to create Gemini API client")
|
||||
}
|
||||
|
||||
m := &EmbeddingModel{
|
||||
Config: cfg,
|
||||
Client: client,
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
var _ embeddingmodels.EmbeddingModel = EmbeddingModel{}
|
||||
|
||||
type EmbeddingModel struct {
|
||||
Client *genai.Client
|
||||
Config
|
||||
}
|
||||
|
||||
// Returns the embedding model kind
|
||||
func (m EmbeddingModel) EmbeddingModelKind() string {
|
||||
return EmbeddingModelKind
|
||||
}
|
||||
|
||||
func (m EmbeddingModel) ToConfig() embeddingmodels.EmbeddingModelConfig {
|
||||
return m.Config
|
||||
}
|
||||
|
||||
func (m EmbeddingModel) EmbedParameters(ctx context.Context, parameters []string) ([][]float32, error) {
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
|
||||
}
|
||||
|
||||
contents := convertStringsToContents(parameters)
|
||||
|
||||
embedConfig := &genai.EmbedContentConfig{
|
||||
TaskType: "SEMANTIC_SIMILARITY",
|
||||
}
|
||||
|
||||
if m.Dimension > 0 {
|
||||
embedConfig.OutputDimensionality = genai.Ptr(m.Dimension)
|
||||
}
|
||||
|
||||
result, err := m.Client.Models.EmbedContent(ctx, m.Model, contents, embedConfig)
|
||||
if err != nil {
|
||||
logger.ErrorContext(ctx, "Error calling EmbedContent for model %s: %v", m.Model, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
embeddings := make([][]float32, 0, len(result.Embeddings))
|
||||
for _, embedding := range result.Embeddings {
|
||||
embeddings = append(embeddings, embedding.Values)
|
||||
}
|
||||
|
||||
logger.InfoContext(ctx, "Successfully embedded %d text parameters using model %s", len(parameters), m.Model)
|
||||
|
||||
return embeddings, nil
|
||||
}
|
||||
|
||||
// convertStringsToContents takes a slice of strings and converts it into a slice of *genai.Content objects.
|
||||
func convertStringsToContents(texts []string) []*genai.Content {
|
||||
contents := make([]*genai.Content, 0, len(texts))
|
||||
|
||||
for _, text := range texts {
|
||||
content := genai.NewContentFromText(text, "")
|
||||
contents = append(contents, content)
|
||||
}
|
||||
return contents
|
||||
}
|
||||
130
internal/embeddingmodels/gemini/gemini_test.go
Normal file
130
internal/embeddingmodels/gemini/gemini_test.go
Normal file
@@ -0,0 +1,130 @@
|
||||
// Copyright 2026 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 gemini_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYamlGemini(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.EmbeddingModelConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
embeddingModels:
|
||||
my-gemini-model:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
`,
|
||||
want: map[string]embeddingmodels.EmbeddingModelConfig{
|
||||
"my-gemini-model": gemini.Config{
|
||||
Name: "my-gemini-model",
|
||||
Kind: gemini.EmbeddingModelKind,
|
||||
Model: "text-embedding-004",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "full example with optional fields",
|
||||
in: `
|
||||
embeddingModels:
|
||||
complex-gemini:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
apiKey: "test-api-key"
|
||||
dimension: 768
|
||||
`,
|
||||
want: map[string]embeddingmodels.EmbeddingModelConfig{
|
||||
"complex-gemini": gemini.Config{
|
||||
Name: "complex-gemini",
|
||||
Kind: gemini.EmbeddingModelKind,
|
||||
Model: "text-embedding-004",
|
||||
ApiKey: "test-api-key",
|
||||
Dimension: 768,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Models server.EmbeddingModelConfigs `yaml:"embeddingModels"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if !cmp.Equal(tc.want, got.Models) {
|
||||
t.Fatalf("incorrect parse: %v", cmp.Diff(tc.want, got.Models))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
func TestFailParseFromYamlGemini(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "missing required model field",
|
||||
in: `
|
||||
embeddingModels:
|
||||
bad-model:
|
||||
kind: gemini
|
||||
`,
|
||||
// Removed the specific model name from the prefix to match your output
|
||||
err: "unable to parse as \"gemini\": Key: 'Config.Model' Error:Field validation for 'Model' failed on the 'required' tag",
|
||||
},
|
||||
{
|
||||
desc: "unknown field",
|
||||
in: `
|
||||
embeddingModels:
|
||||
bad-field:
|
||||
kind: gemini
|
||||
model: text-embedding-004
|
||||
invalid_param: true
|
||||
`,
|
||||
// Updated to match the specific line-starting format of your error output
|
||||
err: "unable to parse as \"gemini\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | kind: gemini\n 3 | model: text-embedding-004",
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Models server.EmbeddingModelConfigs `yaml:"embeddingModels"`
|
||||
}{}
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err == nil {
|
||||
t.Fatalf("expect parsing to fail")
|
||||
}
|
||||
if err.Error() != tc.err {
|
||||
t.Fatalf("unexpected error:\ngot: %q\nwant: %q", err.Error(), tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -276,7 +276,7 @@ func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, tools
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(nil, nil, tools, toolsets, prompts, promptsets)
|
||||
resourceManager := resources.NewResourceManager(nil, nil, nil, tools, toolsets, prompts, promptsets)
|
||||
|
||||
server := Server{
|
||||
version: fakeVersionString,
|
||||
|
||||
@@ -21,6 +21,8 @@ import (
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth/google"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -38,6 +40,8 @@ type ServerConfig struct {
|
||||
SourceConfigs SourceConfigs
|
||||
// AuthServiceConfigs defines what sources of authentication are available for tools.
|
||||
AuthServiceConfigs AuthServiceConfigs
|
||||
// EmbeddingModelConfigs defines a models used to embed parameters.
|
||||
EmbeddingModelConfigs EmbeddingModelConfigs
|
||||
// ToolConfigs defines what tools are available.
|
||||
ToolConfigs ToolConfigs
|
||||
// ToolsetConfigs defines what tools are available.
|
||||
@@ -205,6 +209,50 @@ func (c *AuthServiceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(i
|
||||
return nil
|
||||
}
|
||||
|
||||
// EmbeddingModelConfigs is a type used to allow unmarshal of the embedding model config map
|
||||
type EmbeddingModelConfigs map[string]embeddingmodels.EmbeddingModelConfig
|
||||
|
||||
// validate interface
|
||||
var _ yaml.InterfaceUnmarshalerContext = &EmbeddingModelConfigs{}
|
||||
|
||||
func (c *EmbeddingModelConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
|
||||
*c = make(EmbeddingModelConfigs)
|
||||
// Parse the 'kind' fields for each embedding model
|
||||
var raw map[string]util.DelayedUnmarshaler
|
||||
if err := unmarshal(&raw); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for name, u := range raw {
|
||||
// Unmarshal to a general type that ensure it capture all fields
|
||||
var v map[string]any
|
||||
if err := u.Unmarshal(&v); err != nil {
|
||||
return fmt.Errorf("unable to unmarshal embedding model %q: %w", name, err)
|
||||
}
|
||||
|
||||
kind, ok := v["kind"]
|
||||
if !ok {
|
||||
return fmt.Errorf("missing 'kind' field for embedding model %q", name)
|
||||
}
|
||||
|
||||
dec, err := util.NewStrictDecoder(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error creating decoder: %w", err)
|
||||
}
|
||||
switch kind {
|
||||
case gemini.EmbeddingModelKind:
|
||||
actual := gemini.Config{Name: name}
|
||||
if err := dec.DecodeContext(ctx, &actual); err != nil {
|
||||
return fmt.Errorf("unable to parse as %q: %w", kind, err)
|
||||
}
|
||||
(*c)[name] = actual
|
||||
default:
|
||||
return fmt.Errorf("%q is not a valid kind of auth source", kind)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ToolConfigs is a type used to allow unmarshal of the tool configs
|
||||
type ToolConfigs map[string]tools.ToolConfig
|
||||
|
||||
|
||||
@@ -1107,7 +1107,7 @@ func TestStdioSession(t *testing.T) {
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(nil, nil, toolsMap, toolsets, promptsMap, promptsets)
|
||||
resourceManager := resources.NewResourceManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
|
||||
|
||||
server := &Server{
|
||||
version: fakeVersionString,
|
||||
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
@@ -25,30 +26,33 @@ import (
|
||||
|
||||
// ResourceManager contains available resources for the server. Should be initialized with NewResourceManager().
|
||||
type ResourceManager struct {
|
||||
mu sync.RWMutex
|
||||
sources map[string]sources.Source
|
||||
authServices map[string]auth.AuthService
|
||||
tools map[string]tools.Tool
|
||||
toolsets map[string]tools.Toolset
|
||||
prompts map[string]prompts.Prompt
|
||||
promptsets map[string]prompts.Promptset
|
||||
mu sync.RWMutex
|
||||
sources map[string]sources.Source
|
||||
authServices map[string]auth.AuthService
|
||||
embeddingModels map[string]embeddingmodels.EmbeddingModel
|
||||
tools map[string]tools.Tool
|
||||
toolsets map[string]tools.Toolset
|
||||
prompts map[string]prompts.Prompt
|
||||
promptsets map[string]prompts.Promptset
|
||||
}
|
||||
|
||||
func NewResourceManager(
|
||||
sourcesMap map[string]sources.Source,
|
||||
authServicesMap map[string]auth.AuthService,
|
||||
embeddingModelsMap map[string]embeddingmodels.EmbeddingModel,
|
||||
toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset,
|
||||
promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset,
|
||||
|
||||
) *ResourceManager {
|
||||
resourceMgr := &ResourceManager{
|
||||
mu: sync.RWMutex{},
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
prompts: promptsMap,
|
||||
promptsets: promptsetsMap,
|
||||
mu: sync.RWMutex{},
|
||||
sources: sourcesMap,
|
||||
authServices: authServicesMap,
|
||||
embeddingModels: embeddingModelsMap,
|
||||
tools: toolsMap,
|
||||
toolsets: toolsetsMap,
|
||||
prompts: promptsMap,
|
||||
promptsets: promptsetsMap,
|
||||
}
|
||||
|
||||
return resourceMgr
|
||||
@@ -68,6 +72,13 @@ func (r *ResourceManager) GetAuthService(authServiceName string) (auth.AuthServi
|
||||
return authService, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetEmbeddingModel(embeddingModelName string) (embeddingmodels.EmbeddingModel, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
model, ok := r.embeddingModels[embeddingModelName]
|
||||
return model, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetTool(toolName string) (tools.Tool, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
@@ -96,11 +107,12 @@ func (r *ResourceManager) GetPromptset(promptsetName string) (prompts.Promptset,
|
||||
return promptset, ok
|
||||
}
|
||||
|
||||
func (r *ResourceManager) SetResources(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset, promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset) {
|
||||
func (r *ResourceManager) SetResources(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset, promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.sources = sourcesMap
|
||||
r.authServices = authServicesMap
|
||||
r.embeddingModels = embeddingModelsMap
|
||||
r.tools = toolsMap
|
||||
r.toolsets = toolsetsMap
|
||||
r.prompts = promptsMap
|
||||
@@ -117,6 +129,16 @@ func (r *ResourceManager) GetAuthServiceMap() map[string]auth.AuthService {
|
||||
return copiedMap
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetEmbeddingModelMap() map[string]embeddingmodels.EmbeddingModel {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
copiedMap := make(map[string]embeddingmodels.EmbeddingModel, len(r.embeddingModels))
|
||||
for k, v := range r.embeddingModels {
|
||||
copiedMap[k] = v
|
||||
}
|
||||
return copiedMap
|
||||
}
|
||||
|
||||
func (r *ResourceManager) GetToolsMap() map[string]tools.Tool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
@@ -36,6 +37,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
newAuth := map[string]auth.AuthService{"example-auth": nil}
|
||||
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
|
||||
newTools := map[string]tools.Tool{"example-tool": nil}
|
||||
newToolsets := map[string]tools.Toolset{
|
||||
"example-toolset": {
|
||||
@@ -54,7 +56,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
Prompts: []*prompts.Prompt{},
|
||||
},
|
||||
}
|
||||
resMgr := resources.NewResourceManager(newSources, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
resMgr := resources.NewResourceManager(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
|
||||
gotSource, _ := resMgr.GetSource("example-source")
|
||||
if diff := cmp.Diff(gotSource, newSources["example-source"]); diff != "" {
|
||||
@@ -95,7 +97,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resMgr.SetResources(updateSource, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
resMgr.SetResources(updateSource, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
gotSource, _ = resMgr.GetSource("example-source2")
|
||||
if diff := cmp.Diff(gotSource, updateSource["example-source2"]); diff != "" {
|
||||
t.Errorf("error updating server, sources (-want +got):\n%s", diff)
|
||||
|
||||
@@ -30,6 +30,7 @@ import (
|
||||
"github.com/go-chi/cors"
|
||||
"github.com/go-chi/httplog/v2"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||
@@ -56,6 +57,7 @@ type Server struct {
|
||||
func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
map[string]sources.Source,
|
||||
map[string]auth.AuthService,
|
||||
map[string]embeddingmodels.EmbeddingModel,
|
||||
map[string]tools.Tool,
|
||||
map[string]tools.Toolset,
|
||||
map[string]prompts.Prompt,
|
||||
@@ -91,7 +93,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return s, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
sourcesMap[name] = s
|
||||
}
|
||||
@@ -119,7 +121,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return a, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
authServicesMap[name] = a
|
||||
}
|
||||
@@ -129,6 +131,34 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d authServices: %s", len(authServicesMap), strings.Join(authServiceNames, ", ")))
|
||||
|
||||
// Initialize and validate embedding models from configs.
|
||||
embeddingModelsMap := make(map[string]embeddingmodels.EmbeddingModel)
|
||||
for name, ec := range cfg.EmbeddingModelConfigs {
|
||||
em, err := func() (embeddingmodels.EmbeddingModel, error) {
|
||||
_, span := instrumentation.Tracer.Start(
|
||||
ctx,
|
||||
"toolbox/server/embeddingmodel/init",
|
||||
trace.WithAttributes(attribute.String("model_kind", ec.EmbeddingModelConfigKind())),
|
||||
trace.WithAttributes(attribute.String("model_name", name)),
|
||||
)
|
||||
defer span.End()
|
||||
em, err := ec.Initialize(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize embedding model %q: %w", name, err)
|
||||
}
|
||||
return em, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
embeddingModelsMap[name] = em
|
||||
}
|
||||
embeddingModelNames := make([]string, 0, len(embeddingModelsMap))
|
||||
for name := range embeddingModelsMap {
|
||||
embeddingModelNames = append(embeddingModelNames, name)
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d embeddingModels: %s", len(embeddingModelsMap), strings.Join(embeddingModelNames, ", ")))
|
||||
|
||||
// initialize and validate the tools from configs
|
||||
toolsMap := make(map[string]tools.Tool)
|
||||
for name, tc := range cfg.ToolConfigs {
|
||||
@@ -147,7 +177,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return t, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
toolsMap[name] = t
|
||||
}
|
||||
@@ -184,7 +214,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return t, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
toolsetsMap[name] = t
|
||||
}
|
||||
@@ -216,7 +246,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return p, nil
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
promptsMap[name] = p
|
||||
}
|
||||
@@ -253,7 +283,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
return p, err
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, nil, nil, err
|
||||
return nil, nil, nil, nil, nil, nil, nil, err
|
||||
}
|
||||
promptsetsMap[name] = p
|
||||
}
|
||||
@@ -267,7 +297,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
|
||||
}
|
||||
l.InfoContext(ctx, fmt.Sprintf("Initialized %d promptsets: %s", len(promptsetsMap), strings.Join(promptsetNames, ", ")))
|
||||
|
||||
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
|
||||
}
|
||||
|
||||
// NewServer returns a Server object based on provided Config.
|
||||
@@ -320,7 +350,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
httpLogger := httplog.NewLogger("httplog", httpOpts)
|
||||
r.Use(httplog.RequestLogger(httpLogger))
|
||||
|
||||
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
|
||||
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to initialize configs: %w", err)
|
||||
}
|
||||
@@ -330,7 +360,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
|
||||
sseManager := newSseManager(ctx)
|
||||
|
||||
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
|
||||
|
||||
s := &Server{
|
||||
version: cfg.Version,
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/auth"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/log"
|
||||
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
@@ -144,6 +145,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
},
|
||||
}
|
||||
newAuth := map[string]auth.AuthService{"example-auth": nil}
|
||||
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
|
||||
newTools := map[string]tools.Tool{"example-tool": nil}
|
||||
newToolsets := map[string]tools.Toolset{
|
||||
"example-toolset": {
|
||||
@@ -162,7 +164,7 @@ func TestUpdateServer(t *testing.T) {
|
||||
Prompts: []*prompts.Prompt{},
|
||||
},
|
||||
}
|
||||
s.ResourceMgr.SetResources(newSources, newAuth, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
s.ResourceMgr.SetResources(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
|
||||
if err != nil {
|
||||
t.Errorf("error updating server: %s", err)
|
||||
}
|
||||
|
||||
@@ -333,7 +333,7 @@ func TestInvoke(t *testing.T) {
|
||||
{Name: "prompt", Value: "How many accounts who have region in Prague are eligible for loans?"},
|
||||
}
|
||||
|
||||
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil)
|
||||
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
// Invoke the tool
|
||||
result, err := tool.Invoke(ctx, resourceMgr, params, "") // No accessToken needed for ADC client
|
||||
|
||||
Reference in New Issue
Block a user