Compare commits

..

1 Commits

Author SHA1 Message Date
Yuan Teoh
1a989bab10 docs(sources/dgraph): add best effort maintenance notes 2026-01-15 17:23:17 -08:00
41 changed files with 913 additions and 3059 deletions

View File

@@ -87,7 +87,7 @@ steps:
- "CLOUD_SQL_POSTGRES_REGION=$_REGION"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv:
["CLOUD_SQL_POSTGRES_USER", "CLOUD_SQL_POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
["CLOUD_SQL_POSTGRES_USER", "CLOUD_SQL_POSTGRES_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -134,7 +134,7 @@ steps:
- "ALLOYDB_POSTGRES_DATABASE=$_DATABASE_NAME"
- "ALLOYDB_POSTGRES_REGION=$_REGION"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -293,7 +293,7 @@ steps:
.ci/test_with_coverage.sh \
"Cloud Healthcare API" \
cloudhealthcare \
cloudhealthcare
cloudhealthcare || echo "Integration tests failed."
- id: "postgres"
name: golang:1
@@ -305,7 +305,7 @@ steps:
- "POSTGRES_HOST=$_POSTGRES_HOST"
- "POSTGRES_PORT=$_POSTGRES_PORT"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv: ["POSTGRES_USER", "POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
secretEnv: ["POSTGRES_USER", "POSTGRES_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -964,13 +964,6 @@ steps:
availableSecrets:
secretManager:
# Common secrets
- versionName: projects/$PROJECT_ID/secrets/client_id/versions/latest
env: CLIENT_ID
- versionName: projects/$PROJECT_ID/secrets/api_key/versions/latest
env: API_KEY
# Resource-specific secrets
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest
env: CLOUD_SQL_POSTGRES_USER
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_pass/versions/latest
@@ -987,6 +980,8 @@ availableSecrets:
env: POSTGRES_USER
- versionName: projects/$PROJECT_ID/secrets/postgres_pass/versions/latest
env: POSTGRES_PASS
- versionName: projects/$PROJECT_ID/secrets/client_id/versions/latest
env: CLIENT_ID
- versionName: projects/$PROJECT_ID/secrets/neo4j_user/versions/latest
env: NEO4J_USER
- versionName: projects/$PROJECT_ID/secrets/neo4j_pass/versions/latest

View File

@@ -77,24 +77,11 @@ run_orch_test() {
setup_orch_table
cd "$orch_dir"
echo "--- Debugging npm config for $orch_name ---"
npm config list
echo "--- Active Registry for $orch_name ---"
npm config get registry
echo "--- Inspecting .npmrc files ---"
[ -f ".npmrc" ] && echo "Local .npmrc content:" && cat .npmrc
[ -f "$HOME/.npmrc" ] && echo "Global .npmrc content:" && cat "$HOME/.npmrc"
export GPKG_DEBUG=1
echo "Installing dependencies for $orch_name..."
if [ -f "package-lock.json" ]; then
npm ci --loglevel verbose
npm ci
else
npm install --loglevel verbose
npm install
fi
cd ..

View File

@@ -386,7 +386,6 @@ func NewCommand(opts ...Option) *Command {
// TODO: Insecure by default. Might consider updating this for v1.0.0
flags.StringSliceVar(&cmd.cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
flags.StringSliceVar(&cmd.cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
flags.StringSliceVar(&cmd.cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
// wrap RunE command so that we have access to original Command object
cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) }

View File

@@ -70,9 +70,6 @@ func withDefaults(c server.ServerConfig) server.ServerConfig {
if c.AllowedHosts == nil {
c.AllowedHosts = []string{"*"}
}
if c.UserAgentMetadata == nil {
c.UserAgentMetadata = []string{}
}
return c
}
@@ -233,13 +230,6 @@ func TestServerConfigFlags(t *testing.T) {
AllowedHosts: []string{"http://foo.com", "http://bar.com"},
}),
},
{
desc: "user agent metadata",
args: []string{"--user-agent-metadata", "foo,bar"},
want: withDefaults(server.ServerConfig{
UserAgentMetadata: []string{"foo", "bar"},
}),
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {

File diff suppressed because it is too large Load Diff

View File

@@ -11,7 +11,7 @@
"author": "",
"license": "ISC",
"dependencies": {
"@google/adk": "^0.2.4",
"@google/adk": "^0.1.3",
"@toolbox-sdk/adk": "^0.1.5"
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -11,7 +11,7 @@
"author": "",
"license": "ISC",
"dependencies": {
"@llamaindex/google": "^0.4.0",
"@llamaindex/google": "^0.3.20",
"@llamaindex/workflow": "^1.1.22",
"@toolbox-sdk/core": "^0.1.2",
"llamaindex": "^0.12.0"

View File

@@ -207,7 +207,6 @@ You can connect to Toolbox Cloud Run instances directly through the SDK.
{{< tab header="Python" lang="python" >}}
import asyncio
from toolbox_core import ToolboxClient, auth_methods
from toolbox_core.protocol import Protocol
# Replace with the Cloud Run service URL generated in the previous step
URL = "https://cloud-run-url.app"
@@ -218,7 +217,6 @@ async def main():
async with ToolboxClient(
URL,
client_headers={"Authorization": auth_token_provider},
protocol=Protocol.TOOLBOX,
) as toolbox:
toolset = await toolbox.load_toolset()
# ...
@@ -283,5 +281,3 @@ contain the specific error message needed to diagnose the problem.
Manager, it means the Toolbox service account is missing permissions.
- Ensure the `toolbox-identity` service account has the **Secret Manager
Secret Accessor** (`roles/secretmanager.secretAccessor`) IAM role.
- **Cloud Run Connections via IAP:** Currently we do not support Cloud Run connections via [IAP](https://docs.cloud.google.com/iap/docs/concepts-overview). Please disable IAP if you are using it.

View File

@@ -27,7 +27,6 @@ description: >
| | `--ui` | Launches the Toolbox UI web server. | |
| | `--allowed-origins` | Specifies a list of origins permitted to access this server for CORs access. | `*` |
| | `--allowed-hosts` | Specifies a list of hosts permitted to access this server to prevent DNS rebinding attacks. | `*` |
| | `--user-agent-extra` | Appends additional metadata to the User-Agent. | |
| `-v` | `--version` | version for toolbox | |
## Examples

View File

@@ -7,6 +7,17 @@ description: >
---
{{< notice note >}}
**⚠️ Best Effort Maintenance**
This integration is maintained on a best-effort basis by the project
team/community. While we strive to address issues and provide workarounds when
resources are available, there are no guaranteed response times or code fixes.
The automated integration tests for this module are currently non-functional or
failing.
{{< /notice >}}
## About
[Dgraph][dgraph-docs] is an open-source graph database. It is designed for

View File

@@ -12,9 +12,6 @@ aliases:
The `cloud-gemini-data-analytics-query` tool allows you to send natural language questions to the Gemini Data Analytics API and receive structured responses containing SQL queries, natural language answers, and explanations. For details on defining data agent context for database data sources, see the official [documentation](https://docs.cloud.google.com/gemini/docs/conversational-analytics-api/data-agent-authored-context-databases).
> [!NOTE]
> Only `alloydb`, `spannerReference`, and `cloudSqlReference` are supported as [datasource references](https://clouddocs.devsite.corp.google.com/gemini/docs/conversational-analytics-api/reference/rest/v1beta/projects.locations.dataAgents#DatasourceReferences).
## Example
```yaml

View File

@@ -9,6 +9,17 @@ aliases:
- /resources/tools/dgraph-dql
---
{{< notice note >}}
**⚠️ Best Effort Maintenance**
This integration is maintained on a best-effort basis by the project
team/community. While we strive to address issues and provide workarounds when
resources are available, there are no guaranteed response times or code fixes.
The automated integration tests for this module are currently non-functional or
failing.
{{< /notice >}}
## About
A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph

View File

@@ -30,10 +30,6 @@ following config for example:
- name: userNames
type: array
description: The user names to be set.
items:
name: userName # the item name doesn't matter but it has to exist
type: string
description: username
```
If the input is an array of strings `["Alice", "Sid", "Bob"]`, The final command

View File

@@ -64,14 +64,12 @@ type ServerConfig struct {
Stdio bool
// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.
DisableReload bool
// UI indicates if Toolbox UI endpoints (/ui) are available.
// UI indicates if Toolbox UI endpoints (/ui) are available
UI bool
// Specifies a list of origins permitted to access this server.
AllowedOrigins []string
// Specifies a list of hosts permitted to access this server.
// Specifies a list of hosts permitted to access this server
AllowedHosts []string
// UserAgentMetadata specifies additional metadata to append to the User-Agent string.
UserAgentMetadata []string
}
type logFormat string

View File

@@ -183,13 +183,6 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
embeddingModels := resourceMgr.GetEmbeddingModelMap()
params, err = tool.EmbedParams(ctx, params, embeddingModels)
if err != nil {
err = fmt.Errorf("error embedding parameters: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {

View File

@@ -183,13 +183,6 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
embeddingModels := resourceMgr.GetEmbeddingModelMap()
params, err = tool.EmbedParams(ctx, params, embeddingModels)
if err != nil {
err = fmt.Errorf("error embedding parameters: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {

View File

@@ -176,13 +176,6 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
embeddingModels := resourceMgr.GetEmbeddingModelMap()
params, err = tool.EmbedParams(ctx, params, embeddingModels)
if err != nil {
err = fmt.Errorf("error embedding parameters: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {

View File

@@ -176,13 +176,6 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
embeddingModels := resourceMgr.GetEmbeddingModelMap()
params, err = tool.EmbedParams(ctx, params, embeddingModels)
if err != nil {
err = fmt.Errorf("error embedding parameters: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {

View File

@@ -64,11 +64,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
map[string]prompts.Promptset,
error,
) {
metadataStr := cfg.Version
if len(cfg.UserAgentMetadata) > 0 {
metadataStr += "+" + strings.Join(cfg.UserAgentMetadata, "+")
}
ctx = util.WithUserAgent(ctx, metadataStr)
ctx = util.WithUserAgent(ctx, cfg.Version)
instrumentation, err := util.InstrumentationFromContext(ctx)
if err != nil {
panic(err)

View File

@@ -103,12 +103,10 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", pageURLKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.FHIRFetchPage(ctx, url, tokenStr)
}

View File

@@ -131,12 +131,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", patientIDKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
var opts []googleapi.CallOption

View File

@@ -161,12 +161,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
var summary bool

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDataset(tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDICOMStore(storeID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDICOMStoreMetrics(storeID, tokenStr)
}

View File

@@ -130,12 +130,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", idKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRResource(storeID, resType, resID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRStore(storeID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRStoreMetrics(storeID, tokenStr)
}

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.ListDICOMStores(tokenStr)
}

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.ListFHIRStores(tokenStr)
}

View File

@@ -127,12 +127,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
study, ok := params.AsMap()[studyInstanceUIDKey].(string)
if !ok {

View File

@@ -140,12 +140,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{sopInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})

View File

@@ -138,12 +138,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{seriesInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})

View File

@@ -133,12 +133,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{studyInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey})
if err != nil {

View File

@@ -147,20 +147,12 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
defer teardownTable2(t)
// Set up table for semanti search
vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
defer tearDownVectorTable(t)
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, AlloyDBPostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, AlloyDBPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
// Add semantic search tool config
insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, AlloyDBPostgresToolKind, insertStmt, searchStmt)
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)

View File

@@ -112,7 +112,8 @@ func TestHealthcareToolEndpoints(t *testing.T) {
fhirStoreID := "fhir-store-" + uuid.New().String()
dicomStoreID := "dicom-store-" + uuid.New().String()
patient1ID, patient2ID := setupHealthcareResources(t, healthcareService, healthcareDataset, fhirStoreID, dicomStoreID)
patient1ID, patient2ID, teardown := setupHealthcareResources(t, healthcareService, healthcareDataset, fhirStoreID, dicomStoreID)
defer teardown(t)
toolsFile := getToolsConfig(sourceConfig)
toolsFile = addClientAuthSourceConfig(t, toolsFile)
@@ -172,8 +173,10 @@ func TestHealthcareToolWithStoreRestriction(t *testing.T) {
disallowedFHIRStoreID := "fhir-store-disallowed-" + uuid.New().String()
disallowedDICOMStoreID := "dicom-store-disallowed-" + uuid.New().String()
setupHealthcareResources(t, healthcareService, healthcareDataset, allowedFHIRStoreID, allowedDICOMStoreID)
setupHealthcareResources(t, healthcareService, healthcareDataset, disallowedFHIRStoreID, disallowedDICOMStoreID)
_, _, teardownAllowedStores := setupHealthcareResources(t, healthcareService, healthcareDataset, allowedFHIRStoreID, allowedDICOMStoreID)
defer teardownAllowedStores(t)
_, _, teardownDisallowedStores := setupHealthcareResources(t, healthcareService, healthcareDataset, disallowedFHIRStoreID, disallowedDICOMStoreID)
defer teardownDisallowedStores(t)
// Configure source with dataset restriction.
sourceConfig["allowedFhirStores"] = []string{allowedFHIRStoreID}
@@ -254,7 +257,7 @@ func newHealthcareService(ctx context.Context) (*healthcare.Service, error) {
return healthcareService, nil
}
func setupHealthcareResources(t *testing.T, service *healthcare.Service, datasetID, fhirStoreID, dicomStoreID string) (string, string) {
func setupHealthcareResources(t *testing.T, service *healthcare.Service, datasetID, fhirStoreID, dicomStoreID string) (string, string, func(*testing.T)) {
datasetName := fmt.Sprintf("projects/%s/locations/%s/datasets/%s", healthcareProject, healthcareRegion, datasetID)
var err error
@@ -263,24 +266,12 @@ func setupHealthcareResources(t *testing.T, service *healthcare.Service, dataset
if fhirStore, err = service.Projects.Locations.Datasets.FhirStores.Create(datasetName, fhirStore).FhirStoreId(fhirStoreID).Do(); err != nil {
t.Fatalf("failed to create fhir store: %v", err)
}
// Register cleanup
t.Cleanup(func() {
if _, err := service.Projects.Locations.Datasets.FhirStores.Delete(fhirStore.Name).Do(); err != nil {
t.Logf("failed to delete fhir store: %v", err)
}
})
// Create DICOM store
dicomStore := &healthcare.DicomStore{}
if dicomStore, err = service.Projects.Locations.Datasets.DicomStores.Create(datasetName, dicomStore).DicomStoreId(dicomStoreID).Do(); err != nil {
t.Fatalf("failed to create dicom store: %v", err)
}
// Register cleanup
t.Cleanup(func() {
if _, err := service.Projects.Locations.Datasets.DicomStores.Delete(dicomStore.Name).Do(); err != nil {
t.Logf("failed to delete dicom store: %v", err)
}
})
// Create Patient 1
patient1Body := bytes.NewBuffer([]byte(`{
@@ -326,7 +317,15 @@ func setupHealthcareResources(t *testing.T, service *healthcare.Service, dataset
createFHIRResource(t, service, fhirStore.Name, "Observation", observation2Body)
}
return patient1ID, patient2ID
teardown := func(t *testing.T) {
if _, err := service.Projects.Locations.Datasets.FhirStores.Delete(fhirStore.Name).Do(); err != nil {
t.Logf("failed to delete fhir store: %v", err)
}
if _, err := service.Projects.Locations.Datasets.DicomStores.Delete(dicomStore.Name).Do(); err != nil {
t.Logf("failed to delete dicom store: %v", err)
}
}
return patient1ID, patient2ID, teardown
}
func getToolsConfig(sourceConfig map[string]any) map[string]any {

View File

@@ -132,20 +132,12 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
defer teardownTable2(t)
// Set up table for semantic search
vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
defer tearDownVectorTable(t)
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
// Add semantic search tool config
insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, CloudSQLPostgresToolKind, insertStmt, searchStmt)
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil {
@@ -194,7 +186,6 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
tests.RunPostgresListDatabaseStatsTest(t, ctx, pool)
tests.RunPostgresListRolesTest(t, ctx, pool)
tests.RunPostgresListStoredProcedureTest(t, ctx, pool)
tests.RunSemanticSearchToolInvokeTest(t, "null", "", "The quick brown fox")
}
// Test connection with different IP type

View File

@@ -1,251 +0,0 @@
// Copyright 2026 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 tests contains end to end tests meant to verify the Toolbox Server
// works as expected when executed as a binary.
package tests
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"strings"
"testing"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
"github.com/jackc/pgx/v5/pgxpool"
)
var apiKey = os.Getenv("API_KEY")
// AddSemanticSearchConfig adds embedding models and semantic search tools to the config
// with configurable tool kind and SQL statements.
func AddSemanticSearchConfig(t *testing.T, config map[string]any, toolKind, insertStmt, searchStmt string) map[string]any {
config["embeddingModels"] = map[string]any{
"gemini_model": map[string]any{
"kind": "gemini",
"model": "gemini-embedding-001",
"apiKey": apiKey,
"dimension": 768,
},
}
tools, ok := config["tools"].(map[string]any)
if !ok {
t.Fatalf("unable to get tools from config")
}
tools["insert_docs"] = map[string]any{
"kind": toolKind,
"source": "my-instance",
"description": "Stores content and its vector embedding into the documents table.",
"statement": insertStmt,
"parameters": []any{
map[string]any{
"name": "content",
"type": "string",
"description": "The text content associated with the vector.",
},
map[string]any{
"name": "text_to_embed",
"type": "string",
"description": "The text content used to generate the vector.",
"embeddedBy": "gemini_model",
},
},
}
tools["search_docs"] = map[string]any{
"kind": toolKind,
"source": "my-instance",
"description": "Finds the most semantically similar document to the query vector.",
"statement": searchStmt,
"parameters": []any{
map[string]any{
"name": "query",
"type": "string",
"description": "The text content to search for.",
"embeddedBy": "gemini_model",
},
},
}
config["tools"] = tools
return config
}
// RunSemanticSearchToolInvokeTest runs the insert_docs and search_docs tools
// via both HTTP and MCP endpoints and verifies the output.
func RunSemanticSearchToolInvokeTest(t *testing.T, insertWant, mcpInsertWant, searchWant string) {
// Initialize MCP session once for the MCP test cases
sessionId := RunInitialize(t, "2024-11-05")
tcs := []struct {
name string
api string
isMcp bool
requestBody interface{}
want string
}{
{
name: "HTTP invoke insert_docs",
api: "http://127.0.0.1:5000/api/tool/insert_docs/invoke",
isMcp: false,
requestBody: `{"content": "The quick brown fox jumps over the lazy dog", "text_to_embed": "The quick brown fox jumps over the lazy dog"}`,
want: insertWant,
},
{
name: "HTTP invoke search_docs",
api: "http://127.0.0.1:5000/api/tool/search_docs/invoke",
isMcp: false,
requestBody: `{"query": "fast fox jumping"}`,
want: searchWant,
},
{
name: "MCP invoke insert_docs",
api: "http://127.0.0.1:5000/mcp",
isMcp: true,
requestBody: jsonrpc.JSONRPCRequest{
Jsonrpc: "2.0",
Id: "mcp-insert-docs",
Request: jsonrpc.Request{
Method: "tools/call",
},
Params: map[string]any{
"name": "insert_docs",
"arguments": map[string]any{
"content": "The quick brown fox jumps over the lazy dog",
"text_to_embed": "The quick brown fox jumps over the lazy dog",
},
},
},
want: mcpInsertWant,
},
{
name: "MCP invoke search_docs",
api: "http://127.0.0.1:5000/mcp",
isMcp: true,
requestBody: jsonrpc.JSONRPCRequest{
Jsonrpc: "2.0",
Id: "mcp-search-docs",
Request: jsonrpc.Request{
Method: "tools/call",
},
Params: map[string]any{
"name": "search_docs",
"arguments": map[string]any{
"query": "fast fox jumping",
},
},
},
want: searchWant,
},
}
for _, tc := range tcs {
t.Run(tc.name, func(t *testing.T) {
var bodyReader io.Reader
headers := map[string]string{}
// Prepare Request Body and Headers
if tc.isMcp {
reqBytes, err := json.Marshal(tc.requestBody)
if err != nil {
t.Fatalf("failed to marshal mcp request: %v", err)
}
bodyReader = bytes.NewBuffer(reqBytes)
if sessionId != "" {
headers["Mcp-Session-Id"] = sessionId
}
} else {
bodyReader = bytes.NewBufferString(tc.requestBody.(string))
}
// Send Request
resp, respBody := RunRequest(t, http.MethodPost, tc.api, bodyReader, headers)
if resp.StatusCode != http.StatusOK {
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
}
// Normalize Response to get the actual tool result string
var got string
if tc.isMcp {
var mcpResp struct {
Result struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
} `json:"result"`
}
if err := json.Unmarshal(respBody, &mcpResp); err != nil {
t.Fatalf("error parsing mcp response: %s", err)
}
if len(mcpResp.Result.Content) > 0 {
got = mcpResp.Result.Content[0].Text
}
} else {
var httpResp map[string]interface{}
if err := json.Unmarshal(respBody, &httpResp); err != nil {
t.Fatalf("error parsing http response: %s", err)
}
if res, ok := httpResp["result"].(string); ok {
got = res
}
}
if !strings.Contains(got, tc.want) {
t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
}
})
}
}
// SetupPostgresVectorTable sets up the vector extension and a vector table
func SetupPostgresVectorTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(*testing.T)) {
t.Helper()
if _, err := pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
t.Fatalf("failed to create vector extension: %v", err)
}
tableName := "vector_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
createTableStmt := fmt.Sprintf(`CREATE TABLE %s (
id SERIAL PRIMARY KEY,
content TEXT,
embedding vector(768)
)`, tableName)
if _, err := pool.Exec(ctx, createTableStmt); err != nil {
t.Fatalf("failed to create table %s: %v", tableName, err)
}
return tableName, func(t *testing.T) {
if _, err := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)); err != nil {
t.Errorf("failed to drop table %s: %v", tableName, err)
}
}
}
func GetPostgresVectorSearchStmts(vectorTableName string) (string, string) {
insertStmt := fmt.Sprintf("INSERT INTO %s (content, embedding) VALUES ($1, $2)", vectorTableName)
searchStmt := fmt.Sprintf("SELECT id, content, embedding <-> $1 AS distance FROM %s ORDER BY distance LIMIT 1", vectorTableName)
return insertStmt, searchStmt
}

View File

@@ -111,10 +111,6 @@ func TestPostgres(t *testing.T) {
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
defer teardownTable2(t)
// Set up table for semantic search
vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
defer tearDownVectorTable(t)
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, PostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
@@ -122,10 +118,6 @@ func TestPostgres(t *testing.T) {
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, PostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
// Add semantic search tool config
insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, PostgresToolKind, insertStmt, searchStmt)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil {
t.Fatalf("command initialization returned an error: %s", err)
@@ -173,5 +165,4 @@ func TestPostgres(t *testing.T) {
tests.RunPostgresListDatabaseStatsTest(t, ctx, pool)
tests.RunPostgresListRolesTest(t, ctx, pool)
tests.RunPostgresListStoredProcedureTest(t, ctx, pool)
tests.RunSemanticSearchToolInvokeTest(t, "null", "", "The quick brown fox")
}

View File

@@ -1240,10 +1240,7 @@ func RunPostgresListTablesTest(t *testing.T, tableNameParam, tableNameAuth, user
var filteredGot []any
for _, item := range got {
if tableMap, ok := item.(map[string]interface{}); ok {
name, _ := tableMap["object_name"].(string)
// Only keep the table if it matches expected test tables
if name == tableNameParam || name == tableNameAuth {
if schema, ok := tableMap["schema_name"]; ok && schema == "public" {
filteredGot = append(filteredGot, item)
}
}