Skip to content

Commit 9c62f31

Browse files
duwenxin99duwenxin99
authored
feat: Add embeddingModel support (#2121)
First part of the implementation to support semantic search in tools. Second part: #2151
1 parent 731a32e commit 9c62f31

17 files changed

Lines changed: 634 additions & 71 deletions

File tree

‎cmd/root.go‎

Lines changed: 36 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
"github.com/fsnotify/fsnotify"
3434
yaml "github.com/goccy/go-yaml"
3535
"github.com/googleapis/genai-toolbox/internal/auth"
36+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
3637
"github.com/googleapis/genai-toolbox/internal/log"
3738
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
3839
"github.com/googleapis/genai-toolbox/internal/prompts"
@@ -385,12 +386,13 @@ func NewCommand(opts ...Option) *Command {
385386
}
386387

387388
type ToolsFile struct {
388-
Sources server.SourceConfigs `yaml:"sources"`
389-
AuthSources server.AuthServiceConfigs `yaml:"authSources"` // Deprecated: Kept for compatibility.
390-
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
391-
Tools server.ToolConfigs `yaml:"tools"`
392-
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
393-
Prompts server.PromptConfigs `yaml:"prompts"`
389+
Sources server.SourceConfigs `yaml:"sources"`
390+
AuthSources server.AuthServiceConfigs `yaml:"authSources"` // Deprecated: Kept for compatibility.
391+
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
392+
EmbeddingModels server.EmbeddingModelConfigs `yaml:"embeddingModels"`
393+
Tools server.ToolConfigs `yaml:"tools"`
394+
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
395+
Prompts server.PromptConfigs `yaml:"prompts"`
394396
}
395397

396398
// parseEnv replaces environment variables ${ENV_NAME} with their values.
@@ -439,11 +441,12 @@ func parseToolsFile(ctx context.Context, raw []byte) (ToolsFile, error) {
439441
// All resource names (sources, authServices, tools, toolsets) must be unique across all files.
440442
func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
441443
merged := ToolsFile{
442-
Sources: make(server.SourceConfigs),
443-
AuthServices: make(server.AuthServiceConfigs),
444-
Tools: make(server.ToolConfigs),
445-
Toolsets: make(server.ToolsetConfigs),
446-
Prompts: make(server.PromptConfigs),
444+
Sources: make(server.SourceConfigs),
445+
AuthServices: make(server.AuthServiceConfigs),
446+
EmbeddingModels: make(server.EmbeddingModelConfigs),
447+
Tools: make(server.ToolConfigs),
448+
Toolsets: make(server.ToolsetConfigs),
449+
Prompts: make(server.PromptConfigs),
447450
}
448451

449452
var conflicts []string
@@ -479,6 +482,15 @@ func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
479482
}
480483
}
481484

