mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 16:38:15 -05:00
feat: Add Redis Source and Tool (#519)
1. Added Redis Source and Tool 2. Moved some integration test helpers from tools.go to common.go 3. Make auth integration test want an input variable
This commit is contained in:
Wenxin Du
@@ -332,7 +332,22 @@ steps:
|
||||
- -c
|
||||
- |
|
||||
./couchbase.test -test.v
|
||||
|
||||
|
||||
- id: "redis"
|
||||
name : golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
entrypoint: /bin/bash
|
||||
env:
|
||||
- "GOPATH=/gopath"
|
||||
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
|
||||
secretEnv: ["REDIS_ADDRESS", "REDIS_PASS", "CLIENT_ID"]
|
||||
volumes:
|
||||
- name: "go"
|
||||
path: "/gopath"
|
||||
args:
|
||||
- -c
|
||||
- |
|
||||
./redis.test -test.v
|
||||
|
||||
availableSecrets:
|
||||
secretManager:
|
||||
@@ -380,6 +395,10 @@ availableSecrets:
|
||||
env: COUCHBASE_USER
|
||||
- versionName: projects/$PROJECT_ID/secrets/couchbase_pass/versions/latest
|
||||
env: COUCHBASE_PASS
|
||||
- versionName: projects/$PROJECT_ID/secrets/memorystore_redis_address/versions/latest
|
||||
env: REDIS_ADDRESS
|
||||
- versionName: projects/$PROJECT_ID/secrets/memorystore_redis_pass/versions/latest
|
||||
env: REDIS_PASS
|
||||
|
||||
|
||||
options:
|
||||
|
||||
@@ -53,6 +53,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/neo4j"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgresexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgressql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/redis"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spannerexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql"
|
||||
@@ -72,6 +73,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/mysql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/neo4j"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
|
||||
)
|
||||
|
||||
88
docs/en/resources/sources/redis.md
Normal file
88
docs/en/resources/sources/redis.md
Normal file
@@ -0,0 +1,88 @@
|
||||
---
|
||||
title: "Redis"
|
||||
linkTitle: "Redis"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
Redis is an open-source, in-memory data structure store.
|
||||
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
Redis is an open-source, in-memory data structure store, used as a database, cache, and message broker. It supports data structures such as strings, hashes, lists, sets, sorted sets with range queries, bitmaps, hyperloglogs, and geospatial indexes with radius queries.
|
||||
|
||||
If you are new to Redis, you can find installation and getting started guides on the [official Redis website](https://redis.io/docs/getting-started/).
|
||||
|
||||
## Requirements
|
||||
|
||||
### Redis
|
||||
|
||||
[AUTH string][auth] is a password for connection to Redis. If you have the `requirepass` directive set in your Redis configuration, incoming client connections must authenticate in order to connect.
|
||||
|
||||
Specify your AUTH string in the password field:
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
address: 127.0.0.1
|
||||
username: ${MY_USER_NAME}
|
||||
password: ${MY_AUTH_STRING} # Omit this field if you don't have a password.
|
||||
# database: 0
|
||||
# clusterEnabled: false
|
||||
# useGCPIAM: false
|
||||
```
|
||||
|
||||
{{< notice tip >}}
|
||||
Use environment variable replacement with the format ${ENV_NAME}
|
||||
instead of hardcoding your secrets into the configuration file.
|
||||
{{< /notice >}}
|
||||
|
||||
### Memorystore For Redis
|
||||
|
||||
Memorystore standalone instances support authentication using an [AUTH][auth]
|
||||
string.
|
||||
|
||||
Here is an example tools.yaml config with [AUTH][auth] enabled:
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-redis-cluster-instance:
|
||||
kind: memorystore-redis
|
||||
address: 127.0.0.1
|
||||
password: ${MY_AUTH_STRING}
|
||||
# useGCPIAM: false
|
||||
# clusterEnabled: false
|
||||
```
|
||||
|
||||
Memorystore Redis Cluster supports IAM authentication instead. Grant your account the
|
||||
required [IAM role][iam] and make sure to set `useGCPIAM` to `true`.
|
||||
|
||||
Here is an example tools.yaml config for Memorystore Redis Cluster instances
|
||||
using IAM authentication:
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-redis-cluster-instance:
|
||||
kind: memorystore-redis
|
||||
address: 127.0.0.1
|
||||
useGCPIAM: true
|
||||
clusterEnabled: true
|
||||
```
|
||||
|
||||
[iam]: https://cloud.google.com/memorystore/docs/cluster/about-iam-auth
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|----------------|:--------:|:------------:|---------------------------------------------------------------------------------------------------------------------------------|
|
||||
| kind | string | true | Must be "memorystore-redis". |
|
||||
| address | string | true | Primary endpoint for the Memorystore Redis instance to connect to. |
|
||||
| username | string | false | If you are using a non-default user, specify the user name here. If you are using Memorystore for Redis, leave this field blank |
|
||||
| password | string | false | If you have [Redis AUTH][auth] enabled, specify the AUTH string here |
|
||||
| database | int | false | The Redis database to connect to. Not applicable for cluster enabled instances. The default database is `0`. |
|
||||
| clusterEnabled | bool | false | Set it to `true` if using a Redis Cluster instance. Defaults to `false`. |
|
||||
| useGCPIAM | string | false | Set it to `true` if you are using GCP's IAM authentication. Defaults to `false`. |
|
||||
|
||||
[auth]: https://cloud.google.com/memorystore/docs/redis/about-redis-auth
|
||||
57
docs/en/resources/tools/redis.md
Normal file
57
docs/en/resources/tools/redis.md
Normal file
@@ -0,0 +1,57 @@
|
||||
---
|
||||
title: "redis"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
A "redis" tool executes a set of pre-defined Redis commands against a Redis instance.
|
||||
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
A redis tool executes a series of pre-defined Redis commands against a
|
||||
Redis source.
|
||||
|
||||
The specified Redis commands are executed sequentially. Each command is
|
||||
represented as a string list, where the first element is the command name (e.g., SET,
|
||||
GET, HGETALL) and subsequent elements are its arguments.
|
||||
|
||||
### Dynamic Command Parameters
|
||||
|
||||
Command arguments can be templated using the `$variableName` annotation. The
|
||||
array type parameters will be expanded once into multiple arguments. Take the
|
||||
following config for example:
|
||||
|
||||
```yaml
|
||||
commands:
|
||||
- [SADD, userNames, $userNames] # Array will be flattened into multiple arguments.
|
||||
parameters:
|
||||
- name: userNames
|
||||
type: array
|
||||
description: The user names to be set.
|
||||
```
|
||||
|
||||
If the input is an array of strings `["Alice", "Sid", "Bob"]`, The final command
|
||||
to be executed after argument expansion will be `[SADD, userNames, Alice, Sid, Bob]`.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
user_data_tool:
|
||||
kind: redis
|
||||
source: my-redis-instance
|
||||
description: |
|
||||
Use this tool to interact with user data stored in Redis.
|
||||
It can set, retrieve, and delete user-specific information.
|
||||
commands:
|
||||
- [SADD, userNames, $userNames] # Array will be flattened into multiple arguments.
|
||||
- [GET, $userId]
|
||||
parameters:
|
||||
- name: userId
|
||||
type: string
|
||||
description: The unique identifier for the user.
|
||||
- name: userNames
|
||||
type: array
|
||||
description: The user names to be set.
|
||||
```
|
||||
5
go.mod
5
go.mod
@@ -23,8 +23,10 @@ require (
|
||||
github.com/google/go-cmp v0.7.0
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jackc/pgx/v5 v5.7.5
|
||||
github.com/json-iterator/go v1.1.12
|
||||
github.com/microsoft/go-mssqldb v1.8.2
|
||||
github.com/neo4j/neo4j-go-driver/v5 v5.28.1
|
||||
github.com/redis/go-redis/v9 v9.9.0
|
||||
github.com/spf13/cobra v1.9.1
|
||||
go.opentelemetry.io/contrib/propagators/autoprop v0.61.0
|
||||
go.opentelemetry.io/otel v1.36.0
|
||||
@@ -66,6 +68,7 @@ require (
|
||||
github.com/couchbase/goprotostellar v1.0.2 // indirect
|
||||
github.com/couchbase/tools-common/errors v1.0.0 // indirect
|
||||
github.com/couchbaselabs/gocbconnstr/v2 v2.0.0-20240607131231-fb385523de28 // indirect
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect
|
||||
github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
|
||||
@@ -96,6 +99,8 @@ require (
|
||||
github.com/klauspost/cpuid/v2 v2.2.5 // indirect
|
||||
github.com/leodido/go-urn v1.4.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
|
||||
github.com/modern-go/reflect2 v1.0.2 // indirect
|
||||
github.com/ncruces/go-strftime v0.1.9 // indirect
|
||||
github.com/pierrec/lz4/v4 v4.1.18 // indirect
|
||||
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
|
||||
|
||||
16
go.sum
16
go.sum
@@ -677,6 +677,10 @@ github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2
|
||||
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
|
||||
github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
|
||||
github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
|
||||
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
|
||||
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
|
||||
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
|
||||
github.com/cenkalti/backoff/v5 v5.0.2 h1:rIfFVxEf1QsI7E1ZHfp/B4DF/6QBAUhmgkxc0H7Zss8=
|
||||
github.com/cenkalti/backoff/v5 v5.0.2/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
|
||||
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
|
||||
@@ -728,6 +732,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
|
||||
github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
|
||||
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
@@ -879,6 +885,7 @@ github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
|
||||
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
|
||||
github.com/google/martian v2.1.0+incompatible h1:/CP5g8u/VJHijgedC/Legn3BAbAaWPgecwXBIDzw5no=
|
||||
github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXiIaQGbYVAs8BPL6v8lEs=
|
||||
github.com/google/martian/v3 v3.0.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=
|
||||
@@ -971,6 +978,8 @@ github.com/jackc/puddle v1.3.0 h1:eHK/5clGOatcjX3oWGBO/MpxpbHzSwud5EWTSCI+MX0=
|
||||
github.com/jackc/puddle v1.3.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
|
||||
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
|
||||
github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
|
||||
github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
|
||||
github.com/jung-kurt/gofpdf v1.0.0/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=
|
||||
@@ -1009,6 +1018,11 @@ github.com/microsoft/go-mssqldb v1.8.2 h1:236sewazvC8FvG6Dr3bszrVhMkAl4KYImryLkR
|
||||
github.com/microsoft/go-mssqldb v1.8.2/go.mod h1:vp38dT33FGfVotRiTmDo3bFyaHq+p3LektQrjTULowo=
|
||||
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
|
||||
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
|
||||
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
|
||||
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
|
||||
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
|
||||
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
|
||||
github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=
|
||||
github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/neo4j/neo4j-go-driver/v5 v5.28.1 h1:RKWQW7wTgYAY2fU9S+9LaJ9OwRPbRc0I17tlT7nDmAY=
|
||||
@@ -1035,6 +1049,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH
|
||||
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
|
||||
github.com/prometheus/client_model v0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
|
||||
github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w=
|
||||
github.com/redis/go-redis/v9 v9.9.0 h1:URbPQ4xVQSQhZ27WMQVmZSo3uT3pL+4IdHVcYq2nVfM=
|
||||
github.com/redis/go-redis/v9 v9.9.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
|
||||
@@ -70,7 +70,6 @@ func TestParseFromYamlMssql(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestFailParseFromYaml(t *testing.T) {
|
||||
|
||||
152
internal/sources/redis/redis.go
Normal file
152
internal/sources/redis/redis.go
Normal file
@@ -0,0 +1,152 @@
|
||||
// Copyright 2025 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 redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
const SourceKind string = "redis"
|
||||
|
||||
// validate interface
|
||||
var _ sources.SourceConfig = Config{}
|
||||
|
||||
func init() {
|
||||
if !sources.Register(SourceKind, newConfig) {
|
||||
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Address []string `yaml:"address" validate:"required"`
|
||||
Username string `yaml:"username"`
|
||||
Password string `yaml:"password"`
|
||||
Database int `yaml:"database"`
|
||||
UseGCPIAM bool `yaml:"useGCPIAM"`
|
||||
ClusterEnabled bool `yaml:"clusterEnabled"`
|
||||
}
|
||||
|
||||
func (r Config) SourceConfigKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
// RedisClient is an interface for `redis.Client` and `redis.ClusterClient
|
||||
type RedisClient interface {
|
||||
Do(context.Context, ...any) *redis.Cmd
|
||||
}
|
||||
|
||||
var _ RedisClient = (*redis.Client)(nil)
|
||||
var _ RedisClient = (*redis.ClusterClient)(nil)
|
||||
|
||||
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
|
||||
client, err := initRedisClient(ctx, r)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error initializing Redis client: %s", err)
|
||||
}
|
||||
s := &Source{
|
||||
Name: r.Name,
|
||||
Kind: SourceKind,
|
||||
Client: client,
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func initRedisClient(ctx context.Context, r Config) (RedisClient, error) {
|
||||
var authFn func(ctx context.Context) (username string, password string, err error)
|
||||
if r.UseGCPIAM {
|
||||
// Pass in an access token getter fn for IAM auth
|
||||
authFn = func(ctx context.Context) (username string, password string, err error) {
|
||||
token, err := sources.GetIAMAccessToken(ctx)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return "default", token, nil
|
||||
}
|
||||
}
|
||||
|
||||
var client RedisClient
|
||||
var err error
|
||||
if r.ClusterEnabled {
|
||||
// Create a new Redis Cluster client
|
||||
clusterClient := redis.NewClusterClient(&redis.ClusterOptions{
|
||||
Addrs: r.Address,
|
||||
// PoolSize applies per cluster node and not for the whole cluster.
|
||||
PoolSize: 10,
|
||||
ConnMaxIdleTime: 60 * time.Second,
|
||||
MinIdleConns: 1,
|
||||
CredentialsProviderContext: authFn,
|
||||
Username: r.Username,
|
||||
Password: r.Password,
|
||||
})
|
||||
err = clusterClient.ForEachShard(ctx, func(ctx context.Context, shard *redis.Client) error {
|
||||
return shard.Ping(ctx).Err()
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to connect to redis cluster: %s", err)
|
||||
}
|
||||
client = clusterClient
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// Create a new Redis client
|
||||
standaloneClient := redis.NewClient(&redis.Options{
|
||||
Addr: r.Address[0],
|
||||
PoolSize: 10,
|
||||
ConnMaxIdleTime: 60 * time.Second,
|
||||
MinIdleConns: 1,
|
||||
DB: r.Database,
|
||||
CredentialsProviderContext: authFn,
|
||||
Username: r.Username,
|
||||
Password: r.Password,
|
||||
})
|
||||
_, err = standaloneClient.Ping(ctx).Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to connect to redis: %s", err)
|
||||
}
|
||||
client = standaloneClient
|
||||
return client, nil
|
||||
}
|
||||
|
||||
var _ sources.Source = &Source{}
|
||||
|
||||
type Source struct {
|
||||
Name string `yaml:"name"`
|
||||
Kind string `yaml:"kind"`
|
||||
Client RedisClient
|
||||
}
|
||||
|
||||
func (s *Source) SourceKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
func (s *Source) RedisClient() RedisClient {
|
||||
return s.Client
|
||||
}
|
||||
157
internal/sources/redis/redis_test.go
Normal file
157
internal/sources/redis/redis_test.go
Normal file
@@ -0,0 +1,157 @@
|
||||
// Copyright 2025 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 redis_test
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/redis"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYamlRedis(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.SourceConfigs
|
||||
}{
|
||||
{
|
||||
desc: "default setting",
|
||||
in: `
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
address:
|
||||
- 127.0.0.1
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-redis-instance": redis.Config{
|
||||
Name: "my-redis-instance",
|
||||
Kind: redis.SourceKind,
|
||||
Address: []string{"127.0.0.1"},
|
||||
ClusterEnabled: false,
|
||||
UseGCPIAM: false,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "advanced example",
|
||||
in: `
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
address:
|
||||
- 127.0.0.1
|
||||
password: my-pass
|
||||
database: 1
|
||||
useGCPIAM: true
|
||||
clusterEnabled: true
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-redis-instance": redis.Config{
|
||||
Name: "my-redis-instance",
|
||||
Kind: redis.SourceKind,
|
||||
Address: []string{"127.0.0.1"},
|
||||
Password: "my-pass",
|
||||
Database: 1,
|
||||
ClusterEnabled: true,
|
||||
UseGCPIAM: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
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 TestFailParseFromYaml(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "invalid database",
|
||||
in: `
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
project: my-project
|
||||
address:
|
||||
- 127.0.0.1
|
||||
password: my-pass
|
||||
database: data
|
||||
`,
|
||||
err: "cannot unmarshal string into Go struct field .Sources of type int",
|
||||
},
|
||||
{
|
||||
desc: "extra field",
|
||||
in: `
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
project: my-project
|
||||
address:
|
||||
- 127.0.0.1
|
||||
password: my-pass
|
||||
database: 1
|
||||
`,
|
||||
err: "unable to parse source \"my-redis-instance\" as \"redis\": [6:1] unknown field \"project\"",
|
||||
},
|
||||
{
|
||||
desc: "missing required field",
|
||||
in: `
|
||||
sources:
|
||||
my-redis-instance:
|
||||
kind: redis
|
||||
`,
|
||||
err: "unable to parse source \"my-redis-instance\" as \"redis\": Key: 'Config.Address' Error:Field validation for 'Address' failed on the 'required' tag",
|
||||
},
|
||||
}
|
||||
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")
|
||||
}
|
||||
errStr := err.Error()
|
||||
if !strings.Contains(errStr, tc.err) {
|
||||
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -85,3 +85,20 @@ func GetIAMPrincipalEmailFromADC(ctx context.Context) (string, error) {
|
||||
email := strings.TrimSuffix(emailValue.(string), ".gserviceaccount.com")
|
||||
return email, nil
|
||||
}
|
||||
|
||||
func GetIAMAccessToken(ctx context.Context) (string, error) {
|
||||
creds, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to find default credentials (run 'gcloud auth application-default login'?): %w", err)
|
||||
}
|
||||
|
||||
token, err := creds.TokenSource.Token() // This gets an oauth2.Token
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to get token from token source: %w", err)
|
||||
}
|
||||
|
||||
if !token.Valid() {
|
||||
return "", fmt.Errorf("retrieved token is invalid or expired")
|
||||
}
|
||||
return token.AccessToken, nil
|
||||
}
|
||||
|
||||
210
internal/tools/redis/redis.go
Normal file
210
internal/tools/redis/redis.go
Normal file
@@ -0,0 +1,210 @@
|
||||
// Copyright 2025 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 redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
redissrc "github.com/googleapis/genai-toolbox/internal/sources/redis"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
const kind string = "redis"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
RedisClient() redissrc.RedisClient
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &redissrc.Source{}
|
||||
|
||||
var compatibleSources = [...]string{redissrc.SourceKind}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description" validate:"required"`
|
||||
Commands [][]string `yaml:"commands" validate:"required"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
Parameters tools.Parameters `yaml:"parameters"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
Description: cfg.Description,
|
||||
InputSchema: cfg.Parameters.McpManifest(),
|
||||
}
|
||||
|
||||
// finish tool setup
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: cfg.Parameters,
|
||||
Commands: cfg.Commands,
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
Client: s.RedisClient(),
|
||||
manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
mcpManifest: mcpManifest,
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
Name string `yaml:"name"`
|
||||
Kind string `yaml:"kind"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
Parameters tools.Parameters `yaml:"parameters"`
|
||||
|
||||
Client redissrc.RedisClient
|
||||
Commands [][]string
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
|
||||
cmds, err := replaceCommandsParams(t.Commands, t.Parameters, params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error replacing commands' parameters: %s", err)
|
||||
}
|
||||
|
||||
// Execute commands
|
||||
responses := make([]*redis.Cmd, len(cmds))
|
||||
for i, cmd := range cmds {
|
||||
responses[i] = t.Client.Do(ctx, cmd...)
|
||||
}
|
||||
// Parse responses
|
||||
out := make([]any, len(t.Commands))
|
||||
for i, resp := range responses {
|
||||
if err := resp.Err(); err != nil {
|
||||
// Add error from each command to `errSum`
|
||||
errString := fmt.Sprintf("error from executing command at index %d: %s", i, err)
|
||||
out[i] = errString
|
||||
continue
|
||||
}
|
||||
val, err := resp.Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting result: %s", err)
|
||||
}
|
||||
// If result is a map, convert map[any]any to map[string]any
|
||||
// Because the Go's built-in json/encoding marshalling doesn't support
|
||||
// map[any]any as an input
|
||||
var strMap map[string]any
|
||||
var json = jsoniter.ConfigCompatibleWithStandardLibrary
|
||||
mapStr, err := json.Marshal(val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error marshalling result: %s", err)
|
||||
}
|
||||
err = json.Unmarshal(mapStr, &strMap)
|
||||
if err != nil {
|
||||
// result is not a map
|
||||
out[i] = val
|
||||
continue
|
||||
}
|
||||
out[i] = strMap
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
// replaceCommandsParams is a helper function to replace parameters in the commands
|
||||
func replaceCommandsParams(commands [][]string, params tools.Parameters, paramValues tools.ParamValues) ([][]any, error) {
|
||||
paramMap := paramValues.AsMapWithDollarPrefix()
|
||||
typeMap := make(map[string]string, len(params))
|
||||
for _, p := range params {
|
||||
placeholder := "$" + p.GetName()
|
||||
typeMap[placeholder] = p.GetType()
|
||||
}
|
||||
newCommands := make([][]any, len(commands))
|
||||
for i, cmd := range commands {
|
||||
newCmd := make([]any, len(cmd))
|
||||
for j, part := range cmd {
|
||||
v, ok := paramMap[part]
|
||||
if !ok {
|
||||
// Command part is not a Parameter placeholder
|
||||
newCmd[j] = part
|
||||
continue
|
||||
}
|
||||
if typeMap[part] == "array" {
|
||||
for _, item := range v.([]any) {
|
||||
// Nested arrays will only be expanded once
|
||||
// e.g., [A, [B, C]] --> ["A", "[B C]"]
|
||||
newCmd = append(newCmd, fmt.Sprintf("%s", item))
|
||||
}
|
||||
continue
|
||||
}
|
||||
newCmd[j] = fmt.Sprintf("%s", v)
|
||||
}
|
||||
newCommands[i] = newCmd
|
||||
}
|
||||
return newCommands, nil
|
||||
}
|
||||
85
internal/tools/redis/redis_test.go
Normal file
85
internal/tools/redis/redis_test.go
Normal file
@@ -0,0 +1,85 @@
|
||||
// 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 redis_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/redis"
|
||||
)
|
||||
|
||||
func TestParseFromYamlRedis(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
redis_tool:
|
||||
kind: redis
|
||||
source: my-redis-instance
|
||||
description: some description
|
||||
commands:
|
||||
- [SET, greeting, "hello, {{.name}}"]
|
||||
- [GET, id]
|
||||
parameters:
|
||||
- name: name
|
||||
type: string
|
||||
description: user name
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"redis_tool": redis.Config{
|
||||
Name: "redis_tool",
|
||||
Kind: "redis",
|
||||
Source: "my-redis-instance",
|
||||
Description: "some description",
|
||||
AuthRequired: []string{},
|
||||
Commands: [][]string{{"SET", "greeting", "hello, {{.name}}"}, {"GET", "id"}},
|
||||
Parameters: []tools.Parameter{
|
||||
tools.NewStringParameter("name", "user name"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -447,3 +447,90 @@ func SetupMySQLTable(t *testing.T, ctx context.Context, pool *sql.DB, create_sta
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GetRedisWants return the expected wants for redis
|
||||
func GetRedisValkeyWants() (string, string, string, string) {
|
||||
select1Want := "[\"PONG\"]"
|
||||
failInvocationWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"\"error from executing command at index 0: ERR unknown command 'SELEC 1;', with args beginning with: \""}]}}`
|
||||
invokeParamWant := "[{\"id\":\"1\",\"name\":\"Alice\"},{\"id\":\"3\",\"name\":\"Sid\"}]"
|
||||
mcpInvokeParamWant := `{"jsonrpc":"2.0","id":"my-param-tool","result":{"content":[{"type":"text","text":"{\"id\":\"1\",\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":\"3\",\"name\":\"Sid\"}"}]}}`
|
||||
return select1Want, failInvocationWant, invokeParamWant, mcpInvokeParamWant
|
||||
}
|
||||
|
||||
func GetRedisValkeyToolsConfig(sourceConfig map[string]any, toolKind string) map[string]any {
|
||||
toolsFile := map[string]any{
|
||||
"sources": map[string]any{
|
||||
"my-instance": sourceConfig,
|
||||
},
|
||||
"authServices": map[string]any{
|
||||
"my-google-auth": map[string]any{
|
||||
"kind": "google",
|
||||
"clientId": ClientId,
|
||||
},
|
||||
},
|
||||
"tools": map[string]any{
|
||||
"my-simple-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Simple tool to test end to end functionality.",
|
||||
"commands": [][]string{{"PING"}},
|
||||
},
|
||||
"my-param-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test invocation with params.",
|
||||
"commands": [][]string{{"HGETALL", "row1"}, {"HGETALL", "row3"}},
|
||||
"parameters": []any{
|
||||
map[string]any{
|
||||
"name": "id",
|
||||
"type": "integer",
|
||||
"description": "user ID",
|
||||
},
|
||||
map[string]any{
|
||||
"name": "name",
|
||||
"type": "string",
|
||||
"description": "user name",
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-auth-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test authenticated parameters.",
|
||||
// statement to auto-fill authenticated parameter
|
||||
"commands": [][]string{{"HGETALL", "$email"}},
|
||||
"parameters": []map[string]any{
|
||||
{
|
||||
"name": "email",
|
||||
"type": "string",
|
||||
"description": "user email",
|
||||
"authServices": []map[string]string{
|
||||
{
|
||||
"name": "my-google-auth",
|
||||
"field": "email",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"my-auth-required-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test auth required invocation.",
|
||||
"commands": [][]string{{"PING"}},
|
||||
"authRequired": []string{
|
||||
"my-google-auth",
|
||||
},
|
||||
},
|
||||
"my-fail-tool": map[string]any{
|
||||
"kind": toolKind,
|
||||
"source": "my-instance",
|
||||
"description": "Tool to test statement with incorrect syntax.",
|
||||
"commands": [][]string{{"SELEC 1;"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
return toolsFile
|
||||
|
||||
}
|
||||
|
||||
129
tests/redis/redis_test.go
Normal file
129
tests/redis/redis_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
// Copyright 2025 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 redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"regexp"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var (
|
||||
REDIS_SOURCE_KIND = "redis"
|
||||
REDIS_TOOL_KIND = "redis"
|
||||
REDIS_ADDRESS = os.Getenv("REDIS_ADDRESS")
|
||||
REDIS_PASS = os.Getenv("REDIS_PASS")
|
||||
)
|
||||
|
||||
func getRedisVars(t *testing.T) map[string]any {
|
||||
switch "" {
|
||||
case REDIS_ADDRESS:
|
||||
t.Fatal("'REDIS_ADDRESS' not set")
|
||||
case REDIS_PASS:
|
||||
t.Fatal("'REDIS_PASS' not set")
|
||||
}
|
||||
return map[string]any{
|
||||
"kind": REDIS_SOURCE_KIND,
|
||||
"address": []string{REDIS_ADDRESS},
|
||||
"password": REDIS_PASS,
|
||||
}
|
||||
}
|
||||
|
||||
func initRedisClient(ctx context.Context, address, pass string) (*redis.Client, error) {
|
||||
// Create a new Redis client
|
||||
standaloneClient := redis.NewClient(&redis.Options{
|
||||
Addr: address,
|
||||
PoolSize: 10,
|
||||
ConnMaxIdleTime: 60 * time.Second,
|
||||
MinIdleConns: 1,
|
||||
Password: pass,
|
||||
})
|
||||
_, err := standaloneClient.Ping(ctx).Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to connect to redis: %s", err)
|
||||
}
|
||||
return standaloneClient, nil
|
||||
}
|
||||
|
||||
func TestRedisToolEndpoints(t *testing.T) {
|
||||
sourceConfig := getRedisVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
||||
defer cancel()
|
||||
|
||||
var args []string
|
||||
|
||||
client, err := initRedisClient(ctx, REDIS_ADDRESS, REDIS_PASS)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create Redis connection: %s", err)
|
||||
}
|
||||
|
||||
// set up data for param tool
|
||||
teardownDB := setupRedisDB(t, ctx, client)
|
||||
defer teardownDB(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := tests.GetRedisValkeyToolsConfig(sourceConfig, REDIS_TOOL_KIND)
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`))
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
}
|
||||
|
||||
tests.RunToolGetTest(t)
|
||||
|
||||
select1Want, failInvocationWant, invokeParamWant, mcpInvokeParamWant := tests.GetRedisValkeyWants()
|
||||
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
|
||||
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
|
||||
}
|
||||
|
||||
func setupRedisDB(t *testing.T, ctx context.Context, client *redis.Client) func(*testing.T) {
|
||||
keys := []string{"row1", "row2", "row3"}
|
||||
commands := [][]any{
|
||||
{"HSET", keys[0], "id", 1, "name", "Alice"},
|
||||
{"HSET", keys[1], "id", 2, "name", "Jane"},
|
||||
{"HSET", keys[2], "id", 3, "name", "Sid"},
|
||||
{"HSET", tests.SERVICE_ACCOUNT_EMAIL, "name", "Alice"},
|
||||
}
|
||||
for _, c := range commands {
|
||||
resp := client.Do(ctx, c...)
|
||||
if err := resp.Err(); err != nil {
|
||||
t.Fatalf("unable to insert test data: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
return func(t *testing.T) {
|
||||
// tear down test
|
||||
_, err := client.Del(ctx, keys...).Result()
|
||||
if err != nil {
|
||||
t.Errorf("Teardown failed: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -76,7 +76,7 @@ func RunToolGetTest(t *testing.T) {
|
||||
}
|
||||
|
||||
// RunToolInvoke runs the tool invoke endpoint
|
||||
func RunToolInvokeTest(t *testing.T, select_1_want, invoke_param_want string) {
|
||||
func RunToolInvokeTest(t *testing.T, select1Want, invokeParamWant string) {
|
||||
// Get ID token
|
||||
idToken, err := GetGoogleIdToken(ClientId)
|
||||
if err != nil {
|
||||
@@ -97,7 +97,7 @@ func RunToolInvokeTest(t *testing.T, select_1_want, invoke_param_want string) {
|
||||
api: "http://127.0.0.1:5000/api/tool/my-simple-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(`{}`)),
|
||||
want: select_1_want,
|
||||
want: select1Want,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
@@ -105,7 +105,7 @@ func RunToolInvokeTest(t *testing.T, select_1_want, invoke_param_want string) {
|
||||
api: "http://127.0.0.1:5000/api/tool/my-param-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(`{"id": 3, "name": "Alice"}`)),
|
||||
want: invoke_param_want,
|
||||
want: invokeParamWant,
|
||||
isErr: false,
|
||||
},
|
||||
{
|
||||
@@ -150,7 +150,7 @@ func RunToolInvokeTest(t *testing.T, select_1_want, invoke_param_want string) {
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(`{}`)),
|
||||
isErr: false,
|
||||
want: select_1_want,
|
||||
want: select1Want,
|
||||
},
|
||||
{
|
||||
name: "Invoke my-auth-required-tool with invalid auth token",
|
||||
@@ -460,7 +460,7 @@ func RunExecuteSqlToolInvokeTest(t *testing.T, createTableStatement string, sele
|
||||
}
|
||||
|
||||
// RunMCPToolCallMethod runs the tool/call for mcp endpoint
|
||||
func RunMCPToolCallMethod(t *testing.T, invoke_param_want, fail_invocation_want string) {
|
||||
func RunMCPToolCallMethod(t *testing.T, invokeParamWant, fail_invocation_want string) {
|
||||
// Test tool invoke endpoint
|
||||
invokeTcs := []struct {
|
||||
name string
|
||||
@@ -487,7 +487,7 @@ func RunMCPToolCallMethod(t *testing.T, invoke_param_want, fail_invocation_want
|
||||
},
|
||||
},
|
||||
},
|
||||
want: invoke_param_want,
|
||||
want: invokeParamWant,
|
||||
},
|
||||
{
|
||||
name: "MCP Invoke invalid tool",
|
||||
|
||||
Reference in New Issue
Block a user