-
Notifications
You must be signed in to change notification settings - Fork 56
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add Spanner source and tool (#90)
Add Spanner source and tool.
Spanner source is initialize with the following config:
```
sources:
my-spanner-source:
kind: spanner
project: my-project-name
instance: my-instance-name
database: my_db
# dialect: postgresql # The default dialect is google_standard_sql.
```
Spanner tool (with gsql dialect) is initialize with the following
config.
```
tools:
get_flight_by_id:
kind: spanner
source: my-cloud-sql-source
description: >
Use this tool to list all airports matching search criteria. Takes
at least one of country, city, name, or all and returns all matching
airports. The agent can decide to return the results directly to
the user.
statement: "SELECT * FROM flights WHERE id = @id"
parameters:
- name: id
type: int
description: 'id' represents the unique ID for each flight.
```
Spanner tool (with postgresql dialect) is initialize with the following
config.
```
tools:
get_flight_by_id:
kind: spanner
source: my-cloud-sql-source
description: >
Use this tool to list all airports matching search criteria. Takes
at least one of country, city, name, or all and returns all matching
airports. The agent can decide to return the results directly to
the user.
statement: "SELECT * FROM flights WHERE id = $1"
parameters:
- name: id
type: int
description: 'id' represents the unique ID for each flight.
```
Note: the only difference in config for both dialects is the sql
statement.
---------
Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>- Loading branch information
Showing
10 changed files
with
2,034 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| // Copyright 2024 Google LLC | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package sources | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "strings" | ||
|
|
||
| "gopkg.in/yaml.v3" | ||
| ) | ||
|
|
||
| // Dialect represents the dialect type of a database. | ||
| type Dialect string | ||
|
|
||
| func (i *Dialect) String() string { | ||
| if string(*i) != "" { | ||
| return strings.ToLower(string(*i)) | ||
| } | ||
| return "google_standard_sql" | ||
| } | ||
|
|
||
| func (i *Dialect) UnmarshalYAML(node *yaml.Node) error { | ||
| var dialect string | ||
| if err := node.Decode(&dialect); err != nil { | ||
| return err | ||
| } | ||
| switch strings.ToLower(dialect) { | ||
| case "google_standard_sql", "postgresql": | ||
| *i = Dialect(strings.ToLower(dialect)) | ||
| return nil | ||
| default: | ||
| return fmt.Errorf(`dialect invalid: must be one of "google_standard_sql", or "postgresql"`) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| // Copyright 2024 Google LLC | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package spanner | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "cloud.google.com/go/spanner" | ||
| "github.com/googleapis/genai-toolbox/internal/sources" | ||
| ) | ||
|
|
||
| const SourceKind string = "spanner" | ||
|
|
||
| // validate interface | ||
| var _ sources.SourceConfig = Config{} | ||
|
|
||
| type Config struct { | ||
| Name string `yaml:"name"` | ||
| Kind string `yaml:"kind"` | ||
| Project string `yaml:"project"` | ||
| Instance string `yaml:"instance"` | ||
| Dialect sources.Dialect `yaml:"dialect"` | ||
| Database string `yaml:"database"` | ||
| } | ||
|
|
||
| func (r Config) SourceConfigKind() string { | ||
| return SourceKind | ||
| } | ||
|
|
||
| func (r Config) Initialize() (sources.Source, error) { | ||
| client, err := initSpannerClient(r.Project, r.Instance, r.Database) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("unable to create client: %w", err) | ||
| } | ||
|
|
||
| s := &Source{ | ||
| Name: r.Name, | ||
| Kind: SourceKind, | ||
| Client: client, | ||
| Dialect: r.Dialect.String(), | ||
| } | ||
| return s, nil | ||
| } | ||
|
|
||
| var _ sources.Source = &Source{} | ||
|
|
||
| type Source struct { | ||
| Name string `yaml:"name"` | ||
| Kind string `yaml:"kind"` | ||
| Client *spanner.Client | ||
| Dialect string | ||
| } | ||
|
|
||
| func (s *Source) SourceKind() string { | ||
| return SourceKind | ||
| } | ||
|
|
||
| func (s *Source) SpannerClient() *spanner.Client { | ||
| return s.Client | ||
| } | ||
|
|
||
| func (s *Source) DatabaseDialect() string { | ||
| return s.Dialect | ||
| } | ||
|
|
||
| func initSpannerClient(project, instance, dbname string) (*spanner.Client, error) { | ||
| // Configure the connection to the database | ||
| db := fmt.Sprintf("projects/%s/instances/%s/databases/%s", project, instance, dbname) | ||
|
|
||
| // Configure session pool to automatically clean inactive transactions | ||
| sessionPoolConfig := spanner.SessionPoolConfig{ | ||
| TrackSessionHandles: true, | ||
| InactiveTransactionRemovalOptions: spanner.InactiveTransactionRemovalOptions{ | ||
| ActionOnInactiveTransaction: spanner.WarnAndClose, | ||
| }, | ||
| } | ||
|
|
||
| // Create spanner client | ||
| ctx := context.Background() | ||
| client, err := spanner.NewClientWithConfig(ctx, db, spanner.ClientConfig{SessionPoolConfig: sessionPoolConfig}) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("unable to create new client: %w", err) | ||
| } | ||
| defer client.Close() | ||
|
|
||
| return client, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,148 @@ | ||
| // Copyright 2024 Google LLC | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package spanner_test | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| "github.com/google/go-cmp/cmp" | ||
| "github.com/googleapis/genai-toolbox/internal/server" | ||
| "github.com/googleapis/genai-toolbox/internal/sources" | ||
| "github.com/googleapis/genai-toolbox/internal/sources/spanner" | ||
| "github.com/googleapis/genai-toolbox/internal/testutils" | ||
| "gopkg.in/yaml.v3" | ||
| ) | ||
|
|
||
| func TestParseFromYamlSpannerDb(t *testing.T) { | ||
| tcs := []struct { | ||
| desc string | ||
| in string | ||
| want server.SourceConfigs | ||
| }{ | ||
| { | ||
| desc: "basic example", | ||
| in: ` | ||
| sources: | ||
| my-spanner-instance: | ||
| kind: spanner | ||
| project: my-project | ||
| instance: my-instance | ||
| database: my_db | ||
| `, | ||
| want: map[string]sources.SourceConfig{ | ||
| "my-spanner-instance": spanner.Config{ | ||
| Name: "my-spanner-instance", | ||
| Kind: spanner.SourceKind, | ||
| Project: "my-project", | ||
| Instance: "my-instance", | ||
| Dialect: "google_standard_sql", | ||
| Database: "my_db", | ||
| }, | ||
| }, | ||
| }, | ||
| { | ||
| desc: "gsql dialect", | ||
| in: ` | ||
| sources: | ||
| my-spanner-instance: | ||
| kind: spanner | ||
| project: my-project | ||
| instance: my-instance | ||
| dialect: Google_standard_sql | ||
| database: my_db | ||
| `, | ||
| want: map[string]sources.SourceConfig{ | ||
| "my-spanner-instance": spanner.Config{ | ||
| Name: "my-spanner-instance", | ||
| Kind: spanner.SourceKind, | ||
| Project: "my-project", | ||
| Instance: "my-instance", | ||
| Dialect: "google_standard_sql", | ||
| Database: "my_db", | ||
| }, | ||
| }, | ||
| }, | ||
| { | ||
| desc: "postgresql dialect", | ||
| in: ` | ||
| sources: | ||
| my-spanner-instance: | ||
| kind: spanner | ||
| project: my-project | ||
| instance: my-instance | ||
| dialect: postgresql | ||
| database: my_db | ||
| `, | ||
| want: map[string]sources.SourceConfig{ | ||
| "my-spanner-instance": spanner.Config{ | ||
| Name: "my-spanner-instance", | ||
| Kind: spanner.SourceKind, | ||
| Project: "my-project", | ||
| Instance: "my-instance", | ||
| Dialect: "postgresql", | ||
| Database: "my_db", | ||
| }, | ||
| }, | ||
| }, | ||
| } | ||
| for _, tc := range tcs { | ||
| t.Run(tc.desc, func(t *testing.T) { | ||
| got := struct { | ||
| Sources server.SourceConfigs `yaml:"sources"` | ||
| }{} | ||
| // 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.Sources) { | ||
| t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources) | ||
| } | ||
| }) | ||
| } | ||
|
|
||
| } | ||
|
|
||
| func FailParseFromYamlSpanner(t *testing.T) { | ||
| tcs := []struct { | ||
| desc string | ||
| in string | ||
| }{ | ||
| { | ||
| desc: "invalid dialect", | ||
| in: ` | ||
| sources: | ||
| my-spanner-instance: | ||
| kind: spanner | ||
| project: my-project | ||
| instance: my-instance | ||
| dialect: fail | ||
| database: my_db | ||
| `, | ||
| }, | ||
| } | ||
| for _, tc := range tcs { | ||
| t.Run(tc.desc, func(t *testing.T) { | ||
| got := struct { | ||
| Sources server.SourceConfigs `yaml:"sources"` | ||
| }{} | ||
| // Parse contents | ||
| err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) | ||
| if err == nil { | ||
| t.Fatalf("expect parsing to fail: %s", err) | ||
| } | ||
| }) | ||
| } | ||
| } |
Oops, something went wrong.