mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-04-09 03:02:26 -04:00
890914aae0
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>
101 lines
2.6 KiB
Go
101 lines
2.6 KiB
Go
// 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
|
|
}
|