485+
// Check for conflicts and merge embeddingModels
486+
for name, model := range file.EmbeddingModels {
487+
if _, exists := merged.EmbeddingModels[name]; exists {
488+
conflicts = append(conflicts, fmt.Sprintf("embedding model '%s' (file #%d)", name, fileIndex+1))
489+
} else {
490+
merged.EmbeddingModels[name] = model
491+
}
492+
}
493+
482494
// Check for conflicts and merge tools
483495
for name, tool := range file.Tools {
484496
if _, exists := merged.Tools[name]; exists {
@@ -583,22 +595,22 @@ func handleDynamicReload(ctx context.Context, toolsFile ToolsFile, s *server.Ser
583595
panic(err)
584596
}
585597

586-
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
598+
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
587599
if err != nil {
588600
errMsg := fmt.Errorf("unable to validate reloaded edits: %w", err)
589601
logger.WarnContext(ctx, errMsg.Error())
590602
return err
591603
}
592604

593-
s.ResourceMgr.SetResources(sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
605+
s.ResourceMgr.SetResources(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
594606

595607
return nil
596608
}
597609

598610
// validateReloadEdits checks that the reloaded tools file configs can initialized without failing
599611
func validateReloadEdits(
600612
ctx context.Context, toolsFile ToolsFile,
601-
) (map[string]sources.Source, map[string]auth.AuthService, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
613+
) (map[string]sources.Source, map[string]auth.AuthService, map[string]embeddingmodels.EmbeddingModel, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
602614
) {
603615
logger, err := util.LoggerFromContext(ctx)
604616
if err != nil {
@@ -616,22 +628,23 @@ func validateReloadEdits(
616628
defer span.End()
617629

618630
reloadedConfig := server.ServerConfig{
619-
Version: versionString,
620-
SourceConfigs: toolsFile.Sources,
621-
AuthServiceConfigs: toolsFile.AuthServices,
622-
ToolConfigs: toolsFile.Tools,
623-
ToolsetConfigs: toolsFile.Toolsets,
624-
PromptConfigs: toolsFile.Prompts,
631+
Version: versionString,
632+
SourceConfigs: toolsFile.Sources,
633+
AuthServiceConfigs: toolsFile.AuthServices,
634+
EmbeddingModelConfigs: toolsFile.EmbeddingModels,
635+
ToolConfigs: toolsFile.Tools,
636+
ToolsetConfigs: toolsFile.Toolsets,
637+
PromptConfigs: toolsFile.Prompts,
625638
}
626639

627-
sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
640+
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
628641
if err != nil {
629642
errMsg := fmt.Errorf("unable to initialize reloaded configs: %w", err)
630643
logger.WarnContext(ctx, errMsg.Error())
631-
return nil, nil, nil, nil, nil, nil, err
644+
return nil, nil, nil, nil, nil, nil, nil, err
632645
}
633646

634-
return sourcesMap, authServicesMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
647+
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
635648
}
636649

637650
// watchChanges checks for changes in the provided yaml tools file(s) or folder.

‎cmd/root_test.go‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
"github.com/google/go-cmp/cmp"
3333

3434
"github.com/googleapis/genai-toolbox/internal/auth/google"
35+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels/gemini"
3536
"github.com/googleapis/genai-toolbox/internal/log"
3637
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
3738
"github.com/googleapis/genai-toolbox/internal/prompts"
@@ -1830,9 +1831,10 @@ func TestFileLoadingErrors(t *testing.T) {
18301831

18311832
func TestMergeToolsFiles(t *testing.T) {
18321833
file1 := ToolsFile{
1833-
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
1834-
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}},
1835-
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}},
1834+
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
1835+
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}},
1836+
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}},
1837+
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
18361838
}
18371839
file2 := ToolsFile{
18381840
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
@@ -1854,11 +1856,12 @@ func TestMergeToolsFiles(t *testing.T) {
18541856
name: "merge two distinct files",
18551857
files: []ToolsFile{file1, file2},
18561858
want: ToolsFile{
1857-
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
1858-
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
1859-
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}, "tool2": http.Config{Name: "tool2"}},
1860-
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}, "set2": tools.ToolsetConfig{Name: "set2"}},
1861-
Prompts: server.PromptConfigs{},
1859+
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
1860+
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
1861+
Tools: server.ToolConfigs{"tool1": http.Config{Name: "tool1"}, "tool2": http.Config{Name: "tool2"}},
1862+
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}, "set2": tools.ToolsetConfig{Name: "set2"}},
1863+
Prompts: server.PromptConfigs{},
1864+
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
18621865
},
18631866
wantErr: false,
18641867
},
@@ -1871,22 +1874,24 @@ func TestMergeToolsFiles(t *testing.T) {
18711874
name: "merge single file",
18721875
files: []ToolsFile{file1},
18731876
want: ToolsFile{
1874-
Sources: file1.Sources,
1875-
AuthServices: make(server.AuthServiceConfigs),
1876-
Tools: file1.Tools,
1877-
Toolsets: file1.Toolsets,
1878-
Prompts: server.PromptConfigs{},
1877+
Sources: file1.Sources,
1878+
AuthServices: make(server.AuthServiceConfigs),
1879+
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
1880+
Tools: file1.Tools,
1881+
Toolsets: file1.Toolsets,
1882+
Prompts: server.PromptConfigs{},
18791883
},
18801884
},
18811885
{
18821886
name: "merge empty list",
18831887
files: []ToolsFile{},
18841888
want: ToolsFile{
1885-
Sources: make(server.SourceConfigs),
1886-
AuthServices: make(server.AuthServiceConfigs),
1887-
Tools: make(server.ToolConfigs),
1888-
Toolsets: make(server.ToolsetConfigs),
1889-
Prompts: server.PromptConfigs{},
1889+
Sources: make(server.SourceConfigs),
1890+
AuthServices: make(server.AuthServiceConfigs),
1891+
EmbeddingModels: make(server.EmbeddingModelConfigs),
1892+
Tools: make(server.ToolConfigs),
1893+
Toolsets: make(server.ToolsetConfigs),
1894+
Prompts: server.PromptConfigs{},
18901895
},
18911896
},
18921897
}

