This is page 10 of 45. Use http://codebase.md/googleapis/genai-toolbox?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .ci
│ ├── continuous.release.cloudbuild.yaml
│ ├── generate_release_table.sh
│ ├── integration.cloudbuild.yaml
│ ├── quickstart_test
│ │ ├── go.integration.cloudbuild.yaml
│ │ ├── js.integration.cloudbuild.yaml
│ │ ├── py.integration.cloudbuild.yaml
│ │ ├── run_go_tests.sh
│ │ ├── run_js_tests.sh
│ │ ├── run_py_tests.sh
│ │ └── setup_hotels_sample.sql
│ ├── test_with_coverage.sh
│ └── versioned.release.cloudbuild.yaml
├── .github
│ ├── auto-label.yaml
│ ├── blunderbuss.yml
│ ├── CODEOWNERS
│ ├── header-checker-lint.yml
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ ├── config.yml
│ │ ├── feature_request.yml
│ │ └── question.yml
│ ├── label-sync.yml
│ ├── labels.yaml
│ ├── PULL_REQUEST_TEMPLATE.md
│ ├── release-please.yml
│ ├── renovate.json5
│ ├── sync-repo-settings.yaml
│ └── workflows
│ ├── cloud_build_failure_reporter.yml
│ ├── deploy_dev_docs.yaml
│ ├── deploy_previous_version_docs.yaml
│ ├── deploy_versioned_docs.yaml
│ ├── docs_deploy.yaml
│ ├── docs_preview_clean.yaml
│ ├── docs_preview_deploy.yaml
│ ├── lint.yaml
│ ├── schedule_reporter.yml
│ ├── sync-labels.yaml
│ └── tests.yaml
├── .gitignore
├── .gitmodules
├── .golangci.yaml
├── .hugo
│ ├── archetypes
│ │ └── default.md
│ ├── assets
│ │ ├── icons
│ │ │ └── logo.svg
│ │ └── scss
│ │ ├── _styles_project.scss
│ │ └── _variables_project.scss
│ ├── go.mod
│ ├── go.sum
│ ├── hugo.toml
│ ├── layouts
│ │ ├── _default
│ │ │ └── home.releases.releases
│ │ ├── index.llms-full.txt
│ │ ├── index.llms.txt
│ │ ├── partials
│ │ │ ├── hooks
│ │ │ │ └── head-end.html
│ │ │ ├── navbar-version-selector.html
│ │ │ ├── page-meta-links.html
│ │ │ └── td
│ │ │ └── render-heading.html
│ │ ├── robot.txt
│ │ └── shortcodes
│ │ ├── include.html
│ │ ├── ipynb.html
│ │ └── regionInclude.html
│ ├── package-lock.json
│ ├── package.json
│ └── static
│ ├── favicons
│ │ ├── android-chrome-192x192.png
│ │ ├── android-chrome-512x512.png
│ │ ├── apple-touch-icon.png
│ │ ├── favicon-16x16.png
│ │ ├── favicon-32x32.png
│ │ └── favicon.ico
│ └── js
│ └── w3.js
├── CHANGELOG.md
├── cmd
│ ├── options_test.go
│ ├── options.go
│ ├── root_test.go
│ ├── root.go
│ └── version.txt
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── DEVELOPER.md
├── Dockerfile
├── docs
│ └── en
│ ├── _index.md
│ ├── about
│ │ ├── _index.md
│ │ └── faq.md
│ ├── concepts
│ │ ├── _index.md
│ │ └── telemetry
│ │ ├── index.md
│ │ ├── telemetry_flow.png
│ │ └── telemetry_traces.png
│ ├── getting-started
│ │ ├── _index.md
│ │ ├── colab_quickstart.ipynb
│ │ ├── configure.md
│ │ ├── introduction
│ │ │ ├── _index.md
│ │ │ └── architecture.png
│ │ ├── local_quickstart_go.md
│ │ ├── local_quickstart_js.md
│ │ ├── local_quickstart.md
│ │ ├── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── quickstart
│ │ ├── go
│ │ │ ├── genAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── genkit
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── langchain
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── openAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ └── quickstart_test.go
│ │ ├── golden.txt
│ │ ├── js
│ │ │ ├── genAI
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── genkit
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── langchain
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── llamaindex
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ └── quickstart.test.js
│ │ ├── python
│ │ │ ├── __init__.py
│ │ │ ├── adk
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── core
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── langchain
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── llamaindex
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ └── quickstart_test.py
│ │ └── shared
│ │ ├── cloud_setup.md
│ │ ├── configure_toolbox.md
│ │ └── database_setup.md
│ ├── how-to
│ │ ├── _index.md
│ │ ├── connect_via_geminicli.md
│ │ ├── connect_via_mcp.md
│ │ ├── connect-ide
│ │ │ ├── _index.md
│ │ │ ├── alloydb_pg_admin_mcp.md
│ │ │ ├── alloydb_pg_mcp.md
│ │ │ ├── bigquery_mcp.md
│ │ │ ├── cloud_sql_mssql_admin_mcp.md
│ │ │ ├── cloud_sql_mssql_mcp.md
│ │ │ ├── cloud_sql_mysql_admin_mcp.md
│ │ │ ├── cloud_sql_mysql_mcp.md
│ │ │ ├── cloud_sql_pg_admin_mcp.md
│ │ │ ├── cloud_sql_pg_mcp.md
│ │ │ ├── firestore_mcp.md
│ │ │ ├── looker_mcp.md
│ │ │ ├── mssql_mcp.md
│ │ │ ├── mysql_mcp.md
│ │ │ ├── neo4j_mcp.md
│ │ │ ├── postgres_mcp.md
│ │ │ ├── spanner_mcp.md
│ │ │ └── sqlite_mcp.md
│ │ ├── deploy_docker.md
│ │ ├── deploy_gke.md
│ │ ├── deploy_toolbox.md
│ │ ├── export_telemetry.md
│ │ └── toolbox-ui
│ │ ├── edit-headers.gif
│ │ ├── edit-headers.png
│ │ ├── index.md
│ │ ├── optional-param-checked.png
│ │ ├── optional-param-unchecked.png
│ │ ├── run-tool.gif
│ │ ├── tools.png
│ │ └── toolsets.png
│ ├── reference
│ │ ├── _index.md
│ │ ├── cli.md
│ │ └── prebuilt-tools.md
│ ├── resources
│ │ ├── _index.md
│ │ ├── authServices
│ │ │ ├── _index.md
│ │ │ └── google.md
│ │ ├── sources
│ │ │ ├── _index.md
│ │ │ ├── alloydb-admin.md
│ │ │ ├── alloydb-pg.md
│ │ │ ├── bigquery.md
│ │ │ ├── bigtable.md
│ │ │ ├── cassandra.md
│ │ │ ├── clickhouse.md
│ │ │ ├── cloud-monitoring.md
│ │ │ ├── cloud-sql-admin.md
│ │ │ ├── cloud-sql-mssql.md
│ │ │ ├── cloud-sql-mysql.md
│ │ │ ├── cloud-sql-pg.md
│ │ │ ├── couchbase.md
│ │ │ ├── dataplex.md
│ │ │ ├── dgraph.md
│ │ │ ├── firebird.md
│ │ │ ├── firestore.md
│ │ │ ├── http.md
│ │ │ ├── looker.md
│ │ │ ├── mongodb.md
│ │ │ ├── mssql.md
│ │ │ ├── mysql.md
│ │ │ ├── neo4j.md
│ │ │ ├── oceanbase.md
│ │ │ ├── oracle.md
│ │ │ ├── postgres.md
│ │ │ ├── redis.md
│ │ │ ├── spanner.md
│ │ │ ├── sqlite.md
│ │ │ ├── tidb.md
│ │ │ ├── trino.md
│ │ │ ├── valkey.md
│ │ │ └── yugabytedb.md
│ │ └── tools
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── alloydb-create-cluster.md
│ │ │ ├── alloydb-create-instance.md
│ │ │ ├── alloydb-create-user.md
│ │ │ ├── alloydb-get-cluster.md
│ │ │ ├── alloydb-get-instance.md
│ │ │ ├── alloydb-get-user.md
│ │ │ ├── alloydb-list-clusters.md
│ │ │ ├── alloydb-list-instances.md
│ │ │ ├── alloydb-list-users.md
│ │ │ └── alloydb-wait-for-operation.md
│ │ ├── alloydbainl
│ │ │ ├── _index.md
│ │ │ └── alloydb-ai-nl.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── bigquery-analyze-contribution.md
│ │ │ ├── bigquery-conversational-analytics.md
│ │ │ ├── bigquery-execute-sql.md
│ │ │ ├── bigquery-forecast.md
│ │ │ ├── bigquery-get-dataset-info.md
│ │ │ ├── bigquery-get-table-info.md
│ │ │ ├── bigquery-list-dataset-ids.md
│ │ │ ├── bigquery-list-table-ids.md
│ │ │ ├── bigquery-search-catalog.md
│ │ │ └── bigquery-sql.md
│ │ ├── bigtable
│ │ │ ├── _index.md
│ │ │ └── bigtable-sql.md
│ │ ├── cassandra
│ │ │ ├── _index.md
│ │ │ └── cassandra-cql.md
│ │ ├── clickhouse
│ │ │ ├── _index.md
│ │ │ ├── clickhouse-execute-sql.md
│ │ │ ├── clickhouse-list-databases.md
│ │ │ ├── clickhouse-list-tables.md
│ │ │ └── clickhouse-sql.md
│ │ ├── cloudmonitoring
│ │ │ ├── _index.md
│ │ │ └── cloud-monitoring-query-prometheus.md
│ │ ├── cloudsql
│ │ │ ├── _index.md
│ │ │ ├── cloudsqlcreatedatabase.md
│ │ │ ├── cloudsqlcreateusers.md
│ │ │ ├── cloudsqlgetinstances.md
│ │ │ ├── cloudsqllistdatabases.md
│ │ │ ├── cloudsqllistinstances.md
│ │ │ ├── cloudsqlmssqlcreateinstance.md
│ │ │ ├── cloudsqlmysqlcreateinstance.md
│ │ │ ├── cloudsqlpgcreateinstances.md
│ │ │ └── cloudsqlwaitforoperation.md
│ │ ├── couchbase
│ │ │ ├── _index.md
│ │ │ └── couchbase-sql.md
│ │ ├── dataform
│ │ │ ├── _index.md
│ │ │ └── dataform-compile-local.md
│ │ ├── dataplex
│ │ │ ├── _index.md
│ │ │ ├── dataplex-lookup-entry.md
│ │ │ ├── dataplex-search-aspect-types.md
│ │ │ └── dataplex-search-entries.md
│ │ ├── dgraph
│ │ │ ├── _index.md
│ │ │ └── dgraph-dql.md
│ │ ├── firebird
│ │ │ ├── _index.md
│ │ │ ├── firebird-execute-sql.md
│ │ │ └── firebird-sql.md
│ │ ├── firestore
│ │ │ ├── _index.md
│ │ │ ├── firestore-add-documents.md
│ │ │ ├── firestore-delete-documents.md
│ │ │ ├── firestore-get-documents.md
│ │ │ ├── firestore-get-rules.md
│ │ │ ├── firestore-list-collections.md
│ │ │ ├── firestore-query-collection.md
│ │ │ ├── firestore-query.md
│ │ │ ├── firestore-update-document.md
│ │ │ └── firestore-validate-rules.md
│ │ ├── http
│ │ │ ├── _index.md
│ │ │ └── http.md
│ │ ├── looker
│ │ │ ├── _index.md
│ │ │ ├── looker-add-dashboard-element.md
│ │ │ ├── looker-conversational-analytics.md
│ │ │ ├── looker-get-dashboards.md
│ │ │ ├── looker-get-dimensions.md
│ │ │ ├── looker-get-explores.md
│ │ │ ├── looker-get-filters.md
│ │ │ ├── looker-get-looks.md
│ │ │ ├── looker-get-measures.md
│ │ │ ├── looker-get-models.md
│ │ │ ├── looker-get-parameters.md
│ │ │ ├── looker-health-analyze.md
│ │ │ ├── looker-health-pulse.md
│ │ │ ├── looker-health-vacuum.md
│ │ │ ├── looker-make-dashboard.md
│ │ │ ├── looker-make-look.md
│ │ │ ├── looker-query-sql.md
│ │ │ ├── looker-query-url.md
│ │ │ ├── looker-query.md
│ │ │ └── looker-run-look.md
│ │ ├── mongodb
│ │ │ ├── _index.md
│ │ │ ├── mongodb-aggregate.md
│ │ │ ├── mongodb-delete-many.md
│ │ │ ├── mongodb-delete-one.md
│ │ │ ├── mongodb-find-one.md
│ │ │ ├── mongodb-find.md
│ │ │ ├── mongodb-insert-many.md
│ │ │ ├── mongodb-insert-one.md
│ │ │ ├── mongodb-update-many.md
│ │ │ └── mongodb-update-one.md
│ │ ├── mssql
│ │ │ ├── _index.md
│ │ │ ├── mssql-execute-sql.md
│ │ │ ├── mssql-list-tables.md
│ │ │ └── mssql-sql.md
│ │ ├── mysql
│ │ │ ├── _index.md
│ │ │ ├── mysql-execute-sql.md
│ │ │ ├── mysql-list-active-queries.md
│ │ │ ├── mysql-list-table-fragmentation.md
│ │ │ ├── mysql-list-tables-missing-unique-indexes.md
│ │ │ ├── mysql-list-tables.md
│ │ │ └── mysql-sql.md
│ │ ├── neo4j
│ │ │ ├── _index.md
│ │ │ ├── neo4j-cypher.md
│ │ │ ├── neo4j-execute-cypher.md
│ │ │ └── neo4j-schema.md
│ │ ├── oceanbase
│ │ │ ├── _index.md
│ │ │ ├── oceanbase-execute-sql.md
│ │ │ └── oceanbase-sql.md
│ │ ├── oracle
│ │ │ ├── _index.md
│ │ │ ├── oracle-execute-sql.md
│ │ │ └── oracle-sql.md
│ │ ├── postgres
│ │ │ ├── _index.md
│ │ │ ├── postgres-execute-sql.md
│ │ │ ├── postgres-list-active-queries.md
│ │ │ ├── postgres-list-available-extensions.md
│ │ │ ├── postgres-list-installed-extensions.md
│ │ │ ├── postgres-list-tables.md
│ │ │ └── postgres-sql.md
│ │ ├── redis
│ │ │ ├── _index.md
│ │ │ └── redis.md
│ │ ├── spanner
│ │ │ ├── _index.md
│ │ │ ├── spanner-execute-sql.md
│ │ │ ├── spanner-list-tables.md
│ │ │ └── spanner-sql.md
│ │ ├── sqlite
│ │ │ ├── _index.md
│ │ │ ├── sqlite-execute-sql.md
│ │ │ └── sqlite-sql.md
│ │ ├── tidb
│ │ │ ├── _index.md
│ │ │ ├── tidb-execute-sql.md
│ │ │ └── tidb-sql.md
│ │ ├── trino
│ │ │ ├── _index.md
│ │ │ ├── trino-execute-sql.md
│ │ │ └── trino-sql.md
│ │ ├── utility
│ │ │ ├── _index.md
│ │ │ └── wait.md
│ │ ├── valkey
│ │ │ ├── _index.md
│ │ │ └── valkey.md
│ │ └── yuagbytedb
│ │ ├── _index.md
│ │ └── yugabytedb-sql.md
│ ├── samples
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── ai-nl
│ │ │ │ ├── alloydb_ai_nl.ipynb
│ │ │ │ └── index.md
│ │ │ └── mcp_quickstart.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── colab_quickstart_bigquery.ipynb
│ │ │ ├── local_quickstart.md
│ │ │ └── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── looker
│ │ ├── _index.md
│ │ ├── looker_gemini_oauth
│ │ │ ├── _index.md
│ │ │ ├── authenticated.png
│ │ │ ├── authorize.png
│ │ │ └── registration.png
│ │ ├── looker_gemini.md
│ │ └── looker_mcp_inspector
│ │ ├── _index.md
│ │ ├── inspector_tools.png
│ │ └── inspector.png
│ └── sdks
│ ├── _index.md
│ ├── go-sdk.md
│ ├── js-sdk.md
│ └── python-sdk.md
├── gemini-extension.json
├── go.mod
├── go.sum
├── internal
│ ├── auth
│ │ ├── auth.go
│ │ └── google
│ │ └── google.go
│ ├── log
│ │ ├── handler.go
│ │ ├── log_test.go
│ │ ├── log.go
│ │ └── logger.go
│ ├── prebuiltconfigs
│ │ ├── prebuiltconfigs_test.go
│ │ ├── prebuiltconfigs.go
│ │ └── tools
│ │ ├── alloydb-postgres-admin.yaml
│ │ ├── alloydb-postgres-observability.yaml
│ │ ├── alloydb-postgres.yaml
│ │ ├── bigquery.yaml
│ │ ├── clickhouse.yaml
│ │ ├── cloud-sql-mssql-admin.yaml
│ │ ├── cloud-sql-mssql-observability.yaml
│ │ ├── cloud-sql-mssql.yaml
│ │ ├── cloud-sql-mysql-admin.yaml
│ │ ├── cloud-sql-mysql-observability.yaml
│ │ ├── cloud-sql-mysql.yaml
│ │ ├── cloud-sql-postgres-admin.yaml
│ │ ├── cloud-sql-postgres-observability.yaml
│ │ ├── cloud-sql-postgres.yaml
│ │ ├── dataplex.yaml
│ │ ├── firestore.yaml
│ │ ├── looker-conversational-analytics.yaml
│ │ ├── looker.yaml
│ │ ├── mssql.yaml
│ │ ├── mysql.yaml
│ │ ├── neo4j.yaml
│ │ ├── oceanbase.yaml
│ │ ├── postgres.yaml
│ │ ├── spanner-postgres.yaml
│ │ ├── spanner.yaml
│ │ └── sqlite.yaml
│ ├── server
│ │ ├── api_test.go
│ │ ├── api.go
│ │ ├── common_test.go
│ │ ├── config.go
│ │ ├── mcp
│ │ │ ├── jsonrpc
│ │ │ │ ├── jsonrpc_test.go
│ │ │ │ └── jsonrpc.go
│ │ │ ├── mcp.go
│ │ │ ├── util
│ │ │ │ └── lifecycle.go
│ │ │ ├── v20241105
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ ├── v20250326
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ └── v20250618
│ │ │ ├── method.go
│ │ │ └── types.go
│ │ ├── mcp_test.go
│ │ ├── mcp.go
│ │ ├── server_test.go
│ │ ├── server.go
│ │ ├── static
│ │ │ ├── assets
│ │ │ │ └── mcptoolboxlogo.png
│ │ │ ├── css
│ │ │ │ └── style.css
│ │ │ ├── index.html
│ │ │ ├── js
│ │ │ │ ├── auth.js
│ │ │ │ ├── loadTools.js
│ │ │ │ ├── mainContent.js
│ │ │ │ ├── navbar.js
│ │ │ │ ├── runTool.js
│ │ │ │ ├── toolDisplay.js
│ │ │ │ ├── tools.js
│ │ │ │ └── toolsets.js
│ │ │ ├── tools.html
│ │ │ └── toolsets.html
│ │ ├── web_test.go
│ │ └── web.go
│ ├── sources
│ │ ├── alloydbadmin
│ │ │ ├── alloydbadmin_test.go
│ │ │ └── alloydbadmin.go
│ │ ├── alloydbpg
│ │ │ ├── alloydb_pg_test.go
│ │ │ └── alloydb_pg.go
│ │ ├── bigquery
│ │ │ ├── bigquery_test.go
│ │ │ └── bigquery.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ ├── cassandra_test.go
│ │ │ └── cassandra.go
│ │ ├── clickhouse
│ │ │ ├── clickhouse_test.go
│ │ │ └── clickhouse.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloud_monitoring_test.go
│ │ │ └── cloud_monitoring.go
│ │ ├── cloudsqladmin
│ │ │ ├── cloud_sql_admin_test.go
│ │ │ └── cloud_sql_admin.go
│ │ ├── cloudsqlmssql
│ │ │ ├── cloud_sql_mssql_test.go
│ │ │ └── cloud_sql_mssql.go
│ │ ├── cloudsqlmysql
│ │ │ ├── cloud_sql_mysql_test.go
│ │ │ └── cloud_sql_mysql.go
│ │ ├── cloudsqlpg
│ │ │ ├── cloud_sql_pg_test.go
│ │ │ └── cloud_sql_pg.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataplex
│ │ │ ├── dataplex_test.go
│ │ │ └── dataplex.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── dialect.go
│ │ ├── firebird
│ │ │ ├── firebird_test.go
│ │ │ └── firebird.go
│ │ ├── firestore
│ │ │ ├── firestore_test.go
│ │ │ └── firestore.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── ip_type.go
│ │ ├── looker
│ │ │ ├── looker_test.go
│ │ │ └── looker.go
│ │ ├── mongodb
│ │ │ ├── mongodb_test.go
│ │ │ └── mongodb.go
│ │ ├── mssql
│ │ │ ├── mssql_test.go
│ │ │ └── mssql.go
│ │ ├── mysql
│ │ │ ├── mysql_test.go
│ │ │ └── mysql.go
│ │ ├── neo4j
│ │ │ ├── neo4j_test.go
│ │ │ └── neo4j.go
│ │ ├── oceanbase
│ │ │ ├── oceanbase_test.go
│ │ │ └── oceanbase.go
│ │ ├── oracle
│ │ │ └── oracle.go
│ │ ├── postgres
│ │ │ ├── postgres_test.go
│ │ │ └── postgres.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── sources.go
│ │ ├── spanner
│ │ │ ├── spanner_test.go
│ │ │ └── spanner.go
│ │ ├── sqlite
│ │ │ ├── sqlite_test.go
│ │ │ └── sqlite.go
│ │ ├── tidb
│ │ │ ├── tidb_test.go
│ │ │ └── tidb.go
│ │ ├── trino
│ │ │ ├── trino_test.go
│ │ │ └── trino.go
│ │ ├── util.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedb
│ │ ├── yugabytedb_test.go
│ │ └── yugabytedb.go
│ ├── telemetry
│ │ ├── instrumentation.go
│ │ └── telemetry.go
│ ├── testutils
│ │ └── testutils.go
│ ├── tools
│ │ ├── alloydb
│ │ │ ├── alloydbcreatecluster
│ │ │ │ ├── alloydbcreatecluster_test.go
│ │ │ │ └── alloydbcreatecluster.go
│ │ │ ├── alloydbcreateinstance
│ │ │ │ ├── alloydbcreateinstance_test.go
│ │ │ │ └── alloydbcreateinstance.go
│ │ │ ├── alloydbcreateuser
│ │ │ │ ├── alloydbcreateuser_test.go
│ │ │ │ └── alloydbcreateuser.go
│ │ │ ├── alloydbgetcluster
│ │ │ │ ├── alloydbgetcluster_test.go
│ │ │ │ └── alloydbgetcluster.go
│ │ │ ├── alloydbgetinstance
│ │ │ │ ├── alloydbgetinstance_test.go
│ │ │ │ └── alloydbgetinstance.go
│ │ │ ├── alloydbgetuser
│ │ │ │ ├── alloydbgetuser_test.go
│ │ │ │ └── alloydbgetuser.go
│ │ │ ├── alloydblistclusters
│ │ │ │ ├── alloydblistclusters_test.go
│ │ │ │ └── alloydblistclusters.go
│ │ │ ├── alloydblistinstances
│ │ │ │ ├── alloydblistinstances_test.go
│ │ │ │ └── alloydblistinstances.go
│ │ │ ├── alloydblistusers
│ │ │ │ ├── alloydblistusers_test.go
│ │ │ │ └── alloydblistusers.go
│ │ │ └── alloydbwaitforoperation
│ │ │ ├── alloydbwaitforoperation_test.go
│ │ │ └── alloydbwaitforoperation.go
│ │ ├── alloydbainl
│ │ │ ├── alloydbainl_test.go
│ │ │ └── alloydbainl.go
│ │ ├── bigquery
│ │ │ ├── bigqueryanalyzecontribution
│ │ │ │ ├── bigqueryanalyzecontribution_test.go
│ │ │ │ └── bigqueryanalyzecontribution.go
│ │ │ ├── bigquerycommon
│ │ │ │ ├── table_name_parser_test.go
│ │ │ │ ├── table_name_parser.go
│ │ │ │ └── util.go
│ │ │ ├── bigqueryconversationalanalytics
│ │ │ │ ├── bigqueryconversationalanalytics_test.go
│ │ │ │ └── bigqueryconversationalanalytics.go
│ │ │ ├── bigqueryexecutesql
│ │ │ │ ├── bigqueryexecutesql_test.go
│ │ │ │ └── bigqueryexecutesql.go
│ │ │ ├── bigqueryforecast
│ │ │ │ ├── bigqueryforecast_test.go
│ │ │ │ └── bigqueryforecast.go
│ │ │ ├── bigquerygetdatasetinfo
│ │ │ │ ├── bigquerygetdatasetinfo_test.go
│ │ │ │ └── bigquerygetdatasetinfo.go
│ │ │ ├── bigquerygettableinfo
│ │ │ │ ├── bigquerygettableinfo_test.go
│ │ │ │ └── bigquerygettableinfo.go
│ │ │ ├── bigquerylistdatasetids
│ │ │ │ ├── bigquerylistdatasetids_test.go
│ │ │ │ └── bigquerylistdatasetids.go
│ │ │ ├── bigquerylisttableids
│ │ │ │ ├── bigquerylisttableids_test.go
│ │ │ │ └── bigquerylisttableids.go
│ │ │ ├── bigquerysearchcatalog
│ │ │ │ ├── bigquerysearchcatalog_test.go
│ │ │ │ └── bigquerysearchcatalog.go
│ │ │ └── bigquerysql
│ │ │ ├── bigquerysql_test.go
│ │ │ └── bigquerysql.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ └── cassandracql
│ │ │ ├── cassandracql_test.go
│ │ │ └── cassandracql.go
│ │ ├── clickhouse
│ │ │ ├── clickhouseexecutesql
│ │ │ │ ├── clickhouseexecutesql_test.go
│ │ │ │ └── clickhouseexecutesql.go
│ │ │ ├── clickhouselistdatabases
│ │ │ │ ├── clickhouselistdatabases_test.go
│ │ │ │ └── clickhouselistdatabases.go
│ │ │ ├── clickhouselisttables
│ │ │ │ ├── clickhouselisttables_test.go
│ │ │ │ └── clickhouselisttables.go
│ │ │ └── clickhousesql
│ │ │ ├── clickhousesql_test.go
│ │ │ └── clickhousesql.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloudmonitoring_test.go
│ │ │ └── cloudmonitoring.go
│ │ ├── cloudsql
│ │ │ ├── cloudsqlcreatedatabase
│ │ │ │ ├── cloudsqlcreatedatabase_test.go
│ │ │ │ └── cloudsqlcreatedatabase.go
│ │ │ ├── cloudsqlcreateusers
│ │ │ │ ├── cloudsqlcreateusers_test.go
│ │ │ │ └── cloudsqlcreateusers.go
│ │ │ ├── cloudsqlgetinstances
│ │ │ │ ├── cloudsqlgetinstances_test.go
│ │ │ │ └── cloudsqlgetinstances.go
│ │ │ ├── cloudsqllistdatabases
│ │ │ │ ├── cloudsqllistdatabases_test.go
│ │ │ │ └── cloudsqllistdatabases.go
│ │ │ ├── cloudsqllistinstances
│ │ │ │ ├── cloudsqllistinstances_test.go
│ │ │ │ └── cloudsqllistinstances.go
│ │ │ └── cloudsqlwaitforoperation
│ │ │ ├── cloudsqlwaitforoperation_test.go
│ │ │ └── cloudsqlwaitforoperation.go
│ │ ├── cloudsqlmssql
│ │ │ └── cloudsqlmssqlcreateinstance
│ │ │ ├── cloudsqlmssqlcreateinstance_test.go
│ │ │ └── cloudsqlmssqlcreateinstance.go
│ │ ├── cloudsqlmysql
│ │ │ └── cloudsqlmysqlcreateinstance
│ │ │ ├── cloudsqlmysqlcreateinstance_test.go
│ │ │ └── cloudsqlmysqlcreateinstance.go
│ │ ├── cloudsqlpg
│ │ │ └── cloudsqlpgcreateinstances
│ │ │ ├── cloudsqlpgcreateinstances_test.go
│ │ │ └── cloudsqlpgcreateinstances.go
│ │ ├── common_test.go
│ │ ├── common.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataform
│ │ │ └── dataformcompilelocal
│ │ │ ├── dataformcompilelocal_test.go
│ │ │ └── dataformcompilelocal.go
│ │ ├── dataplex
│ │ │ ├── dataplexlookupentry
│ │ │ │ ├── dataplexlookupentry_test.go
│ │ │ │ └── dataplexlookupentry.go
│ │ │ ├── dataplexsearchaspecttypes
│ │ │ │ ├── dataplexsearchaspecttypes_test.go
│ │ │ │ └── dataplexsearchaspecttypes.go
│ │ │ └── dataplexsearchentries
│ │ │ ├── dataplexsearchentries_test.go
│ │ │ └── dataplexsearchentries.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── firebird
│ │ │ ├── firebirdexecutesql
│ │ │ │ ├── firebirdexecutesql_test.go
│ │ │ │ └── firebirdexecutesql.go
│ │ │ └── firebirdsql
│ │ │ ├── firebirdsql_test.go
│ │ │ └── firebirdsql.go
│ │ ├── firestore
│ │ │ ├── firestoreadddocuments
│ │ │ │ ├── firestoreadddocuments_test.go
│ │ │ │ └── firestoreadddocuments.go
│ │ │ ├── firestoredeletedocuments
│ │ │ │ ├── firestoredeletedocuments_test.go
│ │ │ │ └── firestoredeletedocuments.go
│ │ │ ├── firestoregetdocuments
│ │ │ │ ├── firestoregetdocuments_test.go
│ │ │ │ └── firestoregetdocuments.go
│ │ │ ├── firestoregetrules
│ │ │ │ ├── firestoregetrules_test.go
│ │ │ │ └── firestoregetrules.go
│ │ │ ├── firestorelistcollections
│ │ │ │ ├── firestorelistcollections_test.go
│ │ │ │ └── firestorelistcollections.go
│ │ │ ├── firestorequery
│ │ │ │ ├── firestorequery_test.go
│ │ │ │ └── firestorequery.go
│ │ │ ├── firestorequerycollection
│ │ │ │ ├── firestorequerycollection_test.go
│ │ │ │ └── firestorequerycollection.go
│ │ │ ├── firestoreupdatedocument
│ │ │ │ ├── firestoreupdatedocument_test.go
│ │ │ │ └── firestoreupdatedocument.go
│ │ │ ├── firestorevalidaterules
│ │ │ │ ├── firestorevalidaterules_test.go
│ │ │ │ └── firestorevalidaterules.go
│ │ │ └── util
│ │ │ ├── converter_test.go
│ │ │ ├── converter.go
│ │ │ ├── validator_test.go
│ │ │ └── validator.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── http_method.go
│ │ ├── looker
│ │ │ ├── lookeradddashboardelement
│ │ │ │ ├── lookeradddashboardelement_test.go
│ │ │ │ └── lookeradddashboardelement.go
│ │ │ ├── lookercommon
│ │ │ │ ├── lookercommon_test.go
│ │ │ │ └── lookercommon.go
│ │ │ ├── lookerconversationalanalytics
│ │ │ │ ├── lookerconversationalanalytics_test.go
│ │ │ │ └── lookerconversationalanalytics.go
│ │ │ ├── lookergetdashboards
│ │ │ │ ├── lookergetdashboards_test.go
│ │ │ │ └── lookergetdashboards.go
│ │ │ ├── lookergetdimensions
│ │ │ │ ├── lookergetdimensions_test.go
│ │ │ │ └── lookergetdimensions.go
│ │ │ ├── lookergetexplores
│ │ │ │ ├── lookergetexplores_test.go
│ │ │ │ └── lookergetexplores.go
│ │ │ ├── lookergetfilters
│ │ │ │ ├── lookergetfilters_test.go
│ │ │ │ └── lookergetfilters.go
│ │ │ ├── lookergetlooks
│ │ │ │ ├── lookergetlooks_test.go
│ │ │ │ └── lookergetlooks.go
│ │ │ ├── lookergetmeasures
│ │ │ │ ├── lookergetmeasures_test.go
│ │ │ │ └── lookergetmeasures.go
│ │ │ ├── lookergetmodels
│ │ │ │ ├── lookergetmodels_test.go
│ │ │ │ └── lookergetmodels.go
│ │ │ ├── lookergetparameters
│ │ │ │ ├── lookergetparameters_test.go
│ │ │ │ └── lookergetparameters.go
│ │ │ ├── lookerhealthanalyze
│ │ │ │ ├── lookerhealthanalyze_test.go
│ │ │ │ └── lookerhealthanalyze.go
│ │ │ ├── lookerhealthpulse
│ │ │ │ ├── lookerhealthpulse_test.go
│ │ │ │ └── lookerhealthpulse.go
│ │ │ ├── lookerhealthvacuum
│ │ │ │ ├── lookerhealthvacuum_test.go
│ │ │ │ └── lookerhealthvacuum.go
│ │ │ ├── lookermakedashboard
│ │ │ │ ├── lookermakedashboard_test.go
│ │ │ │ └── lookermakedashboard.go
│ │ │ ├── lookermakelook
│ │ │ │ ├── lookermakelook_test.go
│ │ │ │ └── lookermakelook.go
│ │ │ ├── lookerquery
│ │ │ │ ├── lookerquery_test.go
│ │ │ │ └── lookerquery.go
│ │ │ ├── lookerquerysql
│ │ │ │ ├── lookerquerysql_test.go
│ │ │ │ └── lookerquerysql.go
│ │ │ ├── lookerqueryurl
│ │ │ │ ├── lookerqueryurl_test.go
│ │ │ │ └── lookerqueryurl.go
│ │ │ └── lookerrunlook
│ │ │ ├── lookerrunlook_test.go
│ │ │ └── lookerrunlook.go
│ │ ├── mongodb
│ │ │ ├── mongodbaggregate
│ │ │ │ ├── mongodbaggregate_test.go
│ │ │ │ └── mongodbaggregate.go
│ │ │ ├── mongodbdeletemany
│ │ │ │ ├── mongodbdeletemany_test.go
│ │ │ │ └── mongodbdeletemany.go
│ │ │ ├── mongodbdeleteone
│ │ │ │ ├── mongodbdeleteone_test.go
│ │ │ │ └── mongodbdeleteone.go
│ │ │ ├── mongodbfind
│ │ │ │ ├── mongodbfind_test.go
│ │ │ │ └── mongodbfind.go
│ │ │ ├── mongodbfindone
│ │ │ │ ├── mongodbfindone_test.go
│ │ │ │ └── mongodbfindone.go
│ │ │ ├── mongodbinsertmany
│ │ │ │ ├── mongodbinsertmany_test.go
│ │ │ │ └── mongodbinsertmany.go
│ │ │ ├── mongodbinsertone
│ │ │ │ ├── mongodbinsertone_test.go
│ │ │ │ └── mongodbinsertone.go
│ │ │ ├── mongodbupdatemany
│ │ │ │ ├── mongodbupdatemany_test.go
│ │ │ │ └── mongodbupdatemany.go
│ │ │ └── mongodbupdateone
│ │ │ ├── mongodbupdateone_test.go
│ │ │ └── mongodbupdateone.go
│ │ ├── mssql
│ │ │ ├── mssqlexecutesql
│ │ │ │ ├── mssqlexecutesql_test.go
│ │ │ │ └── mssqlexecutesql.go
│ │ │ ├── mssqllisttables
│ │ │ │ ├── mssqllisttables_test.go
│ │ │ │ └── mssqllisttables.go
│ │ │ └── mssqlsql
│ │ │ ├── mssqlsql_test.go
│ │ │ └── mssqlsql.go
│ │ ├── mysql
│ │ │ ├── mysqlcommon
│ │ │ │ └── mysqlcommon.go
│ │ │ ├── mysqlexecutesql
│ │ │ │ ├── mysqlexecutesql_test.go
│ │ │ │ └── mysqlexecutesql.go
│ │ │ ├── mysqllistactivequeries
│ │ │ │ ├── mysqllistactivequeries_test.go
│ │ │ │ └── mysqllistactivequeries.go
│ │ │ ├── mysqllisttablefragmentation
│ │ │ │ ├── mysqllisttablefragmentation_test.go
│ │ │ │ └── mysqllisttablefragmentation.go
│ │ │ ├── mysqllisttables
│ │ │ │ ├── mysqllisttables_test.go
│ │ │ │ └── mysqllisttables.go
│ │ │ ├── mysqllisttablesmissinguniqueindexes
│ │ │ │ ├── mysqllisttablesmissinguniqueindexes_test.go
│ │ │ │ └── mysqllisttablesmissinguniqueindexes.go
│ │ │ └── mysqlsql
│ │ │ ├── mysqlsql_test.go
│ │ │ └── mysqlsql.go
│ │ ├── neo4j
│ │ │ ├── neo4jcypher
│ │ │ │ ├── neo4jcypher_test.go
│ │ │ │ └── neo4jcypher.go
│ │ │ ├── neo4jexecutecypher
│ │ │ │ ├── classifier
│ │ │ │ │ ├── classifier_test.go
│ │ │ │ │ └── classifier.go
│ │ │ │ ├── neo4jexecutecypher_test.go
│ │ │ │ └── neo4jexecutecypher.go
│ │ │ └── neo4jschema
│ │ │ ├── cache
│ │ │ │ ├── cache_test.go
│ │ │ │ └── cache.go
│ │ │ ├── helpers
│ │ │ │ ├── helpers_test.go
│ │ │ │ └── helpers.go
│ │ │ ├── neo4jschema_test.go
│ │ │ ├── neo4jschema.go
│ │ │ └── types
│ │ │ └── types.go
│ │ ├── oceanbase
│ │ │ ├── oceanbaseexecutesql
│ │ │ │ ├── oceanbaseexecutesql_test.go
│ │ │ │ └── oceanbaseexecutesql.go
│ │ │ └── oceanbasesql
│ │ │ ├── oceanbasesql_test.go
│ │ │ └── oceanbasesql.go
│ │ ├── oracle
│ │ │ ├── oracleexecutesql
│ │ │ │ └── oracleexecutesql.go
│ │ │ └── oraclesql
│ │ │ └── oraclesql.go
│ │ ├── parameters_test.go
│ │ ├── parameters.go
│ │ ├── postgres
│ │ │ ├── postgresexecutesql
│ │ │ │ ├── postgresexecutesql_test.go
│ │ │ │ └── postgresexecutesql.go
│ │ │ ├── postgreslistactivequeries
│ │ │ │ ├── postgreslistactivequeries_test.go
│ │ │ │ └── postgreslistactivequeries.go
│ │ │ ├── postgreslistavailableextensions
│ │ │ │ ├── postgreslistavailableextensions_test.go
│ │ │ │ └── postgreslistavailableextensions.go
│ │ │ ├── postgreslistinstalledextensions
│ │ │ │ ├── postgreslistinstalledextensions_test.go
│ │ │ │ └── postgreslistinstalledextensions.go
│ │ │ ├── postgreslisttables
│ │ │ │ ├── postgreslisttables_test.go
│ │ │ │ └── postgreslisttables.go
│ │ │ └── postgressql
│ │ │ ├── postgressql_test.go
│ │ │ └── postgressql.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── spanner
│ │ │ ├── spannerexecutesql
│ │ │ │ ├── spannerexecutesql_test.go
│ │ │ │ └── spannerexecutesql.go
│ │ │ ├── spannerlisttables
│ │ │ │ ├── spannerlisttables_test.go
│ │ │ │ └── spannerlisttables.go
│ │ │ └── spannersql
│ │ │ ├── spanner_test.go
│ │ │ └── spannersql.go
│ │ ├── sqlite
│ │ │ ├── sqliteexecutesql
│ │ │ │ ├── sqliteexecutesql_test.go
│ │ │ │ └── sqliteexecutesql.go
│ │ │ └── sqlitesql
│ │ │ ├── sqlitesql_test.go
│ │ │ └── sqlitesql.go
│ │ ├── tidb
│ │ │ ├── tidbexecutesql
│ │ │ │ ├── tidbexecutesql_test.go
│ │ │ │ └── tidbexecutesql.go
│ │ │ └── tidbsql
│ │ │ ├── tidbsql_test.go
│ │ │ └── tidbsql.go
│ │ ├── tools_test.go
│ │ ├── tools.go
│ │ ├── toolsets.go
│ │ ├── trino
│ │ │ ├── trinoexecutesql
│ │ │ │ ├── trinoexecutesql_test.go
│ │ │ │ └── trinoexecutesql.go
│ │ │ └── trinosql
│ │ │ ├── trinosql_test.go
│ │ │ └── trinosql.go
│ │ ├── utility
│ │ │ └── wait
│ │ │ ├── wait_test.go
│ │ │ └── wait.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedbsql
│ │ ├── yugabytedbsql_test.go
│ │ └── yugabytedbsql.go
│ └── util
│ └── util.go
├── LICENSE
├── logo.png
├── main.go
├── MCP-TOOLBOX-EXTENSION.md
├── README.md
└── tests
├── alloydb
│ ├── alloydb_integration_test.go
│ └── alloydb_wait_for_operation_test.go
├── alloydbainl
│ └── alloydb_ai_nl_integration_test.go
├── alloydbpg
│ └── alloydb_pg_integration_test.go
├── auth.go
├── bigquery
│ └── bigquery_integration_test.go
├── bigtable
│ └── bigtable_integration_test.go
├── cassandra
│ └── cassandra_integration_test.go
├── clickhouse
│ └── clickhouse_integration_test.go
├── cloudmonitoring
│ └── cloud_monitoring_integration_test.go
├── cloudsql
│ ├── cloud_sql_create_database_test.go
│ ├── cloud_sql_create_users_test.go
│ ├── cloud_sql_get_instances_test.go
│ ├── cloud_sql_list_databases_test.go
│ ├── cloudsql_list_instances_test.go
│ └── cloudsql_wait_for_operation_test.go
├── cloudsqlmssql
│ ├── cloud_sql_mssql_create_instance_integration_test.go
│ └── cloud_sql_mssql_integration_test.go
├── cloudsqlmysql
│ ├── cloud_sql_mysql_create_instance_integration_test.go
│ └── cloud_sql_mysql_integration_test.go
├── cloudsqlpg
│ ├── cloud_sql_pg_create_instances_test.go
│ └── cloud_sql_pg_integration_test.go
├── common.go
├── couchbase
│ └── couchbase_integration_test.go
├── dataform
│ └── dataform_integration_test.go
├── dataplex
│ └── dataplex_integration_test.go
├── dgraph
│ └── dgraph_integration_test.go
├── firebird
│ └── firebird_integration_test.go
├── firestore
│ └── firestore_integration_test.go
├── http
│ └── http_integration_test.go
├── looker
│ └── looker_integration_test.go
├── mongodb
│ └── mongodb_integration_test.go
├── mssql
│ └── mssql_integration_test.go
├── mysql
│ └── mysql_integration_test.go
├── neo4j
│ └── neo4j_integration_test.go
├── oceanbase
│ └── oceanbase_integration_test.go
├── option.go
├── oracle
│ └── oracle_integration_test.go
├── postgres
│ └── postgres_integration_test.go
├── redis
│ └── redis_test.go
├── server.go
├── source.go
├── spanner
│ └── spanner_integration_test.go
├── sqlite
│ └── sqlite_integration_test.go
├── tidb
│ └── tidb_integration_test.go
├── tool.go
├── trino
│ └── trino_integration_test.go
├── utility
│ └── wait_integration_test.go
├── valkey
│ └── valkey_test.go
└── yugabytedb
└── yugabytedb_integration_test.go
```
# Files
--------------------------------------------------------------------------------
/docs/en/getting-started/quickstart/js/genAI/quickstart.js:
--------------------------------------------------------------------------------
```javascript
1 | import { GoogleGenAI } from "@google/genai";
2 | import { ToolboxClient } from "@toolbox-sdk/core";
3 |
4 |
5 | const TOOLBOX_URL = "http://127.0.0.1:5000"; // Update if needed
6 | const GOOGLE_API_KEY = process.env.GOOGLE_API_KEY || 'your-api-key'; // Replace it with your API key
7 |
8 | const prompt = `
9 | You're a helpful hotel assistant. You handle hotel searching, booking, and
10 | cancellations. When the user searches for a hotel, you MUST use the available tools to find information. Mention its name, id,
11 | location and price tier. Always mention hotel id while performing any
12 | searches. This is very important for any operations. For any bookings or
13 | cancellations, please provide the appropriate confirmation. Be sure to
14 | update checkin or checkout dates if mentioned by the user.
15 | Don't ask for confirmations from the user.
16 | `;
17 |
18 | const queries = [
19 | "Find hotels in Basel with Basel in its name.",
20 | "Can you book the Hilton Basel for me?",
21 | "Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.",
22 | "My check in dates would be from April 10, 2024 to April 19, 2024.",
23 | ];
24 |
25 | function mapZodTypeToOpenAPIType(zodTypeName) {
26 |
27 | console.log(zodTypeName)
28 | const typeMap = {
29 | 'ZodString': 'string',
30 | 'ZodNumber': 'number',
31 | 'ZodBoolean': 'boolean',
32 | 'ZodArray': 'array',
33 | 'ZodObject': 'object',
34 | };
35 | return typeMap[zodTypeName] || 'string';
36 | }
37 |
38 | export async function main() {
39 |
40 | const toolboxClient = new ToolboxClient(TOOLBOX_URL);
41 | const toolboxTools = await toolboxClient.loadToolset("my-toolset");
42 |
43 | const geminiTools = [{
44 | functionDeclarations: toolboxTools.map(tool => {
45 |
46 | const schema = tool.getParamSchema();
47 | const properties = {};
48 | const required = [];
49 |
50 |
51 | for (const [key, param] of Object.entries(schema.shape)) {
52 | properties[key] = {
53 | type: mapZodTypeToOpenAPIType(param.constructor.name),
54 | description: param.description || '',
55 | };
56 | required.push(key)
57 | }
58 |
59 | return {
60 | name: tool.getName(),
61 | description: tool.getDescription(),
62 | parameters: { type: 'object', properties, required },
63 | };
64 | })
65 | }];
66 |
67 |
68 | const genAI = new GoogleGenAI({ apiKey: GOOGLE_API_KEY });
69 |
70 | const chat = genAI.chats.create({
71 | model: "gemini-2.5-flash",
72 | config: {
73 | systemInstruction: prompt,
74 | tools: geminiTools,
75 | }
76 | });
77 |
78 | for (const query of queries) {
79 |
80 | let currentResult = await chat.sendMessage({ message: query });
81 |
82 | let finalResponseGiven = false
83 | while (!finalResponseGiven) {
84 |
85 | const response = currentResult;
86 | const functionCalls = response.functionCalls || [];
87 |
88 | if (functionCalls.length === 0) {
89 | console.log(response.text)
90 | finalResponseGiven = true;
91 | } else {
92 | const toolResponses = [];
93 | for (const call of functionCalls) {
94 | const toolName = call.name
95 | const toolToExecute = toolboxTools.find(t => t.getName() === toolName);
96 |
97 | if (toolToExecute) {
98 | try {
99 | const functionResult = await toolToExecute(call.args);
100 | toolResponses.push({
101 | functionResponse: { name: call.name, response: { result: functionResult } }
102 | });
103 | } catch (e) {
104 | console.error(`Error executing tool '${toolName}':`, e);
105 | toolResponses.push({
106 | functionResponse: { name: call.name, response: { error: e.message } }
107 | });
108 | }
109 | }
110 | }
111 |
112 | currentResult = await chat.sendMessage({ message: toolResponses });
113 | }
114 | }
115 |
116 | }
117 | }
118 |
119 | main();
```
--------------------------------------------------------------------------------
/internal/sources/alloydbadmin/alloydbadmin.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 | package alloydbadmin
15 |
16 | import (
17 | "context"
18 | "fmt"
19 | "net/http"
20 |
21 | "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/googleapis/genai-toolbox/internal/util"
24 | "go.opentelemetry.io/otel/trace"
25 | "golang.org/x/oauth2"
26 | "golang.org/x/oauth2/google"
27 | alloydbrestapi "google.golang.org/api/alloydb/v1"
28 | "google.golang.org/api/option"
29 | )
30 |
31 | const SourceKind string = "alloydb-admin"
32 |
33 | type userAgentRoundTripper struct {
34 | userAgent string
35 | next http.RoundTripper
36 | }
37 |
38 | func (rt *userAgentRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
39 | newReq := *req
40 | newReq.Header = make(http.Header)
41 | for k, v := range req.Header {
42 | newReq.Header[k] = v
43 | }
44 | ua := newReq.Header.Get("User-Agent")
45 | if ua == "" {
46 | newReq.Header.Set("User-Agent", rt.userAgent)
47 | } else {
48 | newReq.Header.Set("User-Agent", ua+" "+rt.userAgent)
49 | }
50 | return rt.next.RoundTrip(&newReq)
51 | }
52 |
53 | // validate interface
54 | var _ sources.SourceConfig = Config{}
55 |
56 | func init() {
57 | if !sources.Register(SourceKind, newConfig) {
58 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
59 | }
60 | }
61 |
62 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
63 | actual := Config{Name: name}
64 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
65 | return nil, err
66 | }
67 | return actual, nil
68 | }
69 |
70 | type Config struct {
71 | Name string `yaml:"name" validate:"required"`
72 | Kind string `yaml:"kind" validate:"required"`
73 | UseClientOAuth bool `yaml:"useClientOAuth"`
74 | }
75 |
76 | func (r Config) SourceConfigKind() string {
77 | return SourceKind
78 | }
79 |
80 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
81 | ua, err := util.UserAgentFromContext(ctx)
82 | if err != nil {
83 | fmt.Printf("Error in User Agent retrieval: %s", err)
84 | }
85 |
86 | var client *http.Client
87 | if r.UseClientOAuth {
88 | client = &http.Client{
89 | Transport: &userAgentRoundTripper{
90 | userAgent: ua,
91 | next: http.DefaultTransport,
92 | },
93 | }
94 | } else {
95 | // Use Application Default Credentials
96 | creds, err := google.FindDefaultCredentials(ctx, alloydbrestapi.CloudPlatformScope)
97 | if err != nil {
98 | return nil, fmt.Errorf("failed to find default credentials: %w", err)
99 | }
100 | baseClient := oauth2.NewClient(ctx, creds.TokenSource)
101 | baseClient.Transport = &userAgentRoundTripper{
102 | userAgent: ua,
103 | next: baseClient.Transport,
104 | }
105 | client = baseClient
106 | }
107 |
108 | service, err := alloydbrestapi.NewService(ctx, option.WithHTTPClient(client))
109 | if err != nil {
110 | return nil, fmt.Errorf("error creating new alloydb service: %w", err)
111 | }
112 |
113 | s := &Source{
114 | Name: r.Name,
115 | Kind: SourceKind,
116 | BaseURL: "https://alloydb.googleapis.com",
117 | Service: service,
118 | UseClientOAuth: r.UseClientOAuth,
119 | }
120 |
121 | return s, nil
122 | }
123 |
124 | var _ sources.Source = &Source{}
125 |
126 | type Source struct {
127 | Name string `yaml:"name"`
128 | Kind string `yaml:"kind"`
129 | BaseURL string
130 | Service *alloydbrestapi.Service
131 | UseClientOAuth bool
132 | }
133 |
134 | func (s *Source) SourceKind() string {
135 | return SourceKind
136 | }
137 |
138 | func (s *Source) GetService(ctx context.Context, accessToken string) (*alloydbrestapi.Service, error) {
139 | if s.UseClientOAuth {
140 | token := &oauth2.Token{AccessToken: accessToken}
141 | client := oauth2.NewClient(ctx, oauth2.StaticTokenSource(token))
142 | service, err := alloydbrestapi.NewService(ctx, option.WithHTTPClient(client))
143 | if err != nil {
144 | return nil, fmt.Errorf("error creating new alloydb service: %w", err)
145 | }
146 | return service, nil
147 | }
148 | return s.Service, nil
149 | }
150 |
151 | func (s *Source) UseClientAuthorization() bool {
152 | return s.UseClientOAuth
153 | }
154 |
```
--------------------------------------------------------------------------------
/internal/sources/firestore/firestore.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestore
16 |
17 | import (
18 | "context"
19 | "fmt"
20 |
21 | "cloud.google.com/go/firestore"
22 | "github.com/goccy/go-yaml"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | "github.com/googleapis/genai-toolbox/internal/util"
25 | "go.opentelemetry.io/otel/trace"
26 | "google.golang.org/api/firebaserules/v1"
27 | "google.golang.org/api/option"
28 | )
29 |
30 | const SourceKind string = "firestore"
31 |
32 | // validate interface
33 | var _ sources.SourceConfig = Config{}
34 |
35 | func init() {
36 | if !sources.Register(SourceKind, newConfig) {
37 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
38 | }
39 | }
40 |
41 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
42 | actual := Config{Name: name}
43 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
44 | return nil, err
45 | }
46 | return actual, nil
47 | }
48 |
49 | type Config struct {
50 | // Firestore configs
51 | Name string `yaml:"name" validate:"required"`
52 | Kind string `yaml:"kind" validate:"required"`
53 | Project string `yaml:"project" validate:"required"`
54 | Database string `yaml:"database"` // Optional, defaults to "(default)"
55 | }
56 |
57 | func (r Config) SourceConfigKind() string {
58 | // Returns Firestore source kind
59 | return SourceKind
60 | }
61 |
62 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
63 | // Initializes a Firestore source
64 | client, err := initFirestoreConnection(ctx, tracer, r.Name, r.Project, r.Database)
65 | if err != nil {
66 | return nil, err
67 | }
68 |
69 | // Initialize Firebase Rules client
70 | rulesClient, err := initFirebaseRulesConnection(ctx, r.Project)
71 | if err != nil {
72 | return nil, fmt.Errorf("failed to initialize Firebase Rules client: %w", err)
73 | }
74 |
75 | s := &Source{
76 | Name: r.Name,
77 | Kind: SourceKind,
78 | Client: client,
79 | RulesClient: rulesClient,
80 | ProjectId: r.Project,
81 | DatabaseId: r.Database,
82 | }
83 | return s, nil
84 | }
85 |
86 | var _ sources.Source = &Source{}
87 |
88 | type Source struct {
89 | // Firestore struct with client
90 | Name string `yaml:"name"`
91 | Kind string `yaml:"kind"`
92 | Client *firestore.Client
93 | RulesClient *firebaserules.Service
94 | ProjectId string `yaml:"projectId"`
95 | DatabaseId string `yaml:"databaseId"`
96 | }
97 |
98 | func (s *Source) SourceKind() string {
99 | // Returns Firestore source kind
100 | return SourceKind
101 | }
102 |
103 | func (s *Source) FirestoreClient() *firestore.Client {
104 | return s.Client
105 | }
106 |
107 | func (s *Source) FirebaseRulesClient() *firebaserules.Service {
108 | return s.RulesClient
109 | }
110 |
111 | func (s *Source) GetProjectId() string {
112 | return s.ProjectId
113 | }
114 |
115 | func (s *Source) GetDatabaseId() string {
116 | return s.DatabaseId
117 | }
118 |
119 | func initFirestoreConnection(
120 | ctx context.Context,
121 | tracer trace.Tracer,
122 | name string,
123 | project string,
124 | database string,
125 | ) (*firestore.Client, error) {
126 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
127 | defer span.End()
128 |
129 | userAgent, err := util.UserAgentFromContext(ctx)
130 | if err != nil {
131 | return nil, err
132 | }
133 |
134 | // If database is not specified, use the default database
135 | if database == "" {
136 | database = "(default)"
137 | }
138 |
139 | // Create the Firestore client
140 | client, err := firestore.NewClientWithDatabase(ctx, project, database, option.WithUserAgent(userAgent))
141 | if err != nil {
142 | return nil, fmt.Errorf("failed to create Firestore client for project %q and database %q: %w", project, database, err)
143 | }
144 |
145 | return client, nil
146 | }
147 |
148 | func initFirebaseRulesConnection(
149 | ctx context.Context,
150 | project string,
151 | ) (*firebaserules.Service, error) {
152 | // Create the Firebase Rules client
153 | rulesClient, err := firebaserules.NewService(ctx)
154 | if err != nil {
155 | return nil, fmt.Errorf("failed to create Firebase Rules client for project %q: %w", project, err)
156 | }
157 |
158 | return rulesClient, nil
159 | }
160 |
```
--------------------------------------------------------------------------------
/internal/tools/couchbase/couchbase_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2024 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package couchbase_test
16 |
17 | import (
18 | "testing"
19 |
20 | "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
21 |
22 | yaml "github.com/goccy/go-yaml"
23 | "github.com/google/go-cmp/cmp"
24 | "github.com/googleapis/genai-toolbox/internal/server"
25 | "github.com/googleapis/genai-toolbox/internal/testutils"
26 | "github.com/googleapis/genai-toolbox/internal/tools"
27 | )
28 |
29 | func TestParseFromYamlCouchbase(t *testing.T) {
30 | tcs := []struct {
31 | desc string
32 | in string
33 | want server.ToolConfigs
34 | }{
35 | {
36 | desc: "basic example",
37 | in: `
38 | tools:
39 | example_tool:
40 | kind: couchbase-sql
41 | source: my-couchbase-instance
42 | description: some tool description
43 | statement: |
44 | select * from hotel WHERE name = $hotel;
45 | parameters:
46 | - name: hotel
47 | type: string
48 | description: hotel parameter description
49 | `,
50 | want: server.ToolConfigs{
51 | "example_tool": couchbase.Config{
52 | Name: "example_tool",
53 | Kind: "couchbase-sql",
54 | AuthRequired: []string{},
55 | Source: "my-couchbase-instance",
56 | Description: "some tool description",
57 | Statement: "select * from hotel WHERE name = $hotel;\n",
58 | Parameters: []tools.Parameter{
59 | tools.NewStringParameter("hotel", "hotel parameter description"),
60 | },
61 | },
62 | },
63 | },
64 | }
65 | for _, tc := range tcs {
66 | t.Run(tc.desc, func(t *testing.T) {
67 | got := struct {
68 | Tools server.ToolConfigs `yaml:"tools"`
69 | }{}
70 |
71 | // Create a context with a logger
72 | ctx, err := testutils.ContextWithNewLogger()
73 | if err != nil {
74 | t.Fatalf("unable to create context with logger: %s", err)
75 | }
76 |
77 | // Parse contents with context
78 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
79 | if err != nil {
80 | t.Fatalf("unable to unmarshal: %s", err)
81 | }
82 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
83 | t.Fatalf("incorrect parse: diff %v", diff)
84 | }
85 | })
86 | }
87 | }
88 |
89 | func TestParseFromYamlWithTemplateMssql(t *testing.T) {
90 | ctx, err := testutils.ContextWithNewLogger()
91 | if err != nil {
92 | t.Fatalf("unexpected error: %s", err)
93 | }
94 | tcs := []struct {
95 | desc string
96 | in string
97 | want server.ToolConfigs
98 | }{
99 | {
100 | desc: "basic example",
101 | in: `
102 | tools:
103 | example_tool:
104 | kind: couchbase-sql
105 | source: my-couchbase-instance
106 | description: some tool description
107 | statement: |
108 | select * from {{.tableName}} WHERE name = $hotel;
109 | parameters:
110 | - name: hotel
111 | type: string
112 | description: hotel parameter description
113 | templateParameters:
114 | - name: tableName
115 | type: string
116 | description: The table to select hotels from.
117 | `,
118 | want: server.ToolConfigs{
119 | "example_tool": couchbase.Config{
120 | Name: "example_tool",
121 | Kind: "couchbase-sql",
122 | AuthRequired: []string{},
123 | Source: "my-couchbase-instance",
124 | Description: "some tool description",
125 | Statement: "select * from {{.tableName}} WHERE name = $hotel;\n",
126 | Parameters: []tools.Parameter{
127 | tools.NewStringParameter("hotel", "hotel parameter description"),
128 | },
129 | TemplateParameters: []tools.Parameter{
130 | tools.NewStringParameter("tableName", "The table to select hotels from."),
131 | },
132 | },
133 | },
134 | },
135 | }
136 | for _, tc := range tcs {
137 | t.Run(tc.desc, func(t *testing.T) {
138 | got := struct {
139 | Tools server.ToolConfigs `yaml:"tools"`
140 | }{}
141 | // Parse contents
142 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
143 | if err != nil {
144 | t.Fatalf("unable to unmarshal: %s", err)
145 | }
146 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
147 | t.Fatalf("incorrect parse: diff %v", diff)
148 | }
149 | })
150 | }
151 | }
152 |
```
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.yml:
--------------------------------------------------------------------------------
```yaml
1 | # Copyright 2025 Google LLC
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # https://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | name: 🐞 Bug Report
16 | description: File a report for unexpected or undesired behavior.
17 | title: "<brief summary of what bug or error was observed>"
18 | labels: ["type: bug"]
19 | type: "bug"
20 |
21 | body:
22 | - type: markdown
23 | attributes:
24 | value: |
25 | Thanks for helping us improve! 🙏 Please answer these questions and provide as much information as possible about your problem.
26 |
27 | - id: preamble
28 | type: checkboxes
29 | attributes:
30 | label: Prerequisites
31 | description: |
32 | Please run through the following list and make sure you've tried the usual "quick fixes":
33 | - Search the [current open issues](https://github.com/googleapis/genai-toolbox/issues)
34 | - Update to the [latest version of Toolbox](https://github.com/googleapis/genai-toolbox/releases)
35 | options:
36 | - label: "I've searched the current open issues"
37 | required: true
38 | - label: "I've updated to the latest version of Toolbox"
39 |
40 | - type: input
41 | id: version
42 | attributes:
43 | label: Toolbox version
44 | description: |
45 | What version of Toolbox are you using (`toolbox --version`)? e.g.
46 | - toolbox version 0.3.0
47 | - us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:0.3.0
48 | placeholder: ex. toolbox version 0.3.0
49 | validations:
50 | required: true
51 |
52 | - type: textarea
53 | id: environment
54 | attributes:
55 | label: Environment
56 | description: "Let us know some details about the environment in which you are seeing the bug!"
57 | value: |
58 | 1. OS type and version: (output of `uname -a`)
59 | 2. How are you running Toolbox:
60 | - As a downloaded binary (e.g. from `curl -O https://storage.googleapis.com/genai-toolbox/v$VERSION/linux/amd64/toolbox`)
61 | - As a container (e.g. from `us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:$VERSION`)
62 | - Compiled from source (include the command used to build)
63 |
64 | - type: textarea
65 | id: client
66 | attributes:
67 | label: Client
68 | description: "How are you connecting to Toolbox?"
69 | value: |
70 | 1. Client: <name and link to the client are you using>
71 | 2. Version: <what exact version of the client are you using>
72 | 3. Example: If possible, please include your code of configuration:
73 |
74 | ```python
75 | # Code goes here!
76 | ```
77 |
78 | - id: expected-behavior
79 | type: textarea
80 | attributes:
81 | label: Expected Behavior
82 | description: |
83 | Please enter a detailed description of the behavior you expected, and any information about what behavior you
84 | noticed and why it is defective or unintentional.
85 | validations:
86 | required: true
87 |
88 | - id: current-behavior
89 | type: textarea
90 | attributes:
91 | label: Current Behavior
92 | description: "Please enter a detailed description of the behavior you encountered instead."
93 | validations:
94 | required: true
95 |
96 | - type: textarea
97 | id: repro
98 | attributes:
99 | label: Steps to reproduce?
100 | description: |
101 | How can we reproduce this bug? Please walk us through it step by step,
102 | with as much relevant detail as possible. A 'minimal' reproduction is
103 | preferred, which means removing as much of the examples as possible so
104 | only the minimum required to run and reproduce the bug is left.
105 | value: |
106 | 1. ?
107 | 2. ?
108 | 3. ?
109 | ...
110 | validations:
111 | required: true
112 |
113 | - type: textarea
114 | id: additional-details
115 | attributes:
116 | label: Additional Details
117 | description: |
118 | Any other information you want us to know? Things such as tools config,
119 | server logs, etc. can be included here.
120 |
```
--------------------------------------------------------------------------------
/internal/server/mcp/mcp.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package mcp
16 |
17 | import (
18 | "context"
19 | "encoding/json"
20 | "fmt"
21 | "net/http"
22 | "slices"
23 |
24 | "github.com/googleapis/genai-toolbox/internal/auth"
25 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
26 | mcputil "github.com/googleapis/genai-toolbox/internal/server/mcp/util"
27 | v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
28 | v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
29 | v20250618 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250618"
30 | "github.com/googleapis/genai-toolbox/internal/tools"
31 | )
32 |
33 | // LATEST_PROTOCOL_VERSION is the latest version of the MCP protocol supported.
34 | // Update the version used in InitializeResponse when this value is updated.
35 | const LATEST_PROTOCOL_VERSION = v20250618.PROTOCOL_VERSION
36 |
37 | // SUPPORTED_PROTOCOL_VERSIONS is the MCP protocol versions that are supported.
38 | var SUPPORTED_PROTOCOL_VERSIONS = []string{
39 | v20241105.PROTOCOL_VERSION,
40 | v20250326.PROTOCOL_VERSION,
41 | v20250618.PROTOCOL_VERSION,
42 | }
43 |
44 | // InitializeResponse runs capability negotiation and protocol version agreement.
45 | // This is the Initialization phase of the lifecycle for MCP client-server connections.
46 | // Always start with the latest protocol version supported.
47 | func InitializeResponse(ctx context.Context, id jsonrpc.RequestId, body []byte, toolboxVersion string) (any, string, error) {
48 | var req mcputil.InitializeRequest
49 | if err := json.Unmarshal(body, &req); err != nil {
50 | err = fmt.Errorf("invalid mcp initialize request: %w", err)
51 | return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), "", err
52 | }
53 |
54 | var protocolVersion string
55 | v := req.Params.ProtocolVersion
56 | if slices.Contains(SUPPORTED_PROTOCOL_VERSIONS, v) {
57 | protocolVersion = v
58 | } else {
59 | protocolVersion = LATEST_PROTOCOL_VERSION
60 | }
61 |
62 | toolsListChanged := false
63 | result := mcputil.InitializeResult{
64 | ProtocolVersion: protocolVersion,
65 | Capabilities: mcputil.ServerCapabilities{
66 | Tools: &mcputil.ListChanged{
67 | ListChanged: &toolsListChanged,
68 | },
69 | },
70 | ServerInfo: mcputil.Implementation{
71 | BaseMetadata: mcputil.BaseMetadata{
72 | Name: mcputil.SERVER_NAME,
73 | },
74 | Version: toolboxVersion,
75 | },
76 | }
77 | res := jsonrpc.JSONRPCResponse{
78 | Jsonrpc: jsonrpc.JSONRPC_VERSION,
79 | Id: id,
80 | Result: result,
81 | }
82 |
83 | return res, protocolVersion, nil
84 | }
85 |
86 | // NotificationHandler process notifications request. It MUST NOT send a response.
87 | // Currently Toolbox does not process any notifications.
88 | func NotificationHandler(ctx context.Context, body []byte) error {
89 | var notification jsonrpc.JSONRPCNotification
90 | if err := json.Unmarshal(body, ¬ification); err != nil {
91 | return fmt.Errorf("invalid notification request: %w", err)
92 | }
93 | return nil
94 | }
95 |
96 | // ProcessMethod returns a response for the request.
97 | // This is the Operation phase of the lifecycle for MCP client-server connections.
98 | func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, toolset tools.Toolset, tools map[string]tools.Tool, authServices map[string]auth.AuthService, body []byte, header http.Header) (any, error) {
99 | switch mcpVersion {
100 | case v20250618.PROTOCOL_VERSION:
101 | return v20250618.ProcessMethod(ctx, id, method, toolset, tools, authServices, body, header)
102 | case v20250326.PROTOCOL_VERSION:
103 | return v20250326.ProcessMethod(ctx, id, method, toolset, tools, authServices, body, header)
104 | default:
105 | return v20241105.ProcessMethod(ctx, id, method, toolset, tools, authServices, body, header)
106 | }
107 | }
108 |
109 | // VerifyProtocolVersion verifies if the version string is valid.
110 | func VerifyProtocolVersion(version string) bool {
111 | return slices.Contains(SUPPORTED_PROTOCOL_VERSIONS, version)
112 | }
113 |
```
--------------------------------------------------------------------------------
/.ci/continuous.release.cloudbuild.yaml:
--------------------------------------------------------------------------------
```yaml
1 | # Copyright 2024 Google LLC
2 | # Licensed under the Apache License, Version 2.0 (the "License");
3 | # you may not use this file except in compliance with the License.
4 | # You may obtain a copy of the License at
5 | #
6 | # http://www.apache.org/licenses/LICENSE-2.0
7 | #
8 | # Unless required by applicable law or agreed to in writing, software
9 | # distributed under the License is distributed on an "AS IS" BASIS,
10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 | # See the License for the specific language governing permissions and
12 | # limitations under the License.
13 |
14 | steps:
15 | - id: "build-docker"
16 | name: "gcr.io/cloud-builders/docker"
17 | waitFor: ['-']
18 | script: |
19 | #!/usr/bin/env bash
20 | docker buildx create --name container-builder --driver docker-container --bootstrap --use
21 | docker buildx build --platform linux/amd64,linux/arm64 --build-arg COMMIT_SHA=$(git rev-parse --short HEAD) -t ${_DOCKER_URI}:$REF_NAME --push .
22 |
23 | - id: "install-dependencies"
24 | name: golang:1
25 | waitFor: ['-']
26 | env:
27 | - 'GOPATH=/gopath'
28 | volumes:
29 | - name: 'go'
30 | path: '/gopath'
31 | script: |
32 | go get -d ./...
33 |
34 | - id: "build-linux-amd64"
35 | name: golang:1
36 | waitFor:
37 | - "install-dependencies"
38 | env:
39 | - 'GOPATH=/gopath'
40 | volumes:
41 | - name: 'go'
42 | path: '/gopath'
43 | script: |
44 | #!/usr/bin/env bash
45 | CGO_ENABLED=0 GOOS=linux GOARCH=amd64 \
46 | go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse --short HEAD)" -o toolbox.linux.amd64
47 |
48 | - id: "store-linux-amd64"
49 | name: "gcr.io/cloud-builders/gcloud:latest"
50 | waitFor:
51 | - "build-linux-amd64"
52 | script: |
53 | #!/usr/bin/env bash
54 | gcloud storage cp toolbox.linux.amd64 gs://$_BUCKET_NAME/$REF_NAME/linux/amd64/toolbox
55 |
56 | - id: "build-darwin-arm64"
57 | name: golang:1
58 | waitFor:
59 | - "install-dependencies"
60 | env:
61 | - 'GOPATH=/gopath'
62 | volumes:
63 | - name: 'go'
64 | path: '/gopath'
65 | script: |
66 | #!/usr/bin/env bash
67 | CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 \
68 | go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse --short HEAD)" -o toolbox.darwin.arm64
69 |
70 | - id: "store-darwin-arm64"
71 | name: "gcr.io/cloud-builders/gcloud:latest"
72 | waitFor:
73 | - "build-darwin-arm64"
74 | script: |
75 | #!/usr/bin/env bash
76 | gcloud storage cp toolbox.darwin.arm64 gs://$_BUCKET_NAME/$REF_NAME/darwin/arm64/toolbox
77 |
78 | - id: "build-darwin-amd64"
79 | name: golang:1
80 | waitFor:
81 | - "install-dependencies"
82 | env:
83 | - 'GOPATH=/gopath'
84 | volumes:
85 | - name: 'go'
86 | path: '/gopath'
87 | script: |
88 | #!/usr/bin/env bash
89 | CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 \
90 | go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse --short HEAD)" -o toolbox.darwin.amd64
91 |
92 | - id: "store-darwin-amd64"
93 | name: "gcr.io/cloud-builders/gcloud:latest"
94 | waitFor:
95 | - "build-darwin-amd64"
96 | script: |
97 | #!/usr/bin/env bash
98 | gcloud storage cp toolbox.darwin.amd64 gs://$_BUCKET_NAME/$REF_NAME/darwin/amd64/toolbox
99 |
100 | - id: "build-windows-amd64"
101 | name: golang:1
102 | waitFor:
103 | - "install-dependencies"
104 | env:
105 | - 'GOPATH=/gopath'
106 | volumes:
107 | - name: 'go'
108 | path: '/gopath'
109 | script: |
110 | #!/usr/bin/env bash
111 | CGO_ENABLED=0 GOOS=windows GOARCH=amd64 \
112 | go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse --short HEAD)" -o toolbox.windows.amd64
113 |
114 | - id: "store-windows-amd64"
115 | name: "gcr.io/cloud-builders/gcloud:latest"
116 | waitFor:
117 | - "build-windows-amd64"
118 | script: |
119 | #!/usr/bin/env bash
120 | gcloud storage cp toolbox.windows.amd64 gs://$_BUCKET_NAME/$REF_NAME/windows/amd64/toolbox.exe
121 |
122 | options:
123 | automapSubstitutions: true
124 | dynamicSubstitutions: true
125 | logging: CLOUD_LOGGING_ONLY # Necessary for custom service account
126 | machineType: 'E2_HIGHCPU_32'
127 |
128 | substitutions:
129 | _REGION: us-central1
130 | _AR_HOSTNAME: ${_REGION}-docker.pkg.dev
131 | _AR_REPO_NAME: toolbox-dev
132 | _BUCKET_NAME: genai-toolbox-dev
133 | _DOCKER_URI: ${_AR_HOSTNAME}/${PROJECT_ID}/${_AR_REPO_NAME}/toolbox
134 |
```
--------------------------------------------------------------------------------
/internal/sources/cloudsqladmin/cloud_sql_admin.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 | package cloudsqladmin
15 |
16 | import (
17 | "context"
18 | "fmt"
19 | "net/http"
20 |
21 | "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/googleapis/genai-toolbox/internal/util"
24 | "go.opentelemetry.io/otel/trace"
25 | "golang.org/x/oauth2"
26 | "golang.org/x/oauth2/google"
27 | "google.golang.org/api/option"
28 | sqladmin "google.golang.org/api/sqladmin/v1"
29 | )
30 |
31 | const SourceKind string = "cloud-sql-admin"
32 |
33 | type userAgentRoundTripper struct {
34 | userAgent string
35 | next http.RoundTripper
36 | }
37 |
38 | func (rt *userAgentRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
39 | newReq := *req
40 | newReq.Header = make(http.Header)
41 | for k, v := range req.Header {
42 | newReq.Header[k] = v
43 | }
44 | ua := newReq.Header.Get("User-Agent")
45 | if ua == "" {
46 | newReq.Header.Set("User-Agent", rt.userAgent)
47 | } else {
48 | newReq.Header.Set("User-Agent", ua+" "+rt.userAgent)
49 | }
50 | return rt.next.RoundTrip(&newReq)
51 | }
52 |
53 | // validate interface
54 | var _ sources.SourceConfig = Config{}
55 |
56 | func init() {
57 | if !sources.Register(SourceKind, newConfig) {
58 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
59 | }
60 | }
61 |
62 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
63 | actual := Config{Name: name}
64 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
65 | return nil, err
66 | }
67 | return actual, nil
68 | }
69 |
70 | type Config struct {
71 | Name string `yaml:"name" validate:"required"`
72 | Kind string `yaml:"kind" validate:"required"`
73 | UseClientOAuth bool `yaml:"useClientOAuth"`
74 | }
75 |
76 | func (r Config) SourceConfigKind() string {
77 | return SourceKind
78 | }
79 |
80 | // Initialize initializes a CloudSQL Admin Source instance.
81 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
82 | ua, err := util.UserAgentFromContext(ctx)
83 | if err != nil {
84 | return nil, fmt.Errorf("error in User Agent retrieval: %s", err)
85 | }
86 |
87 | var client *http.Client
88 | if r.UseClientOAuth {
89 | client = &http.Client{
90 | Transport: &userAgentRoundTripper{
91 | userAgent: ua,
92 | next: http.DefaultTransport,
93 | },
94 | }
95 | } else {
96 | // Use Application Default Credentials
97 | creds, err := google.FindDefaultCredentials(ctx, sqladmin.SqlserviceAdminScope)
98 | if err != nil {
99 | return nil, fmt.Errorf("failed to find default credentials: %w", err)
100 | }
101 | baseClient := oauth2.NewClient(ctx, creds.TokenSource)
102 | baseClient.Transport = &userAgentRoundTripper{
103 | userAgent: ua,
104 | next: baseClient.Transport,
105 | }
106 | client = baseClient
107 | }
108 |
109 | service, err := sqladmin.NewService(ctx, option.WithHTTPClient(client))
110 | if err != nil {
111 | return nil, fmt.Errorf("error creating new sqladmin service: %w", err)
112 | }
113 |
114 | s := &Source{
115 | Name: r.Name,
116 | Kind: SourceKind,
117 | BaseURL: "https://sqladmin.googleapis.com",
118 | Service: service,
119 | UseClientOAuth: r.UseClientOAuth,
120 | }
121 | return s, nil
122 | }
123 |
124 | var _ sources.Source = &Source{}
125 |
126 | type Source struct {
127 | Name string `yaml:"name"`
128 | Kind string `yaml:"kind"`
129 | BaseURL string
130 | Service *sqladmin.Service
131 | UseClientOAuth bool
132 | }
133 |
134 | func (s *Source) SourceKind() string {
135 | return SourceKind
136 | }
137 |
138 | func (s *Source) GetService(ctx context.Context, accessToken string) (*sqladmin.Service, error) {
139 | if s.UseClientOAuth {
140 | token := &oauth2.Token{AccessToken: accessToken}
141 | client := oauth2.NewClient(ctx, oauth2.StaticTokenSource(token))
142 | service, err := sqladmin.NewService(ctx, option.WithHTTPClient(client))
143 | if err != nil {
144 | return nil, fmt.Errorf("error creating new sqladmin service: %w", err)
145 | }
146 | return service, nil
147 | }
148 | return s.Service, nil
149 | }
150 |
151 | func (s *Source) UseClientAuthorization() bool {
152 | return s.UseClientOAuth
153 | }
154 |
```
--------------------------------------------------------------------------------
/internal/tools/mongodb/mongodbupdateone/mongodbupdateone_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package mongodbupdateone_test
16 |
17 | import (
18 | "strings"
19 | "testing"
20 |
21 | "github.com/googleapis/genai-toolbox/internal/tools"
22 | "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdateone"
23 |
24 | yaml "github.com/goccy/go-yaml"
25 | "github.com/google/go-cmp/cmp"
26 | "github.com/googleapis/genai-toolbox/internal/server"
27 | "github.com/googleapis/genai-toolbox/internal/testutils"
28 | )
29 |
30 | func TestParseFromYamlMongoQuery(t *testing.T) {
31 | ctx, err := testutils.ContextWithNewLogger()
32 | if err != nil {
33 | t.Fatalf("unexpected error: %s", err)
34 | }
35 | tcs := []struct {
36 | desc string
37 | in string
38 | want server.ToolConfigs
39 | }{
40 | {
41 | desc: "basic example",
42 | in: `
43 | tools:
44 | example_tool:
45 | kind: mongodb-update-one
46 | source: my-instance
47 | description: some description
48 | database: test_db
49 | collection: test_coll
50 | filterPayload: |
51 | { name: {{json .name}} }
52 | filterParams:
53 | - name: name
54 | type: string
55 | description: small description
56 | updatePayload: |
57 | { $set : { item: {{json .item}} } }
58 | updateParams:
59 | - name: item
60 | type: string
61 | description: small description
62 | canonical: true
63 | upsert: true
64 | `,
65 | want: server.ToolConfigs{
66 | "example_tool": mongodbupdateone.Config{
67 | Name: "example_tool",
68 | Kind: "mongodb-update-one",
69 | Source: "my-instance",
70 | AuthRequired: []string{},
71 | Database: "test_db",
72 | Collection: "test_coll",
73 | Canonical: true,
74 | FilterPayload: "{ name: {{json .name}} }\n",
75 | FilterParams: tools.Parameters{
76 | &tools.StringParameter{
77 | CommonParameter: tools.CommonParameter{
78 | Name: "name",
79 | Type: "string",
80 | Desc: "small description",
81 | },
82 | },
83 | },
84 | UpdatePayload: "{ $set : { item: {{json .item}} } }\n",
85 | UpdateParams: tools.Parameters{
86 | &tools.StringParameter{
87 | CommonParameter: tools.CommonParameter{
88 | Name: "item",
89 | Type: "string",
90 | Desc: "small description",
91 | },
92 | },
93 | },
94 | Upsert: true,
95 | Description: "some description",
96 | },
97 | },
98 | },
99 | }
100 | for _, tc := range tcs {
101 | t.Run(tc.desc, func(t *testing.T) {
102 | got := struct {
103 | Tools server.ToolConfigs `yaml:"tools"`
104 | }{}
105 | // Parse contents
106 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
107 | if err != nil {
108 | t.Fatalf("unable to unmarshal: %s", err)
109 | }
110 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
111 | t.Fatalf("incorrect parse: diff %v", diff)
112 | }
113 | })
114 | }
115 |
116 | }
117 |
118 | func TestFailParseFromYamlMongoQuery(t *testing.T) {
119 | ctx, err := testutils.ContextWithNewLogger()
120 | if err != nil {
121 | t.Fatalf("unexpected error: %s", err)
122 | }
123 | tcs := []struct {
124 | desc string
125 | in string
126 | err string
127 | }{
128 | {
129 | desc: "Invalid method",
130 | in: `
131 | tools:
132 | example_tool:
133 | kind: mongodb-update-one
134 | source: my-instance
135 | description: some description
136 | collection: test_coll
137 | filterPayload: |
138 | { name : {{json .name}} }`,
139 | err: `unable to parse tool "example_tool" as kind "mongodb-update-one"`,
140 | },
141 | }
142 | for _, tc := range tcs {
143 | t.Run(tc.desc, func(t *testing.T) {
144 | got := struct {
145 | Tools server.ToolConfigs `yaml:"tools"`
146 | }{}
147 | // Parse contents
148 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
149 | if err == nil {
150 | t.Fatalf("expect parsing to fail")
151 | }
152 | errStr := err.Error()
153 | if !strings.Contains(errStr, tc.err) {
154 | t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err)
155 | }
156 | })
157 | }
158 |
159 | }
160 |
```
--------------------------------------------------------------------------------
/internal/sources/cassandra/cassandra_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package cassandra_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/sources/cassandra"
24 | "github.com/googleapis/genai-toolbox/internal/testutils"
25 | )
26 |
27 | func TestParseFromYamlCassandra(t *testing.T) {
28 | tcs := []struct {
29 | desc string
30 | in string
31 | want server.SourceConfigs
32 | }{
33 | {
34 | desc: "basic example (without optional fields)",
35 | in: `
36 | sources:
37 | my-cassandra-instance:
38 | kind: cassandra
39 | hosts:
40 | - "my-host1"
41 | - "my-host2"
42 | `,
43 | want: server.SourceConfigs{
44 | "my-cassandra-instance": cassandra.Config{
45 | Name: "my-cassandra-instance",
46 | Kind: cassandra.SourceKind,
47 | Hosts: []string{"my-host1", "my-host2"},
48 | Username: "",
49 | Password: "",
50 | ProtoVersion: 0,
51 | CAPath: "",
52 | CertPath: "",
53 | KeyPath: "",
54 | Keyspace: "",
55 | EnableHostVerification: false,
56 | },
57 | },
58 | },
59 | {
60 | desc: "with optional fields",
61 | in: `
62 | sources:
63 | my-cassandra-instance:
64 | kind: cassandra
65 | hosts:
66 | - "my-host1"
67 | - "my-host2"
68 | username: "user"
69 | password: "pass"
70 | keyspace: "example_keyspace"
71 | protoVersion: 4
72 | caPath: "path/to/ca.crt"
73 | certPath: "path/to/cert"
74 | keyPath: "path/to/key"
75 | enableHostVerification: true
76 | `,
77 | want: server.SourceConfigs{
78 | "my-cassandra-instance": cassandra.Config{
79 | Name: "my-cassandra-instance",
80 | Kind: cassandra.SourceKind,
81 | Hosts: []string{"my-host1", "my-host2"},
82 | Username: "user",
83 | Password: "pass",
84 | Keyspace: "example_keyspace",
85 | ProtoVersion: 4,
86 | CAPath: "path/to/ca.crt",
87 | CertPath: "path/to/cert",
88 | KeyPath: "path/to/key",
89 | EnableHostVerification: true,
90 | },
91 | },
92 | },
93 | }
94 | for _, tc := range tcs {
95 | t.Run(tc.desc, func(t *testing.T) {
96 | got := struct {
97 | Sources server.SourceConfigs `yaml:"sources"`
98 | }{}
99 | // Parse contents
100 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
101 | if err != nil {
102 | t.Fatalf("unable to unmarshal: %s", err)
103 | }
104 | if !cmp.Equal(tc.want, got.Sources) {
105 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
106 | }
107 | })
108 | }
109 |
110 | }
111 |
112 | func TestFailParseFromYaml(t *testing.T) {
113 | tcs := []struct {
114 | desc string
115 | in string
116 | err string
117 | }{
118 | {
119 | desc: "extra field",
120 | in: `
121 | sources:
122 | my-cassandra-instance:
123 | kind: cassandra
124 | hosts:
125 | - "my-host"
126 | foo: bar
127 | `,
128 | err: "unable to parse source \"my-cassandra-instance\" as \"cassandra\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | hosts:\n 3 | - my-host\n 4 | kind: cassandra",
129 | },
130 | {
131 | desc: "missing required field",
132 | in: `
133 | sources:
134 | my-cassandra-instance:
135 | kind: cassandra
136 | `,
137 | err: "unable to parse source \"my-cassandra-instance\" as \"cassandra\": Key: 'Config.Hosts' Error:Field validation for 'Hosts' failed on the 'required' tag",
138 | },
139 | }
140 |
141 | for _, tc := range tcs {
142 | t.Run(tc.desc, func(t *testing.T) {
143 | got := struct {
144 | Sources server.SourceConfigs `yaml:"sources"`
145 | }{}
146 | // Parse contents
147 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
148 | if err == nil {
149 | t.Fatalf("expect parsing to fail")
150 | }
151 | errStr := err.Error()
152 | if errStr != tc.err {
153 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
154 | }
155 | })
156 | }
157 |
158 | }
159 |
```
--------------------------------------------------------------------------------
/internal/sources/yugabytedb/yugabytedb.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package yugabytedb
16 |
17 | import (
18 | "context"
19 | "fmt"
20 |
21 | "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/yugabyte/pgx/v5/pgxpool"
24 | "go.opentelemetry.io/otel/trace"
25 | )
26 |
27 | const SourceKind string = "yugabytedb"
28 |
29 | // validate interface
30 | var _ sources.SourceConfig = Config{}
31 |
32 | func init() {
33 | if !sources.Register(SourceKind, newConfig) {
34 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
35 | }
36 | }
37 |
38 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
39 | actual := Config{Name: name}
40 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
41 | return nil, err
42 | }
43 | return actual, nil
44 | }
45 |
46 | type Config struct {
47 | Name string `yaml:"name" validate:"required"`
48 | Kind string `yaml:"kind" validate:"required"`
49 | Host string `yaml:"host" validate:"required"`
50 | Port string `yaml:"port" validate:"required"`
51 | User string `yaml:"user" validate:"required"`
52 | Password string `yaml:"password" validate:"required"`
53 | Database string `yaml:"database" validate:"required"`
54 | LoadBalance string `yaml:"loadBalance"`
55 | TopologyKeys string `yaml:"topologyKeys"`
56 | YBServersRefreshInterval string `yaml:"ybServersRefreshInterval"`
57 | FallBackToTopologyKeysOnly string `yaml:"fallbackToTopologyKeysOnly"`
58 | FailedHostReconnectDelaySeconds string `yaml:"failedHostReconnectDelaySecs"`
59 | }
60 |
61 | func (r Config) SourceConfigKind() string {
62 | return SourceKind
63 | }
64 |
65 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
66 | pool, err := initYugabyteDBConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.LoadBalance, r.TopologyKeys, r.YBServersRefreshInterval, r.FallBackToTopologyKeysOnly, r.FailedHostReconnectDelaySeconds)
67 | if err != nil {
68 | return nil, fmt.Errorf("unable to create pool: %w", err)
69 | }
70 |
71 | err = pool.Ping(ctx)
72 | if err != nil {
73 | return nil, fmt.Errorf("unable to connect successfully: %w", err)
74 | }
75 |
76 | s := &Source{
77 | Name: r.Name,
78 | Kind: SourceKind,
79 | Pool: pool,
80 | }
81 | return s, nil
82 | }
83 |
84 | var _ sources.Source = &Source{}
85 |
86 | type Source struct {
87 | Name string `yaml:"name"`
88 | Kind string `yaml:"kind"`
89 | Pool *pgxpool.Pool
90 | }
91 |
92 | func (s *Source) SourceKind() string {
93 | return SourceKind
94 | }
95 |
96 | func (s *Source) YugabyteDBPool() *pgxpool.Pool {
97 | return s.Pool
98 | }
99 |
100 | func initYugabyteDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, loadBalance, topologyKeys, refreshInterval, explicitFallback, failedHostTTL string) (*pgxpool.Pool, error) {
101 | //nolint:all // Reassigned ctx
102 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
103 | defer span.End()
104 | // urlExample := "postgres://username:password@localhost:5433/database_name"
105 | i := fmt.Sprintf("postgres://%s:%s@%s:%s/%s", user, pass, host, port, dbname)
106 | if loadBalance == "true" {
107 | i = fmt.Sprintf("%s?load_balance=%s", i, loadBalance)
108 | if topologyKeys != "" {
109 | i = fmt.Sprintf("%s&topology_keys=%s", i, topologyKeys)
110 | if explicitFallback == "true" {
111 | i = fmt.Sprintf("%s&fallback_to_topology_keys_only=%s", i, explicitFallback)
112 | }
113 | }
114 | if refreshInterval != "" {
115 | i = fmt.Sprintf("%s&yb_servers_refresh_interval=%s", i, refreshInterval)
116 | }
117 | if failedHostTTL != "" {
118 | i = fmt.Sprintf("%s&failed_host_reconnect_delay_secs=%s", i, failedHostTTL)
119 | }
120 | }
121 | pool, err := pgxpool.New(ctx, i)
122 | if err != nil {
123 | return nil, fmt.Errorf("unable to create connection pool: %w", err)
124 | }
125 |
126 | return pool, nil
127 | }
128 |
```
--------------------------------------------------------------------------------
/internal/sources/bigquery/bigquery_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package bigquery_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
24 | "github.com/googleapis/genai-toolbox/internal/testutils"
25 | )
26 |
27 | func TestParseFromYamlBigQuery(t *testing.T) {
28 | tcs := []struct {
29 | desc string
30 | in string
31 | want server.SourceConfigs
32 | }{
33 | {
34 | desc: "basic example",
35 | in: `
36 | sources:
37 | my-instance:
38 | kind: bigquery
39 | project: my-project
40 | `,
41 | want: server.SourceConfigs{
42 | "my-instance": bigquery.Config{
43 | Name: "my-instance",
44 | Kind: bigquery.SourceKind,
45 | Project: "my-project",
46 | Location: "",
47 | WriteMode: "",
48 | },
49 | },
50 | },
51 | {
52 | desc: "all fields specified",
53 | in: `
54 | sources:
55 | my-instance:
56 | kind: bigquery
57 | project: my-project
58 | location: asia
59 | writeMode: blocked
60 | `,
61 | want: server.SourceConfigs{
62 | "my-instance": bigquery.Config{
63 | Name: "my-instance",
64 | Kind: bigquery.SourceKind,
65 | Project: "my-project",
66 | Location: "asia",
67 | WriteMode: "blocked",
68 | UseClientOAuth: false,
69 | },
70 | },
71 | },
72 | {
73 | desc: "use client auth example",
74 | in: `
75 | sources:
76 | my-instance:
77 | kind: bigquery
78 | project: my-project
79 | location: us
80 | useClientOAuth: true
81 | `,
82 | want: server.SourceConfigs{
83 | "my-instance": bigquery.Config{
84 | Name: "my-instance",
85 | Kind: bigquery.SourceKind,
86 | Project: "my-project",
87 | Location: "us",
88 | UseClientOAuth: true,
89 | },
90 | },
91 | },
92 | {
93 | desc: "with allowed datasets example",
94 | in: `
95 | sources:
96 | my-instance:
97 | kind: bigquery
98 | project: my-project
99 | location: us
100 | allowedDatasets:
101 | - my_dataset
102 | `,
103 | want: server.SourceConfigs{
104 | "my-instance": bigquery.Config{
105 | Name: "my-instance",
106 | Kind: bigquery.SourceKind,
107 | Project: "my-project",
108 | Location: "us",
109 | AllowedDatasets: []string{"my_dataset"},
110 | },
111 | },
112 | },
113 | }
114 | for _, tc := range tcs {
115 | t.Run(tc.desc, func(t *testing.T) {
116 | got := struct {
117 | Sources server.SourceConfigs `yaml:"sources"`
118 | }{}
119 | // Parse contents
120 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
121 | if err != nil {
122 | t.Fatalf("unable to unmarshal: %s", err)
123 | }
124 | if !cmp.Equal(tc.want, got.Sources) {
125 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
126 | }
127 | })
128 | }
129 |
130 | }
131 |
132 | func TestFailParseFromYaml(t *testing.T) {
133 | tcs := []struct {
134 | desc string
135 | in string
136 | err string
137 | }{
138 | {
139 | desc: "extra field",
140 | in: `
141 | sources:
142 | my-instance:
143 | kind: bigquery
144 | project: my-project
145 | location: us
146 | foo: bar
147 | `,
148 | err: "unable to parse source \"my-instance\" as \"bigquery\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | kind: bigquery\n 3 | location: us\n 4 | project: my-project",
149 | },
150 | {
151 | desc: "missing required field",
152 | in: `
153 | sources:
154 | my-instance:
155 | kind: bigquery
156 | location: us
157 | `,
158 | err: "unable to parse source \"my-instance\" as \"bigquery\": Key: 'Config.Project' Error:Field validation for 'Project' failed on the 'required' tag",
159 | },
160 | }
161 | for _, tc := range tcs {
162 | t.Run(tc.desc, func(t *testing.T) {
163 | got := struct {
164 | Sources server.SourceConfigs `yaml:"sources"`
165 | }{}
166 | // Parse contents
167 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
168 | if err == nil {
169 | t.Fatalf("expect parsing to fail")
170 | }
171 | errStr := err.Error()
172 | if errStr != tc.err {
173 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
174 | }
175 | })
176 | }
177 | }
178 |
```
--------------------------------------------------------------------------------
/docs/en/getting-started/quickstart/shared/database_setup.md:
--------------------------------------------------------------------------------
```markdown
1 | <!-- This file has been used in local_quickstart.md, local_quickstart_go.md & local_quickstart_js.md -->
2 | <!-- [START database_setup] -->
3 | In this section, we will create a database, insert some data that needs to be
4 | accessed by our agent, and create a database user for Toolbox to connect with.
5 |
6 | 1. Connect to postgres using the `psql` command:
7 |
8 | ```bash
9 | psql -h 127.0.0.1 -U postgres
10 | ```
11 |
12 | Here, `postgres` denotes the default postgres superuser.
13 |
14 | {{< notice info >}}
15 |
16 | #### **Having trouble connecting?**
17 |
18 | * **Password Prompt:** If you are prompted for a password for the `postgres`
19 | user and do not know it (or a blank password doesn't work), your PostgreSQL
20 | installation might require a password or a different authentication method.
21 | * **`FATAL: role "postgres" does not exist`:** This error means the default
22 | `postgres` superuser role isn't available under that name on your system.
23 | * **`Connection refused`:** Ensure your PostgreSQL server is actually running.
24 | You can typically check with `sudo systemctl status postgresql` and start it
25 | with `sudo systemctl start postgresql` on Linux systems.
26 |
27 | <br/>
28 |
29 | #### **Common Solution**
30 |
31 | For password issues or if the `postgres` role seems inaccessible directly, try
32 | switching to the `postgres` operating system user first. This user often has
33 | permission to connect without a password for local connections (this is called
34 | peer authentication).
35 |
36 | ```bash
37 | sudo -i -u postgres
38 | psql -h 127.0.0.1
39 | ```
40 |
41 | Once you are in the `psql` shell using this method, you can proceed with the
42 | database creation steps below. Afterwards, type `\q` to exit `psql`, and then
43 | `exit` to return to your normal user shell.
44 |
45 | If desired, once connected to `psql` as the `postgres` OS user, you can set a
46 | password for the `postgres` *database* user using: `ALTER USER postgres WITH
47 | PASSWORD 'your_chosen_password';`. This would allow direct connection with `-U
48 | postgres` and a password next time.
49 | {{< /notice >}}
50 |
51 | 1. Create a new database and a new user:
52 |
53 | {{< notice tip >}}
54 | For a real application, it's best to follow the principle of least permission
55 | and only grant the privileges your application needs.
56 | {{< /notice >}}
57 |
58 | ```sql
59 | CREATE USER toolbox_user WITH PASSWORD 'my-password';
60 |
61 | CREATE DATABASE toolbox_db;
62 | GRANT ALL PRIVILEGES ON DATABASE toolbox_db TO toolbox_user;
63 |
64 | ALTER DATABASE toolbox_db OWNER TO toolbox_user;
65 | ```
66 |
67 | 1. End the database session:
68 |
69 | ```bash
70 | \q
71 | ```
72 |
73 | (If you used `sudo -i -u postgres` and then `psql`, remember you might also
74 | need to type `exit` after `\q` to leave the `postgres` user's shell
75 | session.)
76 |
77 | 1. Connect to your database with your new user:
78 |
79 | ```bash
80 | psql -h 127.0.0.1 -U toolbox_user -d toolbox_db
81 | ```
82 |
83 | 1. Create a table using the following command:
84 |
85 | ```sql
86 | CREATE TABLE hotels(
87 | id INTEGER NOT NULL PRIMARY KEY,
88 | name VARCHAR NOT NULL,
89 | location VARCHAR NOT NULL,
90 | price_tier VARCHAR NOT NULL,
91 | checkin_date DATE NOT NULL,
92 | checkout_date DATE NOT NULL,
93 | booked BIT NOT NULL
94 | );
95 | ```
96 |
97 | 1. Insert data into the table.
98 |
99 | ```sql
100 | INSERT INTO hotels(id, name, location, price_tier, checkin_date, checkout_date, booked)
101 | VALUES
102 | (1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-22', '2024-04-20', B'0'),
103 | (2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', B'0'),
104 | (3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', B'0'),
105 | (4, 'Radisson Blu Lucerne', 'Lucerne', 'Midscale', '2024-04-24', '2024-04-05', B'0'),
106 | (5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-23', '2024-04-01', B'0'),
107 | (6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', B'0'),
108 | (7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-27', '2024-04-02', B'0'),
109 | (8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-24', '2024-04-09', B'0'),
110 | (9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', B'0'),
111 | (10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', B'0');
112 | ```
113 |
114 | 1. End the database session:
115 |
116 | ```bash
117 | \q
118 | ```
119 | <!-- [END database_setup] -->
```
--------------------------------------------------------------------------------
/tests/valkey/valkey_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package valkey
16 |
17 | import (
18 | "context"
19 | "log"
20 | "os"
21 | "regexp"
22 | "testing"
23 | "time"
24 |
25 | "github.com/googleapis/genai-toolbox/internal/testutils"
26 | "github.com/googleapis/genai-toolbox/tests"
27 | "github.com/valkey-io/valkey-go"
28 | )
29 |
30 | var (
31 | ValkeySourceKind = "valkey"
32 | ValkeyToolKind = "valkey"
33 | ValkeyAddress = os.Getenv("VALKEY_ADDRESS")
34 | )
35 |
36 | func getValkeyVars(t *testing.T) map[string]any {
37 | switch "" {
38 | case ValkeyAddress:
39 | t.Fatal("'VALKEY_ADDRESS' not set")
40 | }
41 | return map[string]any{
42 | "kind": ValkeySourceKind,
43 | "address": []string{ValkeyAddress},
44 | "disableCache": true,
45 | }
46 | }
47 |
48 | func initValkeyClient(ctx context.Context, addr []string) (valkey.Client, error) {
49 | // Pass in an access token getter fn for IAM auth
50 | client, err := valkey.NewClient(valkey.ClientOption{
51 | InitAddress: addr,
52 | ForceSingleClient: true,
53 | DisableCache: true,
54 | })
55 |
56 | if err != nil {
57 | log.Fatalf("error creating client: %v", err)
58 | }
59 |
60 | // Ping the server to check connectivity (using Do)
61 | pingCmd := client.B().Ping().Build()
62 | _, err = client.Do(ctx, pingCmd).ToString()
63 | if err != nil {
64 | log.Fatalf("Failed to execute PING command: %v", err)
65 | }
66 | log.Println("Successfully connected to Valkey")
67 | return client, nil
68 | }
69 |
70 | func TestValkeyToolEndpoints(t *testing.T) {
71 | sourceConfig := getValkeyVars(t)
72 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
73 | defer cancel()
74 |
75 | var args []string
76 |
77 | client, err := initValkeyClient(ctx, []string{ValkeyAddress})
78 | if err != nil {
79 | t.Fatalf("unable to create Valkey connection: %s", err)
80 | }
81 |
82 | // set up data for param tool
83 | teardownDB := setupValkeyDB(t, ctx, client)
84 | defer teardownDB(t)
85 |
86 | // Write config into a file and pass it to command
87 | toolsFile := tests.GetRedisValkeyToolsConfig(sourceConfig, ValkeyToolKind)
88 |
89 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
90 | if err != nil {
91 | t.Fatalf("command initialization returned an error: %s", err)
92 | }
93 | defer cleanup()
94 |
95 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
96 | defer cancel()
97 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
98 | if err != nil {
99 | t.Logf("toolbox command logs: \n%s", out)
100 | t.Fatalf("toolbox didn't start successfully: %s", err)
101 | }
102 |
103 | // Get configs for tests
104 | select1Want, mcpMyFailToolWant, invokeParamWant, invokeIdNullWant, nullWant, mcpSelect1Want, mcpInvokeParamWant := tests.GetRedisValkeyWants()
105 |
106 | // Run tests
107 | tests.RunToolGetTest(t)
108 | tests.RunToolInvokeTest(t, select1Want,
109 | tests.WithMyToolId3NameAliceWant(invokeParamWant),
110 | tests.WithMyArrayToolWant(invokeParamWant),
111 | tests.WithMyToolById4Want(invokeIdNullWant),
112 | tests.WithNullWant(nullWant),
113 | )
114 | tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want,
115 | tests.WithMcpMyToolId3NameAliceWant(mcpInvokeParamWant),
116 | )
117 | }
118 |
119 | func setupValkeyDB(t *testing.T, ctx context.Context, client valkey.Client) func(*testing.T) {
120 | keys := []string{"row1", "row2", "row3", "row4", "null"}
121 | commands := [][]string{
122 | {"HSET", keys[0], "name", "Alice", "id", "1"},
123 | {"HSET", keys[1], "name", "Jane", "id", "2"},
124 | {"HSET", keys[2], "name", "Sid", "id", "3"},
125 | {"HSET", keys[3], "name", "", "id", "4"},
126 | {"SET", keys[4], "null"},
127 | {"HSET", tests.ServiceAccountEmail, "name", "Alice"},
128 | }
129 | builtCmds := make(valkey.Commands, len(commands))
130 |
131 | for i, cmd := range commands {
132 | builtCmds[i] = client.B().Arbitrary(cmd...).Build()
133 | }
134 |
135 | responses := client.DoMulti(ctx, builtCmds...)
136 | for _, resp := range responses {
137 | if err := resp.Error(); err != nil {
138 | t.Fatalf("unable to insert test data: %s", err)
139 | }
140 | }
141 |
142 | return func(t *testing.T) {
143 | // tear down test
144 | _, err := client.Do(ctx, client.B().Del().Key(keys...).Build()).AsInt64()
145 | if err != nil {
146 | t.Errorf("Teardown failed: %s", err)
147 | }
148 | }
149 |
150 | }
151 |
```
--------------------------------------------------------------------------------
/internal/tools/bigtable/bigtable_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package bigtable_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools"
25 | "github.com/googleapis/genai-toolbox/internal/tools/bigtable"
26 | )
27 |
28 | func TestParseFromYamlBigtable(t *testing.T) {
29 | ctx, err := testutils.ContextWithNewLogger()
30 | if err != nil {
31 | t.Fatalf("unexpected error: %s", err)
32 | }
33 | tcs := []struct {
34 | desc string
35 | in string
36 | want server.ToolConfigs
37 | }{
38 | {
39 | desc: "basic example",
40 | in: `
41 | tools:
42 | example_tool:
43 | kind: bigtable-sql
44 | source: my-pg-instance
45 | description: some description
46 | statement: |
47 | SELECT * FROM SQL_STATEMENT;
48 | parameters:
49 | - name: country
50 | type: string
51 | description: some description
52 | `,
53 | want: server.ToolConfigs{
54 | "example_tool": bigtable.Config{
55 | Name: "example_tool",
56 | Kind: "bigtable-sql",
57 | Source: "my-pg-instance",
58 | Description: "some description",
59 | Statement: "SELECT * FROM SQL_STATEMENT;\n",
60 | AuthRequired: []string{},
61 | Parameters: []tools.Parameter{
62 | tools.NewStringParameter("country", "some description"),
63 | },
64 | },
65 | },
66 | },
67 | }
68 | for _, tc := range tcs {
69 | t.Run(tc.desc, func(t *testing.T) {
70 | got := struct {
71 | Tools server.ToolConfigs `yaml:"tools"`
72 | }{}
73 | // Parse contents
74 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
75 | if err != nil {
76 | t.Fatalf("unable to unmarshal: %s", err)
77 | }
78 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
79 | t.Fatalf("incorrect parse: diff %v", diff)
80 | }
81 | })
82 | }
83 |
84 | }
85 |
86 | func TestParseFromYamlWithTemplateBigtable(t *testing.T) {
87 | ctx, err := testutils.ContextWithNewLogger()
88 | if err != nil {
89 | t.Fatalf("unexpected error: %s", err)
90 | }
91 | tcs := []struct {
92 | desc string
93 | in string
94 | want server.ToolConfigs
95 | }{
96 | {
97 | desc: "basic example",
98 | in: `
99 | tools:
100 | example_tool:
101 | kind: bigtable-sql
102 | source: my-pg-instance
103 | description: some description
104 | statement: |
105 | SELECT * FROM SQL_STATEMENT;
106 | parameters:
107 | - name: country
108 | type: string
109 | description: some description
110 | templateParameters:
111 | - name: tableName
112 | type: string
113 | description: The table to select hotels from.
114 | - name: fieldArray
115 | type: array
116 | description: The columns to return for the query.
117 | items:
118 | name: column
119 | type: string
120 | description: A column name that will be returned from the query.
121 | `,
122 | want: server.ToolConfigs{
123 | "example_tool": bigtable.Config{
124 | Name: "example_tool",
125 | Kind: "bigtable-sql",
126 | Source: "my-pg-instance",
127 | Description: "some description",
128 | Statement: "SELECT * FROM SQL_STATEMENT;\n",
129 | AuthRequired: []string{},
130 | Parameters: []tools.Parameter{
131 | tools.NewStringParameter("country", "some description"),
132 | },
133 | TemplateParameters: []tools.Parameter{
134 | tools.NewStringParameter("tableName", "The table to select hotels from."),
135 | tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
136 | },
137 | },
138 | },
139 | },
140 | }
141 | for _, tc := range tcs {
142 | t.Run(tc.desc, func(t *testing.T) {
143 | got := struct {
144 | Tools server.ToolConfigs `yaml:"tools"`
145 | }{}
146 | // Parse contents
147 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
148 | if err != nil {
149 | t.Fatalf("unable to unmarshal: %s", err)
150 | }
151 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
152 | t.Fatalf("incorrect parse: diff %v", diff)
153 | }
154 | })
155 | }
156 |
157 | }
158 |
```
--------------------------------------------------------------------------------
/internal/server/mcp/util/lifecycle.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package util
16 |
17 | import (
18 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
19 | )
20 |
21 | const (
22 | // SERVER_NAME is the server name used in Implementation.
23 | SERVER_NAME = "Toolbox"
24 | // methods that are supported
25 | INITIALIZE = "initialize"
26 | )
27 |
28 | /* Initialization */
29 |
30 | // Params to define MCP Client during initialize request.
31 | type InitializeParams struct {
32 | // The latest version of the Model Context Protocol that the client supports.
33 | // The client MAY decide to support older versions as well.
34 | ProtocolVersion string `json:"protocolVersion"`
35 | Capabilities ClientCapabilities `json:"capabilities"`
36 | ClientInfo Implementation `json:"clientInfo"`
37 | }
38 |
39 | // InitializeRequest is sent from the client to the server when it first
40 | // connects, asking it to begin initialization.
41 | type InitializeRequest struct {
42 | jsonrpc.Request
43 | Params InitializeParams `json:"params"`
44 | }
45 |
46 | // InitializeResult is sent after receiving an initialize request from the
47 | // client.
48 | type InitializeResult struct {
49 | jsonrpc.Result
50 | // The version of the Model Context Protocol that the server wants to use.
51 | // This may not match the version that the client requested. If the client cannot
52 | // support this version, it MUST disconnect.
53 | ProtocolVersion string `json:"protocolVersion"`
54 | Capabilities ServerCapabilities `json:"capabilities"`
55 | ServerInfo Implementation `json:"serverInfo"`
56 | // Instructions describing how to use the server and its features.
57 | //
58 | // This can be used by clients to improve the LLM's understanding of
59 | // available tools, resources, etc. It can be thought of like a "hint" to the model.
60 | // For example, this information MAY be added to the system prompt.
61 | Instructions string `json:"instructions,omitempty"`
62 | }
63 |
64 | // InitializedNotification is sent from the client to the server after
65 | // initialization has finished.
66 | type InitializedNotification struct {
67 | jsonrpc.Notification
68 | }
69 |
70 | // ListChange represents whether the server supports notification for changes to the capabilities.
71 | type ListChanged struct {
72 | ListChanged *bool `json:"listChanged,omitempty"`
73 | }
74 |
75 | // ClientCapabilities represents capabilities a client may support. Known
76 | // capabilities are defined here, in this schema, but this is not a closed set: any
77 | // client can define its own, additional capabilities.
78 | type ClientCapabilities struct {
79 | // Experimental, non-standard capabilities that the client supports.
80 | Experimental map[string]interface{} `json:"experimental,omitempty"`
81 | // Present if the client supports listing roots.
82 | Roots *ListChanged `json:"roots,omitempty"`
83 | // Present if the client supports sampling from an LLM.
84 | Sampling struct{} `json:"sampling,omitempty"`
85 | }
86 |
87 | // ServerCapabilities represents capabilities that a server may support. Known
88 | // capabilities are defined here, in this schema, but this is not a closed set: any
89 | // server can define its own, additional capabilities.
90 | type ServerCapabilities struct {
91 | Tools *ListChanged `json:"tools,omitempty"`
92 | }
93 |
94 | // Base interface for metadata with name (identifier) and title (display name) properties.
95 | type BaseMetadata struct {
96 | // Intended for programmatic or logical use, but used as a display name in past specs
97 | // or fallback (if title isn't present).
98 | Name string `json:"name"`
99 | // Intended for UI and end-user contexts — optimized to be human-readable and easily understood,
100 | //even by those unfamiliar with domain-specific terminology.
101 | //
102 | // If not provided, the name should be used for display (except for Tool,
103 | // where `annotations.title` should be given precedence over using `name`,
104 | // if present).
105 | Title string `json:"title,omitempty"`
106 | }
107 |
108 | // Implementation describes the name and version of an MCP implementation.
109 | type Implementation struct {
110 | BaseMetadata
111 | Version string `json:"version"`
112 | }
113 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/util/validator.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package util
16 |
17 | import (
18 | "fmt"
19 | "regexp"
20 | "strings"
21 | )
22 |
23 | // Regular expressions for validating Firestore paths
24 | var (
25 | // Pattern to detect absolute paths (those starting with "projects/")
26 | absolutePathRegex = regexp.MustCompile(`^projects/[^/]+/databases/[^/]+/documents/`)
27 | )
28 |
29 | // PathType represents the type of Firestore path
30 | type PathType int
31 |
32 | const (
33 | CollectionPath PathType = iota
34 | DocumentPath
35 | )
36 |
37 | // ValidateCollectionPath validates that a path is a valid Firestore collection path.
38 | // Collection paths must have an odd number of segments (collection/doc/collection)
39 | func ValidateCollectionPath(path string) error {
40 | return validatePath(path, CollectionPath)
41 | }
42 |
43 | // ValidateDocumentPath validates that a path is a valid Firestore document path.
44 | // Document paths must have an even number of segments (collection/doc or collection/doc/collection/doc)
45 | func ValidateDocumentPath(path string) error {
46 | return validatePath(path, DocumentPath)
47 | }
48 |
49 | // validatePath is the common validation function for both collection and document paths
50 | func validatePath(path string, pathType PathType) error {
51 | pathTypeName := "document"
52 | if pathType == CollectionPath {
53 | pathTypeName = "collection"
54 | }
55 |
56 | // Check for empty path
57 | if path == "" {
58 | return fmt.Errorf("%s path cannot be empty", pathTypeName)
59 | }
60 |
61 | // Check if it's an absolute path
62 | if absolutePathRegex.MatchString(path) {
63 | example := "users/userId"
64 | if pathType == CollectionPath {
65 | example = "users"
66 | }
67 | return fmt.Errorf("path must be relative (e.g., '%s'), not absolute (matching pattern: ^projects/[^/]+/databases/[^/]+/documents/)", example)
68 | }
69 |
70 | // Split the path using strings.Split to preserve empty segments
71 | segments := strings.Split(path, "/")
72 |
73 | // Check for empty result
74 | if len(segments) == 0 {
75 | return fmt.Errorf("%s path cannot be empty or contain only slashes", pathTypeName)
76 | }
77 |
78 | // Check segment count based on path type
79 | segmentCount := len(segments)
80 | if pathType == CollectionPath && segmentCount%2 == 0 {
81 | // Collection paths must have an odd number of segments
82 | return fmt.Errorf("invalid collection path: must have an odd number of segments (e.g., 'collection' or 'collection/doc/subcollection'), got %d segments", segmentCount)
83 | } else if pathType == DocumentPath && segmentCount%2 != 0 {
84 | // Document paths must have an even number of segments
85 | return fmt.Errorf("invalid document path: must have an even number of segments (e.g., 'collection/doc'), got %d segments", segmentCount)
86 | }
87 |
88 | // Validate each segment
89 | for i, segment := range segments {
90 | isCollectionSegment := (i % 2) == 0
91 | if err := validateSegment(segment, isCollectionSegment); err != nil {
92 | return fmt.Errorf("invalid segment at position %d (%s): %w", i+1, segment, err)
93 | }
94 | }
95 |
96 | return nil
97 | }
98 |
99 | // validateSegment validates a single path segment
100 | func validateSegment(segment string, isCollection bool) error {
101 | segmentType := "document ID"
102 | if isCollection {
103 | segmentType = "collection ID"
104 | }
105 |
106 | // Check for empty segment
107 | if segment == "" {
108 | return fmt.Errorf("segment cannot be empty")
109 | }
110 |
111 | // Check for whitespace-only segment
112 | if strings.TrimSpace(segment) == "" {
113 | return fmt.Errorf("segment cannot be only whitespace")
114 | }
115 |
116 | // Check for single or double period
117 | if segment == "." || segment == ".." {
118 | return fmt.Errorf("segment cannot be '.' or '..'")
119 | }
120 |
121 | // Check for reserved prefix
122 | if strings.HasPrefix(segment, "__") {
123 | return fmt.Errorf("%s cannot start with '__' (reserved prefix)", segmentType)
124 | }
125 |
126 | return nil
127 | }
128 |
129 | // IsAbsolutePath checks if a path is an absolute Firestore path
130 | func IsAbsolutePath(path string) bool {
131 | return absolutePathRegex.MatchString(path)
132 | }
133 |
134 | // IsRelativePath checks if a path is a relative Firestore path
135 | func IsRelativePath(path string) bool {
136 | return path != "" && !IsAbsolutePath(path)
137 | }
138 |
```
--------------------------------------------------------------------------------
/internal/tools/bigquery/bigquerysql/bigquerysql_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package bigquerysql_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools"
25 | "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
26 | )
27 |
28 | func TestParseFromYamlBigQuery(t *testing.T) {
29 | ctx, err := testutils.ContextWithNewLogger()
30 | if err != nil {
31 | t.Fatalf("unexpected error: %s", err)
32 | }
33 | tcs := []struct {
34 | desc string
35 | in string
36 | want server.ToolConfigs
37 | }{
38 | {
39 | desc: "basic example",
40 | in: `
41 | tools:
42 | example_tool:
43 | kind: bigquery-sql
44 | source: my-instance
45 | description: some description
46 | statement: |
47 | SELECT * FROM SQL_STATEMENT;
48 | parameters:
49 | - name: country
50 | type: string
51 | description: some description
52 | `,
53 | want: server.ToolConfigs{
54 | "example_tool": bigquerysql.Config{
55 | Name: "example_tool",
56 | Kind: "bigquery-sql",
57 | Source: "my-instance",
58 | Description: "some description",
59 | Statement: "SELECT * FROM SQL_STATEMENT;\n",
60 | AuthRequired: []string{},
61 | Parameters: []tools.Parameter{
62 | tools.NewStringParameter("country", "some description"),
63 | },
64 | },
65 | },
66 | },
67 | }
68 | for _, tc := range tcs {
69 | t.Run(tc.desc, func(t *testing.T) {
70 | got := struct {
71 | Tools server.ToolConfigs `yaml:"tools"`
72 | }{}
73 | // Parse contents
74 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
75 | if err != nil {
76 | t.Fatalf("unable to unmarshal: %s", err)
77 | }
78 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
79 | t.Fatalf("incorrect parse: diff %v", diff)
80 | }
81 | })
82 | }
83 |
84 | }
85 |
86 | func TestParseFromYamlWithTemplateBigQuery(t *testing.T) {
87 | ctx, err := testutils.ContextWithNewLogger()
88 | if err != nil {
89 | t.Fatalf("unexpected error: %s", err)
90 | }
91 | tcs := []struct {
92 | desc string
93 | in string
94 | want server.ToolConfigs
95 | }{
96 | {
97 | desc: "basic example",
98 | in: `
99 | tools:
100 | example_tool:
101 | kind: bigquery-sql
102 | source: my-instance
103 | description: some description
104 | statement: |
105 | SELECT * FROM SQL_STATEMENT;
106 | parameters:
107 | - name: country
108 | type: string
109 | description: some description
110 | templateParameters:
111 | - name: tableName
112 | type: string
113 | description: The table to select hotels from.
114 | - name: fieldArray
115 | type: array
116 | description: The columns to return for the query.
117 | items:
118 | name: column
119 | type: string
120 | description: A column name that will be returned from the query.
121 | `,
122 | want: server.ToolConfigs{
123 | "example_tool": bigquerysql.Config{
124 | Name: "example_tool",
125 | Kind: "bigquery-sql",
126 | Source: "my-instance",
127 | Description: "some description",
128 | Statement: "SELECT * FROM SQL_STATEMENT;\n",
129 | AuthRequired: []string{},
130 | Parameters: []tools.Parameter{
131 | tools.NewStringParameter("country", "some description"),
132 | },
133 | TemplateParameters: []tools.Parameter{
134 | tools.NewStringParameter("tableName", "The table to select hotels from."),
135 | tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
136 | },
137 | },
138 | },
139 | },
140 | }
141 | for _, tc := range tcs {
142 | t.Run(tc.desc, func(t *testing.T) {
143 | got := struct {
144 | Tools server.ToolConfigs `yaml:"tools"`
145 | }{}
146 | // Parse contents
147 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
148 | if err != nil {
149 | t.Fatalf("unable to unmarshal: %s", err)
150 | }
151 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
152 | t.Fatalf("incorrect parse: diff %v", diff)
153 | }
154 | })
155 | }
156 |
157 | }
158 |
```
--------------------------------------------------------------------------------
/internal/sources/redis/redis.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 | package redis
15 |
16 | import (
17 | "context"
18 | "fmt"
19 | "time"
20 |
21 | "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/redis/go-redis/v9"
24 | "go.opentelemetry.io/otel/trace"
25 | )
26 |
27 | const SourceKind string = "redis"
28 |
29 | // validate interface
30 | var _ sources.SourceConfig = Config{}
31 |
32 | func init() {
33 | if !sources.Register(SourceKind, newConfig) {
34 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
35 | }
36 | }
37 |
38 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
39 | actual := Config{Name: name}
40 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
41 | return nil, err
42 | }
43 | return actual, nil
44 | }
45 |
46 | type Config struct {
47 | Name string `yaml:"name" validate:"required"`
48 | Kind string `yaml:"kind" validate:"required"`
49 | Address []string `yaml:"address" validate:"required"`
50 | Username string `yaml:"username"`
51 | Password string `yaml:"password"`
52 | Database int `yaml:"database"`
53 | UseGCPIAM bool `yaml:"useGCPIAM"`
54 | ClusterEnabled bool `yaml:"clusterEnabled"`
55 | }
56 |
57 | func (r Config) SourceConfigKind() string {
58 | return SourceKind
59 | }
60 |
61 | // RedisClient is an interface for `redis.Client` and `redis.ClusterClient
62 | type RedisClient interface {
63 | Do(context.Context, ...any) *redis.Cmd
64 | }
65 |
66 | var _ RedisClient = (*redis.Client)(nil)
67 | var _ RedisClient = (*redis.ClusterClient)(nil)
68 |
69 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
70 | client, err := initRedisClient(ctx, r)
71 | if err != nil {
72 | return nil, fmt.Errorf("error initializing Redis client: %s", err)
73 | }
74 | s := &Source{
75 | Name: r.Name,
76 | Kind: SourceKind,
77 | Client: client,
78 | }
79 | return s, nil
80 | }
81 |
82 | func initRedisClient(ctx context.Context, r Config) (RedisClient, error) {
83 | var authFn func(ctx context.Context) (username string, password string, err error)
84 | if r.UseGCPIAM {
85 | // Pass in an access token getter fn for IAM auth
86 | authFn = func(ctx context.Context) (username string, password string, err error) {
87 | token, err := sources.GetIAMAccessToken(ctx)
88 | if err != nil {
89 | return "", "", err
90 | }
91 | return "default", token, nil
92 | }
93 | }
94 |
95 | var client RedisClient
96 | var err error
97 | if r.ClusterEnabled {
98 | // Create a new Redis Cluster client
99 | clusterClient := redis.NewClusterClient(&redis.ClusterOptions{
100 | Addrs: r.Address,
101 | // PoolSize applies per cluster node and not for the whole cluster.
102 | PoolSize: 10,
103 | ConnMaxIdleTime: 60 * time.Second,
104 | MinIdleConns: 1,
105 | CredentialsProviderContext: authFn,
106 | Username: r.Username,
107 | Password: r.Password,
108 | })
109 | err = clusterClient.ForEachShard(ctx, func(ctx context.Context, shard *redis.Client) error {
110 | return shard.Ping(ctx).Err()
111 | })
112 | if err != nil {
113 | return nil, fmt.Errorf("unable to connect to redis cluster: %s", err)
114 | }
115 | client = clusterClient
116 | return client, nil
117 | }
118 |
119 | // Create a new Redis client
120 | standaloneClient := redis.NewClient(&redis.Options{
121 | Addr: r.Address[0],
122 | PoolSize: 10,
123 | ConnMaxIdleTime: 60 * time.Second,
124 | MinIdleConns: 1,
125 | DB: r.Database,
126 | CredentialsProviderContext: authFn,
127 | Username: r.Username,
128 | Password: r.Password,
129 | })
130 | _, err = standaloneClient.Ping(ctx).Result()
131 | if err != nil {
132 | return nil, fmt.Errorf("unable to connect to redis: %s", err)
133 | }
134 | client = standaloneClient
135 | return client, nil
136 | }
137 |
138 | var _ sources.Source = &Source{}
139 |
140 | type Source struct {
141 | Name string `yaml:"name"`
142 | Kind string `yaml:"kind"`
143 | Client RedisClient
144 | }
145 |
146 | func (s *Source) SourceKind() string {
147 | return SourceKind
148 | }
149 |
150 | func (s *Source) RedisClient() RedisClient {
151 | return s.Client
152 | }
153 |
```
--------------------------------------------------------------------------------
/internal/sources/cloudsqlmssql/cloud_sql_mssql.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package cloudsqlmssql
16 |
17 | import (
18 | "context"
19 | "database/sql"
20 | "fmt"
21 | "net/url"
22 | "slices"
23 |
24 | "cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
25 | "github.com/goccy/go-yaml"
26 | "github.com/googleapis/genai-toolbox/internal/sources"
27 | "github.com/googleapis/genai-toolbox/internal/util"
28 | "go.opentelemetry.io/otel/trace"
29 | )
30 |
31 | const SourceKind string = "cloud-sql-mssql"
32 |
33 | // validate interface
34 | var _ sources.SourceConfig = Config{}
35 |
36 | func init() {
37 | if !sources.Register(SourceKind, newConfig) {
38 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
39 | }
40 | }
41 |
42 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
43 | actual := Config{Name: name, IPType: "public"} // Default IPType
44 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
45 | return nil, err
46 | }
47 | return actual, nil
48 | }
49 |
50 | type Config struct {
51 | // Cloud SQL MSSQL configs
52 | Name string `yaml:"name" validate:"required"`
53 | Kind string `yaml:"kind" validate:"required"`
54 | Project string `yaml:"project" validate:"required"`
55 | Region string `yaml:"region" validate:"required"`
56 | Instance string `yaml:"instance" validate:"required"`
57 | IPAddress string `yaml:"ipAddress" validate:"required"`
58 | IPType sources.IPType `yaml:"ipType" validate:"required"`
59 | User string `yaml:"user" validate:"required"`
60 | Password string `yaml:"password" validate:"required"`
61 | Database string `yaml:"database" validate:"required"`
62 | }
63 |
64 | func (r Config) SourceConfigKind() string {
65 | // Returns Cloud SQL MSSQL source kind
66 | return SourceKind
67 | }
68 |
69 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
70 | // Initializes a Cloud SQL MSSQL source
71 | db, err := initCloudSQLMssqlConnection(ctx, tracer, r.Name, r.Project, r.Region, r.Instance, r.IPAddress, r.IPType.String(), r.User, r.Password, r.Database)
72 | if err != nil {
73 | return nil, fmt.Errorf("unable to create db connection: %w", err)
74 | }
75 |
76 | // Verify db connection
77 | err = db.PingContext(ctx)
78 | if err != nil {
79 | return nil, fmt.Errorf("unable to connect successfully: %w", err)
80 | }
81 |
82 | s := &Source{
83 | Name: r.Name,
84 | Kind: SourceKind,
85 | Db: db,
86 | }
87 | return s, nil
88 | }
89 |
90 | var _ sources.Source = &Source{}
91 |
92 | type Source struct {
93 | // Cloud SQL MSSQL struct with connection pool
94 | Name string `yaml:"name"`
95 | Kind string `yaml:"kind"`
96 | Db *sql.DB
97 | }
98 |
99 | func (s *Source) SourceKind() string {
100 | // Returns Cloud SQL MSSQL source kind
101 | return SourceKind
102 | }
103 |
104 | func (s *Source) MSSQLDB() *sql.DB {
105 | // Returns a Cloud SQL MSSQL database connection pool
106 | return s.Db
107 | }
108 |
109 | func initCloudSQLMssqlConnection(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipAddress, ipType, user, pass, dbname string) (*sql.DB, error) {
110 | //nolint:all // Reassigned ctx
111 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
112 | defer span.End()
113 |
114 | userAgent, err := util.UserAgentFromContext(ctx)
115 | if err != nil {
116 | return nil, err
117 | }
118 |
119 | // Create dsn
120 | query := url.Values{}
121 | query.Add("app name", userAgent)
122 | query.Add("database", dbname)
123 | query.Add("cloudsql", fmt.Sprintf("%s:%s:%s", project, region, instance))
124 |
125 | url := &url.URL{
126 | Scheme: "sqlserver",
127 | User: url.UserPassword(user, pass),
128 | Host: ipAddress,
129 | RawQuery: query.Encode(),
130 | }
131 |
132 | // Get dial options
133 | opts, err := sources.GetCloudSQLOpts(ipType, userAgent, false)
134 | if err != nil {
135 | return nil, err
136 | }
137 |
138 | // Register sql server driver
139 | if !slices.Contains(sql.Drivers(), "cloudsql-sqlserver-driver") {
140 | _, err := mssql.RegisterDriver("cloudsql-sqlserver-driver", opts...)
141 | if err != nil {
142 | return nil, err
143 | }
144 | }
145 |
146 | // Open database connection
147 | db, err := sql.Open(
148 | "cloudsql-sqlserver-driver",
149 | url.String(),
150 | )
151 | if err != nil {
152 | return nil, err
153 | }
154 | return db, nil
155 | }
156 |
```
--------------------------------------------------------------------------------
/internal/tools/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package clickhouse
16 |
17 | import (
18 | "context"
19 | "database/sql"
20 | "fmt"
21 |
22 | yaml "github.com/goccy/go-yaml"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | "github.com/googleapis/genai-toolbox/internal/tools"
25 | )
26 |
27 | type compatibleSource interface {
28 | ClickHousePool() *sql.DB
29 | }
30 |
31 | var compatibleSources = []string{"clickhouse"}
32 |
33 | const listDatabasesKind string = "clickhouse-list-databases"
34 |
35 | func init() {
36 | if !tools.Register(listDatabasesKind, newListDatabasesConfig) {
37 | panic(fmt.Sprintf("tool kind %q already registered", listDatabasesKind))
38 | }
39 | }
40 |
41 | func newListDatabasesConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
42 | actual := Config{Name: name}
43 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
44 | return nil, err
45 | }
46 | return actual, nil
47 | }
48 |
49 | type Config struct {
50 | Name string `yaml:"name" validate:"required"`
51 | Kind string `yaml:"kind" validate:"required"`
52 | Source string `yaml:"source" validate:"required"`
53 | Description string `yaml:"description" validate:"required"`
54 | AuthRequired []string `yaml:"authRequired"`
55 | Parameters tools.Parameters `yaml:"parameters"`
56 | }
57 |
58 | var _ tools.ToolConfig = Config{}
59 |
60 | func (cfg Config) ToolConfigKind() string {
61 | return listDatabasesKind
62 | }
63 |
64 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
65 | rawS, ok := srcs[cfg.Source]
66 | if !ok {
67 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
68 | }
69 |
70 | s, ok := rawS.(compatibleSource)
71 | if !ok {
72 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", listDatabasesKind, compatibleSources)
73 | }
74 |
75 | allParameters, paramManifest, _ := tools.ProcessParameters(nil, cfg.Parameters)
76 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters)
77 |
78 | t := Tool{
79 | Name: cfg.Name,
80 | Kind: listDatabasesKind,
81 | Parameters: cfg.Parameters,
82 | AllParams: allParameters,
83 | AuthRequired: cfg.AuthRequired,
84 | Pool: s.ClickHousePool(),
85 | manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
86 | mcpManifest: mcpManifest,
87 | }
88 | return t, nil
89 | }
90 |
91 | var _ tools.Tool = Tool{}
92 |
93 | type Tool struct {
94 | Name string `yaml:"name"`
95 | Kind string `yaml:"kind"`
96 | AuthRequired []string `yaml:"authRequired"`
97 | Parameters tools.Parameters `yaml:"parameters"`
98 | AllParams tools.Parameters `yaml:"allParams"`
99 |
100 | Pool *sql.DB
101 | manifest tools.Manifest
102 | mcpManifest tools.McpManifest
103 | }
104 |
105 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, token tools.AccessToken) (any, error) {
106 | // Query to list all databases
107 | query := "SHOW DATABASES"
108 |
109 | results, err := t.Pool.QueryContext(ctx, query)
110 | if err != nil {
111 | return nil, fmt.Errorf("unable to execute query: %w", err)
112 | }
113 | defer results.Close()
114 |
115 | var databases []map[string]any
116 | for results.Next() {
117 | var dbName string
118 | err := results.Scan(&dbName)
119 | if err != nil {
120 | return nil, fmt.Errorf("unable to parse row: %w", err)
121 | }
122 | databases = append(databases, map[string]any{
123 | "name": dbName,
124 | })
125 | }
126 |
127 | if err := results.Err(); err != nil {
128 | return nil, fmt.Errorf("errors encountered by results.Scan: %w", err)
129 | }
130 |
131 | return databases, nil
132 | }
133 |
134 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
135 | return tools.ParseParams(t.AllParams, data, claims)
136 | }
137 |
138 | func (t Tool) Manifest() tools.Manifest {
139 | return t.manifest
140 | }
141 |
142 | func (t Tool) McpManifest() tools.McpManifest {
143 | return t.mcpManifest
144 | }
145 |
146 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
147 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
148 | }
149 |
150 | func (t Tool) RequiresClientAuthorization() bool {
151 | return false
152 | }
153 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestoreadddocuments_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreadddocuments"
25 | )
26 |
27 | func TestParseFromYamlFirestoreAddDocuments(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | add_docs_tool:
42 | kind: firestore-add-documents
43 | source: my-firestore-instance
44 | description: Add documents to Firestore collections
45 | `,
46 | want: server.ToolConfigs{
47 | "add_docs_tool": firestoreadddocuments.Config{
48 | Name: "add_docs_tool",
49 | Kind: "firestore-add-documents",
50 | Source: "my-firestore-instance",
51 | Description: "Add documents to Firestore collections",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_add_docs:
61 | kind: firestore-add-documents
62 | source: prod-firestore
63 | description: Add documents with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_add_docs": firestoreadddocuments.Config{
70 | Name: "secure_add_docs",
71 | Kind: "firestore-add-documents",
72 | Source: "prod-firestore",
73 | Description: "Add documents with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | add_user_docs:
104 | kind: firestore-add-documents
105 | source: users-firestore
106 | description: Add user documents
107 | authRequired:
108 | - user-auth
109 | add_product_docs:
110 | kind: firestore-add-documents
111 | source: products-firestore
112 | description: Add product documents
113 | add_order_docs:
114 | kind: firestore-add-documents
115 | source: orders-firestore
116 | description: Add order documents
117 | authRequired:
118 | - user-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "add_user_docs": firestoreadddocuments.Config{
123 | Name: "add_user_docs",
124 | Kind: "firestore-add-documents",
125 | Source: "users-firestore",
126 | Description: "Add user documents",
127 | AuthRequired: []string{"user-auth"},
128 | },
129 | "add_product_docs": firestoreadddocuments.Config{
130 | Name: "add_product_docs",
131 | Kind: "firestore-add-documents",
132 | Source: "products-firestore",
133 | Description: "Add product documents",
134 | AuthRequired: []string{},
135 | },
136 | "add_order_docs": firestoreadddocuments.Config{
137 | Name: "add_order_docs",
138 | Kind: "firestore-add-documents",
139 | Source: "orders-firestore",
140 | Description: "Add order documents",
141 | AuthRequired: []string{"user-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/server/mcp/v20241105/types.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package v20241105
16 |
17 | import (
18 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
19 | "github.com/googleapis/genai-toolbox/internal/tools"
20 | )
21 |
22 | // SERVER_NAME is the server name used in Implementation.
23 | const SERVER_NAME = "Toolbox"
24 |
25 | // PROTOCOL_VERSION is the version of the MCP protocol in this package.
26 | const PROTOCOL_VERSION = "2024-11-05"
27 |
28 | // methods that are supported.
29 | const (
30 | PING = "ping"
31 | TOOLS_LIST = "tools/list"
32 | TOOLS_CALL = "tools/call"
33 | )
34 |
35 | /* Empty result */
36 |
37 | // EmptyResult represents a response that indicates success but carries no data.
38 | type EmptyResult jsonrpc.Result
39 |
40 | /* Pagination */
41 |
42 | // Cursor is an opaque token used to represent a cursor for pagination.
43 | type Cursor string
44 |
45 | type PaginatedRequest struct {
46 | jsonrpc.Request
47 | Params struct {
48 | // An opaque token representing the current pagination position.
49 | // If provided, the server should return results starting after this cursor.
50 | Cursor Cursor `json:"cursor,omitempty"`
51 | } `json:"params,omitempty"`
52 | }
53 |
54 | type PaginatedResult struct {
55 | jsonrpc.Result
56 | // An opaque token representing the pagination position after the last returned result.
57 | // If present, there may be more results available.
58 | NextCursor Cursor `json:"nextCursor,omitempty"`
59 | }
60 |
61 | /* Tools */
62 |
63 | // Sent from the client to request a list of tools the server has.
64 | type ListToolsRequest struct {
65 | PaginatedRequest
66 | }
67 |
68 | // The server's response to a tools/list request from the client.
69 | type ListToolsResult struct {
70 | PaginatedResult
71 | Tools []tools.McpManifest `json:"tools"`
72 | }
73 |
74 | // Used by the client to invoke a tool provided by the server.
75 | type CallToolRequest struct {
76 | jsonrpc.Request
77 | Params struct {
78 | Name string `json:"name"`
79 | Arguments map[string]any `json:"arguments,omitempty"`
80 | } `json:"params,omitempty"`
81 | }
82 |
83 | // The sender or recipient of messages and data in a conversation.
84 | type Role string
85 |
86 | const (
87 | RoleUser Role = "user"
88 | RoleAssistant Role = "assistant"
89 | )
90 |
91 | // Base for objects that include optional annotations for the client.
92 | // The client can use annotations to inform how objects are used or displayed
93 | type Annotated struct {
94 | Annotations *struct {
95 | // Describes who the intended customer of this object or data is.
96 | // It can include multiple entries to indicate content useful for multiple
97 | // audiences (e.g., `["user", "assistant"]`).
98 | Audience []Role `json:"audience,omitempty"`
99 | // Describes how important this data is for operating the server.
100 | //
101 | // A value of 1 means "most important," and indicates that the data is
102 | // effectively required, while 0 means "least important," and indicates that
103 | // the data is entirely optional.
104 | //
105 | // @TJS-type number
106 | // @minimum 0
107 | // @maximum 1
108 | Priority float64 `json:"priority,omitempty"`
109 | } `json:"annotations,omitempty"`
110 | }
111 |
112 | // TextContent represents text provided to or from an LLM.
113 | type TextContent struct {
114 | Annotated
115 | Type string `json:"type"`
116 | // The text content of the message.
117 | Text string `json:"text"`
118 | }
119 |
120 | // The server's response to a tool call.
121 | //
122 | // Any errors that originate from the tool SHOULD be reported inside the result
123 | // object, with `isError` set to true, _not_ as an MCP protocol-level error
124 | // response. Otherwise, the LLM would not be able to see that an error occurred
125 | // and self-correct.
126 | //
127 | // However, any errors in _finding_ the tool, an error indicating that the
128 | // server does not support tool calls, or any other exceptional conditions,
129 | // should be reported as an MCP error response.
130 | type CallToolResult struct {
131 | jsonrpc.Result
132 | // Could be either a TextContent, ImageContent, or EmbeddedResources
133 | // For Toolbox, we will only be sending TextContent
134 | Content []TextContent `json:"content"`
135 | // Whether the tool call ended in an error.
136 | // If not set, this is assumed to be false (the call was successful).
137 | IsError bool `json:"isError,omitempty"`
138 | }
139 |
```