This is page 35 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
├── 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
├── 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/samples/bigquery/local_quickstart.md:
--------------------------------------------------------------------------------
```markdown
1 | ---
2 | title: "Quickstart (Local with BigQuery)"
3 | type: docs
4 | weight: 1
5 | description: >
6 | How to get started running Toolbox locally with Python, BigQuery, and
7 | LangGraph, LlamaIndex, or ADK.
8 | ---
9 |
10 | [](https://colab.research.google.com/github/googleapis/genai-toolbox/blob/main/docs/en/samples/bigquery/colab_quickstart_bigquery.ipynb)
12 |
13 | ## Before you begin
14 |
15 | This guide assumes you have already done the following:
16 |
17 | 1. Installed [Python 3.9+][install-python] (including [pip][install-pip] and
18 | your preferred virtual environment tool for managing dependencies e.g.
19 | [venv][install-venv]).
20 | 1. Installed and configured the [Google Cloud SDK (gcloud CLI)][install-gcloud].
21 | 1. Authenticated with Google Cloud for Application Default Credentials (ADC):
22 |
23 | ```bash
24 | gcloud auth login --update-adc
25 | ```
26 |
27 | 1. Set your default Google Cloud project (replace `YOUR_PROJECT_ID` with your
28 | actual project ID):
29 |
30 | ```bash
31 | gcloud config set project YOUR_PROJECT_ID
32 | export GOOGLE_CLOUD_PROJECT=YOUR_PROJECT_ID
33 | ```
34 |
35 | Toolbox and the client libraries will use this project for BigQuery, unless
36 | overridden in configurations.
37 | 1. [Enabled the BigQuery API][enable-bq-api] in your Google Cloud project.
38 | 1. Installed the BigQuery client library for Python:
39 |
40 | ```bash
41 | pip install google-cloud-bigquery
42 | ```
43 |
44 | 1. Completed setup for usage with an LLM model such as
45 | {{< tabpane text=true persist=header >}}
46 | {{% tab header="Core" lang="en" %}}
47 |
48 | - [langchain-vertexai](https://python.langchain.com/docs/integrations/llms/google_vertex_ai_palm/#setup)
49 | package.
50 |
51 | - [langchain-google-genai](https://python.langchain.com/docs/integrations/chat/google_generative_ai/#setup)
52 | package.
53 |
54 | - [langchain-anthropic](https://python.langchain.com/docs/integrations/chat/anthropic/#setup)
55 | package.
56 | {{% /tab %}}
57 | {{% tab header="LangChain" lang="en" %}}
58 | - [langchain-vertexai](https://python.langchain.com/docs/integrations/llms/google_vertex_ai_palm/#setup)
59 | package.
60 |
61 | - [langchain-google-genai](https://python.langchain.com/docs/integrations/chat/google_generative_ai/#setup)
62 | package.
63 |
64 | - [langchain-anthropic](https://python.langchain.com/docs/integrations/chat/anthropic/#setup)
65 | package.
66 | {{% /tab %}}
67 | {{% tab header="LlamaIndex" lang="en" %}}
68 | - [llama-index-llms-google-genai](https://pypi.org/project/llama-index-llms-google-genai/)
69 | package.
70 |
71 | - [llama-index-llms-anthropic](https://docs.llamaindex.ai/en/stable/examples/llm/anthropic)
72 | package.
73 | {{% /tab %}}
74 | {{% tab header="ADK" lang="en" %}}
75 | - [google-adk](https://pypi.org/project/google-adk/) package.
76 | {{% /tab %}}
77 | {{< /tabpane >}}
78 |
79 | [install-python]: https://wiki.python.org/moin/BeginnersGuide/Download
80 | [install-pip]: https://pip.pypa.io/en/stable/installation/
81 | [install-venv]:
82 | https://packaging.python.org/en/latest/tutorials/installing-packages/#creating-virtual-environments
83 | [install-gcloud]: https://cloud.google.com/sdk/docs/install
84 | [enable-bq-api]:
85 | https://cloud.google.com/bigquery/docs/quickstarts/query-public-dataset-console#before-you-begin
86 |
87 | ## Step 1: Set up your BigQuery Dataset and Table
88 |
89 | In this section, we will create a BigQuery dataset and a table, then insert some
90 | data that needs to be accessed by our agent. BigQuery operations are performed
91 | against your configured Google Cloud project.
92 |
93 | 1. Create a new BigQuery dataset (replace `YOUR_DATASET_NAME` with your desired
94 | dataset name, e.g., `toolbox_ds`, and optionally specify a location like `US`
95 | or `EU`):
96 |
97 | ```bash
98 | export BQ_DATASET_NAME="YOUR_DATASET_NAME" # e.g., toolbox_ds
99 | export BQ_LOCATION="US" # e.g., US, EU, asia-northeast1
100 |
101 | bq --location=$BQ_LOCATION mk $BQ_DATASET_NAME
102 | ```
103 |
104 | You can also do this through the [Google Cloud
105 | Console](https://console.cloud.google.com/bigquery).
106 |
107 | {{< notice tip >}}
108 | For a real application, ensure that the service account or user running Toolbox
109 | has the necessary IAM permissions (e.g., BigQuery Data Editor, BigQuery User)
110 | on the dataset or project. For this local quickstart with user credentials,
111 | your own permissions will apply.
112 | {{< /notice >}}
113 |
114 | 1. The hotels table needs to be defined in your new dataset for use with the bq
115 | query command. First, create a file named `create_hotels_table.sql` with the
116 | following content:
117 |
118 | ```sql
119 | CREATE TABLE IF NOT EXISTS `YOUR_PROJECT_ID.YOUR_DATASET_NAME.hotels` (
120 | id INT64 NOT NULL,
121 | name STRING NOT NULL,
122 | location STRING NOT NULL,
123 | price_tier STRING NOT NULL,
124 | checkin_date DATE NOT NULL,
125 | checkout_date DATE NOT NULL,
126 | booked BOOLEAN NOT NULL
127 | );
128 | ```
129 |
130 | > **Note:** Replace `YOUR_PROJECT_ID` and `YOUR_DATASET_NAME` in the SQL
131 | > with your actual project ID and dataset name.
132 |
133 | Then run the command below to execute the sql query:
134 |
135 | ```bash
136 | bq query --project_id=$GOOGLE_CLOUD_PROJECT --dataset_id=$BQ_DATASET_NAME --use_legacy_sql=false < create_hotels_table.sql
137 | ```
138 |
139 | 1. Next, populate the hotels table with some initial data. To do this, create a
140 | file named `insert_hotels_data.sql` and add the following SQL INSERT
141 | statement to it.
142 |
143 | ```sql
144 | INSERT INTO `YOUR_PROJECT_ID.YOUR_DATASET_NAME.hotels` (id, name, location, price_tier, checkin_date, checkout_date, booked)
145 | VALUES
146 | (1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-20', '2024-04-22', FALSE),
147 | (2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', FALSE),
148 | (3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', FALSE),
149 | (4, 'Radisson Blu Lucerne', 'Lucerne', 'Midscale', '2024-04-05', '2024-04-24', FALSE),
150 | (5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-01', '2024-04-23', FALSE),
151 | (6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', FALSE),
152 | (7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-02', '2024-04-27', FALSE),
153 | (8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-09', '2024-04-24', FALSE),
154 | (9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', FALSE),
155 | (10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', FALSE);
156 | ```
157 |
158 | > **Note:** Replace `YOUR_PROJECT_ID` and `YOUR_DATASET_NAME` in the SQL
159 | > with your actual project ID and dataset name.
160 |
161 | Then run the command below to execute the sql query:
162 |
163 | ```bash
164 | bq query --project_id=$GOOGLE_CLOUD_PROJECT --dataset_id=$BQ_DATASET_NAME --use_legacy_sql=false < insert_hotels_data.sql
165 | ```
166 |
167 | ## Step 2: Install and configure Toolbox
168 |
169 | In this section, we will download Toolbox, configure our tools in a `tools.yaml`
170 | to use BigQuery, and then run the Toolbox server.
171 |
172 | 1. Download the latest version of Toolbox as a binary:
173 |
174 | {{< notice tip >}}
175 | Select the
176 | [correct binary](https://github.com/googleapis/genai-toolbox/releases)
177 | corresponding to your OS and CPU architecture.
178 | {{< /notice >}}
179 | <!-- {x-release-please-start-version} -->
180 | ```bash
181 | export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
182 | curl -O https://storage.googleapis.com/genai-toolbox/v0.17.0/$OS/toolbox
183 | ```
184 | <!-- {x-release-please-end} -->
185 |
186 | 1. Make the binary executable:
187 |
188 | ```bash
189 | chmod +x toolbox
190 | ```
191 |
192 | 1. Write the following into a `tools.yaml` file. You must replace the
193 | `YOUR_PROJECT_ID` and `YOUR_DATASET_NAME` placeholder in the config with your
194 | actual BigQuery project and dataset name. The `location` field is optional;
195 | if not specified, it defaults to 'us'. The table name `hotels` is used
196 | directly in the statements.
197 |
198 | {{< notice tip >}}
199 | Authentication with BigQuery is handled via Application Default Credentials
200 | (ADC). Ensure you have run `gcloud auth application-default login`.
201 | {{< /notice >}}
202 |
203 | ```yaml
204 | sources:
205 | my-bigquery-source:
206 | kind: bigquery
207 | project: YOUR_PROJECT_ID
208 | location: us
209 | tools:
210 | search-hotels-by-name:
211 | kind: bigquery-sql
212 | source: my-bigquery-source
213 | description: Search for hotels based on name.
214 | parameters:
215 | - name: name
216 | type: string
217 | description: The name of the hotel.
218 | statement: SELECT * FROM `YOUR_DATASET_NAME.hotels` WHERE LOWER(name) LIKE LOWER(CONCAT('%', @name, '%'));
219 | search-hotels-by-location:
220 | kind: bigquery-sql
221 | source: my-bigquery-source
222 | description: Search for hotels based on location.
223 | parameters:
224 | - name: location
225 | type: string
226 | description: The location of the hotel.
227 | statement: SELECT * FROM `YOUR_DATASET_NAME.hotels` WHERE LOWER(location) LIKE LOWER(CONCAT('%', @location, '%'));
228 | book-hotel:
229 | kind: bigquery-sql
230 | source: my-bigquery-source
231 | description: >-
232 | Book a hotel by its ID. If the hotel is successfully booked, returns a NULL, raises an error if not.
233 | parameters:
234 | - name: hotel_id
235 | type: integer
236 | description: The ID of the hotel to book.
237 | statement: UPDATE `YOUR_DATASET_NAME.hotels` SET booked = TRUE WHERE id = @hotel_id;
238 | update-hotel:
239 | kind: bigquery-sql
240 | source: my-bigquery-source
241 | description: >-
242 | Update a hotel's check-in and check-out dates by its ID. Returns a message indicating whether the hotel was successfully updated or not.
243 | parameters:
244 | - name: checkin_date
245 | type: string
246 | description: The new check-in date of the hotel.
247 | - name: checkout_date
248 | type: string
249 | description: The new check-out date of the hotel.
250 | - name: hotel_id
251 | type: integer
252 | description: The ID of the hotel to update.
253 | statement: >-
254 | UPDATE `YOUR_DATASET_NAME.hotels` SET checkin_date = PARSE_DATE('%Y-%m-%d', @checkin_date), checkout_date = PARSE_DATE('%Y-%m-%d', @checkout_date) WHERE id = @hotel_id;
255 | cancel-hotel:
256 | kind: bigquery-sql
257 | source: my-bigquery-source
258 | description: Cancel a hotel by its ID.
259 | parameters:
260 | - name: hotel_id
261 | type: integer
262 | description: The ID of the hotel to cancel.
263 | statement: UPDATE `YOUR_DATASET_NAME.hotels` SET booked = FALSE WHERE id = @hotel_id;
264 | ```
265 |
266 | **Important Note on `toolsets`**: The `tools.yaml` content above does not
267 | include a `toolsets` section. The Python agent examples in Step 3 (e.g.,
268 | `await toolbox_client.load_toolset("my-toolset")`) rely on a toolset named
269 | `my-toolset`. To make those examples work, you will need to add a `toolsets`
270 | section to your `tools.yaml` file, for example:
271 |
272 | ```yaml
273 | # Add this to your tools.yaml if using load_toolset("my-toolset")
274 | # Ensure it's at the same indentation level as 'sources:' and 'tools:'
275 | toolsets:
276 | my-toolset:
277 | - search-hotels-by-name
278 | - search-hotels-by-location
279 | - book-hotel
280 | - update-hotel
281 | - cancel-hotel
282 | ```
283 |
284 | Alternatively, you can modify the agent code to load tools individually
285 | (e.g., using `await toolbox_client.load_tool("search-hotels-by-name")`).
286 |
287 | For more info on tools, check out the [Resources](../../resources/) section
288 | of the docs.
289 |
290 | 1. Run the Toolbox server, pointing to the `tools.yaml` file created earlier:
291 |
292 | ```bash
293 | ./toolbox --tools-file "tools.yaml"
294 | ```
295 |
296 | {{< notice note >}}
297 | Toolbox enables dynamic reloading by default. To disable, use the
298 | `--disable-reload` flag.
299 | {{< /notice >}}
300 |
301 | ## Step 3: Connect your agent to Toolbox
302 |
303 | In this section, we will write and run an agent that will load the Tools
304 | from Toolbox.
305 |
306 | {{< notice tip>}} If you prefer to experiment within a Google Colab environment,
307 | you can connect to a
308 | [local runtime](https://research.google.com/colaboratory/local-runtimes.html).
309 | {{< /notice >}}
310 |
311 | 1. In a new terminal, install the SDK package.
312 |
313 | {{< tabpane persist=header >}}
314 | {{< tab header="Core" lang="bash" >}}
315 |
316 | pip install toolbox-core
317 | {{< /tab >}}
318 | {{< tab header="Langchain" lang="bash" >}}
319 |
320 | pip install toolbox-langchain
321 | {{< /tab >}}
322 | {{< tab header="LlamaIndex" lang="bash" >}}
323 |
324 | pip install toolbox-llamaindex
325 | {{< /tab >}}
326 | {{< tab header="ADK" lang="bash" >}}
327 |
328 | pip install google-adk
329 | {{< /tab >}}
330 |
331 | {{< /tabpane >}}
332 |
333 | 1. Install other required dependencies:
334 |
335 | {{< tabpane persist=header >}}
336 | {{< tab header="Core" lang="bash" >}}
337 |
338 | # TODO(developer): replace with correct package if needed
339 |
340 | pip install langgraph langchain-google-vertexai
341 |
342 | # pip install langchain-google-genai
343 |
344 | # pip install langchain-anthropic
345 |
346 | {{< /tab >}}
347 | {{< tab header="Langchain" lang="bash" >}}
348 |
349 | # TODO(developer): replace with correct package if needed
350 |
351 | pip install langgraph langchain-google-vertexai
352 |
353 | # pip install langchain-google-genai
354 |
355 | # pip install langchain-anthropic
356 |
357 | {{< /tab >}}
358 | {{< tab header="LlamaIndex" lang="bash" >}}
359 |
360 | # TODO(developer): replace with correct package if needed
361 |
362 | pip install llama-index-llms-google-genai
363 |
364 | # pip install llama-index-llms-anthropic
365 |
366 | {{< /tab >}}
367 | {{< tab header="ADK" lang="bash" >}}
368 | pip install toolbox-core
369 | {{< /tab >}}
370 | {{< /tabpane >}}
371 |
372 | 1. Create a new file named `hotel_agent.py` and copy the following
373 | code to create an agent:
374 | {{< tabpane persist=header >}}
375 | {{< tab header="Core" lang="python" >}}
376 |
377 | import asyncio
378 |
379 | from google import genai
380 | from google.genai.types import (
381 | Content,
382 | FunctionDeclaration,
383 | GenerateContentConfig,
384 | Part,
385 | Tool,
386 | )
387 |
388 | from toolbox_core import ToolboxClient
389 |
390 | prompt = """
391 | You're a helpful hotel assistant. You handle hotel searching, booking and
392 | cancellations. When the user searches for a hotel, mention it's name, id,
393 | location and price tier. Always mention hotel id while performing any
394 | searches. This is very important for any operations. For any bookings or
395 | cancellations, please provide the appropriate confirmation. Be sure to
396 | update checkin or checkout dates if mentioned by the user.
397 | Don't ask for confirmations from the user.
398 | """
399 |
400 | queries = [
401 | "Find hotels in Basel with Basel in it's name.",
402 | "Please book the hotel Hilton Basel for me.",
403 | "This is too expensive. Please cancel it.",
404 | "Please book Hyatt Regency for me",
405 | "My check in dates for my booking would be from April 10, 2024 to April 19, 2024.",
406 | ]
407 |
408 | async def run_application():
409 | async with ToolboxClient("<http://127.0.0.1:5000>") as toolbox_client:
410 |
411 | # The toolbox_tools list contains Python callables (functions/methods) designed for LLM tool-use
412 | # integration. While this example uses Google's genai client, these callables can be adapted for
413 | # various function-calling or agent frameworks. For easier integration with supported frameworks
414 | # (https://github.com/googleapis/mcp-toolbox-python-sdk/tree/main/packages), use the
415 | # provided wrapper packages, which handle framework-specific boilerplate.
416 | toolbox_tools = await toolbox_client.load_toolset("my-toolset")
417 | genai_client = genai.Client(
418 | vertexai=True, project="project-id", location="us-central1"
419 | )
420 |
421 | genai_tools = [
422 | Tool(
423 | function_declarations=[
424 | FunctionDeclaration.from_callable_with_api_option(callable=tool)
425 | ]
426 | )
427 | for tool in toolbox_tools
428 | ]
429 | history = []
430 | for query in queries:
431 | user_prompt_content = Content(
432 | role="user",
433 | parts=[Part.from_text(text=query)],
434 | )
435 | history.append(user_prompt_content)
436 |
437 | response = genai_client.models.generate_content(
438 | model="gemini-2.0-flash-001",
439 | contents=history,
440 | config=GenerateContentConfig(
441 | system_instruction=prompt,
442 | tools=genai_tools,
443 | ),
444 | )
445 | history.append(response.candidates[0].content)
446 | function_response_parts = []
447 | for function_call in response.function_calls:
448 | fn_name = function_call.name
449 | # The tools are sorted alphabetically
450 | if fn_name == "search-hotels-by-name":
451 | function_result = await toolbox_tools[3](**function_call.args)
452 | elif fn_name == "search-hotels-by-location":
453 | function_result = await toolbox_tools[2](**function_call.args)
454 | elif fn_name == "book-hotel":
455 | function_result = await toolbox_tools[0](**function_call.args)
456 | elif fn_name == "update-hotel":
457 | function_result = await toolbox_tools[4](**function_call.args)
458 | elif fn_name == "cancel-hotel":
459 | function_result = await toolbox_tools[1](**function_call.args)
460 | else:
461 | raise ValueError("Function name not present.")
462 | function_response = {"result": function_result}
463 | function_response_part = Part.from_function_response(
464 | name=function_call.name,
465 | response=function_response,
466 | )
467 | function_response_parts.append(function_response_part)
468 |
469 | if function_response_parts:
470 | tool_response_content = Content(role="tool", parts=function_response_parts)
471 | history.append(tool_response_content)
472 |
473 | response2 = genai_client.models.generate_content(
474 | model="gemini-2.0-flash-001",
475 | contents=history,
476 | config=GenerateContentConfig(
477 | tools=genai_tools,
478 | ),
479 | )
480 | final_model_response_content = response2.candidates[0].content
481 | history.append(final_model_response_content)
482 | print(response2.text)
483 |
484 | asyncio.run(run_application())
485 | {{< /tab >}}
486 | {{< tab header="LangChain" lang="python" >}}
487 |
488 | import asyncio
489 | from langgraph.prebuilt import create_react_agent
490 |
491 | # TODO(developer): replace this with another import if needed
492 |
493 | from langchain_google_vertexai import ChatVertexAI
494 |
495 | # from langchain_google_genai import ChatGoogleGenerativeAI
496 |
497 | # from langchain_anthropic import ChatAnthropic
498 |
499 | from langgraph.checkpoint.memory import MemorySaver
500 |
501 | from toolbox_langchain import ToolboxClient
502 |
503 | prompt = """
504 | You're a helpful hotel assistant. You handle hotel searching, booking and
505 | cancellations. When the user searches for a hotel, mention it's name, id,
506 | location and price tier. Always mention hotel ids while performing any
507 | searches. This is very important for any operations. For any bookings or
508 | cancellations, please provide the appropriate confirmation. Be sure to
509 | update checkin or checkout dates if mentioned by the user.
510 | Don't ask for confirmations from the user.
511 | """
512 |
513 | queries = [
514 | "Find hotels in Basel with Basel in its name.",
515 | "Can you book the Hilton Basel for me?",
516 | "Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.",
517 | "My check in dates would be from April 10, 2024 to April 19, 2024.",
518 | ]
519 |
520 | async def main():
521 | # TODO(developer): replace this with another model if needed
522 | model = ChatVertexAI(model_name="gemini-2.0-flash-001")
523 | # model = ChatGoogleGenerativeAI(model="gemini-2.0-flash-001")
524 | # model = ChatAnthropic(model="claude-3-5-sonnet-20240620")
525 |
526 | # Load the tools from the Toolbox server
527 | client = ToolboxClient("http://127.0.0.1:5000")
528 | tools = await client.aload_toolset()
529 |
530 | agent = create_react_agent(model, tools, checkpointer=MemorySaver())
531 |
532 | config = {"configurable": {"thread_id": "thread-1"}}
533 | for query in queries:
534 | inputs = {"messages": [("user", prompt + query)]}
535 | response = await agent.ainvoke(inputs, stream_mode="values", config=config)
536 | print(response["messages"][-1].content)
537 |
538 | asyncio.run(main())
539 | {{< /tab >}}
540 | {{< tab header="LlamaIndex" lang="python" >}}
541 | import asyncio
542 | import os
543 |
544 | from llama_index.core.agent.workflow import AgentWorkflow
545 |
546 | from llama_index.core.workflow import Context
547 |
548 | # TODO(developer): replace this with another import if needed
549 |
550 | from llama_index.llms.google_genai import GoogleGenAI
551 |
552 | # from llama_index.llms.anthropic import Anthropic
553 |
554 | from toolbox_llamaindex import ToolboxClient
555 |
556 | prompt = """
557 | You're a helpful hotel assistant. You handle hotel searching, booking and
558 | cancellations. When the user searches for a hotel, mention it's name, id,
559 | location and price tier. Always mention hotel ids while performing any
560 | searches. This is very important for any operations. For any bookings or
561 | cancellations, please provide the appropriate confirmation. Be sure to
562 | update checkin or checkout dates if mentioned by the user.
563 | Don't ask for confirmations from the user.
564 | """
565 |
566 | queries = [
567 | "Find hotels in Basel with Basel in it's name.",
568 | "Can you book the Hilton Basel for me?",
569 | "Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.",
570 | "My check in dates would be from April 10, 2024 to April 19, 2024.",
571 | ]
572 |
573 | async def main():
574 | # TODO(developer): replace this with another model if needed
575 | llm = GoogleGenAI(
576 | model="gemini-2.0-flash-001",
577 | vertexai_config={"location": "us-central1"},
578 | )
579 | # llm = GoogleGenAI(
580 | # api_key=os.getenv("GOOGLE_API_KEY"),
581 | # model="gemini-2.0-flash-001",
582 | # )
583 | # llm = Anthropic(
584 | # model="claude-3-7-sonnet-latest",
585 | # api_key=os.getenv("ANTHROPIC_API_KEY")
586 | # )
587 |
588 | # Load the tools from the Toolbox server
589 | client = ToolboxClient("http://127.0.0.1:5000")
590 | tools = await client.aload_toolset()
591 |
592 | agent = AgentWorkflow.from_tools_or_functions(
593 | tools,
594 | llm=llm,
595 | system_prompt=prompt,
596 | )
597 | ctx = Context(agent)
598 | for query in queries:
599 | response = await agent.arun(user_msg=query, ctx=ctx)
600 | print(f"---- {query} ----")
601 | print(str(response))
602 |
603 | asyncio.run(main())
604 | {{< /tab >}}
605 | {{< tab header="ADK" lang="python" >}}
606 | from google.adk.agents import Agent
607 | from google.adk.runners import Runner
608 | from google.adk.sessions import InMemorySessionService
609 | from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
610 | from google.genai import types # For constructing message content
611 | from toolbox_core import ToolboxSyncClient
612 |
613 | import os
614 | os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'
615 |
616 | # TODO(developer): Replace 'YOUR_PROJECT_ID' with your Google Cloud Project ID
617 |
618 | os.environ['GOOGLE_CLOUD_PROJECT'] = 'YOUR_PROJECT_ID'
619 |
620 | # TODO(developer): Replace 'us-central1' with your Google Cloud Location (region)
621 |
622 | os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'
623 |
624 | # --- Load Tools from Toolbox ---
625 |
626 | # TODO(developer): Ensure the Toolbox server is running at <http://127.0.0.1:5000>
627 |
628 | with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
629 | # TODO(developer): Replace "my-toolset" with the actual ID of your toolset as configured in your MCP Toolbox server.
630 | agent_toolset = toolbox_client.load_toolset("my-toolset")
631 |
632 | # --- Define the Agent's Prompt ---
633 | prompt = """
634 | You're a helpful hotel assistant. You handle hotel searching, booking and
635 | cancellations. When the user searches for a hotel, mention it's name, id,
636 | location and price tier. Always mention hotel ids while performing any
637 | searches. This is very important for any operations. For any bookings or
638 | cancellations, please provide the appropriate confirmation. Be sure to
639 | update checkin or checkout dates if mentioned by the user.
640 | Don't ask for confirmations from the user.
641 | """
642 |
643 | # --- Configure the Agent ---
644 |
645 | root_agent = Agent(
646 | model='gemini-2.0-flash-001',
647 | name='hotel_agent',
648 | description='A helpful AI assistant that can search and book hotels.',
649 | instruction=prompt,
650 | tools=agent_toolset, # Pass the loaded toolset
651 | )
652 |
653 | # --- Initialize Services for Running the Agent ---
654 | session_service = InMemorySessionService()
655 | artifacts_service = InMemoryArtifactService()
656 | # Create a new session for the interaction.
657 | session = session_service.create_session(
658 | state={}, app_name='hotel_agent', user_id='123'
659 | )
660 |
661 | runner = Runner(
662 | app_name='hotel_agent',
663 | agent=root_agent,
664 | artifact_service=artifacts_service,
665 | session_service=session_service,
666 | )
667 |
668 | # --- Define Queries and Run the Agent ---
669 | queries = [
670 | "Find hotels in Basel with Basel in it's name.",
671 | "Can you book the Hilton Basel for me?",
672 | "Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.",
673 | "My check in dates would be from April 10, 2024 to April 19, 2024.",
674 | ]
675 |
676 | for query in queries:
677 | content = types.Content(role='user', parts=[types.Part(text=query)])
678 | events = runner.run(session_id=session.id,
679 | user_id='123', new_message=content)
680 |
681 | responses = (
682 | part.text
683 | for event in events
684 | for part in event.content.parts
685 | if part.text is not None
686 | )
687 |
688 | for text in responses:
689 | print(text)
690 | {{< /tab >}}
691 | {{< /tabpane >}}
692 |
693 | {{< tabpane text=true persist=header >}}
694 | {{% tab header="Core" lang="en" %}}
695 | To learn more about the Core SDK, check out the [Toolbox Core SDK
696 | documentation.](https://github.com/googleapis/mcp-toolbox-sdk-python/blob/main/packages/toolbox-core/README.md)
697 | {{% /tab %}}
698 | {{% tab header="Langchain" lang="en" %}}
699 | To learn more about Agents in LangChain, check out the [LangGraph Agent
700 | documentation.](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.chat_agent_executor.create_react_agent)
701 | {{% /tab %}}
702 | {{% tab header="LlamaIndex" lang="en" %}}
703 | To learn more about Agents in LlamaIndex, check out the [LlamaIndex
704 | AgentWorkflow
705 | documentation.](https://docs.llamaindex.ai/en/stable/examples/agent/agent_workflow_basic/)
706 | {{% /tab %}}
707 | {{% tab header="ADK" lang="en" %}}
708 | To learn more about Agents in ADK, check out the [ADK
709 | documentation.](https://google.github.io/adk-docs/)
710 | {{% /tab %}}
711 | {{< /tabpane >}}
712 |
713 | 1. Run your agent, and observe the results:
714 |
715 | ```sh
716 | python hotel_agent.py
717 | ```
718 |
```
--------------------------------------------------------------------------------
/internal/server/mcp_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 server
16 |
17 | import (
18 | "bufio"
19 | "bytes"
20 | "context"
21 | "encoding/json"
22 | "fmt"
23 | "net/http"
24 | "net/http/httptest"
25 | "os"
26 | "reflect"
27 | "strings"
28 | "testing"
29 |
30 | "github.com/googleapis/genai-toolbox/internal/log"
31 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
32 | "github.com/googleapis/genai-toolbox/internal/telemetry"
33 | "github.com/googleapis/genai-toolbox/internal/tools"
34 | )
35 |
36 | const jsonrpcVersion = "2.0"
37 | const protocolVersion20241105 = "2024-11-05"
38 | const protocolVersion20250326 = "2025-03-26"
39 | const protocolVersion20250618 = "2025-06-18"
40 | const serverName = "Toolbox"
41 |
42 | var basicInputSchema = map[string]any{
43 | "type": "object",
44 | "properties": map[string]any{},
45 | "required": []any{},
46 | }
47 |
48 | var tool2InputSchema = map[string]any{
49 | "type": "object",
50 | "properties": map[string]any{
51 | "param1": map[string]any{"type": "integer", "description": "This is the first parameter."},
52 | "param2": map[string]any{"type": "integer", "description": "This is the second parameter."},
53 | },
54 | "required": []any{"param1", "param2"},
55 | }
56 |
57 | var tool3InputSchema = map[string]any{
58 | "type": "object",
59 | "properties": map[string]any{
60 | "my_array": map[string]any{
61 | "type": "array",
62 | "description": "this param is an array of strings",
63 | "items": map[string]any{"type": "string", "description": "string item"},
64 | },
65 | },
66 | "required": []any{"my_array"},
67 | }
68 |
69 | func TestMcpEndpointWithoutInitialized(t *testing.T) {
70 | mockTools := []MockTool{tool1, tool2, tool3, tool4, tool5}
71 | toolsMap, toolsets := setUpResources(t, mockTools)
72 | r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets)
73 | defer shutdown()
74 | ts := runServer(r, false)
75 | defer ts.Close()
76 |
77 | testCases := []struct {
78 | name string
79 | url string
80 | isErr bool
81 | body jsonrpc.JSONRPCRequest
82 | want map[string]any
83 | }{
84 | {
85 | name: "ping",
86 | url: "/",
87 | body: jsonrpc.JSONRPCRequest{
88 | Jsonrpc: jsonrpcVersion,
89 | Id: "ping-test-123",
90 | Request: jsonrpc.Request{
91 | Method: "ping",
92 | },
93 | },
94 | isErr: false,
95 | want: map[string]any{
96 | "jsonrpc": "2.0",
97 | "id": "ping-test-123",
98 | "result": map[string]any{},
99 | },
100 | },
101 | {
102 | name: "tools/list",
103 | url: "/",
104 | body: jsonrpc.JSONRPCRequest{
105 | Jsonrpc: jsonrpcVersion,
106 | Id: "tools-list",
107 | Request: jsonrpc.Request{
108 | Method: "tools/list",
109 | },
110 | },
111 | isErr: false,
112 | want: map[string]any{
113 | "jsonrpc": "2.0",
114 | "id": "tools-list",
115 | "result": map[string]any{
116 | "tools": []any{
117 | map[string]any{
118 | "name": "no_params",
119 | "inputSchema": basicInputSchema,
120 | },
121 | map[string]any{
122 | "name": "some_params",
123 | "inputSchema": tool2InputSchema,
124 | },
125 | map[string]any{
126 | "name": "array_param",
127 | "description": "some description",
128 | "inputSchema": tool3InputSchema,
129 | },
130 | map[string]any{
131 | "name": "unauthorized_tool",
132 | "inputSchema": basicInputSchema,
133 | },
134 | map[string]any{
135 | "name": "require_client_auth_tool",
136 | "inputSchema": basicInputSchema,
137 | },
138 | },
139 | },
140 | },
141 | },
142 | {
143 | name: "missing method",
144 | url: "/",
145 | isErr: true,
146 | body: jsonrpc.JSONRPCRequest{
147 | Jsonrpc: jsonrpcVersion,
148 | Id: "missing-method",
149 | Request: jsonrpc.Request{},
150 | },
151 | want: map[string]any{
152 | "jsonrpc": "2.0",
153 | "id": "missing-method",
154 | "error": map[string]any{
155 | "code": -32601.0,
156 | "message": "method not found",
157 | },
158 | },
159 | },
160 | {
161 | name: "invalid jsonrpc version",
162 | url: "/",
163 | isErr: true,
164 | body: jsonrpc.JSONRPCRequest{
165 | Jsonrpc: "1.0",
166 | Id: "invalid-jsonrpc-version",
167 | Request: jsonrpc.Request{
168 | Method: "foo",
169 | },
170 | },
171 | want: map[string]any{
172 | "jsonrpc": "2.0",
173 | "id": "invalid-jsonrpc-version",
174 | "error": map[string]any{
175 | "code": -32600.0,
176 | "message": "invalid json-rpc version",
177 | },
178 | },
179 | },
180 | {
181 | name: "call tool1 unauthorized tool",
182 | url: "/",
183 | body: jsonrpc.JSONRPCRequest{
184 | Jsonrpc: jsonrpcVersion,
185 | Id: "tools-call-tool1",
186 | Request: jsonrpc.Request{
187 | Method: "tools/call",
188 | },
189 | Params: map[string]any{
190 | "name": "no_params",
191 | },
192 | },
193 | want: map[string]any{
194 | "jsonrpc": "2.0",
195 | "id": "tools-call-tool1",
196 | "result": map[string]any{
197 | "content": []any{
198 | map[string]any{
199 | "type": "text",
200 | "text": `"no_params"`,
201 | },
202 | },
203 | },
204 | },
205 | },
206 | {
207 | name: "call tool4 unauthorized tool",
208 | url: "/",
209 | body: jsonrpc.JSONRPCRequest{
210 | Jsonrpc: jsonrpcVersion,
211 | Id: "tools-call-tool4",
212 | Request: jsonrpc.Request{
213 | Method: "tools/call",
214 | },
215 | Params: map[string]any{
216 | "name": "unauthorized_tool",
217 | },
218 | },
219 | want: map[string]any{
220 | "jsonrpc": "2.0",
221 | "id": "tools-call-tool4",
222 | "error": map[string]any{
223 | "code": -32600.0,
224 | "message": "unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized",
225 | },
226 | },
227 | },
228 | {
229 | name: "call tool5 unauthorized tool",
230 | url: "/",
231 | body: jsonrpc.JSONRPCRequest{
232 | Jsonrpc: jsonrpcVersion,
233 | Id: "tools-call-tool5",
234 | Request: jsonrpc.Request{
235 | Method: "tools/call",
236 | },
237 | Params: map[string]any{
238 | "name": "require_client_auth_tool",
239 | },
240 | },
241 | want: map[string]any{
242 | "jsonrpc": "2.0",
243 | "id": "tools-call-tool5",
244 | "error": map[string]any{
245 | "code": -32600.0,
246 | "message": "missing access token in the 'Authorization' header",
247 | },
248 | },
249 | },
250 | }
251 | for _, tc := range testCases {
252 | t.Run(tc.name, func(t *testing.T) {
253 | reqMarshal, err := json.Marshal(tc.body)
254 | if err != nil {
255 | t.Fatalf("unexpected error during marshaling of body")
256 | }
257 |
258 | resp, body, err := runRequest(ts, http.MethodPost, tc.url, bytes.NewBuffer(reqMarshal), nil)
259 | if err != nil {
260 | t.Fatalf("unexpected error during request: %s", err)
261 | }
262 |
263 | // Notifications don't expect a response.
264 | if tc.want != nil {
265 | if contentType := resp.Header.Get("Content-type"); contentType != "application/json" {
266 | t.Fatalf("unexpected content-type header: want %s, got %s", "application/json", contentType)
267 | }
268 |
269 | var got map[string]any
270 | if err := json.Unmarshal(body, &got); err != nil {
271 | t.Fatalf("unexpected error unmarshalling body: %s", err)
272 | }
273 | if !reflect.DeepEqual(got, tc.want) {
274 | t.Fatalf("unexpected response: got %+v, want %+v", got, tc.want)
275 | }
276 | }
277 | })
278 | }
279 | }
280 |
281 | func runInitializeLifecycle(t *testing.T, ts *httptest.Server, protocolVersion string, initializeWant map[string]any, idHeader bool) string {
282 | initializeRequestBody := map[string]any{
283 | "jsonrpc": jsonrpcVersion,
284 | "id": "mcp-initialize",
285 | "method": "initialize",
286 | "params": map[string]any{
287 | "protocolVersion": protocolVersion,
288 | },
289 | }
290 | reqMarshal, err := json.Marshal(initializeRequestBody)
291 | if err != nil {
292 | t.Fatalf("unexpected error during marshaling of body")
293 | }
294 |
295 | resp, body, err := runRequest(ts, http.MethodPost, "/", bytes.NewBuffer(reqMarshal), nil)
296 | if err != nil {
297 | t.Fatalf("unexpected error during request: %s", err)
298 | }
299 |
300 | if contentType := resp.Header.Get("Content-type"); contentType != "application/json" {
301 | t.Fatalf("unexpected content-type header: want %s, got %s", "application/json", contentType)
302 | }
303 |
304 | sessionId := resp.Header.Get("Mcp-Session-Id")
305 | if idHeader && sessionId == "" {
306 | t.Fatalf("Mcp-Session-Id header is expected")
307 | }
308 |
309 | var got map[string]any
310 | if err := json.Unmarshal(body, &got); err != nil {
311 | t.Fatalf("unexpected error unmarshalling body: %s", err)
312 | }
313 | if !reflect.DeepEqual(got, initializeWant) {
314 | t.Fatalf("unexpected response: got %+v, want %+v", got, initializeWant)
315 | }
316 |
317 | header := map[string]string{}
318 | if sessionId != "" {
319 | header["Mcp-Session-Id"] = sessionId
320 | }
321 |
322 | initializeNotificationBody := map[string]any{
323 | "jsonrpc": jsonrpcVersion,
324 | "method": "notifications/initialized",
325 | }
326 | notiMarshal, err := json.Marshal(initializeNotificationBody)
327 | if err != nil {
328 | t.Fatalf("unexpected error during marshaling of notifications body")
329 | }
330 |
331 | _, _, err = runRequest(ts, http.MethodPost, "/", bytes.NewBuffer(notiMarshal), header)
332 | if err != nil {
333 | t.Fatalf("unexpected error during request: %s", err)
334 | }
335 | return sessionId
336 | }
337 |
338 | func TestMcpEndpoint(t *testing.T) {
339 | mockTools := []MockTool{tool1, tool2, tool3, tool4, tool5}
340 | toolsMap, toolsets := setUpResources(t, mockTools)
341 | r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets)
342 | defer shutdown()
343 | ts := runServer(r, false)
344 | defer ts.Close()
345 |
346 | versTestCases := []struct {
347 | name string
348 | protocol string
349 | idHeader bool
350 | initWant map[string]any
351 | }{
352 | {
353 | name: "version 2024-11-05",
354 | protocol: protocolVersion20241105,
355 | idHeader: false,
356 | initWant: map[string]any{
357 | "jsonrpc": "2.0",
358 | "id": "mcp-initialize",
359 | "result": map[string]any{
360 | "protocolVersion": "2024-11-05",
361 | "capabilities": map[string]any{
362 | "tools": map[string]any{"listChanged": false},
363 | },
364 | "serverInfo": map[string]any{"name": serverName, "version": fakeVersionString},
365 | },
366 | },
367 | },
368 | {
369 | name: "version 2025-03-26",
370 | protocol: protocolVersion20250326,
371 | idHeader: true,
372 | initWant: map[string]any{
373 | "jsonrpc": "2.0",
374 | "id": "mcp-initialize",
375 | "result": map[string]any{
376 | "protocolVersion": "2025-03-26",
377 | "capabilities": map[string]any{
378 | "tools": map[string]any{"listChanged": false},
379 | },
380 | "serverInfo": map[string]any{"name": serverName, "version": fakeVersionString},
381 | },
382 | },
383 | },
384 | {
385 | name: "version 2025-06-18",
386 | protocol: protocolVersion20250618,
387 | idHeader: false,
388 | initWant: map[string]any{
389 | "jsonrpc": "2.0",
390 | "id": "mcp-initialize",
391 | "result": map[string]any{
392 | "protocolVersion": "2025-06-18",
393 | "capabilities": map[string]any{
394 | "tools": map[string]any{"listChanged": false},
395 | },
396 | "serverInfo": map[string]any{"name": serverName, "version": fakeVersionString},
397 | },
398 | },
399 | },
400 | }
401 | for _, vtc := range versTestCases {
402 | t.Run(vtc.name, func(t *testing.T) {
403 | sessionId := runInitializeLifecycle(t, ts, vtc.protocol, vtc.initWant, vtc.idHeader)
404 |
405 | header := map[string]string{}
406 | if sessionId != "" {
407 | header["Mcp-Session-Id"] = sessionId
408 | }
409 |
410 | if vtc.protocol == protocolVersion20250618 {
411 | header["MCP-Protocol-Version"] = vtc.protocol
412 | }
413 |
414 | testCases := []struct {
415 | name string
416 | url string
417 | isErr bool
418 | body any
419 | wantStatusCode int
420 | want map[string]any
421 | }{
422 | {
423 | name: "basic notification",
424 | url: "/",
425 | body: jsonrpc.JSONRPCRequest{
426 | Jsonrpc: jsonrpcVersion,
427 | Request: jsonrpc.Request{
428 | Method: "notification",
429 | },
430 | },
431 | wantStatusCode: http.StatusAccepted,
432 | },
433 | {
434 | name: "ping",
435 | url: "/",
436 | body: jsonrpc.JSONRPCRequest{
437 | Jsonrpc: jsonrpcVersion,
438 | Id: "ping-test-123",
439 | Request: jsonrpc.Request{
440 | Method: "ping",
441 | },
442 | },
443 | wantStatusCode: http.StatusOK,
444 | want: map[string]any{
445 | "jsonrpc": "2.0",
446 | "id": "ping-test-123",
447 | "result": map[string]any{},
448 | },
449 | },
450 | {
451 | name: "tools/list",
452 | url: "/",
453 | body: jsonrpc.JSONRPCRequest{
454 | Jsonrpc: jsonrpcVersion,
455 | Id: "tools-list",
456 | Request: jsonrpc.Request{
457 | Method: "tools/list",
458 | },
459 | },
460 | wantStatusCode: http.StatusOK,
461 | want: map[string]any{
462 | "jsonrpc": "2.0",
463 | "id": "tools-list",
464 | "result": map[string]any{
465 | "tools": []any{
466 | map[string]any{
467 | "name": "no_params",
468 | "inputSchema": basicInputSchema,
469 | },
470 | map[string]any{
471 | "name": "some_params",
472 | "inputSchema": tool2InputSchema,
473 | },
474 | map[string]any{
475 | "name": "array_param",
476 | "description": "some description",
477 | "inputSchema": tool3InputSchema,
478 | },
479 | map[string]any{
480 | "name": "unauthorized_tool",
481 | "inputSchema": basicInputSchema,
482 | },
483 | map[string]any{
484 | "name": "require_client_auth_tool",
485 | "inputSchema": basicInputSchema,
486 | },
487 | },
488 | },
489 | },
490 | },
491 | {
492 | name: "tools/list on tool1_only",
493 | url: "/tool1_only",
494 | body: jsonrpc.JSONRPCRequest{
495 | Jsonrpc: jsonrpcVersion,
496 | Id: "tools-list-tool1",
497 | Request: jsonrpc.Request{
498 | Method: "tools/list",
499 | },
500 | },
501 | wantStatusCode: http.StatusOK,
502 | want: map[string]any{
503 | "jsonrpc": "2.0",
504 | "id": "tools-list-tool1",
505 | "result": map[string]any{
506 | "tools": []any{
507 | map[string]any{
508 | "name": "no_params",
509 | "inputSchema": basicInputSchema,
510 | },
511 | },
512 | },
513 | },
514 | },
515 | {
516 | name: "tools/list on invalid tool set",
517 | url: "/foo",
518 | isErr: true,
519 | body: jsonrpc.JSONRPCRequest{
520 | Jsonrpc: jsonrpcVersion,
521 | Id: "tools-list-invalid-toolset",
522 | Request: jsonrpc.Request{
523 | Method: "tools/list",
524 | },
525 | },
526 | wantStatusCode: http.StatusOK,
527 | want: map[string]any{
528 | "jsonrpc": "2.0",
529 | "id": "tools-list-invalid-toolset",
530 | "error": map[string]any{
531 | "code": -32600.0,
532 | "message": "toolset does not exist",
533 | },
534 | },
535 | },
536 | {
537 | name: "missing method",
538 | url: "/",
539 | isErr: true,
540 | body: jsonrpc.JSONRPCRequest{
541 | Jsonrpc: jsonrpcVersion,
542 | Id: "missing-method",
543 | Request: jsonrpc.Request{},
544 | },
545 | wantStatusCode: http.StatusOK,
546 | want: map[string]any{
547 | "jsonrpc": "2.0",
548 | "id": "missing-method",
549 | "error": map[string]any{
550 | "code": -32601.0,
551 | "message": "method not found",
552 | },
553 | },
554 | },
555 | {
556 | name: "invalid method",
557 | url: "/",
558 | isErr: true,
559 | body: jsonrpc.JSONRPCRequest{
560 | Jsonrpc: jsonrpcVersion,
561 | Id: "invalid-method",
562 | Request: jsonrpc.Request{
563 | Method: "foo",
564 | },
565 | },
566 | wantStatusCode: http.StatusOK,
567 | want: map[string]any{
568 | "jsonrpc": "2.0",
569 | "id": "invalid-method",
570 | "error": map[string]any{
571 | "code": -32601.0,
572 | "message": "invalid method foo",
573 | },
574 | },
575 | },
576 | {
577 | name: "invalid jsonrpc version",
578 | url: "/",
579 | isErr: true,
580 | body: jsonrpc.JSONRPCRequest{
581 | Jsonrpc: "1.0",
582 | Id: "invalid-jsonrpc-version",
583 | Request: jsonrpc.Request{
584 | Method: "foo",
585 | },
586 | },
587 | wantStatusCode: http.StatusOK,
588 | want: map[string]any{
589 | "jsonrpc": "2.0",
590 | "id": "invalid-jsonrpc-version",
591 | "error": map[string]any{
592 | "code": -32600.0,
593 | "message": "invalid json-rpc version",
594 | },
595 | },
596 | },
597 | {
598 | name: "batch requests",
599 | url: "/",
600 | isErr: true,
601 | body: []any{
602 | jsonrpc.JSONRPCRequest{
603 | Jsonrpc: "1.0",
604 | Id: "batch-requests1",
605 | Request: jsonrpc.Request{
606 | Method: "foo",
607 | },
608 | },
609 | jsonrpc.JSONRPCRequest{
610 | Jsonrpc: jsonrpcVersion,
611 | Id: "batch-requests2",
612 | Request: jsonrpc.Request{
613 | Method: "tools/list",
614 | },
615 | },
616 | },
617 | wantStatusCode: http.StatusOK,
618 | want: map[string]any{
619 | "jsonrpc": "2.0",
620 | "error": map[string]any{
621 | "code": -32600.0,
622 | "message": "not supporting batch requests",
623 | },
624 | },
625 | },
626 | {
627 | name: "call tool1 unauthorized tool",
628 | url: "/",
629 | body: jsonrpc.JSONRPCRequest{
630 | Jsonrpc: jsonrpcVersion,
631 | Id: "tools-call-tool1",
632 | Request: jsonrpc.Request{
633 | Method: "tools/call",
634 | },
635 | Params: map[string]any{
636 | "name": "no_params",
637 | },
638 | },
639 | wantStatusCode: http.StatusOK,
640 | want: map[string]any{
641 | "jsonrpc": "2.0",
642 | "id": "tools-call-tool1",
643 | "result": map[string]any{
644 | "content": []any{
645 | map[string]any{
646 | "type": "text",
647 | "text": `"no_params"`,
648 | },
649 | },
650 | },
651 | },
652 | },
653 | {
654 | name: "call tool4 unauthorized tool",
655 | url: "/",
656 | body: jsonrpc.JSONRPCRequest{
657 | Jsonrpc: jsonrpcVersion,
658 | Id: "tools-call-tool4",
659 | Request: jsonrpc.Request{
660 | Method: "tools/call",
661 | },
662 | Params: map[string]any{
663 | "name": "unauthorized_tool",
664 | },
665 | },
666 | wantStatusCode: http.StatusUnauthorized,
667 | want: map[string]any{
668 | "jsonrpc": "2.0",
669 | "id": "tools-call-tool4",
670 | "error": map[string]any{
671 | "code": -32600.0,
672 | "message": "unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized",
673 | },
674 | },
675 | },
676 | {
677 | name: "call tool5 unauthorized tool",
678 | url: "/",
679 | body: jsonrpc.JSONRPCRequest{
680 | Jsonrpc: jsonrpcVersion,
681 | Id: "tools-call-tool5",
682 | Request: jsonrpc.Request{
683 | Method: "tools/call",
684 | },
685 | Params: map[string]any{
686 | "name": "require_client_auth_tool",
687 | },
688 | },
689 | wantStatusCode: http.StatusUnauthorized,
690 | want: map[string]any{
691 | "jsonrpc": "2.0",
692 | "id": "tools-call-tool5",
693 | "error": map[string]any{
694 | "code": -32600.0,
695 | "message": "missing access token in the 'Authorization' header",
696 | },
697 | },
698 | },
699 | }
700 | for _, tc := range testCases {
701 | t.Run(tc.name, func(t *testing.T) {
702 | reqMarshal, err := json.Marshal(tc.body)
703 | if err != nil {
704 | t.Fatalf("unexpected error during marshaling of body")
705 | }
706 |
707 | if vtc.protocol != protocolVersion20241105 && len(header) == 0 {
708 | t.Fatalf("header is missing")
709 | }
710 |
711 | resp, body, err := runRequest(ts, http.MethodPost, tc.url, bytes.NewBuffer(reqMarshal), header)
712 |
713 | if err != nil {
714 | t.Fatalf("unexpected error during request: %s", err)
715 | }
716 |
717 | if resp.StatusCode != tc.wantStatusCode {
718 | t.Errorf("StatusCode mismatch: got %d, want %d", resp.StatusCode, tc.wantStatusCode)
719 | }
720 |
721 | // Notifications don't expect a response.
722 | if tc.want != nil {
723 | if contentType := resp.Header.Get("Content-type"); contentType != "application/json" {
724 | t.Fatalf("unexpected content-type header: want %s, got %s", "application/json", contentType)
725 | }
726 |
727 | var got map[string]any
728 | if err := json.Unmarshal(body, &got); err != nil {
729 | t.Fatalf("unexpected error unmarshalling body: %s", err)
730 | }
731 | // for decode failure, a random uuid is generated in server
732 | if tc.want["id"] == nil {
733 | tc.want["id"] = got["id"]
734 | }
735 | if !reflect.DeepEqual(got, tc.want) {
736 | t.Fatalf("unexpected response: got %+v, want %+v", got, tc.want)
737 | }
738 | }
739 | })
740 | }
741 | })
742 | }
743 | }
744 |
745 | func TestInvalidProtocolVersionHeader(t *testing.T) {
746 | toolsMap, toolsets := map[string]tools.Tool{}, map[string]tools.Toolset{}
747 | r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets)
748 | defer shutdown()
749 | ts := runServer(r, false)
750 | defer ts.Close()
751 |
752 | header := map[string]string{}
753 | header["MCP-Protocol-Version"] = "foo"
754 |
755 | resp, body, err := runRequest(ts, http.MethodPost, "/", nil, header)
756 | if resp.Status != "400 Bad Request" {
757 | t.Fatalf("unexpected status: %s", resp.Status)
758 | }
759 | var got map[string]any
760 | if err := json.Unmarshal(body, &got); err != nil {
761 | t.Fatalf("unexpected error unmarshalling body: %s", err)
762 | }
763 | want := "invalid protocol version: foo"
764 | if got["error"] != want {
765 | t.Fatalf("unexpected error message: got %s, want %s", got["error"], want)
766 | }
767 | if err != nil {
768 | t.Fatalf("unexpected error during request: %s", err)
769 | }
770 | }
771 |
772 | func TestDeleteEndpoint(t *testing.T) {
773 | toolsMap, toolsets := map[string]tools.Tool{}, map[string]tools.Toolset{}
774 | r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets)
775 | defer shutdown()
776 | ts := runServer(r, false)
777 | defer ts.Close()
778 |
779 | resp, _, err := runRequest(ts, http.MethodDelete, "/", nil, nil)
780 | if resp.Status != "200 OK" {
781 | t.Fatalf("unexpected status: %s", resp.Status)
782 | }
783 | if err != nil {
784 | t.Fatalf("unexpected error during request: %s", err)
785 | }
786 | }
787 |
788 | func TestGetEndpoint(t *testing.T) {
789 | toolsMap, toolsets := map[string]tools.Tool{}, map[string]tools.Toolset{}
790 | r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets)
791 | defer shutdown()
792 | ts := runServer(r, false)
793 | defer ts.Close()
794 |
795 | resp, body, err := runRequest(ts, http.MethodGet, "/", nil, nil)
796 | if resp.Status != "405 Method Not Allowed" {
797 | t.Fatalf("unexpected status: %s", resp.Status)
798 | }
799 | var got map[string]any
800 | if err := json.Unmarshal(body, &got); err != nil {
801 | t.Fatalf("unexpected error unmarshalling body: %s", err)
802 | }
803 | want := "toolbox does not support streaming in streamable HTTP transport"
804 | if got["error"] != want {
805 | t.Fatalf("unexpected error message: %s", got["error"])
806 | }
807 | if err != nil {
808 | t.Fatalf("unexpected error during request: %s", err)
809 | }
810 | }
811 |
812 | func TestSseEndpoint(t *testing.T) {
813 | r, shutdown := setUpServer(t, "mcp", nil, nil)
814 | defer shutdown()
815 | ts := runServer(r, false)
816 | defer ts.Close()
817 | if !strings.Contains(ts.URL, "http://127.0.0.1") {
818 | t.Fatalf("unexpected url, got %s", ts.URL)
819 | }
820 | tsPort := strings.TrimPrefix(ts.URL, "http://127.0.0.1:")
821 | tls := runServer(r, true)
822 | defer tls.Close()
823 | if !strings.Contains(tls.URL, "https://127.0.0.1") {
824 | t.Fatalf("unexpected url, got %s", tls.URL)
825 | }
826 | tlsPort := strings.TrimPrefix(tls.URL, "https://127.0.0.1:")
827 |
828 | contentType := "text/event-stream"
829 | cacheControl := "no-cache"
830 | connection := "keep-alive"
831 | accessControlAllowOrigin := "*"
832 |
833 | testCases := []struct {
834 | name string
835 | server *httptest.Server
836 | path string
837 | proto string
838 | event string
839 | }{
840 | {
841 | name: "basic",
842 | server: ts,
843 | path: "/sse",
844 | event: fmt.Sprintf("event: endpoint\ndata: %s/mcp?sessionId=", ts.URL),
845 | },
846 | {
847 | name: "toolset1",
848 | server: ts,
849 | path: "/tool1_only/sse",
850 | event: fmt.Sprintf("event: endpoint\ndata: http://127.0.0.1:%s/mcp/tool1_only?sessionId=", tsPort),
851 | },
852 | {
853 | name: "basic with http proto",
854 | server: ts,
855 | path: "/sse",
856 | proto: "http",
857 | event: fmt.Sprintf("event: endpoint\ndata: http://127.0.0.1:%s/mcp?sessionId=", tsPort),
858 | },
859 | {
860 | name: "basic tls with https proto",
861 | server: ts,
862 | path: "/sse",
863 | proto: "https",
864 | event: fmt.Sprintf("event: endpoint\ndata: https://127.0.0.1:%s/mcp?sessionId=", tsPort),
865 | },
866 | {
867 | name: "basic tls",
868 | server: tls,
869 | path: "/sse",
870 | event: fmt.Sprintf("event: endpoint\ndata: https://127.0.0.1:%s/mcp?sessionId=", tlsPort),
871 | },
872 | }
873 |
874 | for _, tc := range testCases {
875 | t.Run(tc.name, func(t *testing.T) {
876 | resp, err := runSseRequest(tc.server, tc.path, tc.proto)
877 | if err != nil {
878 | t.Fatalf("unable to run sse request: %s", err)
879 | }
880 | defer resp.Body.Close()
881 |
882 | if gotContentType := resp.Header.Get("Content-type"); gotContentType != contentType {
883 | t.Fatalf("unexpected content-type header: want %s, got %s", contentType, gotContentType)
884 | }
885 | if gotCacheControl := resp.Header.Get("Cache-Control"); gotCacheControl != cacheControl {
886 | t.Fatalf("unexpected cache-control header: want %s, got %s", cacheControl, gotCacheControl)
887 | }
888 | if gotConnection := resp.Header.Get("Connection"); gotConnection != connection {
889 | t.Fatalf("unexpected content-type header: want %s, got %s", connection, gotConnection)
890 | }
891 | if gotAccessControlAllowOrigin := resp.Header.Get("Access-Control-Allow-Origin"); gotAccessControlAllowOrigin != accessControlAllowOrigin {
892 | t.Fatalf("unexpected cache-control header: want %s, got %s", accessControlAllowOrigin, gotAccessControlAllowOrigin)
893 | }
894 |
895 | buffer := make([]byte, 1024)
896 | n, err := resp.Body.Read(buffer)
897 | if err != nil {
898 | t.Fatalf("unable to read response: %s", err)
899 | }
900 | endpointEvent := string(buffer[:n])
901 | if !strings.Contains(endpointEvent, tc.event) {
902 | t.Fatalf("unexpected event: got %s, want to contain %s", endpointEvent, tc.event)
903 | }
904 | })
905 | }
906 | }
907 |
908 | func runSseRequest(ts *httptest.Server, path string, proto string) (*http.Response, error) {
909 | req, err := http.NewRequest(http.MethodGet, ts.URL+path, nil)
910 | if err != nil {
911 | return nil, fmt.Errorf("unable to create request: %w", err)
912 | }
913 | if proto != "" {
914 | req.Header.Set("X-Forwarded-Proto", proto)
915 | }
916 | resp, err := ts.Client().Do(req)
917 | if err != nil {
918 | return nil, fmt.Errorf("unable to send request: %w", err)
919 | }
920 | return resp, nil
921 | }
922 |
923 | func TestStdioSession(t *testing.T) {
924 | ctx, cancel := context.WithCancel(context.Background())
925 | defer cancel()
926 |
927 | mockTools := []MockTool{tool1, tool2, tool3}
928 | toolsMap, toolsets := setUpResources(t, mockTools)
929 |
930 | pr, pw, err := os.Pipe()
931 | if err != nil {
932 | t.Fatalf("error with Pipe: %s", err)
933 | }
934 |
935 | testLogger, err := log.NewStdLogger(pw, os.Stderr, "warn")
936 | if err != nil {
937 | t.Fatalf("unable to initialize logger: %s", err)
938 | }
939 |
940 | otelShutdown, err := telemetry.SetupOTel(ctx, fakeVersionString, "", false, "toolbox")
941 | if err != nil {
942 | t.Fatalf("unable to setup otel: %s", err)
943 | }
944 | defer func() {
945 | err := otelShutdown(ctx)
946 | if err != nil {
947 | t.Fatalf("error shutting down OpenTelemetry: %s", err)
948 | }
949 | }()
950 |
951 | instrumentation, err := telemetry.CreateTelemetryInstrumentation(fakeVersionString)
952 | if err != nil {
953 | t.Fatalf("unable to create custom metrics: %s", err)
954 | }
955 |
956 | sseManager := newSseManager(ctx)
957 |
958 | resourceManager := NewResourceManager(nil, nil, toolsMap, toolsets)
959 |
960 | server := &Server{
961 | version: fakeVersionString,
962 | logger: testLogger,
963 | instrumentation: instrumentation,
964 | sseManager: sseManager,
965 | ResourceMgr: resourceManager,
966 | }
967 |
968 | in := bufio.NewReader(pr)
969 | stdioSession := NewStdioSession(server, in, pw)
970 |
971 | // test stdioSession.readLine()
972 | input := "test readLine function\n"
973 | _, err = fmt.Fprintf(pw, "%s", input)
974 | if err != nil {
975 | t.Fatalf("error writing into pipe w: %s", err)
976 | }
977 |
978 | line, err := stdioSession.readLine(ctx)
979 | if err != nil {
980 | t.Fatalf("error with stdioSession.readLine: %s", err)
981 | }
982 | if line != input {
983 | t.Fatalf("unexpected line: got %s, want %s", line, input)
984 | }
985 |
986 | // test stdioSession.write()
987 | write := "test write function"
988 | err = stdioSession.write(ctx, write)
989 | if err != nil {
990 | t.Fatalf("error with stdioSession.write: %s", err)
991 | }
992 |
993 | read, err := in.ReadString('\n')
994 | if err != nil {
995 | t.Fatalf("error reading: %s", err)
996 | }
997 | want := fmt.Sprintf(`"%s"`, write) + "\n"
998 | if read != want {
999 | t.Fatalf("unexpected read: got %s, want %s", read, want)
1000 | }
1001 | }
1002 |
```
--------------------------------------------------------------------------------
/tests/spanner/spanner_integration_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 spanner
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "fmt"
22 | "io"
23 | "net/http"
24 | "os"
25 | "regexp"
26 | "strings"
27 | "testing"
28 | "time"
29 |
30 | "cloud.google.com/go/spanner"
31 | database "cloud.google.com/go/spanner/admin/database/apiv1"
32 | "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
33 | "github.com/google/uuid"
34 | "github.com/googleapis/genai-toolbox/internal/testutils"
35 | "github.com/googleapis/genai-toolbox/internal/tools"
36 | "github.com/googleapis/genai-toolbox/tests"
37 | )
38 |
39 | var (
40 | SpannerSourceKind = "spanner"
41 | SpannerToolKind = "spanner-sql"
42 | SpannerProject = os.Getenv("SPANNER_PROJECT")
43 | SpannerDatabase = os.Getenv("SPANNER_DATABASE")
44 | SpannerInstance = os.Getenv("SPANNER_INSTANCE")
45 | )
46 |
47 | func getSpannerVars(t *testing.T) map[string]any {
48 | switch "" {
49 | case SpannerProject:
50 | t.Fatal("'SPANNER_PROJECT' not set")
51 | case SpannerDatabase:
52 | t.Fatal("'SPANNER_DATABASE' not set")
53 | case SpannerInstance:
54 | t.Fatal("'SPANNER_INSTANCE' not set")
55 | }
56 |
57 | return map[string]any{
58 | "kind": SpannerSourceKind,
59 | "project": SpannerProject,
60 | "instance": SpannerInstance,
61 | "database": SpannerDatabase,
62 | }
63 | }
64 |
65 | func initSpannerClients(ctx context.Context, project, instance, dbname string) (*spanner.Client, *database.DatabaseAdminClient, error) {
66 | // Configure the connection to the database
67 | db := fmt.Sprintf("projects/%s/instances/%s/databases/%s", project, instance, dbname)
68 |
69 | // Configure session pool to automatically clean inactive transactions
70 | sessionPoolConfig := spanner.SessionPoolConfig{
71 | TrackSessionHandles: true,
72 | InactiveTransactionRemovalOptions: spanner.InactiveTransactionRemovalOptions{
73 | ActionOnInactiveTransaction: spanner.WarnAndClose,
74 | },
75 | }
76 |
77 | // Create Spanner client (for queries)
78 | dataClient, err := spanner.NewClientWithConfig(context.Background(), db, spanner.ClientConfig{SessionPoolConfig: sessionPoolConfig})
79 | if err != nil {
80 | return nil, nil, fmt.Errorf("unable to create new Spanner client: %w", err)
81 | }
82 |
83 | // Create Spanner admin client (for creating databases)
84 | adminClient, err := database.NewDatabaseAdminClient(ctx)
85 | if err != nil {
86 | return nil, nil, fmt.Errorf("unable to create new Spanner admin client: %w", err)
87 | }
88 |
89 | return dataClient, adminClient, nil
90 | }
91 |
92 | func TestSpannerToolEndpoints(t *testing.T) {
93 | sourceConfig := getSpannerVars(t)
94 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
95 | defer cancel()
96 |
97 | var args []string
98 |
99 | // Create Spanner client
100 | dataClient, adminClient, err := initSpannerClients(ctx, SpannerProject, SpannerInstance, SpannerDatabase)
101 | if err != nil {
102 | t.Fatalf("unable to create Spanner client: %s", err)
103 | }
104 |
105 | // create table name with UUID
106 | tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
107 | tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
108 | tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
109 |
110 | // set up data for param tool
111 | createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := getSpannerParamToolInfo(tableNameParam)
112 | dbString := fmt.Sprintf(
113 | "projects/%s/instances/%s/databases/%s",
114 | SpannerProject,
115 | SpannerInstance,
116 | SpannerDatabase,
117 | )
118 | teardownTable1 := setupSpannerTable(t, ctx, adminClient, dataClient, createParamTableStmt, insertParamTableStmt, tableNameParam, dbString, paramTestParams)
119 | defer teardownTable1(t)
120 |
121 | // set up data for auth tool
122 | createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := getSpannerAuthToolInfo(tableNameAuth)
123 | teardownTable2 := setupSpannerTable(t, ctx, adminClient, dataClient, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, dbString, authTestParams)
124 | defer teardownTable2(t)
125 |
126 | // set up data for template param tool
127 | createStatementTmpl := fmt.Sprintf("CREATE TABLE %s (id INT64, name STRING(MAX), age INT64) PRIMARY KEY (id)", tableNameTemplateParam)
128 | teardownTableTmpl := setupSpannerTable(t, ctx, adminClient, dataClient, createStatementTmpl, "", tableNameTemplateParam, dbString, nil)
129 | defer teardownTableTmpl(t)
130 |
131 | // Write config into a file and pass it to command
132 | toolsFile := tests.GetToolsConfig(sourceConfig, SpannerToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
133 | toolsFile = addSpannerExecuteSqlConfig(t, toolsFile)
134 | toolsFile = addSpannerReadOnlyConfig(t, toolsFile)
135 | toolsFile = addTemplateParamConfig(t, toolsFile)
136 | toolsFile = addSpannerListTablesConfig(t, toolsFile)
137 |
138 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
139 | if err != nil {
140 | t.Fatalf("command initialization returned an error: %s", err)
141 | }
142 | defer cleanup()
143 |
144 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
145 | defer cancel()
146 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
147 | if err != nil {
148 | t.Logf("toolbox command logs: \n%s", out)
149 | t.Fatalf("toolbox didn't start successfully: %s", err)
150 | }
151 |
152 | // Get configs for tests
153 | select1Want := "[{\"\":\"1\"}]"
154 | invokeParamWant := "[{\"id\":\"1\",\"name\":\"Alice\"},{\"id\":\"3\",\"name\":\"Sid\"}]"
155 | accessSchemaWant := "[{\"schema_name\":\"INFORMATION_SCHEMA\"}]"
156 | toolInvokeMyToolById4Want := `[{"id":"4","name":null}]`
157 | mcpMyFailToolWant := `"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute client: unable to parse row: spanner: code = \"InvalidArgument\", desc = \"Syntax error: Unexpected identifier \\\\\\\"SELEC\\\\\\\" [at 1:1]\\\\nSELEC 1;\\\\n^\"`
158 | mcpMyToolId3NameAliceWant := `{"jsonrpc":"2.0","id":"my-tool","result":{"content":[{"type":"text","text":"{\"id\":\"1\",\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":\"3\",\"name\":\"Sid\"}"}]}}`
159 | mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"\":\"1\"}"}]}}`
160 | tmplSelectAllWwant := "[{\"age\":\"21\",\"id\":\"1\",\"name\":\"Alex\"},{\"age\":\"100\",\"id\":\"2\",\"name\":\"Alice\"}]"
161 | tmplSelectId1Want := "[{\"age\":\"21\",\"id\":\"1\",\"name\":\"Alex\"}]"
162 |
163 | // Run tests
164 | tests.RunToolGetTest(t)
165 | tests.RunToolInvokeTest(t, select1Want,
166 | tests.WithMyToolId3NameAliceWant(invokeParamWant),
167 | tests.WithMyArrayToolWant(invokeParamWant),
168 | tests.WithMyToolById4Want(toolInvokeMyToolById4Want),
169 | )
170 | tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want, tests.WithMcpMyToolId3NameAliceWant(mcpMyToolId3NameAliceWant))
171 | tests.RunToolInvokeWithTemplateParameters(
172 | t, tableNameTemplateParam,
173 | tests.WithSelectAllWant(tmplSelectAllWwant),
174 | tests.WithTmplSelectId1Want(tmplSelectId1Want),
175 | tests.DisableDdlTest(),
176 | )
177 | runSpannerSchemaToolInvokeTest(t, accessSchemaWant)
178 | runSpannerExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam, tableNameAuth)
179 | runSpannerListTablesTest(t, tableNameParam, tableNameAuth, tableNameTemplateParam)
180 | }
181 |
182 | // getSpannerToolInfo returns statements and param for my-tool for spanner-sql kind
183 | func getSpannerParamToolInfo(tableName string) (string, string, string, string, string, string, map[string]any) {
184 | createStatement := fmt.Sprintf("CREATE TABLE %s (id INT64, name STRING(MAX)) PRIMARY KEY (id)", tableName)
185 | insertStatement := fmt.Sprintf("INSERT INTO %s (id, name) VALUES (1, @name1), (2, @name2), (3, @name3), (4, @name4)", tableName)
186 | toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = @id OR name = @name", tableName)
187 | idToolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = @id", tableName)
188 | nameToolStatement := fmt.Sprintf("SELECT * FROM %s WHERE name = @name", tableName)
189 | arrayToolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id IN UNNEST(@idArray) AND name IN UNNEST(@nameArray)", tableName)
190 | params := map[string]any{"name1": "Alice", "name2": "Jane", "name3": "Sid", "name4": nil}
191 | return createStatement, insertStatement, toolStatement, idToolStatement, nameToolStatement, arrayToolStatement, params
192 | }
193 |
194 | // getSpannerAuthToolInfo returns statements and param of my-auth-tool for spanner-sql kind
195 | func getSpannerAuthToolInfo(tableName string) (string, string, string, map[string]any) {
196 | createStatement := fmt.Sprintf("CREATE TABLE %s (id INT64, name STRING(MAX), email STRING(MAX)) PRIMARY KEY (id)", tableName)
197 | insertStatement := fmt.Sprintf("INSERT INTO %s (id, name, email) VALUES (1, @name1, @email1), (2, @name2, @email2)", tableName)
198 | toolStatement := fmt.Sprintf("SELECT name FROM %s WHERE email = @email", tableName)
199 | params := map[string]any{
200 | "name1": "Alice",
201 | "email1": tests.ServiceAccountEmail,
202 | "name2": "Jane",
203 | "email2": "[email protected]",
204 | }
205 | return createStatement, insertStatement, toolStatement, params
206 | }
207 |
208 | // setupSpannerTable creates and inserts data into a table of tool
209 | // compatible with spanner-sql tool
210 | func setupSpannerTable(t *testing.T, ctx context.Context, adminClient *database.DatabaseAdminClient, dataClient *spanner.Client, createStatement, insertStatement, tableName, dbString string, params map[string]any) func(*testing.T) {
211 |
212 | // Create table
213 | op, err := adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{
214 | Database: dbString,
215 | Statements: []string{createStatement},
216 | })
217 | if err != nil {
218 | t.Fatalf("unable to start create table operation %s: %s", tableName, err)
219 | }
220 | err = op.Wait(ctx)
221 | if err != nil {
222 | t.Fatalf("unable to create test table %s: %s", tableName, err)
223 | }
224 |
225 | // Insert test data
226 | if insertStatement != "" {
227 | _, err = dataClient.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
228 | stmt := spanner.Statement{
229 | SQL: insertStatement,
230 | Params: params,
231 | }
232 | _, err := txn.Update(ctx, stmt)
233 | return err
234 | })
235 | if err != nil {
236 | t.Fatalf("unable to insert test data: %s", err)
237 | }
238 | }
239 |
240 | return func(t *testing.T) {
241 | // tear down test
242 | op, err = adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{
243 | Database: dbString,
244 | Statements: []string{fmt.Sprintf("DROP TABLE %s", tableName)},
245 | })
246 | if err != nil {
247 | t.Errorf("unable to start drop %s operation: %s", tableName, err)
248 | return
249 | }
250 |
251 | opErr := op.Wait(ctx)
252 | if opErr != nil {
253 | t.Errorf("Teardown failed: %s", opErr)
254 | }
255 | }
256 | }
257 |
258 | // addSpannerExecuteSqlConfig gets the tools config for `spanner-execute-sql`
259 | func addSpannerExecuteSqlConfig(t *testing.T, config map[string]any) map[string]any {
260 | tools, ok := config["tools"].(map[string]any)
261 | if !ok {
262 | t.Fatalf("unable to get tools from config")
263 | }
264 | tools["my-exec-sql-tool-read-only"] = map[string]any{
265 | "kind": "spanner-execute-sql",
266 | "source": "my-instance",
267 | "description": "Tool to execute sql",
268 | "readOnly": true,
269 | }
270 | tools["my-exec-sql-tool"] = map[string]any{
271 | "kind": "spanner-execute-sql",
272 | "source": "my-instance",
273 | "description": "Tool to execute sql",
274 | }
275 | tools["my-auth-exec-sql-tool"] = map[string]any{
276 | "kind": "spanner-execute-sql",
277 | "source": "my-instance",
278 | "description": "Tool to execute sql",
279 | "authRequired": []string{
280 | "my-google-auth",
281 | },
282 | }
283 | config["tools"] = tools
284 | return config
285 | }
286 |
287 | func addSpannerReadOnlyConfig(t *testing.T, config map[string]any) map[string]any {
288 | tools, ok := config["tools"].(map[string]any)
289 | if !ok {
290 | t.Fatalf("unable to get tools from config")
291 | }
292 | tools["access-schema-read-only"] = map[string]any{
293 | "kind": "spanner-sql",
294 | "source": "my-instance",
295 | "description": "Tool to access information schema in read-only mode.",
296 | "statement": "SELECT schema_name FROM `INFORMATION_SCHEMA`.SCHEMATA WHERE schema_name='INFORMATION_SCHEMA';",
297 | "readOnly": true,
298 | }
299 | tools["access-schema"] = map[string]any{
300 | "kind": "spanner-sql",
301 | "source": "my-instance",
302 | "description": "Tool to access information schema.",
303 | "statement": "SELECT schema_name FROM `INFORMATION_SCHEMA`.SCHEMATA WHERE schema_name='INFORMATION_SCHEMA';",
304 | }
305 | config["tools"] = tools
306 | return config
307 | }
308 |
309 | // addSpannerListTablesConfig adds the spanner-list-tables tool configuration
310 | func addSpannerListTablesConfig(t *testing.T, config map[string]any) map[string]any {
311 | tools, ok := config["tools"].(map[string]any)
312 | if !ok {
313 | t.Fatalf("unable to get tools from config")
314 | }
315 |
316 | // Add spanner-list-tables tool
317 | tools["list-tables-tool"] = map[string]any{
318 | "kind": "spanner-list-tables",
319 | "source": "my-instance",
320 | "description": "Lists tables with their schema information",
321 | }
322 |
323 | config["tools"] = tools
324 | return config
325 | }
326 |
327 | func addTemplateParamConfig(t *testing.T, config map[string]any) map[string]any {
328 | toolsMap, ok := config["tools"].(map[string]any)
329 | if !ok {
330 | t.Fatalf("unable to get tools from config")
331 | }
332 | toolsMap["insert-table-templateParams-tool"] = map[string]any{
333 | "kind": "spanner-sql",
334 | "source": "my-instance",
335 | "description": "Insert tool with template parameters",
336 | "statement": "INSERT INTO {{.tableName}} ({{array .columns}}) VALUES ({{.values}})",
337 | "templateParameters": []tools.Parameter{
338 | tools.NewStringParameter("tableName", "some description"),
339 | tools.NewArrayParameter("columns", "The columns to insert into", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
340 | tools.NewStringParameter("values", "The values to insert as a comma separated string"),
341 | },
342 | }
343 | toolsMap["select-templateParams-tool"] = map[string]any{
344 | "kind": "spanner-sql",
345 | "source": "my-instance",
346 | "description": "Create table tool with template parameters",
347 | "statement": "SELECT * FROM {{.tableName}}",
348 | "templateParameters": []tools.Parameter{
349 | tools.NewStringParameter("tableName", "some description"),
350 | },
351 | }
352 | toolsMap["select-templateParams-combined-tool"] = map[string]any{
353 | "kind": "spanner-sql",
354 | "source": "my-instance",
355 | "description": "Create table tool with template parameters",
356 | "statement": "SELECT * FROM {{.tableName}} WHERE id = @id",
357 | "parameters": []tools.Parameter{tools.NewIntParameter("id", "the id of the user")},
358 | "templateParameters": []tools.Parameter{
359 | tools.NewStringParameter("tableName", "some description"),
360 | },
361 | }
362 | toolsMap["select-fields-templateParams-tool"] = map[string]any{
363 | "kind": "spanner-sql",
364 | "source": "my-instance",
365 | "description": "Create table tool with template parameters",
366 | "statement": "SELECT {{array .fields}} FROM {{.tableName}}",
367 | "templateParameters": []tools.Parameter{
368 | tools.NewStringParameter("tableName", "some description"),
369 | tools.NewArrayParameter("fields", "The fields to select from", tools.NewStringParameter("field", "A field that will be returned from the query.")),
370 | },
371 | }
372 | toolsMap["select-filter-templateParams-combined-tool"] = map[string]any{
373 | "kind": "spanner-sql",
374 | "source": "my-instance",
375 | "description": "Create table tool with template parameters",
376 | "statement": "SELECT * FROM {{.tableName}} WHERE {{.columnFilter}} = @name",
377 | "parameters": []tools.Parameter{tools.NewStringParameter("name", "the name of the user")},
378 | "templateParameters": []tools.Parameter{
379 | tools.NewStringParameter("tableName", "some description"),
380 | tools.NewStringParameter("columnFilter", "some description"),
381 | },
382 | }
383 | config["tools"] = toolsMap
384 | return config
385 | }
386 |
387 | func runSpannerExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWant, tableNameParam, tableNameAuth string) {
388 | // Get ID token
389 | idToken, err := tests.GetGoogleIdToken(tests.ClientId)
390 | if err != nil {
391 | t.Fatalf("error getting Google ID token: %s", err)
392 | }
393 |
394 | // Test tool invoke endpoint
395 | invokeTcs := []struct {
396 | name string
397 | api string
398 | requestHeader map[string]string
399 | requestBody io.Reader
400 | want string
401 | isErr bool
402 | }{
403 | {
404 | name: "invoke my-exec-sql-tool-read-only",
405 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool-read-only/invoke",
406 | requestHeader: map[string]string{},
407 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
408 | want: select1Want,
409 | isErr: false,
410 | },
411 | {
412 | name: "invoke my-exec-sql-tool-read-only with data present in table",
413 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool-read-only/invoke",
414 | requestHeader: map[string]string{},
415 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"SELECT * FROM %s WHERE id = 3 OR name = 'Alice'\"}", tableNameParam))),
416 | want: invokeParamWant,
417 | isErr: false,
418 | },
419 | {
420 | name: "invoke my-exec-sql-tool-read-only create table",
421 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool-read-only/invoke",
422 | requestHeader: map[string]string{},
423 | requestBody: bytes.NewBuffer([]byte(`{"sql":"CREATE TABLE t (id SERIAL PRIMARY KEY, name TEXT)"}`)),
424 | isErr: true,
425 | },
426 | {
427 | name: "invoke my-exec-sql-tool-read-only drop table",
428 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool-read-only/invoke",
429 | requestHeader: map[string]string{},
430 | requestBody: bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)),
431 | isErr: true,
432 | },
433 | {
434 | name: "invoke my-exec-sql-tool-read-only insert entry",
435 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool-read-only/invoke",
436 | requestHeader: map[string]string{},
437 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"INSERT INTO %s (id, name) VALUES (4, 'test_name')\"}", tableNameParam))),
438 | isErr: true,
439 | },
440 | {
441 | name: "invoke my-exec-sql-tool without body",
442 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
443 | requestHeader: map[string]string{},
444 | requestBody: bytes.NewBuffer([]byte(`{}`)),
445 | isErr: true,
446 | },
447 | {
448 | name: "invoke my-exec-sql-tool",
449 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
450 | requestHeader: map[string]string{},
451 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
452 | want: select1Want,
453 | isErr: false,
454 | },
455 | {
456 | name: "invoke my-exec-sql-tool create table",
457 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
458 | requestHeader: map[string]string{},
459 | requestBody: bytes.NewBuffer([]byte(`{"sql":"CREATE TABLE t (id SERIAL PRIMARY KEY, name TEXT)"}`)),
460 | isErr: true,
461 | },
462 | {
463 | name: "invoke my-exec-sql-tool drop table",
464 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
465 | requestHeader: map[string]string{},
466 | requestBody: bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)),
467 | isErr: true,
468 | },
469 | {
470 | name: "invoke my-exec-sql-tool insert entry",
471 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
472 | requestHeader: map[string]string{},
473 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"sql\":\"INSERT INTO %s (id, name) VALUES (5, 'test_name')\"}", tableNameParam))),
474 | want: "null",
475 | isErr: false,
476 | },
477 | {
478 | name: "invoke my-exec-sql-tool without body",
479 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
480 | requestHeader: map[string]string{},
481 | requestBody: bytes.NewBuffer([]byte(`{}`)),
482 | isErr: true,
483 | },
484 | {
485 | name: "Invoke my-auth-exec-sql-tool with auth token",
486 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
487 | requestHeader: map[string]string{"my-google-auth_token": idToken},
488 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
489 | isErr: false,
490 | want: select1Want,
491 | },
492 | {
493 | name: "Invoke my-auth-exec-sql-tool with invalid auth token",
494 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
495 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
496 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
497 | isErr: true,
498 | },
499 | {
500 | name: "Invoke my-auth-exec-sql-tool without auth token",
501 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
502 | requestHeader: map[string]string{},
503 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
504 | isErr: true,
505 | },
506 | }
507 | for _, tc := range invokeTcs {
508 | t.Run(tc.name, func(t *testing.T) {
509 | // Send Tool invocation request
510 | req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
511 | if err != nil {
512 | t.Fatalf("unable to create request: %s", err)
513 | }
514 | req.Header.Add("Content-type", "application/json")
515 | for k, v := range tc.requestHeader {
516 | req.Header.Add(k, v)
517 | }
518 | resp, err := http.DefaultClient.Do(req)
519 | if err != nil {
520 | t.Fatalf("unable to send request: %s", err)
521 | }
522 | defer resp.Body.Close()
523 |
524 | if resp.StatusCode != http.StatusOK {
525 | if tc.isErr {
526 | return
527 | }
528 | bodyBytes, _ := io.ReadAll(resp.Body)
529 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
530 | }
531 |
532 | // Check response body
533 | var body map[string]interface{}
534 | err = json.NewDecoder(resp.Body).Decode(&body)
535 | if err != nil {
536 | t.Fatalf("error parsing response body")
537 | }
538 |
539 | got, ok := body["result"].(string)
540 | if !ok {
541 | t.Fatalf("unable to find result in response body")
542 | }
543 |
544 | if got != tc.want {
545 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
546 | }
547 | })
548 | }
549 | }
550 |
551 | // Helper function to verify table list results
552 | func verifyTableListResult(t *testing.T, body map[string]interface{}, expectedTables []string, expectedSimpleFormat bool) {
553 | // Parse the result
554 | result, ok := body["result"].(string)
555 | if !ok {
556 | t.Fatalf("unable to find result in response body")
557 | }
558 |
559 | var tables []interface{}
560 | err := json.Unmarshal([]byte(result), &tables)
561 | if err != nil {
562 | t.Fatalf("unable to parse result as JSON array: %s", err)
563 | }
564 |
565 | // If we expect specific tables, verify they exist
566 | if len(expectedTables) > 0 {
567 | tableNames := make(map[string]bool)
568 | requiredKeys := []string{"schema_name", "object_name", "object_type", "columns", "constraints", "indexes"}
569 | if expectedSimpleFormat {
570 | requiredKeys = []string{"name"}
571 | }
572 |
573 | for _, table := range tables {
574 | tableMap, ok := table.(map[string]interface{})
575 | if !ok {
576 | continue
577 | }
578 |
579 | // Parse object_details JSON string into map[string]interface{}
580 | if objectDetailsStr, ok := tableMap["object_details"].(string); ok {
581 | var objectDetails map[string]interface{}
582 | if err := json.Unmarshal([]byte(objectDetailsStr), &objectDetails); err != nil {
583 | t.Errorf("failed to parse object_details JSON: %v for %v", err, objectDetailsStr)
584 | continue
585 | }
586 |
587 | for _, reqKey := range requiredKeys {
588 | if _, hasKey := objectDetails[reqKey]; !hasKey {
589 | t.Errorf("missing required key '%s', for object_details: %v", reqKey, objectDetails)
590 | }
591 | }
592 | }
593 |
594 | if name, ok := tableMap["object_name"].(string); ok {
595 | tableNames[name] = true
596 | }
597 | }
598 |
599 | for _, expected := range expectedTables {
600 | if !tableNames[expected] {
601 | t.Errorf("expected table %s not found in results", expected)
602 | }
603 | }
604 | }
605 | }
606 |
607 | // runSpannerListTablesTest tests the spanner-list-tables tool
608 | func runSpannerListTablesTest(t *testing.T, tableNameParam, tableNameAuth, tableNameTemplateParam string) {
609 | invokeTcs := []struct {
610 | name string
611 | requestBody io.Reader
612 | expectedTables []string // empty means don't check specific tables
613 | useSimpleFormat bool
614 | }{
615 | {
616 | name: "list all tables with detailed format",
617 | requestBody: bytes.NewBuffer([]byte(`{}`)),
618 | expectedTables: []string{tableNameParam, tableNameAuth, tableNameTemplateParam},
619 | },
620 | {
621 | name: "list tables with simple format",
622 | requestBody: bytes.NewBuffer([]byte(`{"output_format": "simple"}`)),
623 | expectedTables: []string{tableNameParam, tableNameAuth, tableNameTemplateParam},
624 | useSimpleFormat: true,
625 | },
626 | {
627 | name: "list specific tables",
628 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth))),
629 | expectedTables: []string{tableNameParam, tableNameAuth},
630 | },
631 | {
632 | name: "list non-existent table",
633 | requestBody: bytes.NewBuffer([]byte(`{"table_names": "non_existent_table_xyz"}`)),
634 | expectedTables: []string{},
635 | },
636 | }
637 |
638 | for _, tc := range invokeTcs {
639 | t.Run(tc.name, func(t *testing.T) {
640 | // Use RunRequest helper function from tests package
641 | url := "http://127.0.0.1:5000/api/tool/list-tables-tool/invoke"
642 | headers := map[string]string{}
643 |
644 | resp, respBody := tests.RunRequest(t, http.MethodPost, url, tc.requestBody, headers)
645 |
646 | if resp.StatusCode != http.StatusOK {
647 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
648 | }
649 |
650 | // Check response body
651 | var body map[string]interface{}
652 | err := json.Unmarshal(respBody, &body)
653 | if err != nil {
654 | t.Fatalf("error parsing response body: %s", err)
655 | }
656 |
657 | verifyTableListResult(t, body, tc.expectedTables, tc.useSimpleFormat)
658 | })
659 | }
660 | }
661 |
662 | func runSpannerSchemaToolInvokeTest(t *testing.T, accessSchemaWant string) {
663 | invokeTcs := []struct {
664 | name string
665 | api string
666 | requestHeader map[string]string
667 | requestBody io.Reader
668 | want string
669 | isErr bool
670 | }{
671 | {
672 | name: "invoke list-tables-read-only",
673 | api: "http://127.0.0.1:5000/api/tool/access-schema-read-only/invoke",
674 | requestHeader: map[string]string{},
675 | requestBody: bytes.NewBuffer([]byte(`{}`)),
676 | want: accessSchemaWant,
677 | isErr: false,
678 | },
679 | {
680 | name: "invoke list-tables",
681 | api: "http://127.0.0.1:5000/api/tool/access-schema/invoke",
682 | requestHeader: map[string]string{},
683 | requestBody: bytes.NewBuffer([]byte(`{}`)),
684 | isErr: true,
685 | },
686 | }
687 | for _, tc := range invokeTcs {
688 | t.Run(tc.name, func(t *testing.T) {
689 | // Send Tool invocation request
690 | req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
691 | if err != nil {
692 | t.Fatalf("unable to create request: %s", err)
693 | }
694 | req.Header.Add("Content-type", "application/json")
695 | for k, v := range tc.requestHeader {
696 | req.Header.Add(k, v)
697 | }
698 | resp, err := http.DefaultClient.Do(req)
699 | if err != nil {
700 | t.Fatalf("unable to send request: %s", err)
701 | }
702 | defer resp.Body.Close()
703 |
704 | if resp.StatusCode != http.StatusOK {
705 | if tc.isErr {
706 | return
707 | }
708 | bodyBytes, _ := io.ReadAll(resp.Body)
709 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
710 | }
711 |
712 | // Check response body
713 | var body map[string]interface{}
714 | err = json.NewDecoder(resp.Body).Decode(&body)
715 | if err != nil {
716 | t.Fatalf("error parsing response body")
717 | }
718 |
719 | got, ok := body["result"].(string)
720 | if !ok {
721 | t.Fatalf("unable to find result in response body")
722 | }
723 |
724 | if got != tc.want {
725 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
726 | }
727 | })
728 | }
729 | }
730 |
```