‎docs/en/resources/embeddingModels/_index.md‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
---
2+
title: "EmbeddingModels"
3+
type: docs
4+
weight: 2
5+
description: >
6+
EmbeddingModels represent services that transform text into vector embeddings for semantic search.
7+
---
8+
9+
EmbeddingModels represent services that generate vector representations of text
10+
data. In the MCP Toolbox, these models enable **Semantic Queries**,
11+
allowing [Tools](../tools/) to automatically convert human-readable text into
12+
numerical vectors before using them in a query.
13+
14+
This is primarily used in two scenarios:
15+
16+
- **Vector Ingestion**: Converting a text parameter into a vector string during
17+
an `INSERT` operation.
18+
19+
- **Semantic Search**: Converting a natural language query into a vector to
20+
perform similarity searches.
21+
22+
## Example
23+
24+
The following configuration defines an embedding model and applies it to
25+
specific tool parameters.
26+
27+
{{< notice tip >}}
28+
Use environment variable replacement with the format ${ENV_NAME}
29+
instead of hardcoding your API keys into the configuration file.
30+
{{< /notice >}}
31+
32+
### Step 1 - Define an Embedding Model
33+
34+
Define an embedding model in the `embeddingModels` section:
35+
36+
```yaml
37+
embeddingModels:
38+
gemini-model: # Name of the embedding model
39+
kind: gemini
40+
model: gemini-embedding-001
41+
apiKey: ${GOOGLE_API_KEY}
42+
dimension: 768
43+
44+
```
45+
46+
### Step 2 - Embed Tool Parameters
47+
48+
Use the defined embedding model, embed your query parameters using the
49+
`embeddedBy` field. Only string-typed
50+
parameters can be embedded:
51+
52+
```yaml
53+
tools:
54+
# Vector ingestion tool
55+
insert_embedding:
56+
kind: postgres-sql
57+
source: my-pg-instance
58+
statement: |
59+
INSERT INTO documents (content, embedding)
60+
VALUES ($1, $2);
61+
parameters:
62+
- name: content
63+
type: string
64+
- name: vector_string
65+
type: string
66+
description: The text to be vectorized and stored.
67+
embeddedBy: gemini-model # refers to the name of a defined embedding model
68+
69+
# Semantic search tool
70+
search_embedding:
71+
kind: postgres-sql
72+
source: my-pg-instance
73+
statement: |
74+
SELECT id, content, embedding <-> $1 AS distance
75+
FROM documents
76+
ORDER BY distance LIMIT 1
77+
parameters:
78+
- name: semantic_search_string
79+
type: string
80+
description: The search query that will be converted to a vector.
81+
embeddedBy: gemini-model # refers to the name of a defined embedding model
82+
```
83+
84+
## Kinds of Embedding Models

