Skip to content

Commit 17b70cc

Browse files
duwenxin99duwenxin99
authored
feat(tools/postgressql): Add Parameter embeddedBy config support (googleapis#2151)
Add parameter `embeddedBy` field to support vector embedding & semantic search. Major change in `internal/util/parameters/parameters.go` This PR only adds vector formatter for the postgressql tool. Other tools requiring vector formatting may not work with embeddedBy. Second part of the Semantic Search support. First part: googleapis#2121
1 parent 001d634 commit 17b70cc

186 files changed

Lines changed: 1044 additions & 3 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cmd/root.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,11 +484,11 @@ func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
484484
}
485485

486486
// Check for conflicts and merge embeddingModels
487-
for name, model := range file.EmbeddingModels {
487+
for name, em := range file.EmbeddingModels {
488488
if _, exists := merged.EmbeddingModels[name]; exists {
489489
conflicts = append(conflicts, fmt.Sprintf("embedding model '%s' (file #%d)", name, fileIndex+1))
490490
} else {
491-
merged.EmbeddingModels[name] = model
491+
merged.EmbeddingModels[name] = em
492492
}
493493
}
494494

‎internal/embeddingmodels/embeddingmodels.go‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,11 @@
1414

1515
package embeddingmodels
1616

17-
import "context"
17+
import (
18+
"context"
19+
"strconv"
20+
"strings"
21+
)
1822

1923
// EmbeddingModelConfig is the interface for configuring embedding models.
2024
type EmbeddingModelConfig interface {
@@ -27,3 +31,29 @@ type EmbeddingModel interface {
2731
ToConfig() EmbeddingModelConfig
2832
EmbedParameters(context.Context, []string) ([][]float32, error)
2933
}
34+
35+
type VectorFormatter func(vectorFloats []float32) any
36+
37+
// FormatVectorForPgvector converts a slice of floats into a PostgreSQL vector literal string: '[x, y, z]'
38+
func FormatVectorForPgvector(vectorFloats []float32) any {
39+
if len(vectorFloats) == 0 {
40+
return "[]"
41+
}
42+
43+
// Pre-allocate the builder.
44+
var b strings.Builder
45+
b.Grow(len(vectorFloats) * 10)
46+
47+
b.WriteByte('[')
48+
for i, f := range vectorFloats {
49+
if i > 0 {
50+
b.WriteString(", ")
51+
}
52+
b.Write(strconv.AppendFloat(nil, float64(f), 'g', -1, 32))
53+
}
54+
b.WriteByte(']')
55+
56+
return b.String()
57+
}
58+
59+
var _ VectorFormatter = FormatVectorForPgvector

‎internal/server/api.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,14 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
246246
}
247247
s.logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
248248

249+
params, err = tool.EmbedParams(ctx, params, s.ResourceMgr.GetEmbeddingModelMap())
250+
if err != nil {
251+
err = fmt.Errorf("error embedding parameters: %w", err)
252+
s.logger.DebugContext(ctx, err.Error())
253+
_ = render.Render(w, r, newErrResponse(err, http.StatusBadRequest))
254+
return
255+
}
256+
249257
res, err := tool.Invoke(ctx, s.ResourceMgr, params, accessToken)
250258

251259
// Determine what error to return to the users.

‎internal/server/common_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"testing"
2525

2626
"github.com/go-chi/chi/v5"
27+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2728
"github.com/googleapis/genai-toolbox/internal/log"
2829
"github.com/googleapis/genai-toolbox/internal/prompts"
2930
"github.com/googleapis/genai-toolbox/internal/server/resources"
@@ -64,6 +65,10 @@ func (t MockTool) ParseParams(data map[string]any, claimsMap map[string]map[stri
6465
return parameters.ParseParams(t.Params, data, claimsMap)
6566
}
6667

68+
func (t MockTool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
69+
return parameters.EmbedParams(ctx, t.Params, paramValues, embeddingModelsMap, nil)
70+
}
71+
6772
func (t MockTool) Manifest() tools.Manifest {
6873
pMs := make([]parameters.ParameterManifest, 0, len(t.Params))
6974
for _, p := range t.Params {

‎internal/tools/alloydb/alloydbcreatecluster/alloydbcreatecluster.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -166,6 +167,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
166167
return parameters.ParseParams(t.AllParams, data, claims)
167168
}
168169

170+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
171+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
172+
}
173+
169174
// Manifest returns the tool's manifest.
170175
func (t Tool) Manifest() tools.Manifest {
171176
return t.manifest

‎internal/tools/alloydb/alloydbcreateinstance/alloydbcreateinstance.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -172,6 +173,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
172173
return parameters.ParseParams(t.AllParams, data, claims)
173174
}
174175

176+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
177+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
178+
}
179+
175180
// Manifest returns the tool's manifest.
176181
func (t Tool) Manifest() tools.Manifest {
177182
return t.manifest

‎internal/tools/alloydb/alloydbcreateuser/alloydbcreateuser.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -177,6 +178,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
177178
return parameters.ParseParams(t.AllParams, data, claims)
178179
}
179180

181+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
182+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
183+
}
184+
180185
// Manifest returns the tool's manifest.
181186
func (t Tool) Manifest() tools.Manifest {
182187
return t.manifest

‎internal/tools/alloydb/alloydbgetcluster/alloydbgetcluster.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -148,6 +149,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
148149
return parameters.ParseParams(t.AllParams, data, claims)
149150
}
150151

152+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
153+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
154+
}
155+
151156
// Manifest returns the tool's manifest.
152157
func (t Tool) Manifest() tools.Manifest {
153158
return t.manifest

‎internal/tools/alloydb/alloydbgetinstance/alloydbgetinstance.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -152,6 +153,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
152153
return parameters.ParseParams(t.AllParams, data, claims)
153154
}
154155

156+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
157+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
158+
}
159+
155160
// Manifest returns the tool's manifest.
156161
func (t Tool) Manifest() tools.Manifest {
157162
return t.manifest

‎internal/tools/alloydb/alloydbgetuser/alloydbgetuser.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"fmt"
2020

2121
yaml "github.com/goccy/go-yaml"
22+
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
2223
"github.com/googleapis/genai-toolbox/internal/sources"
2324
"github.com/googleapis/genai-toolbox/internal/tools"
2425
"github.com/googleapis/genai-toolbox/internal/util/parameters"
@@ -152,6 +153,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any)
152153
return parameters.ParseParams(t.AllParams, data, claims)
153154
}
154155

156+
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
157+
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
158+
}
159+
155160
// Manifest returns the tool's manifest.
156161
func (t Tool) Manifest() tools.Manifest {
157162
return t.manifest

0 commit comments

Comments
 (0)