‎docs/en/resources/embeddingModels/gemini.md‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
---
2+
title: "Gemini Embedding"
3+
type: docs
4+
weight: 1
5+
description: >
6+
Use Google's Gemini models to generate high-performance text embeddings for vector databases.
7+
---
8+
9+
## About
10+
11+
Google Gemini provides state-of-the-art embedding models that convert text into
12+
high-dimensional vectors.
13+
14+
### Authentication
15+
16+
Toolbox uses your [Application Default Credentials
17+
(ADC)][adc] to authorize with the
18+
Gemini API client.
19+
20+
Optionally, you can use an [API key][api-key] obtain an API
21+
Key from the [Google AI Studio][ai-studio].
22+
23+
We recommend using an API key for testing and using application default
24+
credentials for production.
25+
26+
[adc]: https://cloud.google.com/docs/authentication#adc
27+
[api-key]: https://ai.google.dev/gemini-api/docs/api-key#api-keys
28+
[ai-studio]: https://aistudio.google.com/app/apikey
29+
30+
## Behavior
31+
32+
### Automatic Vectorization
33+
34+
When a tool parameter is configured with `embeddedBy: <your-gemini-model-name>`,
35+
the Toolbox intercepts the raw text input from the client and sends it to the
36+
Gemini API. The resulting numerical array is then formatted before being passed
37+
to your database source.
38+
39+
### Dimension Matching
40+
41+
The `dimension` field must match the expected size of your database column
42+
(e.g., a `vector(768)` column in PostgreSQL). This setting is supported by newer
43+
models since 2024 only. You cannot set this value if using the earlier model
44+
(`models/embedding-001`). Check out [available Gemini models][modellist] for more
45+
information.
46+
47+
[modellist]:
48+
https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models
49+
50+
## Example
51+
52+
```yaml
53+
embeddingModels:
54+
gemini-model:
55+
kind: gemini
56+
model: gemini-embedding-001
57+
apiKey: ${GOOGLE_API_KEY}
58+
dimension: 768
59+
```
60+
61+
{{< notice tip >}}
62+
Use environment variable replacement with the format ${ENV_NAME}
63+
instead of hardcoding your secrets into the configuration file.
64+
{{< /notice >}}
65+
66+
## Reference
67+
68+
| **field** | **type** | **required** | **description** |
69+
|-----------|:--------:|:------------:|--------------------------------------------------------------|
70+
| kind | string | true | Must be `gemini`. |
71+
| model | string | true | The Gemini model ID to use (e.g., `gemini-embedding-001`). |
72+
| apiKey | string | false | Your API Key from Google AI Studio. |
73+
| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). |

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ require (
5959
go.opentelemetry.io/otel/trace v1.38.0
6060
golang.org/x/oauth2 v0.33.0
6161
google.golang.org/api v0.256.0
62+
google.golang.org/genai v1.37.0
6263
google.golang.org/genproto v0.0.0-20251022142026-3a174f9686a8
6364
google.golang.org/protobuf v1.36.10
6465
modernc.org/sqlite v1.40.0

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1869,6 +1869,8 @@ google.golang.org/appengine v1.6.1/go.mod h1:i06prIuMbXzDqacNJfV5OdTW448YApPu5ww
18691869
google.golang.org/appengine v1.6.5/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
18701870
google.golang.org/appengine v1.6.6/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
18711871
google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
1872+
google.golang.org/genai v1.37.0 h1:dgp71k1wQ+/+APdZrN3LFgAGnVnr5IdTF1Oj0Dg+BQc=
1873+
google.golang.org/genai v1.37.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
18721874
google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
18731875
google.golang.org/genproto v0.0.0-20190307195333-5fe7a883aa19/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=
18741876
google.golang.org/genproto v0.0.0-20190418145605-e7d98fc518a7/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=

0 commit comments

Comments
 (0)