├── .ci ├── continuous.release.cloudbuild.yaml ├── generate_release_table.sh ├── integration.cloudbuild.yaml └── versioned.release.cloudbuild.yaml ├── .github ├── CODEOWNERS ├── ISSUE_TEMPLATE │ ├── bug_report.yml │ ├── config.yml │ ├── feature_request.yml │ └── question.yml ├── auto-label.yaml ├── blunderbuss.yml ├── header-checker-lint.yaml ├── label-sync.yml ├── labels.yaml ├── release-please.yml ├── renovate.json5 ├── sync-repo-settings.yaml └── workflows │ ├── cloud_build_failure_reporter.yml │ ├── 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 │ │ └── _variables_project.scss ├── go.mod ├── go.sum ├── hugo.toml ├── layouts │ ├── _default │ │ └── _markup │ │ │ └── render-heading.html │ ├── partials │ │ └── page-meta-links.html │ └── robot.txt ├── 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 ├── CHANGELOG.md ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── DEVELOPER.md ├── Dockerfile ├── LICENSE ├── README.md ├── cmd ├── options.go ├── options_test.go ├── root.go ├── root_test.go └── version.txt ├── 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.md │ └── mcp_quickstart │ │ ├── _index.md │ │ ├── inspector.png │ │ └── inspector_tools.png │ ├── how-to │ ├── _index.md │ ├── connect-ide │ │ ├── _index.md │ │ ├── alloydb_pg_mcp.md │ │ ├── bigquery_mcp.md │ │ ├── cloud_sql_mssql_mcp.md │ │ ├── cloud_sql_mysql_mcp.md │ │ ├── cloud_sql_pg_mcp.md │ │ ├── postgres_mcp.md │ │ └── spanner_mcp.md │ ├── connect_via_mcp.md │ ├── deploy_docker.md │ ├── deploy_gke.md │ ├── deploy_toolbox.md │ └── export_telemetry.md │ ├── resources │ ├── _index.md │ ├── authServices │ │ ├── _index.md │ │ └── google.md │ ├── sources │ │ ├── _index.md │ │ ├── alloydb-pg.md │ │ ├── bigquery.md │ │ ├── bigtable.md │ │ ├── cloud-sql-mssql.md │ │ ├── cloud-sql-mysql.md │ │ ├── cloud-sql-pg.md │ │ ├── couchbase.md │ │ ├── dgraph.md │ │ ├── http.md │ │ ├── mssql.md │ │ ├── mysql.md │ │ ├── neo4j.md │ │ ├── postgres.md │ │ ├── spanner.md │ │ └── sqlite.md │ └── tools │ │ ├── _index.md │ │ ├── alloydb-ai-nl.md │ │ ├── bigquery-execute-sql.md │ │ ├── bigquery-get-dataset-info.md │ │ ├── bigquery-get-table-info.md │ │ ├── bigquery-list-dataset-ids.md │ │ ├── bigquery-list-table-ids.md │ │ ├── bigquery-sql.md │ │ ├── bigtable-sql.md │ │ ├── couchbase-sql.md │ │ ├── dgraph-dql.md │ │ ├── http.md │ │ ├── mssql-execute-sql.md │ │ ├── mssql-sql.md │ │ ├── mysql-execute-sql.md │ │ ├── mysql-sql.md │ │ ├── neo4j-cypher.md │ │ ├── postgres-execute-sql.md │ │ ├── postgres-sql.md │ │ ├── spanner-execute-sql.md │ │ ├── spanner-sql.md │ │ └── sqlite-sql.md │ └── samples │ ├── _index.md │ └── bigquery │ ├── _index.md │ ├── colab_quickstart_bigquery.ipynb │ ├── local_quickstart.md │ └── mcp_quickstart │ ├── _index.md │ ├── inspector.png │ └── inspector_tools.png ├── go.mod ├── go.sum ├── internal ├── auth │ ├── auth.go │ └── google │ │ └── google.go ├── log │ ├── handler.go │ ├── log.go │ ├── log_test.go │ └── logger.go ├── prebuiltconfigs │ ├── prebuiltconfigs.go │ ├── prebuiltconfigs_test.go │ └── tools │ │ ├── alloydb-postgres.yaml │ │ ├── bigquery.yaml │ │ ├── cloud-sql-mssql.yaml │ │ ├── cloud-sql-mysql.yaml │ │ ├── cloud-sql-postgres.yaml │ │ ├── postgres.yaml │ │ ├── spanner-postgres.yaml │ │ └── spanner.yaml ├── server │ ├── api.go │ ├── api_test.go │ ├── common_test.go │ ├── config.go │ ├── instrumentation.go │ ├── mcp.go │ ├── mcp │ │ ├── method.go │ │ └── types.go │ ├── mcp_test.go │ ├── server.go │ └── server_test.go ├── sources │ ├── alloydbpg │ │ ├── alloydb_pg.go │ │ └── alloydb_pg_test.go │ ├── bigquery │ │ ├── bigquery.go │ │ └── bigquery_test.go │ ├── bigtable │ │ ├── bigtable.go │ │ └── bigtable_test.go │ ├── cloudsqlmssql │ │ ├── cloud_sql_mssql.go │ │ └── cloud_sql_mssql_test.go │ ├── cloudsqlmysql │ │ ├── cloud_sql_mysql.go │ │ └── cloud_sql_mysql_test.go │ ├── cloudsqlpg │ │ ├── cloud_sql_pg.go │ │ └── cloud_sql_pg_test.go │ ├── couchbase │ │ ├── couchbase.go │ │ └── couchbase_test.go │ ├── dgraph │ │ ├── dgraph.go │ │ └── dgraph_test.go │ ├── dialect.go │ ├── http │ │ ├── http.go │ │ └── http_test.go │ ├── ip_type.go │ ├── mssql │ │ ├── mssql.go │ │ └── mssql_test.go │ ├── mysql │ │ ├── mysql.go │ │ └── mysql_test.go │ ├── neo4j │ │ ├── neo4j.go │ │ └── neo4j_test.go │ ├── postgres │ │ ├── postgres.go │ │ └── postgres_test.go │ ├── sources.go │ ├── spanner │ │ ├── spanner.go │ │ └── spanner_test.go │ ├── sqlite │ │ ├── sqlite.go │ │ └── sqlite_test.go │ └── util.go ├── telemetry │ └── telemetry.go ├── testutils │ └── testutils.go ├── tools │ ├── alloydbainl │ │ ├── alloydbainl.go │ │ └── alloydbainl_test.go │ ├── bigquery │ │ ├── bigquery.go │ │ └── bigquery_test.go │ ├── bigqueryexecutesql │ │ ├── bigqueryexecutesql.go │ │ └── bigqueryexecutesql_test.go │ ├── bigquerygetdatasetinfo │ │ ├── bigquerygetdatasetinfo.go │ │ └── bigquerygetdatasetinfo_test.go │ ├── bigquerygettableinfo │ │ ├── bigquerygettableinfo.go │ │ └── bigquerygettableinfo_test.go │ ├── bigquerylistdatasetids │ │ ├── bigquerylistdatasetids.go │ │ └── bigquerylistdatasetids_test.go │ ├── bigquerylisttableids │ │ ├── bigquerylisttableids.go │ │ └── bigquerylisttableids_test.go │ ├── bigtable │ │ ├── bigtable.go │ │ └── bigtable_test.go │ ├── common.go │ ├── couchbase │ │ ├── couchbase.go │ │ └── couchbase_test.go │ ├── dgraph │ │ ├── dgraph.go │ │ └── dgraph_test.go │ ├── http │ │ ├── http.go │ │ └── http_test.go │ ├── http_method.go │ ├── mssqlexecutesql │ │ ├── mssqlexecutesql.go │ │ └── mssqlexecutesql_test.go │ ├── mssqlsql │ │ ├── mssqlsql.go │ │ └── mssqlsql_test.go │ ├── mysqlexecutesql │ │ ├── mysqlexecutesql.go │ │ └── mysqlexecutesql_test.go │ ├── mysqlsql │ │ ├── mysqlsql.go │ │ └── mysqlsql_test.go │ ├── neo4j │ │ ├── neo4j.go │ │ └── neo4j_test.go │ ├── parameters.go │ ├── parameters_test.go │ ├── postgresexecutesql │ │ ├── postgresexecutesql.go │ │ └── postgresexecutesql_test.go │ ├── postgressql │ │ ├── postgressql.go │ │ └── postgressql_test.go │ ├── spanner │ │ ├── spanner.go │ │ └── spanner_test.go │ ├── spannerexecutesql │ │ ├── spannerexecutesql.go │ │ └── spannerexecutesql_test.go │ ├── sqlitesql │ │ ├── sqlitesql.go │ │ └── sqlitesql_test.go │ ├── tools.go │ └── toolsets.go └── util │ └── util.go ├── logo.png ├── main.go └── tests ├── 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 ├── cloudsqlmssql └── cloud_sql_mssql_integration_test.go ├── cloudsqlmysql └── cloud_sql_mysql_integration_test.go ├── cloudsqlpg └── cloud_sql_pg_integration_test.go ├── common.go ├── couchbase └── couchbase_integration_test.go ├── dgraph └── dgraph_integration_test.go ├── http └── http_integration_test.go ├── mssql └── mssql_integration_test.go ├── mysql └── mysql_integration_test.go ├── neo4j └── neo4j_integration_test.go ├── postgres └── postgres_integration_test.go ├── server.go ├── source.go ├── spanner └── spanner_integration_test.go ├── sqlite └── sqlite_integration_test.go └── tool.go /.ci/generate_release_table.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | 3 | 4 | # Check if VERSION has been set 5 | if [ -z "${VERSION}" ]; then 6 | echo "Error: VERSION env var is not set" >&2 # Print to stderr 7 | exit 1 # Exit with a non-zero status to indicate an error 8 | fi 9 | 10 | 11 | FILES=("linux.amd64" "darwin.arm64" "darwin.amd64" "windows.amd64") 12 | output_string="" 13 | 14 | # Define the descriptions - ensure this array's order matches FILES 15 | DESCRIPTIONS=( 16 | "For **Linux** systems running on **Intel/AMD 64-bit processors**." 17 | "For **macOS** systems running on **Apple Silicon** (M1, M2, M3, etc.) processors." 18 | "For **macOS** systems running on **Intel processors**." 19 | "For **Windows** systems running on **Intel/AMD 64-bit processors**." 20 | ) 21 | 22 | # Write the table header 23 | ROW_FMT="| %-105s | %-120s | %-67s |\n" 24 | output_string+=$(printf "$ROW_FMT" "**OS/Architecture**" "**Description**" "**SHA256 Hash**")$'\n' 25 | output_string+=$(printf "$ROW_FMT" "$(printf -- '-%0.s' {1..105})" "$(printf -- '-%0.s' {1..120})" "$(printf -- '-%0.s' {1..67})")$'\n' 26 | 27 | 28 | # Loop through all files matching the pattern "toolbox.*.*" 29 | for i in "${!FILES[@]}" 30 | do 31 | file_key="${FILES[$i]}" # e.g., "linux.amd64" 32 | description_text="${DESCRIPTIONS[$i]}" 33 | 34 | # Extract OS and ARCH from the filename 35 | OS=$(echo "$file_key" | cut -d '.' -f 1) 36 | ARCH=$(echo "$file_key" | cut -d '.' -f 2) 37 | 38 | # Get release URL 39 | URL="https://storage.googleapis.com/genai-toolbox/$VERSION/$OS/$ARCH/toolbox" 40 | 41 | curl "$URL" --fail --output toolbox || exit 1 42 | 43 | # Calculate the SHA256 checksum of the file 44 | SHA256=$(shasum -a 256 toolbox | awk '{print $1}') 45 | 46 | # Write the table row 47 | output_string+=$(printf "$ROW_FMT" "[$OS/$ARCH]($URL)" "$description_text" "$SHA256")$'\n' 48 | 49 | rm toolbox 50 | done 51 | 52 | printf "$output_string\n" 53 | 54 | -------------------------------------------------------------------------------- /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | # This file controls who is tagged for review for any given pull request. 2 | # 3 | # For syntax help see: 4 | # https://help.github.com/en/github/creating-cloning-and-archiving-repositories/about-code-owners#codeowners-syntax 5 | 6 | * @googleapis/senseai-eco 7 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: false 2 | contact_links: 3 | - name: Google Cloud Support 4 | url: https://cloud.google.com/support/ 5 | about: If you have a support contract with Google, please both open an issue here and open Google Cloud Support portal with a link to the issue. 6 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: ✨ Feature Request 16 | description: Suggest an idea for new or improved behavior. 17 | title: "" 18 | labels: ["type: feature request"] 19 | type: feature 20 | body: 21 | - type: markdown 22 | attributes: 23 | value: | 24 | Thanks for helping us improve! 🙏 Please answer these questions and provide as much information as possible about your feature request. 25 | 26 | - id: preamble 27 | type: checkboxes 28 | attributes: 29 | label: Prerequisites 30 | description: | 31 | Please run through the following list and make sure you've tried the usual "quick fixes": 32 | options: 33 | - label: "Search the [current open issues](https://github.com/googleapis/genai-toolbox/issues)" 34 | required: true 35 | 36 | - type: textarea 37 | id: use-case 38 | attributes: 39 | label: What are you trying to do that currently feels hard or impossible? 40 | description: "A clear and concise description of what the end goal for the feature should be -- avoid generalizing and try to provide a specific use-case." 41 | validations: 42 | required: true 43 | 44 | - type: textarea 45 | id: suggested-solution 46 | attributes: 47 | label: Suggested Solution(s) 48 | description: "If you have a suggestion for how this use-case can be solved, please feel free to include it." 49 | 50 | - type: textarea 51 | id: alternatives-considered 52 | attributes: 53 | label: Alternatives Considered 54 | description: "Are there any workaround or third party tools to replicate this behavior? Why would adding this feature be preferred over them?" 55 | 56 | - type: textarea 57 | id: additional-details 58 | attributes: 59 | label: Additional Details 60 | description: "Any additional information we should know? Please reference it here (issues, PRs, descriptions, or screenshots)" 61 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: 💬 Question 16 | description: Questions on how something works or the best way to do something? 17 | title: "" 18 | labels: ["type: question"] 19 | 20 | body: 21 | - type: markdown 22 | attributes: 23 | value: | 24 | Thanks for helping us improve! 🙏 Please provide as much information as possible about your question. 25 | 26 | - id: preamble 27 | type: checkboxes 28 | attributes: 29 | label: Prerequisites 30 | description: | 31 | Please run through the following list and make sure you've tried the usual "quick fixes": 32 | options: 33 | - label: "Search the [current open issues](https://github.com/googleapis/genai-toolbox/issues)" 34 | required: true 35 | 36 | - type: textarea 37 | id: question 38 | attributes: 39 | label: Question 40 | description: "What's your question? Please provide as much relevant information as possible to reduce turnaround time. Include information like what environment, language, or framework you are using." 41 | placeholder: "Example: How do I connect using private IP with the AlloyDB source?" 42 | validations: 43 | required: true 44 | 45 | - type: textarea 46 | id: code 47 | attributes: 48 | label: Code 49 | description: "Please paste any useful application code that might be relevant to your question. (if your code is in a public repo, feel free to paste a link!)" 50 | 51 | - type: textarea 52 | id: additional-details 53 | attributes: 54 | label: Additional Details 55 | description: "Any other information you want us to know that might be helpful in answering your question? (link issues, PRs, descriptions, or screenshots)." 56 | -------------------------------------------------------------------------------- /.github/auto-label.yaml: -------------------------------------------------------------------------------- 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 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | enabled: false 16 | -------------------------------------------------------------------------------- /.github/blunderbuss.yml: -------------------------------------------------------------------------------- 1 | assign_issues: 2 | - kurtisvg 3 | - Yuan325 4 | - duwenxin99 5 | assign_prs: 6 | - kurtisvg 7 | - Yuan325 8 | - duwenxin99 -------------------------------------------------------------------------------- /.github/header-checker-lint.yaml: -------------------------------------------------------------------------------- 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 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | allowedCopyrightHolders: 16 | - 'Google LLC' 17 | allowedLicenses: 18 | - 'Apache-2.0' 19 | sourceFileExtensions: 20 | - 'go' 21 | - 'yaml' 22 | - 'yml' 23 | -------------------------------------------------------------------------------- /.github/label-sync.yml: -------------------------------------------------------------------------------- 1 | --- 2 | ignored: true 3 | -------------------------------------------------------------------------------- /.github/labels.yaml: -------------------------------------------------------------------------------- 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 | - name: duplicate 16 | color: ededed 17 | description: "" 18 | 19 | - name: 'type: bug' 20 | color: db4437 21 | description: Error or flaw in code with unintended results or allowing sub-optimal 22 | usage patterns. 23 | - name: 'type: cleanup' 24 | color: c5def5 25 | description: An internal cleanup or hygiene concern. 26 | - name: 'type: docs' 27 | color: 0000A0 28 | description: Improvement to the documentation for an API. 29 | - name: 'type: feature request' 30 | color: c5def5 31 | description: ‘Nice-to-have’ improvement, new feature or different behavior or design. 32 | - name: 'type: process' 33 | color: c5def5 34 | description: A process-related concern. May include testing, release, or the like. 35 | - name: 'type: question' 36 | color: c5def5 37 | description: Request for information or clarification. 38 | 39 | - name: 'priority: p0' 40 | color: b60205 41 | description: Highest priority. Critical issue. P0 implies highest priority. 42 | - name: 'priority: p1' 43 | color: ffa03e 44 | description: Important issue which blocks shipping the next release. Will be fixed 45 | prior to next release. 46 | - name: 'priority: p2' 47 | color: fef2c0 48 | description: Moderately-important priority. Fix may not be included in next release. 49 | - name: 'priority: p3' 50 | color: ffffc7 51 | description: Desirable enhancement or fix. May not be included in next release. 52 | 53 | - name: do not merge 54 | color: d93f0b 55 | description: Indicates a pull request not ready for merge, due to either quality 56 | or timing. 57 | 58 | - name: 'autorelease: pending' 59 | color: ededed 60 | description: Release please needs to do its work on this. 61 | - name: 'autorelease: triggered' 62 | color: ededed 63 | description: Release please has triggered a release for this. 64 | - name: 'autorelease: tagged' 65 | color: ededed 66 | description: Release please has completed a release for this. 67 | 68 | - name: 'tests: run' 69 | color: 3DED97 70 | description: Label to trigger Github Action tests. 71 | 72 | - name: 'status: contribution welcome' 73 | color: 8befd7 74 | description: Status - Contributions welcome. 75 | 76 | - name: 'status: awaiting response' 77 | color: 8befd7 78 | description: Status - Awaiting response from author. 79 | 80 | - name: 'status: awaiting codeowners' 81 | color: 8befd7 82 | description: Status - Awaiting response from code owners. 83 | -------------------------------------------------------------------------------- /.github/release-please.yml: -------------------------------------------------------------------------------- 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 | handleGHRelease: true 16 | packageName: genai-toolbox 17 | releaseType: simple 18 | versionFile: "cmd/version.txt" 19 | extraFiles: [ 20 | "README.md", 21 | "docs/en/getting-started/introduction/_index.md", 22 | "docs/en/getting-started/local_quickstart.md", 23 | "docs/en/getting-started/mcp_quickstart/_index.md", 24 | "docs/en/how-to/deploy_gke.md", 25 | "docs/en/samples/bigquery/local_quickstart.md", 26 | "docs/en/samples/bigquery/mcp_quickstart.md", 27 | "docs/en/getting-started/colab_quickstart.ipynb", 28 | "docs/en/samples/bigquery/colab_quickstart_bigquery.ipynb", 29 | ] 30 | -------------------------------------------------------------------------------- /.github/renovate.json5: -------------------------------------------------------------------------------- 1 | { 2 | extends: [ 3 | 'config:recommended', 4 | ':semanticCommitTypeAll(chore)', 5 | ':ignoreUnstable', 6 | ':separateMajorReleases', 7 | ':prConcurrentLimitNone', 8 | ':prHourlyLimitNone', 9 | ':preserveSemverRanges', 10 | ], 11 | minimumReleaseAge: '3', 12 | rebaseWhen: 'conflicted', 13 | dependencyDashboardLabels: [ 14 | 'type: process', 15 | ], 16 | "postUpdateOptions": [ 17 | "gomodTidy" 18 | ], 19 | packageRules: [ 20 | { 21 | groupName: 'GitHub Actions', 22 | matchManagers: [ 23 | 'github-actions', 24 | ], 25 | pinDigests: true, 26 | }, 27 | ], 28 | } 29 | -------------------------------------------------------------------------------- /.github/sync-repo-settings.yaml: -------------------------------------------------------------------------------- 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 | # Synchronize repository settings from a centralized config 16 | # https://github.com/googleapis/repo-automation-bots/tree/main/packages/sync-repo-settings 17 | # Install: https://github.com/apps/sync-repo-settings 18 | 19 | # Disable merge commits 20 | rebaseMergeAllowed: true 21 | squashMergeAllowed: true 22 | mergeCommitAllowed: false 23 | # Enable branch protection 24 | branchProtectionRules: 25 | - pattern: main 26 | isAdminEnforced: true 27 | requiredStatusCheckContexts: 28 | - "cla/google" 29 | - "lint" 30 | - "conventionalcommits.org" 31 | - "header-check" 32 | # - Add required status checks like presubmit tests 33 | - "unit tests (ubuntu-latest)" 34 | - "integration-test-pr (toolbox-testing-438616)" 35 | requiredApprovingReviewCount: 1 36 | requiresCodeOwnerReviews: true 37 | requiresStrictStatusChecks: true 38 | 39 | # Set team access 40 | permissionRules: 41 | - team: senseai-eco 42 | permission: admin 43 | -------------------------------------------------------------------------------- /.github/workflows/docs_deploy.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # 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 | name: "docs" 16 | 17 | permissions: 18 | contents: write 19 | 20 | on: 21 | push: 22 | branches: 23 | - main 24 | paths: 25 | - 'docs/**' 26 | - 'github/workflows/docs**' 27 | - '.hugo/**' 28 | 29 | # Allow triggering manually. 30 | workflow_dispatch: 31 | 32 | jobs: 33 | deploy: 34 | runs-on: ubuntu-24.04 35 | defaults: 36 | run: 37 | working-directory: .hugo 38 | concurrency: 39 | group: ${{ github.workflow }}-${{ github.ref }} 40 | cancel-in-progress: true 41 | steps: 42 | - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4 43 | with: 44 | fetch-depth: 0 # Fetch all history for .GitInfo and .Lastmod 45 | 46 | - name: Setup Hugo 47 | uses: peaceiris/actions-hugo@75d2e84710de30f6ff7268e08f310b60ef14033f # v3 48 | with: 49 | hugo-version: "0.145.0" 50 | extended: true 51 | 52 | - name: Setup Node 53 | uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4 54 | with: 55 | node-version: "22" 56 | 57 | - name: Cache dependencies 58 | uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4 59 | with: 60 | path: ~/.npm 61 | key: ${{ runner.os }}-node-${{ hashFiles('**/package-lock.json') }} 62 | restore-keys: | 63 | ${{ runner.os }}-node- 64 | 65 | - run: npm ci 66 | - run: hugo --minify 67 | env: 68 | HUGO_BASEURL: https://${{ github.repository_owner }}.github.io/${{ github.event.repository.name }}/ 69 | HUGO_RELATIVEURLS: false 70 | 71 | - name: Deploy 72 | uses: peaceiris/actions-gh-pages@4f9cc6602d3f66b9c108549d475ec49e8ef4d45e # v4 73 | with: 74 | github_token: ${{ secrets.GITHUB_TOKEN }} 75 | publish_dir: .hugo/public 76 | # Do not delete previews on each production deploy. 77 | # CSS or JS changes will require manual clean-up. 78 | keep_files: true 79 | commit_message: "deploy: ${{ github.event.head_commit.message }}" 80 | -------------------------------------------------------------------------------- /.github/workflows/docs_preview_clean.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # 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 | name: "docs" 16 | 17 | permissions: 18 | contents: write 19 | pull-requests: write 20 | 21 | # This Workflow depends on 'github.event.number', 22 | # not compatible with branch or manual triggers. 23 | on: 24 | pull_request: 25 | types: 26 | - closed 27 | 28 | jobs: 29 | clean: 30 | if: ${{ !github.event.pull_request.head.repo.fork }} 31 | runs-on: ubuntu-24.04 32 | concurrency: 33 | # Shared concurrency group wih preview staging. 34 | group: "preview-${{ github.event.number }}" 35 | cancel-in-progress: true 36 | steps: 37 | - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4 38 | with: 39 | ref: gh-pages 40 | 41 | - name: Remove Preview 42 | run: | 43 | rm -Rf ./previews/PR-${{ github.event.number }} 44 | git config user.name 'github-actions[bot]' 45 | git config user.email 'github-actions[bot]@users.noreply.github.com' 46 | git add -u previews/PR-${{ github.event.number }} 47 | git commit --message "cleanup: previews/PR-${{ github.event.number }}" 48 | git push 49 | 50 | - name: Comment 51 | uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7 52 | with: 53 | script: | 54 | github.rest.issues.createComment({ 55 | issue_number: context.payload.number, 56 | owner: context.repo.owner, 57 | repo: context.repo.repo, 58 | body: "🧨 Preview deployments removed." 59 | }) -------------------------------------------------------------------------------- /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- 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 | name: lint 16 | on: 17 | pull_request: 18 | pull_request_target: 19 | types: [labeled] 20 | 21 | # Declare default permissions as read only. 22 | permissions: read-all 23 | 24 | jobs: 25 | lint: 26 | if: "${{ github.event.action != 'labeled' || github.event.label.name == 'tests: run' }}" 27 | name: lint 28 | runs-on: ubuntu-latest 29 | concurrency: 30 | group: ${{ github.workflow }}-${{ github.ref }} 31 | cancel-in-progress: true 32 | permissions: 33 | contents: 'read' 34 | issues: 'write' 35 | pull-requests: 'write' 36 | steps: 37 | - name: Remove PR Label 38 | if: "${{ github.event.action == 'labeled' && github.event.label.name == 'tests: run' }}" 39 | uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 40 | with: 41 | github-token: ${{ secrets.GITHUB_TOKEN }} 42 | script: | 43 | try { 44 | await github.rest.issues.removeLabel({ 45 | name: 'tests: run', 46 | owner: context.repo.owner, 47 | repo: context.repo.repo, 48 | issue_number: context.payload.pull_request.number 49 | }); 50 | } catch (e) { 51 | console.log('Failed to remove label. Another job may have already removed it!'); 52 | } 53 | - name: Setup Go 54 | uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 55 | with: 56 | go-version: "1.22" 57 | - name: Checkout code 58 | uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 59 | with: 60 | ref: ${{ github.event.pull_request.head.sha }} 61 | repository: ${{ github.event.pull_request.head.repo.full_name }} 62 | token: ${{ secrets.GITHUB_TOKEN }} 63 | - name: > 64 | Verify go mod tidy. If you're reading this and the check has 65 | failed, run `goimports -w . && go mod tidy && golangci-lint run` 66 | run: | 67 | go mod tidy && git diff --exit-code 68 | - name: golangci-lint 69 | uses: golangci/golangci-lint-action@4afd733a84b1f43292c63897423277bb7f4313a9 # v8.0.0 70 | with: 71 | version: latest 72 | args: --timeout 3m 73 | -------------------------------------------------------------------------------- /.github/workflows/schedule_reporter.yml: -------------------------------------------------------------------------------- 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 | name: Schedule Reporter 16 | 17 | on: 18 | schedule: 19 | - cron: '0 6 * * *' # Runs at 6 AM every morning 20 | 21 | jobs: 22 | run_reporter: 23 | permissions: 24 | issues: 'write' 25 | checks: 'read' 26 | contents: 'read' 27 | uses: ./.github/workflows/cloud_build_failure_reporter.yml 28 | with: 29 | trigger_names: "toolbox-test-nightly,toolbox-test-on-merge" 30 | -------------------------------------------------------------------------------- /.github/workflows/sync-labels.yaml: -------------------------------------------------------------------------------- 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 | name: Sync Labels 16 | on: 17 | push: 18 | branches: 19 | - main 20 | 21 | # Declare default permissions as read only. 22 | permissions: read-all 23 | 24 | jobs: 25 | build: 26 | runs-on: ubuntu-latest 27 | permissions: 28 | contents: 'read' 29 | issues: 'write' 30 | pull-requests: 'write' 31 | steps: 32 | - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 33 | - uses: micnncim/action-label-syncer@3abd5ab72fda571e69fffd97bd4e0033dd5f495c # v1.3.0 34 | env: 35 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 36 | with: 37 | manifest: .github/labels.yaml 38 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # direnv 2 | .envrc 3 | 4 | # vscode 5 | .vscode/ 6 | 7 | # npm 8 | node_modules 9 | 10 | # hugo 11 | .hugo/public/ 12 | .hugo/resources/_gen 13 | .hugo_build.lock 14 | 15 | # coverage 16 | .coverage 17 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "docs2/themes/godocs"] 2 | path = docs2/themes/godocs 3 | url = https://github.com/themefisher/godocs.git 4 | -------------------------------------------------------------------------------- /.golangci.yaml: -------------------------------------------------------------------------------- 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 | version: "2" 16 | linters: 17 | enable: 18 | - errcheck 19 | - govet 20 | - ineffassign 21 | - staticcheck 22 | - unused 23 | exclusions: 24 | presets: 25 | - std-error-handling 26 | issues: 27 | fix: true 28 | formatters: 29 | enable: 30 | - goimports 31 | settings: 32 | gofmt: 33 | rewrite-rules: 34 | - pattern: interface{} 35 | replacement: any 36 | - pattern: a[b:len(a)] 37 | replacement: a[b:] 38 | -------------------------------------------------------------------------------- /.hugo/archetypes/default.md: -------------------------------------------------------------------------------- 1 | +++ 2 | date = '{{ .Date }}' 3 | draft = true 4 | title = '{{ replace .File.ContentBaseName "-" " " | title }}' 5 | +++ 6 | -------------------------------------------------------------------------------- /.hugo/assets/scss/_variables_project.scss: -------------------------------------------------------------------------------- 1 | $primary: #80a7e9; 2 | $secondary: #4484f4; 3 | -------------------------------------------------------------------------------- /.hugo/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/googleapis/genai-toolbox 2 | 3 | go 1.23.2 4 | 5 | require ( 6 | github.com/google/docsy v0.11.0 // indirect 7 | github.com/martignoni/hugo-notice v0.0.0-20240707105359-40327ac00cc4 // indirect 8 | ) 9 | -------------------------------------------------------------------------------- /.hugo/go.sum: -------------------------------------------------------------------------------- 1 | github.com/FortAwesome/Font-Awesome v0.0.0-20240716171331-37eff7fa00de/go.mod h1:IUgezN/MFpCDIlFezw3L8j83oeiIuYoj28Miwr/KUYo= 2 | github.com/google/docsy v0.11.0 h1:QnV40cc28QwS++kP9qINtrIv4hlASruhC/K3FqkHAmM= 3 | github.com/google/docsy v0.11.0/go.mod h1:hGGW0OjNuG5ZbH5JRtALY3yvN8ybbEP/v2iaK4bwOUI= 4 | github.com/martignoni/hugo-notice v0.0.0-20240707105359-40327ac00cc4 h1:lxS0B1ta9/uW+orrnvsGHMCC0TgN4DymEgdlb0PL/uU= 5 | github.com/martignoni/hugo-notice v0.0.0-20240707105359-40327ac00cc4/go.mod h1:MIQPOMgEcbyRC0gNLzQFSgrS+wIy3RuQ/HbaZYtTOKU= 6 | github.com/twbs/bootstrap v5.3.3+incompatible/go.mod h1:fZTSrkpSf0/HkL0IIJzvVspTt1r9zuf7XlZau8kpcY0= 7 | -------------------------------------------------------------------------------- /.hugo/hugo.toml: -------------------------------------------------------------------------------- 1 | title = 'MCP Toolbox for Databases' 2 | relativeURLs = true 3 | 4 | languageCode = 'en-us' 5 | defaultContentLanguage = "en" 6 | defaultContentLanguageInSubdir = false 7 | 8 | enableGitInfo = true 9 | enableRobotsTXT = true 10 | 11 | [languages] 12 | [languages.en] 13 | languageName ="English" 14 | weight = 1 15 | 16 | [module] 17 | proxy = "direct" 18 | [module.hugoVersion] 19 | extended = true 20 | min = "0.73.0" 21 | [[module.mounts]] 22 | source = "../docs/en" 23 | target = 'content' 24 | [[module.imports]] 25 | path = "github.com/google/docsy" 26 | disable = false 27 | [[module.imports]] 28 | path = "github.com/martignoni/hugo-notice" 29 | 30 | [params] 31 | copyright = "Google LLC" 32 | github_repo = "https://github.com/googleapis/genai-toolbox" 33 | github_project_repo = "https://github.com/googleapis/genai-toolbox" 34 | github_subdir = "docs" 35 | offlineSearch = true 36 | [params.ui] 37 | ul_show = 100 38 | showLightDarkModeMenu = true 39 | breadcrumb_disable = true 40 | sidebar_menu_foldable = true 41 | sidebar_menu_compact = false 42 | 43 | [[menu.main]] 44 | name = "GitHub" 45 | weight = 50 46 | url = "https://github.com/googleapis/genai-toolbox" 47 | pre = "" 48 | 49 | [markup.goldmark.renderer] 50 | unsafe= true -------------------------------------------------------------------------------- /.hugo/layouts/_default/_markup/render-heading.html: -------------------------------------------------------------------------------- 1 | {{ template "_default/_markup/td-render-heading.html" . }} 2 | -------------------------------------------------------------------------------- /.hugo/layouts/partials/page-meta-links.html: -------------------------------------------------------------------------------- 1 | {{/* cSpell:ignore querify subdir */ -}} 2 | {{/* Class names ending with `--KIND` are deprecated in favor of `__KIND`, but we're keeping them for a few releases after 0.9.0 */ -}} 3 | 4 | {{ if .File -}} 5 | {{ $path := urls.JoinPath $.Language.Lang $.File.Path -}} 6 | {{ $gh_repo := $.Param "github_repo" -}} 7 | {{ $gh_url := $.Param "github_url" -}} 8 | {{ $gh_subdir := $.Param "github_subdir" | default "" -}} 9 | {{ $gh_project_repo := $.Param "github_project_repo" -}} 10 | {{ $gh_branch := $.Param "github_branch" | default "main" -}} 11 |
12 | {{ if $gh_url -}} 13 | {{ warnf "Warning: use of `github_url` is deprecated. For details, see https://www.docsy.dev/docs/adding-content/repository-links/#github_url-optional" -}} 14 | {{ T "post_edit_this" }} 15 | {{ else if $gh_repo -}} 16 | 17 | {{/* Adjust $path based on path_base_for_github_subdir */ -}} 18 | {{ $ghs_base := $.Param "path_base_for_github_subdir" -}} 19 | {{ $ghs_rename := "" -}} 20 | {{ if reflect.IsMap $ghs_base -}} 21 | {{ $ghs_rename = $ghs_base.to -}} 22 | {{ $ghs_base = $ghs_base.from -}} 23 | {{ end -}} 24 | {{ with $ghs_base -}} 25 | {{ $path = replaceRE . $ghs_rename $path -}} 26 | {{ end -}} 27 | 28 | {{ $gh_repo_path := printf "%s/%s/%s" $gh_branch $gh_subdir $path -}} 29 | {{ $gh_repo_path = replaceRE "//+" "/" $gh_repo_path -}} 30 | 31 | {{ $viewURL := printf "%s/tree/%s" $gh_repo $gh_repo_path -}} 32 | {{ $editURL := printf "%s/edit/%s" $gh_repo $gh_repo_path -}} 33 | {{ $issuesURL := printf "%s/issues/new?title=%s" $gh_repo (safeURL $.Title ) -}} 34 | {{ $newPageStub := resources.Get "stubs/new-page-template.md" -}} 35 | {{ $newPageQS := querify "value" $newPageStub.Content "filename" "change-me.md" | safeURL -}} 36 | {{ $newPageURL := printf "%s/new/%s?%s" $gh_repo (path.Dir $gh_repo_path) $newPageQS -}} 37 | 38 | {{ T "post_view_this" }} 39 | {{ T "post_edit_this" }} 40 | {{ T "post_create_child_page" }} 41 | {{ T "post_create_issue" }} 42 | {{ with $gh_project_repo -}} 43 | {{ $project_issueURL := printf "%s/issues/new" . -}} 44 | {{ T "post_create_project_issue" }} 45 | {{ end -}} 46 | 47 | {{ end -}} 48 | {{ with .CurrentSection.AlternativeOutputFormats.Get "print" -}} 49 | {{ T "print_entire_section" }} 50 | {{ end }} 51 |
52 | {{ end -}} 53 | -------------------------------------------------------------------------------- /.hugo/layouts/robot.txt: -------------------------------------------------------------------------------- 1 | User-agent: * 2 | {{ if eq hugo.Environment "preview" }} 3 | Disallow: /* 4 | {{ end }} -------------------------------------------------------------------------------- /.hugo/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "devDependencies": { 3 | "autoprefixer": "^10.4.20", 4 | "postcss": "^8.4.49", 5 | "postcss-cli": "^11.0.0" 6 | } 7 | } 8 | -------------------------------------------------------------------------------- /.hugo/static/favicons/android-chrome-192x192.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/android-chrome-192x192.png -------------------------------------------------------------------------------- /.hugo/static/favicons/android-chrome-512x512.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/android-chrome-512x512.png -------------------------------------------------------------------------------- /.hugo/static/favicons/apple-touch-icon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/apple-touch-icon.png -------------------------------------------------------------------------------- /.hugo/static/favicons/favicon-16x16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/favicon-16x16.png -------------------------------------------------------------------------------- /.hugo/static/favicons/favicon-32x32.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/favicon-32x32.png -------------------------------------------------------------------------------- /.hugo/static/favicons/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/.hugo/static/favicons/favicon.ico -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to contribute 2 | 3 | We'd love to accept your patches and contributions to this project. 4 | 5 | ## Before you begin 6 | 7 | ### Sign our Contributor License Agreement 8 | 9 | Contributions to this project must be accompanied by a 10 | [Contributor License Agreement](https://cla.developers.google.com/about) (CLA). 11 | You (or your employer) retain the copyright to your contribution; this simply 12 | gives us permission to use and redistribute your contributions as part of the 13 | project. 14 | 15 | If you or your current employer have already signed the Google CLA (even if it 16 | was for a different project), you probably don't need to do it again. 17 | 18 | Visit to see your current agreements or to 19 | sign a new one. 20 | 21 | ### Review our community guidelines 22 | 23 | This project follows 24 | [Google's Open Source Community Guidelines](https://opensource.google/conduct/). 25 | 26 | ## Contribution process 27 | 28 | ### Code reviews 29 | 30 | All submissions, including submissions by project members, require review. We 31 | use GitHub pull requests for this purpose. Consult 32 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 33 | information on using pull requests. -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 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 | # Use the latest stable golang 1.x to compile to a binary 16 | FROM --platform=$BUILDPLATFORM golang:1 as build 17 | 18 | WORKDIR /go/src/genai-toolbox 19 | COPY . . 20 | 21 | ARG TARGETOS 22 | ARG TARGETARCH 23 | ARG METADATA_TAGS=dev 24 | 25 | RUN go get ./... 26 | RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \ 27 | go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.metadataString=container.${METADATA_TAGS}" 28 | 29 | # Final Stage 30 | FROM gcr.io/distroless/static:nonroot 31 | 32 | WORKDIR /app 33 | COPY --from=build --chown=nonroot /go/src/genai-toolbox/genai-toolbox /toolbox 34 | USER nonroot 35 | 36 | ENTRYPOINT ["/toolbox"] 37 | -------------------------------------------------------------------------------- /cmd/options.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 cmd 16 | 17 | import ( 18 | "io" 19 | ) 20 | 21 | // Option is a function that configures a Command. 22 | type Option func(*Command) 23 | 24 | // WithStreams overrides the default writer. 25 | func WithStreams(out, err io.Writer) Option { 26 | return func(c *Command) { 27 | c.outStream = out 28 | c.errStream = err 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /cmd/options_test.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 cmd 16 | 17 | import ( 18 | "errors" 19 | "io" 20 | "testing" 21 | 22 | "github.com/spf13/cobra" 23 | ) 24 | 25 | func TestCommandOptions(t *testing.T) { 26 | w := io.Discard 27 | tcs := []struct { 28 | desc string 29 | isValid func(*Command) error 30 | option Option 31 | }{ 32 | { 33 | desc: "with logger", 34 | isValid: func(c *Command) error { 35 | if c.outStream != w || c.errStream != w { 36 | return errors.New("loggers do not match") 37 | } 38 | return nil 39 | }, 40 | option: WithStreams(w, w), 41 | }, 42 | } 43 | for _, tc := range tcs { 44 | t.Run(tc.desc, func(t *testing.T) { 45 | got, err := invokeProxyWithOption(tc.option) 46 | if err != nil { 47 | t.Fatal(err) 48 | } 49 | if err := tc.isValid(got); err != nil { 50 | t.Errorf("option did not initialize command correctly: %v", err) 51 | } 52 | }) 53 | } 54 | } 55 | 56 | func invokeProxyWithOption(o Option) (*Command, error) { 57 | c := NewCommand(o) 58 | // Keep the test output quiet 59 | c.SilenceUsage = true 60 | c.SilenceErrors = true 61 | // Disable execute behavior 62 | c.RunE = func(*cobra.Command, []string) error { 63 | return nil 64 | } 65 | 66 | err := c.Execute() 67 | return c, err 68 | } 69 | -------------------------------------------------------------------------------- /cmd/version.txt: -------------------------------------------------------------------------------- 1 | 0.6.0 2 | -------------------------------------------------------------------------------- /docs/en/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Documentation" 3 | type: docs 4 | notoc: false 5 | weight: 1 6 | description: > 7 | All of Toolbox's documentation. 8 | --- 9 | 10 | 11 | 12 | 13 | 14 | 15 | -------------------------------------------------------------------------------- /docs/en/about/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "About" 3 | type: docs 4 | weight: 6 5 | description: A list of other information related to Toolbox. 6 | --- 7 | -------------------------------------------------------------------------------- /docs/en/concepts/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Concepts" 3 | type: docs 4 | weight: 2 5 | description: Some core concepts in Toolbox 6 | --- 7 | -------------------------------------------------------------------------------- /docs/en/concepts/telemetry/telemetry_flow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/concepts/telemetry/telemetry_flow.png -------------------------------------------------------------------------------- /docs/en/concepts/telemetry/telemetry_traces.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/concepts/telemetry/telemetry_traces.png -------------------------------------------------------------------------------- /docs/en/getting-started/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Getting Started" 3 | type: docs 4 | weight: 1 5 | description: > 6 | How to get started with Toolbox 7 | --- 8 | -------------------------------------------------------------------------------- /docs/en/getting-started/configure.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Configuration" 3 | type: docs 4 | weight: 4 5 | description: How to configure Toolbox's tools.yaml file. 6 | --- 7 | 8 | The primary way to configure Toolbox is through the `tools.yaml` file. If you 9 | have multiple files, you can tell toolbox which to load with the `--tools-file 10 | tools.yaml` flag. 11 | 12 | You can find more detailed reference documentation to all resource types in the 13 | [Resources](../resources/). 14 | 15 | ### Using Environment Variables 16 | 17 | To avoid hardcoding certain secret fields like passwords, usernames, API keys 18 | etc., you could use environment variables instead with the format `${ENV_NAME}`. 19 | 20 | ```yaml 21 | user: ${USER_NAME} 22 | password: ${PASSWORD} 23 | ``` 24 | 25 | ### Sources 26 | 27 | The `sources` section of your `tools.yaml` defines what data sources your 28 | Toolbox should have access to. Most tools will have at least one source to 29 | execute against. 30 | 31 | ```yaml 32 | sources: 33 | my-pg-source: 34 | kind: postgres 35 | host: 127.0.0.1 36 | port: 5432 37 | database: toolbox_db 38 | user: ${USER_NAME} 39 | password: ${PASSWORD} 40 | ``` 41 | 42 | For more details on configuring different types of sources, see the 43 | [Sources](../resources/sources/). 44 | 45 | ### Tools 46 | 47 | The `tools` section of your `tools.yaml` define your the actions your agent can 48 | take: what kind of tool it is, which source(s) it affects, what parameters it 49 | uses, etc. 50 | 51 | ```yaml 52 | tools: 53 | search-hotels-by-name: 54 | kind: postgres-sql 55 | source: my-pg-source 56 | description: Search for hotels based on name. 57 | parameters: 58 | - name: name 59 | type: string 60 | description: The name of the hotel. 61 | statement: SELECT * FROM hotels WHERE name ILIKE '%' || $1 || '%'; 62 | ``` 63 | 64 | For more details on configuring different types of tools, see the 65 | [Tools](../resources/tools/). 66 | 67 | ### Toolsets 68 | 69 | The `toolsets` section of your `tools.yaml` allows you to define groups of tools 70 | that you want to be able to load together. This can be useful for defining 71 | different sets for different agents or different applications. 72 | 73 | ```yaml 74 | toolsets: 75 | my_first_toolset: 76 | - my_first_tool 77 | - my_second_tool 78 | my_second_toolset: 79 | - my_second_tool 80 | - my_third_tool 81 | ``` 82 | 83 | You can load toolsets by name: 84 | 85 | ```python 86 | # This will load all tools 87 | all_tools = client.load_toolset() 88 | 89 | # This will only load the tools listed in 'my_second_toolset' 90 | my_second_toolset = client.load_toolset("my_second_toolset") 91 | ``` 92 | -------------------------------------------------------------------------------- /docs/en/getting-started/introduction/architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/getting-started/introduction/architecture.png -------------------------------------------------------------------------------- /docs/en/getting-started/mcp_quickstart/inspector.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/getting-started/mcp_quickstart/inspector.png -------------------------------------------------------------------------------- /docs/en/getting-started/mcp_quickstart/inspector_tools.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/getting-started/mcp_quickstart/inspector_tools.png -------------------------------------------------------------------------------- /docs/en/how-to/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "How-to" 3 | type: docs 4 | weight: 3 5 | description: > 6 | List of guides detailing how to do different things with Toolbox. 7 | --- 8 | -------------------------------------------------------------------------------- /docs/en/how-to/connect-ide/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Connect from your IDE" 3 | type: docs 4 | weight: 1 5 | description: > 6 | List of guides detailing how to connect your AI tools (IDEs) to Toolbox using MCP. 7 | aliases: 8 | - /how-to/connect_tools_using_mcp 9 | --- 10 | -------------------------------------------------------------------------------- /docs/en/how-to/deploy_docker.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Deploy using Docker Compose" 3 | type: docs 4 | weight: 4 5 | description: > 6 | How to deploy Toolbox using Docker Compose. 7 | --- 8 | 9 | 10 | 11 | 12 | ## Before you begin 13 | 14 | 1. [Install Docker Compose.](https://docs.docker.com/compose/install/) 15 | 16 | ## Configure `tools.yaml` file 17 | 18 | Create a `tools.yaml` file that contains your configuration for Toolbox. For 19 | details, see the 20 | [configuration](https://github.com/googleapis/genai-toolbox/blob/main/README.md#configuration) 21 | section. 22 | 23 | ## Deploy using Docker Compose 24 | 25 | 1. Create a `docker-compose.yml` file, customizing as needed: 26 | 27 | ```yaml 28 | services: 29 | toolbox: 30 | # TODO: It is recommended to pin to a specific image version instead of latest. 31 | image: us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:latest 32 | hostname: toolbox 33 | platform: linux/amd64 34 | ports: 35 | - "5000:5000" 36 | volumes: 37 | - ./config:/config 38 | command: [ "toolbox", "--tools-file", "/config/tools.yaml", "--address", "0.0.0.0"] 39 | depends_on: 40 | db: 41 | condition: service_healthy 42 | networks: 43 | - tool-network 44 | db: 45 | # TODO: It is recommended to pin to a specific image version instead of latest. 46 | image: postgres 47 | hostname: db 48 | environment: 49 | POSTGRES_USER: toolbox_user 50 | POSTGRES_PASSWORD: my-password 51 | POSTGRES_DB: toolbox_db 52 | ports: 53 | - "5432:5432" 54 | volumes: 55 | - ./db:/var/lib/postgresql/data 56 | # This file can be used to bootstrap your schema if needed. 57 | # See "initialization scripts" on https://hub.docker.com/_/postgres/ for more info 58 | - ./config/init.sql:/docker-entrypoint-initdb.d/init.sql 59 | healthcheck: 60 | test: ["CMD-SHELL", "pg_isready -U toolbox_user -d toolbox_db"] 61 | interval: 10s 62 | timeout: 5s 63 | retries: 5 64 | networks: 65 | - tool-network 66 | networks: 67 | tool-network: 68 | 69 | ``` 70 | 71 | 1. Run the following command to bring up the Toolbox and Postgres instance 72 | 73 | ```bash 74 | docker-compose up -d 75 | ``` 76 | 77 | 78 | {{< notice tip >}} 79 | 80 | You can use this setup quickly set up Toolbox + Postgres to follow along in our 81 | [Quickstart](../getting-started/local_quickstart.md) 82 | 83 | {{< /notice >}} 84 | 85 | 86 | 87 | ## Connecting with Toolbox Client SDK 88 | 89 | Next, we will use Toolbox with the Client SDKs: 90 | 91 | 1. The url for the Toolbox server running using docker-compose will be: 92 | 93 | ``` 94 | http://localhost:5000 95 | ``` 96 | 97 | 1. Import and initialize the client with the URL: 98 | 99 | {{< tabpane persist=header >}} 100 | {{< tab header="LangChain" lang="Python" >}} 101 | from toolbox_langchain import ToolboxClient 102 | 103 | # Replace with the cloud run service URL generated above 104 | async with ToolboxClient("http://$YOUR_URL") as toolbox: 105 | {{< /tab >}} 106 | {{< tab header="Llamaindex" lang="Python" >}} 107 | from toolbox_llamaindex import ToolboxClient 108 | 109 | # Replace with the cloud run service URL generated above 110 | async with ToolboxClient("http://$YOUR_URL") as toolbox: 111 | {{< /tab >}} 112 | {{< /tabpane >}} 113 | 114 | 115 | -------------------------------------------------------------------------------- /docs/en/how-to/export_telemetry.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Export Telemetry" 3 | type: docs 4 | weight: 5 5 | description: > 6 | How to set up and configure Toolbox to use the Otel Collector. 7 | --- 8 | 9 | 10 | ## About 11 | 12 | The [OpenTelemetry Collector][about-collector] offers a vendor-agnostic 13 | implementation of how to receive, process and export telemetry data. It removes 14 | the need to run, operate, and maintain multiple agents/collectors. 15 | 16 | [about-collector]: https://opentelemetry.io/docs/collector/ 17 | 18 | ## Configure the Collector 19 | 20 | To configure the collector, you will have to provide a configuration file. The 21 | configuration file consists of four classes of pipeline component that access 22 | telemetry data. 23 | - `Receivers` 24 | - `Processors` 25 | - `Exporters` 26 | - `Connectors` 27 | 28 | Example of setting up the classes of pipeline components (in this example, we 29 | don't use connectors): 30 | 31 | ```yaml 32 | receivers: 33 | otlp: 34 | protocols: 35 | http: 36 | endpoint: "127.0.0.1:4553" 37 | 38 | exporters: 39 | googlecloud: 40 | project: 41 | 42 | processors: 43 | batch: 44 | send_batch_size: 200 45 | ``` 46 | 47 | After each pipeline component is configured, you will enable it within the 48 | `service` section of the configuration file. 49 | 50 | ```yaml 51 | service: 52 | pipelines: 53 | traces: 54 | receivers: ["otlp"] 55 | processors: ["batch"] 56 | exporters: ["googlecloud"] 57 | ``` 58 | 59 | ## Running the Collector 60 | 61 | There are a couple of steps to run and use a Collector. 62 | 63 | 1. [Install the 64 | Collector](https://opentelemetry.io/docs/collector/installation/) binary. 65 | Pull a binary or Docker image for the OpenTelemetry contrib collector. 66 | 67 | 1. Set up credentials for telemetry backend. 68 | 69 | 1. Set up the Collector config. Below are some examples for setting up the 70 | Collector config: 71 | - [Google Cloud Exporter][google-cloud-exporter] 72 | - [Google Managed Service for Prometheus Exporter][google-prometheus-exporter] 73 | 74 | 1. Run the Collector with the configuration file. 75 | 76 | ```bash 77 | ./otelcol-contrib --config=collector-config.yaml 78 | ``` 79 | 80 | 1. Run toolbox with the `--telemetry-otlp` flag. Configure it to send them to 81 | `http://127.0.0.1:4553` (for HTTP) or the Collector's URL. 82 | 83 | ```bash 84 | ./toolbox --telemetry-otlp=http://127.0.0.1:4553 85 | ``` 86 | 87 | 1. Once telemetry datas are collected, you can view them in your telemetry 88 | backend. If you are using GCP exporters, telemetry will be visible in GCP 89 | dashboard at [Metrics Explorer][metrics-explorer] and [Trace 90 | Explorer][trace-explorer]. 91 | 92 | {{< notice note >}} 93 | If you are exporting to Google Cloud monitoring, we recommend that you use 94 | the Google Cloud Exporter for traces and the Google Managed Service for 95 | Prometheus Exporter for metrics. 96 | {{< /notice >}} 97 | 98 | [google-cloud-exporter]: 99 | https://github.com/open-telemetry/opentelemetry-collector-contrib/tree/main/exporter/googlecloudexporter 100 | [google-prometheus-exporter]: 101 | https://github.com/open-telemetry/opentelemetry-collector-contrib/tree/main/exporter/googlemanagedprometheusexporter#example-configuration 102 | [metrics-explorer]: https://console.cloud.google.com/monitoring/metrics-explorer 103 | [trace-explorer]: https://console.cloud.google.com/traces 104 | -------------------------------------------------------------------------------- /docs/en/resources/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Resources" 3 | type: docs 4 | weight: 4 5 | description: List of reference documentation for resources in Toolbox. 6 | --- 7 | -------------------------------------------------------------------------------- /docs/en/resources/authServices/google.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Google Sign-In" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Use Google Sign-In for Oauth 2.0 flow and token lifecycle. 7 | --- 8 | 9 | ## Getting Started 10 | 11 | Google Sign-In manages the OAuth 2.0 flow and token lifecycle. To integrate the 12 | Google Sign-In workflow to your web app [follow this guide][gsi-setup]. 13 | 14 | After setting up the Google Sign-In workflow, you should have registered your 15 | application and retrieved a [Client ID][client-id]. Configure your auth service 16 | in with the `Client ID`. 17 | 18 | [gsi-setup]: https://developers.google.com/identity/sign-in/web/sign-in 19 | [client-id]: https://developers.google.com/identity/sign-in/web/sign-in#create_authorization_credentials 20 | 21 | ## Behavior 22 | 23 | ### Authorized Invocations 24 | 25 | When using [Authorized Invocations][auth-invoke], a tool will be 26 | considered authorized if it has a valid Oauth 2.0 token that matches the Client 27 | ID. 28 | 29 | [auth-invoke]: ../tools/#authorized-invocations 30 | 31 | ### Authenticated Parameters 32 | 33 | When using [Authenticated Parameters][auth-params], any [claim provided by the 34 | id-token][provided-claims] can be used for the parameter. 35 | 36 | [auth-params]: ../tools/#authenticated-phugarameters 37 | [provided-claims]: 38 | https://developers.google.com/identity/openid-connect/openid-connect#obtaininguserprofileinformation 39 | 40 | ## Example 41 | 42 | ```yaml 43 | authServices: 44 | my-google-auth: 45 | kind: google 46 | clientId: ${YOUR_GOOGLE_CLIENT_ID} 47 | ``` 48 | 49 | {{< notice tip >}} 50 | Use environment variable replacement with the format ${ENV_NAME} 51 | instead of hardcoding your secrets into the configuration file. 52 | {{< /notice >}} 53 | 54 | ## Reference 55 | 56 | | **field** | **type** | **required** | **description** | 57 | |-----------|:--------:|:------------:|------------------------------------------------------------------| 58 | | kind | string | true | Must be "google". | 59 | | clientId | string | true | Client ID of your application from registering your application. | 60 | -------------------------------------------------------------------------------- /docs/en/resources/sources/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Sources" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Sources represent your different data sources that a tool can interact with. 7 | --- 8 | 9 | A Source represents a data sources that a tool can interact with. You can define 10 | Sources as a map in the `sources` section of your `tools.yaml` file. Typically, 11 | a source configuration will contain any information needed to connect with and 12 | interact with the database. 13 | 14 | {{< notice tip >}} 15 | Use environment variable replacement with the format ${ENV_NAME} 16 | instead of hardcoding your secrets into the configuration file. 17 | {{< /notice >}} 18 | 19 | ```yaml 20 | sources: 21 | my-cloud-sql-source: 22 | kind: cloud-sql-postgres 23 | project: my-project-id 24 | region: us-central1 25 | instance: my-instance-name 26 | database: my_db 27 | user: ${USER_NAME} 28 | password: ${PASSWORD} 29 | ``` 30 | 31 | In implementation, each source is a different connection pool or client that used 32 | to connect to the database and execute the tool. 33 | 34 | ## Available Sources 35 | -------------------------------------------------------------------------------- /docs/en/resources/sources/bigtable.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Bigtable" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Bigtable is a low-latency NoSQL database service for machine learning, operational analytics, and user-facing operations. It's a wide-column, key-value store that can scale to billions of rows and thousands of columns. With Bigtable, you can replicate your data to regions across the world for high availability and data resiliency. 7 | 8 | --- 9 | 10 | # Bigtable Source 11 | 12 | [Bigtable][bigtable-docs] is a low-latency NoSQL database service for machine 13 | learning, operational analytics, and user-facing operations. It's a wide-column, 14 | key-value store that can scale to billions of rows and thousands of columns. 15 | With Bigtable, you can replicate your data to regions across the world for high 16 | availability and data resiliency. 17 | 18 | If you are new to Bigtable, you can try to [create an instance and write data 19 | with the cbt CLI][bigtable-quickstart-with-cli]. 20 | 21 | You can use [GoogleSQL statements][bigtable-googlesql] to query your Bigtable 22 | data. GoogleSQL is an ANSI-compliant structured query language (SQL) that is 23 | also implemented for other Google Cloud services. SQL queries are handled by 24 | cluster nodes in the same way as NoSQL data requests. Therefore, the same best 25 | practices apply when creating SQL queries to run against your Bigtable data, 26 | such as avoiding full table scans or complex filters. 27 | 28 | [bigtable-docs]: https://cloud.google.com/bigtable/docs 29 | [bigtable-quickstart-with-cli]: 30 | https://cloud.google.com/bigtable/docs/create-instance-write-data-cbt-cli 31 | 32 | [bigtable-googlesql]: 33 | https://cloud.google.com/bigtable/docs/googlesql-overview 34 | 35 | ## Requirements 36 | 37 | ### IAM Permissions 38 | 39 | Bigtable uses [Identity and Access Management (IAM)][iam-overview] to control 40 | user and group access to Bigtable resources at the project, instance, table, and 41 | backup level. Toolbox will use your [Application Default Credentials (ADC)][adc] 42 | to authorize and authenticate when interacting with [Bigtable][bigtable-docs]. 43 | 44 | In addition to [setting the ADC for your server][set-adc], you need to ensure 45 | the IAM identity has been given the correct IAM permissions for the query 46 | provided. See [Apply IAM roles][grant-permissions] for more information on 47 | applying IAM permissions and roles to an identity. 48 | 49 | [iam-overview]: https://cloud.google.com/bigtable/docs/access-control 50 | [adc]: https://cloud.google.com/docs/authentication#adc 51 | [set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc 52 | [grant-permissions]: https://cloud.google.com/bigtable/docs/access-control#iam-management-instance 53 | 54 | ## Example 55 | 56 | ```yaml 57 | sources: 58 | my-bigtable-source: 59 | kind: "bigtable" 60 | project: "my-project-id" 61 | instance: "test-instance" 62 | ``` 63 | 64 | ## Reference 65 | 66 | | **field** | **type** | **required** | **description** | 67 | |-----------|:--------:|:------------:|-------------------------------------------------------------------------------| 68 | | kind | string | true | Must be "bigtable". | 69 | | project | string | true | Id of the GCP project that the cluster was created in (e.g. "my-project-id"). | 70 | | instance | string | true | Name of the Bigtable instance. | 71 | -------------------------------------------------------------------------------- /docs/en/resources/sources/dgraph.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Dgraph" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Dgraph is fully open-source, built-for-scale graph database for Gen AI workloads 7 | 8 | --- 9 | 10 | ## About 11 | 12 | [Dgraph][dgraph-docs] is an open-source graph database. It is designed for real-time workloads, horizontal scalability, and data flexibility. Implemented as a distributed system, Dgraph processes queries in parallel to deliver the fastest result. 13 | 14 | This source can connect to either a self-managed Dgraph cluster or one hosted on 15 | Dgraph Cloud. If you're new to Dgraph, the fastest way to get started is to 16 | [sign up for Dgraph Cloud][dgraph-login]. 17 | 18 | [dgraph-docs]: https://dgraph.io/docs 19 | [dgraph-login]: https://cloud.dgraph.io/login 20 | 21 | ## Requirements 22 | 23 | ### Database User 24 | 25 | When **connecting to a hosted Dgraph database**, this source uses the API key 26 | for access. If you are using a dedicated environment, you will additionally need 27 | the namespace and user credentials for that namespace. 28 | 29 | For **connecting to a local or self-hosted Dgraph database**, use the namespace 30 | and user credentials for that namespace. 31 | 32 | ## Example 33 | 34 | ```yaml 35 | sources: 36 | my-dgraph-source: 37 | kind: dgraph 38 | dgraphUrl: https://xxxx.cloud.dgraph.io 39 | user: ${USER_NAME} 40 | password: ${PASSWORD} 41 | apiKey: ${API_KEY} 42 | namespace : 0 43 | ``` 44 | 45 | {{< notice tip >}} 46 | Use environment variable replacement with the format ${ENV_NAME} 47 | instead of hardcoding your secrets into the configuration file. 48 | {{< /notice >}} 49 | 50 | ## Reference 51 | 52 | | **Field** | **Type** | **Required** | **Description** | 53 | |-------------|:--------:|:------------:|--------------------------------------------------------------------------------------------------| 54 | | kind | string | true | Must be "dgraph". | 55 | | dgraphUrl | string | true | Connection URI (e.g. "", ""). | 56 | | user | string | false | Name of the Dgraph user to connect as (e.g., "groot"). | 57 | | password | string | false | Password of the Dgraph user (e.g., "password"). | 58 | | apiKey | string | false | API key to connect to a Dgraph Cloud instance. | 59 | | namespace | uint64 | false | Dgraph namespace (not required for Dgraph Cloud Shared Clusters). | 60 | -------------------------------------------------------------------------------- /docs/en/resources/sources/http.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "HTTP" 3 | linkTitle: "HTTP" 4 | type: docs 5 | weight: 1 6 | description: > 7 | The HTTP source enables the Toolbox to retrieve data from a remote server using HTTP requests. 8 | --- 9 | 10 | ## About 11 | 12 | The HTTP Source allows Toolbox to retrieve data from arbitrary HTTP 13 | endpoints. This enables Generative AI applications to access data from web APIs 14 | and other HTTP-accessible resources. 15 | 16 | ## Example 17 | 18 | ```yaml 19 | sources: 20 | my-http-source: 21 | kind: http 22 | baseUrl: https://api.example.com/data 23 | timeout: 10s # default to 30s 24 | headers: 25 | Authorization: Bearer ${API_KEY} 26 | Content-Type: application/json 27 | queryParams: 28 | param1: value1 29 | param2: value2 30 | ``` 31 | 32 | {{< notice tip >}} 33 | Use environment variable replacement with the format ${ENV_NAME} 34 | instead of hardcoding your secrets into the configuration file. 35 | {{< /notice >}} 36 | 37 | ## Reference 38 | 39 | | **field** | **type** | **required** | **description** | 40 | |-------------|:-----------------:|:------------:|-----------------------------------------------------------------------------------------------------------------------------------| 41 | | kind | string | true | Must be "http". | 42 | | baseUrl | string | true | The base URL for the HTTP requests (e.g., `https://api.example.com`). | 43 | | timeout | string | false | The timeout for HTTP requests (e.g., "5s", "1m", refer to [ParseDuration][parse-duration-doc] for more examples). Defaults to 30s. | 44 | | headers | map[string]string | false | Default headers to include in the HTTP requests. | 45 | | queryParams | map[string]string | false | Default query parameters to include in the HTTP requests. | 46 | 47 | [parse-duration-doc]: https://pkg.go.dev/time#ParseDuration 48 | -------------------------------------------------------------------------------- /docs/en/resources/sources/mssql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "SQL Server" 3 | type: docs 4 | weight: 1 5 | description: > 6 | SQL Server is a relational database management system (RDBMS). 7 | 8 | --- 9 | 10 | ## About 11 | 12 | [SQL Server][mssql-docs] is a relational database management system (RDBMS) 13 | developed by Microsoft that allows users to store, retrieve, and manage large 14 | amount of data through a structured format. 15 | 16 | [mssql-docs]: https://www.microsoft.com/en-us/sql-server 17 | 18 | ## Requirements 19 | 20 | ### Database User 21 | 22 | This source only uses standard authentication. You will need to [create a 23 | SQL Server user][mssql-users] to login to the database with. 24 | 25 | [mssql-users]: https://learn.microsoft.com/en-us/sql/relational-databases/security/authentication-access/create-a-database-user?view=sql-server-ver16 26 | 27 | ## Example 28 | 29 | ```yaml 30 | sources: 31 | my-mssql-source: 32 | kind: mssql 33 | host: 127.0.0.1 34 | port: 1433 35 | database: my_db 36 | user: ${USER_NAME} 37 | password: ${PASSWORD} 38 | ``` 39 | 40 | {{< notice tip >}} 41 | Use environment variable replacement with the format ${ENV_NAME} 42 | instead of hardcoding your secrets into the configuration file. 43 | {{< /notice >}} 44 | 45 | ## Reference 46 | 47 | | **field** | **type** | **required** | **description** | 48 | |-----------|:--------:|:------------:|------------------------------------------------------------------------| 49 | | kind | string | true | Must be "mssql". | 50 | | host | string | true | IP address to connect to (e.g. "127.0.0.1"). | 51 | | port | string | true | Port to connect to (e.g. "1433"). | 52 | | database | string | true | Name of the SQL Server database to connect to (e.g. "my_db"). | 53 | | user | string | true | Name of the SQL Server user to connect as (e.g. "my-user"). | 54 | | password | string | true | Password of the SQL Server user (e.g. "my-password"). | 55 | -------------------------------------------------------------------------------- /docs/en/resources/sources/mysql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "MySQL" 3 | type: docs 4 | weight: 1 5 | description: > 6 | MySQL is a relational database management system that stores and manages data. 7 | 8 | --- 9 | 10 | ## About 11 | 12 | [MySQL][mysql-docs] is a relational database management system (RDBMS) that 13 | stores and manages data. It's a popular choice for developers because of its 14 | reliability, performance, and ease of use. 15 | 16 | [mysql-docs]: https://www.mysql.com/ 17 | 18 | ## Requirements 19 | 20 | ### Database User 21 | 22 | This source only uses standard authentication. You will need to [create a 23 | MySQL user][mysql-users] to login to the database with. 24 | 25 | [mysql-users]: https://dev.mysql.com/doc/refman/8.4/en/user-names.html 26 | 27 | ## Example 28 | 29 | ```yaml 30 | sources: 31 | my-mysql-source: 32 | kind: mysql 33 | host: 127.0.0.1 34 | port: 3306 35 | database: my_db 36 | user: ${USER_NAME} 37 | password: ${PASSWORD} 38 | ``` 39 | 40 | {{< notice tip >}} 41 | Use environment variable replacement with the format ${ENV_NAME} 42 | instead of hardcoding your secrets into the configuration file. 43 | {{< /notice >}} 44 | 45 | ## Reference 46 | 47 | | **field** | **type** | **required** | **description** | 48 | |-----------|:--------:|:------------:|---------------------------------------------------------------------------------------------| 49 | | kind | string | true | Must be "mysql". | 50 | | host | string | true | IP address to connect to (e.g. "127.0.0.1"). | 51 | | port | string | true | Port to connect to (e.g. "3306"). | 52 | | database | string | true | Name of the MySQL database to connect to (e.g. "my_db"). | 53 | | user | string | true | Name of the MySQL user to connect as (e.g. "my-mysql-user"). | 54 | | password | string | true | Password of the MySQL user (e.g. "my-password"). | 55 | -------------------------------------------------------------------------------- /docs/en/resources/sources/neo4j.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Neo4j" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Neo4j is a powerful, open source graph database system 7 | 8 | --- 9 | 10 | [Neo4j][neo4j-docs] is a powerful, open source graph database system with over 11 | 15 years of active development that has earned it a strong reputation for 12 | reliability, feature robustness, and performance. 13 | 14 | [neo4j-docs]: https://neo4j.com/docs 15 | 16 | ## Requirements 17 | 18 | ### Database User 19 | 20 | This source only uses standard authentication. You will need to [create a Neo4j 21 | user][neo4j-users] to log in to the database with, or use the default `neo4j` 22 | user if available. 23 | 24 | [neo4j-users]: https://neo4j.com/docs/operations-manual/current/authentication-authorization/manage-users/ 25 | 26 | ## Example 27 | 28 | ```yaml 29 | sources: 30 | my-neo4j-source: 31 | kind: neo4j 32 | uri: neo4j+s://xxxx.databases.neo4j.io:7687 33 | user: ${USER_NAME} 34 | password: ${PASSWORD} 35 | database: "neo4j" 36 | ``` 37 | 38 | {{< notice tip >}} 39 | Use environment variable replacement with the format ${ENV_NAME} 40 | instead of hardcoding your secrets into the configuration file. 41 | {{< /notice >}} 42 | 43 | ## Reference 44 | 45 | | **field** | **type** | **required** | **description** | 46 | |-----------|:--------:|:------------:|----------------------------------------------------------------------| 47 | | kind | string | true | Must be "neo4j". | 48 | | uri | string | true | Connect URI ("bolt://localhost", "neo4j+s://xxx.databases.neo4j.io") | 49 | | user | string | true | Name of the Neo4j user to connect as (e.g. "neo4j"). | 50 | | password | string | true | Password of the Neo4j user (e.g. "my-password"). | 51 | | database | string | true | Name of the Neo4j database to connect to (e.g. "neo4j"). | 52 | -------------------------------------------------------------------------------- /docs/en/resources/sources/postgres.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "PostgreSQL" 3 | type: docs 4 | weight: 1 5 | description: > 6 | PostgreSQL is a powerful, open source object-relational database. 7 | 8 | --- 9 | 10 | ## About 11 | 12 | [PostgreSQL][pg-docs] is a powerful, open source object-relational database 13 | system with over 35 years of active development that has earned it a strong 14 | reputation for reliability, feature robustness, and performance. 15 | 16 | [pg-docs]: https://www.postgresql.org/ 17 | 18 | ## Requirements 19 | 20 | ### Database User 21 | 22 | This source only uses standard authentication. You will need to [create a 23 | PostgreSQL user][pg-users] to login to the database with. 24 | 25 | [pg-users]: https://www.postgresql.org/docs/current/sql-createuser.html 26 | 27 | ## Example 28 | 29 | ```yaml 30 | sources: 31 | my-pg-source: 32 | kind: postgres 33 | host: 127.0.0.1 34 | port: 5432 35 | database: my_db 36 | user: ${USER_NAME} 37 | password: ${PASSWORD} 38 | ``` 39 | 40 | {{< notice tip >}} 41 | Use environment variable replacement with the format ${ENV_NAME} 42 | instead of hardcoding your secrets into the configuration file. 43 | {{< /notice >}} 44 | 45 | ## Reference 46 | 47 | | **field** | **type** | **required** | **description** | 48 | |-----------|:--------:|:------------:|------------------------------------------------------------------------| 49 | | kind | string | true | Must be "postgres". | 50 | | host | string | true | IP address to connect to (e.g. "127.0.0.1") | 51 | | port | string | true | Port to connect to (e.g. "5432") | 52 | | database | string | true | Name of the Postgres database to connect to (e.g. "my_db"). | 53 | | user | string | true | Name of the Postgres user to connect as (e.g. "my-pg-user"). | 54 | | password | string | true | Password of the Postgres user (e.g. "my-password"). | 55 | -------------------------------------------------------------------------------- /docs/en/resources/sources/spanner.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Spanner" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Spanner is a fully managed database service from Google Cloud that combines 7 | relational, key-value, graph, and search capabilities. 8 | 9 | --- 10 | 11 | # Spanner Source 12 | 13 | [Spanner][spanner-docs] is a fully managed, mission-critical database service 14 | that brings together relational, graph, key-value, and search. It offers 15 | transactional consistency at global scale, automatic, synchronous replication 16 | for high availability, and support for two SQL dialects: GoogleSQL (ANSI 2011 17 | with extensions) and PostgreSQL. 18 | 19 | If you are new to Spanner, you can try to [create and query a database using 20 | the Google Cloud console][spanner-quickstart]. 21 | 22 | [spanner-docs]: https://cloud.google.com/spanner/docs 23 | [spanner-quickstart]: 24 | https://cloud.google.com/spanner/docs/create-query-database-console 25 | 26 | ## Requirements 27 | 28 | ### IAM Permissions 29 | 30 | Spanner uses [Identity and Access Management (IAM)][iam-overview] to control 31 | user and group access to Spanner resources at the project, Spanner instance, and 32 | Spanner database levels. Toolbox will use your [Application Default Credentials 33 | (ADC)][adc] to authorize and authenticate when interacting with Spanner. 34 | 35 | In addition to [setting the ADC for your server][set-adc], you need to ensure 36 | the IAM identity has been given the correct IAM permissions for the query 37 | provided. See [Apply IAM roles][grant-permissions] for more information on 38 | applying IAM permissions and roles to an identity. 39 | 40 | [iam-overview]: https://cloud.google.com/spanner/docs/iam 41 | [adc]: https://cloud.google.com/docs/authentication#adc 42 | [set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc 43 | [grant-permissions]: https://cloud.google.com/spanner/docs/grant-permissions 44 | 45 | ## Example 46 | 47 | ```yaml 48 | sources: 49 | my-spanner-source: 50 | kind: "spanner" 51 | project: "my-project-id" 52 | instance: "my-instance" 53 | database: "my_db" 54 | # dialect: "googlesql" 55 | ``` 56 | 57 | ## Reference 58 | 59 | | **field** | **type** | **required** | **description** | 60 | |-----------|:--------:|:------------:|---------------------------------------------------------------------------------------------------------------------| 61 | | kind | string | true | Must be "spanner". | 62 | | project | string | true | Id of the GCP project that the cluster was created in (e.g. "my-project-id"). | 63 | | instance | string | true | Name of the Spanner instance. | 64 | | database | string | true | Name of the database on the Spanner instance | 65 | | dialect | string | false | Name of the dialect type of the Spanner database, must be either `googlesql` or `postgresql`. Default: `googlesql`. | 66 | -------------------------------------------------------------------------------- /docs/en/resources/sources/sqlite.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "SQLite" 3 | linkTitle: "SQLite" 4 | type: docs 5 | weight: 1 6 | description: > 7 | SQLite is a C-language library that implements a small, fast, self-contained, 8 | high-reliability, full-featured, SQL database engine. 9 | --- 10 | 11 | ## About 12 | 13 | [SQLite](https://sqlite.org/) is a software library that provides a relational 14 | database management system. The lite in SQLite means lightweight in terms of 15 | setup, database administration, and required resources. 16 | 17 | SQLite has the following notable characteristics: 18 | - Self-contained with no external dependencies 19 | - Serverless - the SQLite library accesses its storage files directly 20 | - Single database file that can be easily copied or moved 21 | - Zero-configuration - no setup or administration needed 22 | - Transactional with ACID properties 23 | 24 | ## Requirements 25 | 26 | ### Database File 27 | 28 | You need a SQLite database file. This can be: 29 | - An existing database file 30 | - A path where a new database file should be created 31 | - `:memory:` for an in-memory database 32 | 33 | ## Example 34 | 35 | ```yaml 36 | sources: 37 | my-sqlite-db: 38 | kind: "sqlite" 39 | database: "/path/to/database.db" 40 | ``` 41 | 42 | For an in-memory database: 43 | ```yaml 44 | sources: 45 | my-sqlite-memory-db: 46 | kind: "sqlite" 47 | database: ":memory:" 48 | ``` 49 | 50 | ## Reference 51 | 52 | ### Configuration Fields 53 | 54 | | Field | Type | Required | Description | 55 | |-------|------|----------|-------------| 56 | | kind | string | Yes | Must be "sqlite" | 57 | | database | string | Yes | Path to SQLite database file, or ":memory:" for an in-memory database | 58 | 59 | ### Connection Properties 60 | 61 | SQLite connections are configured with these defaults for optimal performance: 62 | - `MaxOpenConns`: 1 (SQLite only supports one writer at a time) 63 | - `MaxIdleConns`: 1 -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-execute-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-execute-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-execute-sql" tool executes a SQL statement against BigQuery. 7 | --- 8 | 9 | ## About 10 | 11 | A `bigquery-execute-sql` tool executes a SQL statement against BigQuery. 12 | It's compatible with the following sources: 13 | 14 | - [bigquery](../sources/bigquery.md) 15 | 16 | `bigquery-execute-sql` takes one input parameter `sql` and runs the sql 17 | statement against the `source`. 18 | 19 | ## Example 20 | 21 | ```yaml 22 | tools: 23 | execute_sql_tool: 24 | kind: bigquery-execute-sql 25 | source: my-bigquery-source 26 | description: Use this tool to execute sql statement. 27 | ``` 28 | 29 | ## Reference 30 | 31 | | **field** | **type** | **required** | **description** | 32 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 33 | | kind | string | true | Must be "bigquery-execute-sql". | 34 | | source | string | true | Name of the source the SQL should execute on. | 35 | | description | string | true | Description of the tool that is passed to the LLM. | 36 | -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-get-dataset-info.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-get-dataset-info" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-get-dataset-info" tool retrieves metadata for a BigQuery dataset. 7 | --- 8 | 9 | ## About 10 | 11 | A `bigquery-get-dataset-info` tool retrieves metadata for a BigQuery dataset. 12 | It's compatible with the following sources: 13 | 14 | - [bigquery](../sources/bigquery.md) 15 | 16 | bigquery-get-dataset-info takes a dataset parameter to specify the dataset 17 | on the given source. 18 | 19 | ## Example 20 | 21 | ```yaml 22 | tools: 23 | bigquery_get_dataset_info: 24 | kind: bigquery-get-dataset-info 25 | source: my-bigquery-source 26 | description: Use this tool to get dataset metadata. 27 | ``` 28 | 29 | ## Reference 30 | 31 | | **field** | **type** | **required** | **description** | 32 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 33 | | kind | string | true | Must be "bigquery-get-dataset-info". | 34 | | source | string | true | Name of the source the SQL should execute on. | 35 | | description | string | true | Description of the tool that is passed to the LLM. | 36 | -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-get-table-info.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-get-table-info" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-get-table-info" tool retrieves metadata for a BigQuery table. 7 | --- 8 | 9 | ## About 10 | 11 | A `bigquery-get-table-info` tool retrieves metadata for a BigQuery table. 12 | It's compatible with the following sources: 13 | 14 | - [bigquery](../sources/bigquery.md) 15 | 16 | bigquery-get-table-info takes dataset and table parameters to specify 17 | the target table. 18 | 19 | ## Example 20 | 21 | ```yaml 22 | tools: 23 | bigquery_get_table_info: 24 | kind: bigquery-get-table-info 25 | source: my-bigquery-source 26 | description: Use this tool to get table metadata. 27 | ``` 28 | 29 | ## Reference 30 | 31 | | **field** | **type** | **required** | **description** | 32 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 33 | | kind | string | true | Must be "bigquery-get-table-info". | 34 | | source | string | true | Name of the source the SQL should execute on. | 35 | | description | string | true | Description of the tool that is passed to the LLM. | 36 | -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-list-dataset-ids.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-list-dataset-ids" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-list-dataset-ids" tool returns all dataset IDs from the source. 7 | --- 8 | 9 | ## About 10 | 11 | A `bigquery-list-dataset-ids` tool returns all dataset IDs from the source. 12 | It's compatible with the following sources: 13 | 14 | - [bigquery](../sources/bigquery.md) 15 | 16 | bigquery-list-dataset-ids requires no input parameters beyond the configured 17 | source. 18 | 19 | ## Example 20 | 21 | ```yaml 22 | tools: 23 | bigquery_list_dataset_ids: 24 | kind: bigquery-list-dataset-ids 25 | source: my-bigquery-source 26 | description: Use this tool to get dataset metadata. 27 | ``` 28 | 29 | ## Reference 30 | 31 | | **field** | **type** | **required** | **description** | 32 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 33 | | kind | string | true | Must be "bigquery-list-dataset-ids". | 34 | | source | string | true | Name of the source the SQL should execute on. | 35 | | description | string | true | Description of the tool that is passed to the LLM. | 36 | -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-list-table-ids.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-list-table-ids" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-list-table-ids" tool returns table IDs in a given BigQuery dataset. 7 | --- 8 | 9 | ## About 10 | 11 | A `bigquery-list-table-ids` tool returns table IDs in a given BigQuery dataset. 12 | It's compatible with the following sources: 13 | 14 | - [bigquery](../sources/bigquery.md) 15 | 16 | bigquery-get-dataset-info takes a dataset parameter to specify the dataset 17 | from which to list table IDs. 18 | 19 | ## Example 20 | 21 | ```yaml 22 | tools: 23 | bigquery_list_table_ids: 24 | kind: bigquery-list-table-ids 25 | source: my-bigquery-source 26 | description: Use this tool to get table metadata. 27 | ``` 28 | 29 | ## Reference 30 | 31 | | **field** | **type** | **required** | **description** | 32 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 33 | | kind | string | true | Must be "bigquery-list-table-ids". | 34 | | source | string | true | Name of the source the SQL should execute on. | 35 | | description | string | true | Description of the tool that is passed to the LLM. | 36 | -------------------------------------------------------------------------------- /docs/en/resources/tools/bigquery-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "bigquery-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "bigquery-sql" tool executes a pre-defined SQL statement. 7 | --- 8 | 9 | ## About 10 | A `bigquery-sql` tool executes a pre-defined SQL statement. It's compatible with 11 | the following sources: 12 | 13 | - [bigquery](../sources/bigquery.md) 14 | 15 | ### GoogleSQL 16 | 17 | BigQuery uses [GoogleSQL][bigquery-googlesql] for querying data. The integration 18 | with Toolbox supports this dialect. The specified SQL statement is executed, and 19 | parameters can be inserted into the query. BigQuery supports both named parameters 20 | (e.g., `@name`) and positional parameters (`?`), but they cannot be mixed in the 21 | same query. 22 | 23 | > **Note:** This tool uses [parameterized queries](https://cloud.google.com/bigquery/docs/parameterized-queries) to prevent SQL injections. Query parameters can be used as substitutes for arbitrary expressions. Parameters cannot be used as substitutes for identifiers, column names, table names, or other parts of the query. 24 | 25 | [bigquery-googlesql]: https://cloud.google.com/bigquery/docs/reference/standard-sql/ 26 | 27 | ## Example 28 | 29 | ```yaml 30 | tools: 31 | # Example: Querying a user table in BigQuery 32 | search_users_bq: 33 | kind: bigquery-sql 34 | source: my-bigquery-source 35 | statement: | 36 | SELECT 37 | id, 38 | name, 39 | email 40 | FROM 41 | `my-project.my-dataset.users` 42 | WHERE 43 | id = @id OR email = @email; 44 | description: | 45 | Use this tool to get information for a specific user. 46 | Takes an id number or a name and returns info on the user. 47 | 48 | Example: 49 | {{ 50 | "id": 123, 51 | "name": "Alice", 52 | }} 53 | parameters: 54 | - name: id 55 | type: integer 56 | description: User ID 57 | - name: email 58 | type: string 59 | description: Email address of the user 60 | ``` 61 | 62 | ## Reference 63 | 64 | | **field** | **type** | **required** | **description** | 65 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 66 | | kind | string | true | Must be "bigquery-sql". | 67 | | source | string | true | Name of the source the GoogleSQL should execute on. | 68 | | description | string | true | Description of the tool that is passed to the LLM. | 69 | | statement | string | true | The GoogleSQL statement to execute. | 70 | | parameters | [parameters](_index#specifying-parameters) | false | List of [parameters](_index#specifying-parameters) that will be inserted into the SQL statement. | 71 | -------------------------------------------------------------------------------- /docs/en/resources/tools/mssql-execute-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "mssql-execute-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "mssql-execute-sql" tool executes a SQL statement against a SQL Server 7 | database. 8 | --- 9 | 10 | ## About 11 | 12 | A `mssql-execute-sql` tool executes a SQL statement against a SQL Server 13 | database. It's compatible with any of the following sources: 14 | 15 | - [cloud-sql-mssql](../sources/cloud-sql-mssql.md) 16 | - [mssql](../sources/mssql.md) 17 | 18 | `mssql-execute-sql` takes one input parameter `sql` and run the sql 19 | statement against the `source`. 20 | 21 | > **Note:** This tool is intended for developer assistant workflows with 22 | > human-in-the-loop and shouldn't be used for production agents. 23 | 24 | ## Example 25 | 26 | ```yaml 27 | tools: 28 | execute_sql_tool: 29 | kind: mssql-execute-sql 30 | source: my-mssql-instance 31 | description: Use this tool to execute sql statement. 32 | ``` 33 | ## Reference 34 | 35 | | **field** | **type** | **required** | **description** | 36 | |-------------|:------------------------------------------:|:------------:|----------------------------------------------------| 37 | | kind | string | true | Must be "mssql-execute-sql". | 38 | | source | string | true | Name of the source the SQL should execute on. | 39 | | description | string | true | Description of the tool that is passed to the LLM. | 40 | -------------------------------------------------------------------------------- /docs/en/resources/tools/mysql-execute-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "mysql-execute-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "mysql-execute-sql" tool executes a SQL statement against a MySQL 7 | database. 8 | --- 9 | 10 | ## About 11 | 12 | A `mysql-execute-sql` tool executes a SQL statement against a MySQL 13 | database. It's compatible with any of the following sources: 14 | 15 | - [cloud-sql-mysql](../sources/cloud-sql-mysql.md) 16 | - [mysql](../sources/mysql.md) 17 | 18 | `mysql-execute-sql` takes one input parameter `sql` and run the sql 19 | statement against the `source`. 20 | 21 | > **Note:** This tool is intended for developer assistant workflows with 22 | > human-in-the-loop and shouldn't be used for production agents. 23 | 24 | ## Example 25 | 26 | ```yaml 27 | tools: 28 | execute_sql_tool: 29 | kind: mysql-execute-sql 30 | source: my-mysql-instance 31 | description: Use this tool to execute sql statement. 32 | ``` 33 | 34 | ## Reference 35 | 36 | | **field** | **type** | **required** | **description** | 37 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 38 | | kind | string | true | Must be "mysql-execute-sql". | 39 | | source | string | true | Name of the source the SQL should execute on. | 40 | | description | string | true | Description of the tool that is passed to the LLM. | 41 | -------------------------------------------------------------------------------- /docs/en/resources/tools/postgres-execute-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "postgres-execute-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "postgres-execute-sql" tool executes a SQL statement against a Postgres 7 | database. 8 | --- 9 | 10 | ## About 11 | 12 | A `postgres-execute-sql` tool executes a SQL statement against a Postgres 13 | database. It's compatible with any of the following sources: 14 | 15 | - [alloydb-postgres](../sources/alloydb-pg.md) 16 | - [cloud-sql-postgres](../sources/cloud-sql-pg.md) 17 | - [postgres](../sources/postgres.md) 18 | 19 | `postgres-execute-sql` takes one input parameter `sql` and run the sql 20 | statement against the `source`. 21 | 22 | > **Note:** This tool is intended for developer assistant workflows with 23 | > human-in-the-loop and shouldn't be used for production agents. 24 | 25 | ## Example 26 | 27 | ```yaml 28 | tools: 29 | execute_sql_tool: 30 | kind: postgres-execute-sql 31 | source: my-pg-instance 32 | description: Use this tool to execute sql statement. 33 | ``` 34 | 35 | ## Reference 36 | 37 | | **field** | **type** | **required** | **description** | 38 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 39 | | kind | string | true | Must be "postgres-execute-sql". | 40 | | source | string | true | Name of the source the SQL should execute on. | 41 | | description | string | true | Description of the tool that is passed to the LLM. | 42 | -------------------------------------------------------------------------------- /docs/en/resources/tools/spanner-execute-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "spanner-execute-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | A "spanner-execute-sql" tool executes a SQL statement against a Spanner 7 | database. 8 | --- 9 | 10 | ## About 11 | 12 | A `spanner-execute-sql` tool executes a SQL statement against a Spanner 13 | database. It's compatible with any of the following sources: 14 | 15 | - [spanner](../sources/spanner.md) 16 | 17 | `spanner-execute-sql` takes one input parameter `sql` and run the sql 18 | statement against the `source`. 19 | 20 | > **Note:** This tool is intended for developer assistant workflows with 21 | > human-in-the-loop and shouldn't be used for production agents. 22 | 23 | ## Example 24 | 25 | ```yaml 26 | tools: 27 | execute_sql_tool: 28 | kind: spanner-execute-sql 29 | source: my-spanner-instance 30 | description: Use this tool to execute sql statement. 31 | ``` 32 | 33 | ## Reference 34 | 35 | | **field** | **type** | **required** | **description** | 36 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 37 | | kind | string | true | Must be "spanner-execute-sql". | 38 | | source | string | true | Name of the source the SQL should execute on. | 39 | | description | string | true | Description of the tool that is passed to the LLM. | 40 | | readOnly | bool | false | When set to `true`, the `statement` is run as a read-only transaction. Default: `false`. | 41 | -------------------------------------------------------------------------------- /docs/en/resources/tools/sqlite-sql.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "sqlite-sql" 3 | type: docs 4 | weight: 1 5 | description: > 6 | Execute SQL statements against a SQLite database. 7 | --- 8 | 9 | ## About 10 | 11 | A `sqlite-sql` tool executes SQL statements against a SQLite database. 12 | It's compatible with any of the following sources: 13 | 14 | - [sqlite](../sources/sqlite.md) 15 | 16 | SQLite uses the `?` placeholder for parameters in SQL statements. Parameters are 17 | bound in the order they are provided. 18 | 19 | The statement field supports any valid SQLite SQL statement, including `SELECT`, 20 | `INSERT`, `UPDATE`, `DELETE`, `CREATE/ALTER/DROP` table statements, and other 21 | DDL statements. 22 | 23 | > **Note:** This tool uses parameterized queries to prevent SQL injections. 24 | > Query parameters can be used as substitutes for arbitrary expressions. 25 | > Parameters cannot be used as substitutes for identifiers, column names, table 26 | > names, or other parts of the query. 27 | 28 | ### Example 29 | 30 | ```yaml 31 | tools: 32 | search-users: 33 | kind: sqlite-sql 34 | source: my-sqlite-db 35 | description: Search users by name and age 36 | parameters: 37 | - name: name 38 | type: string 39 | description: The name to search for 40 | - name: min_age 41 | type: integer 42 | description: Minimum age 43 | statement: SELECT * FROM users WHERE name LIKE ? AND age >= ? 44 | ``` 45 | 46 | ## Reference 47 | 48 | | **field** | **type** | **required** | **description** | 49 | |-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| 50 | | kind | string | Yes | Must be "sqlite-sql" | 51 | | source | string | Yes | Name of a SQLite source configuration | 52 | | description | string | Yes | Description of what the tool does | 53 | | parameters | array | No | List of parameters for the SQL statement | 54 | | statement | string | Yes | The SQL statement to execute | 55 | -------------------------------------------------------------------------------- /docs/en/samples/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "Samples" 3 | type: docs 4 | weight: 5 5 | description: Samples and guides for using Toolbox. 6 | --- 7 | -------------------------------------------------------------------------------- /docs/en/samples/bigquery/_index.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: "BigQuery" 3 | type: docs 4 | weight: 1 5 | description: > 6 | How to get started with Toolbox using BigQuery. 7 | --- 8 | -------------------------------------------------------------------------------- /docs/en/samples/bigquery/mcp_quickstart/inspector.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/samples/bigquery/mcp_quickstart/inspector.png -------------------------------------------------------------------------------- /docs/en/samples/bigquery/mcp_quickstart/inspector_tools.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/docs/en/samples/bigquery/mcp_quickstart/inspector_tools.png -------------------------------------------------------------------------------- /internal/auth/auth.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 auth 16 | 17 | import ( 18 | "context" 19 | "net/http" 20 | ) 21 | 22 | // AuthServiceConfig is the interface for configuring authentication services. 23 | type AuthServiceConfig interface { 24 | AuthServiceConfigKind() string 25 | Initialize() (AuthService, error) 26 | } 27 | 28 | // AuthService is the interface for authentication services. 29 | type AuthService interface { 30 | AuthServiceKind() string 31 | GetName() string 32 | GetClaimsFromHeader(context.Context, http.Header) (map[string]any, error) 33 | } 34 | -------------------------------------------------------------------------------- /internal/auth/google/google.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 google 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | "net/http" 21 | 22 | "github.com/googleapis/genai-toolbox/internal/auth" 23 | "google.golang.org/api/idtoken" 24 | ) 25 | 26 | const AuthServiceKind string = "google" 27 | 28 | // validate interface 29 | var _ auth.AuthServiceConfig = Config{} 30 | 31 | // Auth service configuration 32 | type Config struct { 33 | Name string `yaml:"name" validate:"required"` 34 | Kind string `yaml:"kind" validate:"required"` 35 | ClientID string `yaml:"clientId" validate:"required"` 36 | } 37 | 38 | // Returns the auth service kind 39 | func (cfg Config) AuthServiceConfigKind() string { 40 | return AuthServiceKind 41 | } 42 | 43 | // Initialize a Google auth service 44 | func (cfg Config) Initialize() (auth.AuthService, error) { 45 | a := &AuthService{ 46 | Name: cfg.Name, 47 | Kind: AuthServiceKind, 48 | ClientID: cfg.ClientID, 49 | } 50 | return a, nil 51 | } 52 | 53 | var _ auth.AuthService = AuthService{} 54 | 55 | // struct used to store auth service info 56 | type AuthService struct { 57 | Name string `yaml:"name"` 58 | Kind string `yaml:"kind"` 59 | ClientID string `yaml:"clientId"` 60 | } 61 | 62 | // Returns the auth service kind 63 | func (a AuthService) AuthServiceKind() string { 64 | return AuthServiceKind 65 | } 66 | 67 | // Returns the name of the auth service 68 | func (a AuthService) GetName() string { 69 | return a.Name 70 | } 71 | 72 | // Verifies Google ID token and return claims 73 | func (a AuthService) GetClaimsFromHeader(ctx context.Context, h http.Header) (map[string]any, error) { 74 | if token := h.Get(a.Name + "_token"); token != "" { 75 | payload, err := idtoken.Validate(ctx, token, a.ClientID) 76 | if err != nil { 77 | return nil, fmt.Errorf("Google ID token verification failure: %w", err) //nolint:staticcheck 78 | } 79 | return payload.Claims, nil 80 | } 81 | return nil, nil 82 | } 83 | -------------------------------------------------------------------------------- /internal/log/logger.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 log 16 | 17 | import ( 18 | "context" 19 | ) 20 | 21 | // Logger is the interface used throughout the project for logging. 22 | type Logger interface { 23 | // DebugContext is for reporting additional information about internal operations. 24 | DebugContext(ctx context.Context, format string, args ...interface{}) 25 | // InfoContext is for reporting informational messages. 26 | InfoContext(ctx context.Context, format string, args ...interface{}) 27 | // WarnContext is for reporting warning messages. 28 | WarnContext(ctx context.Context, format string, args ...interface{}) 29 | // ErrorContext is for reporting errors. 30 | ErrorContext(ctx context.Context, format string, args ...interface{}) 31 | } 32 | -------------------------------------------------------------------------------- /internal/prebuiltconfigs/prebuiltconfigs.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 prebuiltconfigs 16 | 17 | import ( 18 | "embed" 19 | "fmt" 20 | "path/filepath" 21 | "strings" 22 | ) 23 | 24 | var ( 25 | //go:embed tools/*.yaml 26 | prebuiltConfigsFS embed.FS 27 | 28 | // Map of sources to their prebuilt tools 29 | prebuiltToolYAMLs map[string][]byte 30 | // List of sources with prebuilt tools 31 | prebuiltToolsSources []string 32 | ) 33 | 34 | func init() { 35 | var err error 36 | prebuiltToolYAMLs, prebuiltToolsSources, err = loadPrebuiltToolYAMLs() 37 | if err != nil { 38 | panic(fmt.Sprintf("Unexpected Error: %v\n", err)) 39 | } 40 | } 41 | 42 | // Get prebuilt tools for a source 43 | func Get(prebuiltSourceConfig string) ([]byte, error) { 44 | content, ok := prebuiltToolYAMLs[prebuiltSourceConfig] 45 | if !ok { 46 | prebuiltHelpSuffix := "no prebuilt configurations found." 47 | if len(prebuiltToolsSources) > 0 { 48 | prebuiltHelpSuffix = fmt.Sprintf("available: %s", strings.Join(prebuiltToolsSources, ", ")) 49 | } 50 | errMsg := fmt.Errorf("prebuilt source tool for '%s' not found. %s", prebuiltSourceConfig, prebuiltHelpSuffix) 51 | return nil, errMsg 52 | } 53 | return content, nil 54 | } 55 | 56 | // Load all available pre built tools 57 | func loadPrebuiltToolYAMLs() (map[string][]byte, []string, error) { 58 | toolYAMLs := make(map[string][]byte) 59 | var sourceTypes []string 60 | entries, err := prebuiltConfigsFS.ReadDir("tools") 61 | if err != nil { 62 | errMsg := fmt.Errorf("failed to read prebuilt tools %w", err) 63 | return nil, nil, errMsg 64 | } 65 | 66 | for _, entry := range entries { 67 | lowerName := strings.ToLower(entry.Name()) 68 | if !entry.IsDir() && (strings.HasSuffix(lowerName, ".yaml")) { 69 | filePathInFS := filepath.Join("tools", entry.Name()) 70 | content, err := prebuiltConfigsFS.ReadFile(filePathInFS) 71 | if err != nil { 72 | errMsg := fmt.Errorf("failed to read a prebuilt tool %w", err) 73 | return nil, nil, errMsg 74 | } 75 | sourceTypeKey := entry.Name()[:len(entry.Name())-len(".yaml")] 76 | 77 | sourceTypes = append(sourceTypes, sourceTypeKey) 78 | toolYAMLs[sourceTypeKey] = content 79 | } 80 | } 81 | if len(toolYAMLs) == 0 { 82 | errMsg := fmt.Errorf("no prebuilt tool configurations were loaded.%w", err) 83 | return nil, nil, errMsg 84 | } 85 | 86 | return toolYAMLs, sourceTypes, nil 87 | } 88 | -------------------------------------------------------------------------------- /internal/prebuiltconfigs/tools/bigquery.yaml: -------------------------------------------------------------------------------- 1 | sources: 2 | bigquery-source: 3 | kind: "bigquery" 4 | project: ${BIGQUERY_PROJECT} 5 | 6 | tools: 7 | execute_sql: 8 | kind: bigquery-execute-sql 9 | source: bigquery-source 10 | description: Use this tool to execute sql statement. 11 | 12 | get_dataset_info: 13 | kind: bigquery-get-dataset-info 14 | source: bigquery-source 15 | description: Use this tool to get dataset metadata. 16 | 17 | get_table_info: 18 | kind: bigquery-get-table-info 19 | source: bigquery-source 20 | description: Use this tool to get table metadata. 21 | 22 | list_dataset_ids: 23 | kind: bigquery-list-dataset-ids 24 | source: bigquery-source 25 | description: Use this tool to list datasets. 26 | 27 | list_table_ids: 28 | kind: bigquery-list-table-ids 29 | source: bigquery-source 30 | description: Use this tool to list tables. 31 | 32 | toolsets: 33 | bigquery-database-tools: 34 | - execute_sql 35 | - get_dataset_info 36 | - get_table_info 37 | - list_dataset_ids 38 | - list_table_ids 39 | -------------------------------------------------------------------------------- /internal/server/mcp/method.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package mcp 16 | 17 | import ( 18 | "context" 19 | "encoding/json" 20 | "fmt" 21 | 22 | "github.com/googleapis/genai-toolbox/internal/tools" 23 | ) 24 | 25 | func Initialize(version string) InitializeResult { 26 | toolsListChanged := false 27 | result := InitializeResult{ 28 | ProtocolVersion: LATEST_PROTOCOL_VERSION, 29 | Capabilities: ServerCapabilities{ 30 | Tools: &ListChanged{ 31 | ListChanged: &toolsListChanged, 32 | }, 33 | }, 34 | ServerInfo: Implementation{ 35 | Name: SERVER_NAME, 36 | Version: version, 37 | }, 38 | } 39 | return result 40 | } 41 | 42 | // ToolsList return a ListToolsResult 43 | func ToolsList(toolset tools.Toolset) ListToolsResult { 44 | mcpManifest := toolset.McpManifest 45 | 46 | result := ListToolsResult{ 47 | Tools: mcpManifest, 48 | } 49 | return result 50 | } 51 | 52 | // ToolCall runs tool invocation and return a CallToolResult 53 | func ToolCall(ctx context.Context, tool tools.Tool, params tools.ParamValues) CallToolResult { 54 | res, err := tool.Invoke(ctx, params) 55 | if err != nil { 56 | text := TextContent{ 57 | Type: "text", 58 | Text: err.Error(), 59 | } 60 | return CallToolResult{Content: []TextContent{text}, IsError: true} 61 | } 62 | 63 | content := make([]TextContent, 0) 64 | for _, d := range res { 65 | text := TextContent{Type: "text"} 66 | dM, err := json.Marshal(d) 67 | if err != nil { 68 | text.Text = fmt.Sprintf("fail to marshal: %s, result: %s", err, d) 69 | } else { 70 | text.Text = string(dM) 71 | } 72 | content = append(content, text) 73 | } 74 | return CallToolResult{Content: content} 75 | } 76 | -------------------------------------------------------------------------------- /internal/server/server_test.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 server_test 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | "io" 21 | "net/http" 22 | "os" 23 | "strings" 24 | "testing" 25 | 26 | "github.com/googleapis/genai-toolbox/internal/log" 27 | "github.com/googleapis/genai-toolbox/internal/server" 28 | "github.com/googleapis/genai-toolbox/internal/telemetry" 29 | ) 30 | 31 | func TestServe(t *testing.T) { 32 | ctx, cancel := context.WithCancel(context.Background()) 33 | defer cancel() 34 | 35 | addr, port := "127.0.0.1", 5000 36 | cfg := server.ServerConfig{ 37 | Version: "0.0.0", 38 | Address: addr, 39 | Port: port, 40 | } 41 | 42 | otelShutdown, err := telemetry.SetupOTel(ctx, "0.0.0", "", false, "toolbox") 43 | if err != nil { 44 | t.Fatalf("unexpected error: %s", err) 45 | } 46 | defer func() { 47 | err := otelShutdown(ctx) 48 | if err != nil { 49 | t.Fatalf("unexpected error: %s", err) 50 | } 51 | }() 52 | 53 | testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info") 54 | if err != nil { 55 | t.Fatalf("unexpected error: %s", err) 56 | } 57 | 58 | s, err := server.NewServer(ctx, cfg, testLogger) 59 | if err != nil { 60 | t.Fatalf("unable to initialize server: %v", err) 61 | } 62 | 63 | err = s.Listen(ctx) 64 | if err != nil { 65 | t.Fatalf("unable to start server: %v", err) 66 | } 67 | 68 | // start server in background 69 | errCh := make(chan error) 70 | go func() { 71 | defer close(errCh) 72 | 73 | err = s.Serve(ctx) 74 | if err != nil { 75 | errCh <- err 76 | } 77 | }() 78 | 79 | url := fmt.Sprintf("http://%s:%d/", addr, port) 80 | resp, err := http.Get(url) 81 | if err != nil { 82 | t.Fatalf("error when sending a request: %s", err) 83 | } 84 | defer resp.Body.Close() 85 | if resp.StatusCode != 200 { 86 | t.Fatalf("response status code is not 200") 87 | } 88 | raw, err := io.ReadAll(resp.Body) 89 | if err != nil { 90 | t.Fatalf("error reading from request body: %s", err) 91 | } 92 | if got := string(raw); strings.Contains(got, "0.0.0") { 93 | t.Fatalf("version missing from output: %q", got) 94 | } 95 | } 96 | -------------------------------------------------------------------------------- /internal/sources/bigquery/bigquery.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package bigquery 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | 21 | bigqueryapi "cloud.google.com/go/bigquery" 22 | "golang.org/x/oauth2/google" 23 | "google.golang.org/api/option" 24 | 25 | "github.com/googleapis/genai-toolbox/internal/sources" 26 | "github.com/googleapis/genai-toolbox/internal/util" 27 | "go.opentelemetry.io/otel/trace" 28 | ) 29 | 30 | const SourceKind string = "bigquery" 31 | 32 | // validate interface 33 | var _ sources.SourceConfig = Config{} 34 | 35 | type Config struct { 36 | // BigQuery configs 37 | Name string `yaml:"name" validate:"required"` 38 | Kind string `yaml:"kind" validate:"required"` 39 | Project string `yaml:"project" validate:"required"` 40 | Location string `yaml:"location"` 41 | } 42 | 43 | func (r Config) SourceConfigKind() string { 44 | // Returns BigQuery source kind 45 | return SourceKind 46 | } 47 | 48 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 49 | // Initializes a BigQuery Google SQL source 50 | client, err := initBigQueryConnection(ctx, tracer, r.Name, r.Project, r.Location) 51 | if err != nil { 52 | return nil, err 53 | } 54 | s := &Source{ 55 | Name: r.Name, 56 | Kind: SourceKind, 57 | Client: client, 58 | Location: r.Location, 59 | } 60 | return s, nil 61 | 62 | } 63 | 64 | var _ sources.Source = &Source{} 65 | 66 | type Source struct { 67 | // BigQuery Google SQL struct with client 68 | Name string `yaml:"name"` 69 | Kind string `yaml:"kind"` 70 | Client *bigqueryapi.Client 71 | Location string `yaml:"location"` 72 | } 73 | 74 | func (s *Source) SourceKind() string { 75 | // Returns BigQuery Google SQL source kind 76 | return SourceKind 77 | } 78 | 79 | func (s *Source) BigQueryClient() *bigqueryapi.Client { 80 | return s.Client 81 | } 82 | 83 | func initBigQueryConnection( 84 | ctx context.Context, 85 | tracer trace.Tracer, 86 | name string, 87 | project string, 88 | location string, 89 | ) (*bigqueryapi.Client, error) { 90 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 91 | defer span.End() 92 | 93 | cred, err := google.FindDefaultCredentials(ctx, bigqueryapi.Scope) 94 | if err != nil { 95 | return nil, fmt.Errorf("failed to find default Google Cloud credentials with scope %q: %w", bigqueryapi.Scope, err) 96 | } 97 | 98 | userAgent, err := util.UserAgentFromContext(ctx) 99 | if err != nil { 100 | return nil, err 101 | } 102 | 103 | client, err := bigqueryapi.NewClient(ctx, project, option.WithUserAgent(userAgent), option.WithCredentials(cred)) 104 | client.Location = location 105 | if err != nil { 106 | return nil, fmt.Errorf("failed to create BigQuery client for project %q: %w", project, err) 107 | } 108 | return client, nil 109 | } 110 | -------------------------------------------------------------------------------- /internal/sources/bigquery/bigquery_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package bigquery_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/sources/bigquery" 24 | "github.com/googleapis/genai-toolbox/internal/testutils" 25 | ) 26 | 27 | func TestParseFromYamlBigQuery(t *testing.T) { 28 | tcs := []struct { 29 | desc string 30 | in string 31 | want server.SourceConfigs 32 | }{ 33 | { 34 | desc: "basic example", 35 | in: ` 36 | sources: 37 | my-instance: 38 | kind: bigquery 39 | project: my-project 40 | location: us 41 | `, 42 | want: server.SourceConfigs{ 43 | "my-instance": bigquery.Config{ 44 | Name: "my-instance", 45 | Kind: bigquery.SourceKind, 46 | Project: "my-project", 47 | Location: "us", 48 | }, 49 | }, 50 | }, 51 | } 52 | for _, tc := range tcs { 53 | t.Run(tc.desc, func(t *testing.T) { 54 | got := struct { 55 | Sources server.SourceConfigs `yaml:"sources"` 56 | }{} 57 | // Parse contents 58 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) 59 | if err != nil { 60 | t.Fatalf("unable to unmarshal: %s", err) 61 | } 62 | if !cmp.Equal(tc.want, got.Sources) { 63 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources) 64 | } 65 | }) 66 | } 67 | 68 | } 69 | 70 | func TestFailParseFromYaml(t *testing.T) { 71 | tcs := []struct { 72 | desc string 73 | in string 74 | err string 75 | }{ 76 | { 77 | desc: "extra field", 78 | in: ` 79 | sources: 80 | my-instance: 81 | kind: bigquery 82 | project: my-project 83 | location: us 84 | foo: bar 85 | `, 86 | err: "unable to parse as \"bigquery\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | kind: bigquery\n 3 | location: us\n 4 | project: my-project", 87 | }, 88 | { 89 | desc: "missing required field", 90 | in: ` 91 | sources: 92 | my-mysql-instance: 93 | kind: bigquery 94 | location: us 95 | `, 96 | err: "unable to parse as \"bigquery\": Key: 'Config.Project' Error:Field validation for 'Project' failed on the 'required' tag", 97 | }, 98 | } 99 | for _, tc := range tcs { 100 | t.Run(tc.desc, func(t *testing.T) { 101 | got := struct { 102 | Sources server.SourceConfigs `yaml:"sources"` 103 | }{} 104 | // Parse contents 105 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) 106 | if err == nil { 107 | t.Fatalf("expect parsing to fail") 108 | } 109 | errStr := err.Error() 110 | if errStr != tc.err { 111 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) 112 | } 113 | }) 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /internal/sources/bigtable/bigtable.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package bigtable 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | 21 | "cloud.google.com/go/bigtable" 22 | "github.com/googleapis/genai-toolbox/internal/sources" 23 | "github.com/googleapis/genai-toolbox/internal/util" 24 | "go.opentelemetry.io/otel/trace" 25 | "google.golang.org/api/option" 26 | ) 27 | 28 | const SourceKind string = "bigtable" 29 | 30 | // validate interface 31 | var _ sources.SourceConfig = Config{} 32 | 33 | type Config struct { 34 | Name string `yaml:"name" validate:"required"` 35 | Kind string `yaml:"kind" validate:"required"` 36 | Project string `yaml:"project" validate:"required"` 37 | Instance string `yaml:"instance" validate:"required"` 38 | } 39 | 40 | func (r Config) SourceConfigKind() string { 41 | return SourceKind 42 | } 43 | 44 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 45 | client, err := initBigtableClient(ctx, tracer, r.Name, r.Project, r.Instance) 46 | if err != nil { 47 | return nil, fmt.Errorf("unable to create client: %w", err) 48 | } 49 | 50 | s := &Source{ 51 | Name: r.Name, 52 | Kind: SourceKind, 53 | Client: client, 54 | } 55 | return s, nil 56 | } 57 | 58 | var _ sources.Source = &Source{} 59 | 60 | type Source struct { 61 | Name string `yaml:"name"` 62 | Kind string `yaml:"kind"` 63 | Client *bigtable.Client 64 | } 65 | 66 | func (s *Source) SourceKind() string { 67 | return SourceKind 68 | } 69 | 70 | func (s *Source) BigtableClient() *bigtable.Client { 71 | return s.Client 72 | } 73 | 74 | func initBigtableClient(ctx context.Context, tracer trace.Tracer, name, project, instance string) (*bigtable.Client, error) { 75 | //nolint:all // Reassigned ctx 76 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 77 | defer span.End() 78 | 79 | // Set up Bigtable data operations client. 80 | poolSize := 10 81 | userAgent, err := util.UserAgentFromContext(ctx) 82 | if err != nil { 83 | return nil, err 84 | } 85 | 86 | client, err := bigtable.NewClient(ctx, project, instance, option.WithUserAgent(userAgent), option.WithGRPCConnectionPool(poolSize)) 87 | 88 | if err != nil { 89 | return nil, fmt.Errorf("unable to create bigtable.NewClient: %w", err) 90 | } 91 | 92 | return client, nil 93 | } 94 | -------------------------------------------------------------------------------- /internal/sources/dialect.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 sources 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | "strings" 21 | ) 22 | 23 | // Dialect represents the dialect type of a database. 24 | type Dialect string 25 | 26 | func (i *Dialect) String() string { 27 | if string(*i) != "" { 28 | return strings.ToLower(string(*i)) 29 | } 30 | return "googlesql" 31 | } 32 | 33 | func (i *Dialect) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error { 34 | var dialect string 35 | if err := unmarshal(&dialect); err != nil { 36 | return err 37 | } 38 | switch strings.ToLower(dialect) { 39 | case "googlesql", "postgresql": 40 | *i = Dialect(strings.ToLower(dialect)) 41 | return nil 42 | default: 43 | return fmt.Errorf(`dialect invalid: must be one of "googlesql", or "postgresql"`) 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /internal/sources/http/http.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | package http 15 | 16 | import ( 17 | "context" 18 | "fmt" 19 | "net/http" 20 | "net/url" 21 | "time" 22 | 23 | "github.com/googleapis/genai-toolbox/internal/sources" 24 | "go.opentelemetry.io/otel/trace" 25 | ) 26 | 27 | const SourceKind string = "http" 28 | 29 | // validate interface 30 | var _ sources.SourceConfig = Config{} 31 | 32 | type Config struct { 33 | Name string `yaml:"name" validate:"required"` 34 | Kind string `yaml:"kind" validate:"required"` 35 | BaseURL string `yaml:"baseUrl"` 36 | Timeout string `yaml:"timeout"` 37 | DefaultHeaders map[string]string `yaml:"headers"` 38 | QueryParams map[string]string `yaml:"queryParams"` 39 | } 40 | 41 | func (r Config) SourceConfigKind() string { 42 | return SourceKind 43 | } 44 | 45 | // DefaultConfig is a helper function that generates the default configuration for an HTTP Tool Config. 46 | func DefaultConfig(name string) Config { 47 | return Config{Name: name, Timeout: "30s"} 48 | } 49 | 50 | // Initialize initializes an HTTP Source instance. 51 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 52 | duration, err := time.ParseDuration(r.Timeout) 53 | if err != nil { 54 | return nil, fmt.Errorf("unable to parse Timeout string as time.Duration: %s", err) 55 | } 56 | client := http.Client{ 57 | Timeout: duration, 58 | } 59 | 60 | // Validate BaseURL 61 | _, err = url.ParseRequestURI(r.BaseURL) 62 | if err != nil { 63 | return nil, fmt.Errorf("failed to parse BaseUrl %v", err) 64 | } 65 | 66 | s := &Source{ 67 | Name: r.Name, 68 | Kind: SourceKind, 69 | BaseURL: r.BaseURL, 70 | DefaultHeaders: r.DefaultHeaders, 71 | QueryParams: r.QueryParams, 72 | Client: &client, 73 | } 74 | return s, nil 75 | 76 | } 77 | 78 | var _ sources.Source = &Source{} 79 | 80 | type Source struct { 81 | Name string `yaml:"name"` 82 | Kind string `yaml:"kind"` 83 | BaseURL string `yaml:"baseUrl"` 84 | DefaultHeaders map[string]string `yaml:"headers"` 85 | QueryParams map[string]string `yaml:"queryParams"` 86 | Client *http.Client 87 | } 88 | 89 | func (s *Source) SourceKind() string { 90 | return SourceKind 91 | } 92 | -------------------------------------------------------------------------------- /internal/sources/ip_type.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 sources 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | "strings" 21 | ) 22 | 23 | type IPType string 24 | 25 | func (i *IPType) String() string { 26 | if string(*i) != "" { 27 | return strings.ToLower(string(*i)) 28 | } 29 | return "public" 30 | } 31 | 32 | func (i *IPType) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error { 33 | var ipType string 34 | if err := unmarshal(&ipType); err != nil { 35 | return err 36 | } 37 | switch strings.ToLower(ipType) { 38 | case "private", "public": 39 | *i = IPType(strings.ToLower(ipType)) 40 | return nil 41 | default: 42 | return fmt.Errorf(`ipType invalid: must be one of "public", or "private"`) 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /internal/sources/mssql/mssql.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 mssql 16 | 17 | import ( 18 | "context" 19 | "database/sql" 20 | "fmt" 21 | 22 | "github.com/googleapis/genai-toolbox/internal/sources" 23 | _ "github.com/microsoft/go-mssqldb" 24 | "go.opentelemetry.io/otel/trace" 25 | ) 26 | 27 | const SourceKind string = "mssql" 28 | 29 | // validate interface 30 | var _ sources.SourceConfig = Config{} 31 | 32 | type Config struct { 33 | // Cloud SQL MSSQL configs 34 | Name string `yaml:"name" validate:"required"` 35 | Kind string `yaml:"kind" validate:"required"` 36 | Host string `yaml:"host" validate:"required"` 37 | Port string `yaml:"port" validate:"required"` 38 | User string `yaml:"user" validate:"required"` 39 | Password string `yaml:"password" validate:"required"` 40 | Database string `yaml:"database" validate:"required"` 41 | } 42 | 43 | func (r Config) SourceConfigKind() string { 44 | // Returns Cloud SQL MSSQL source kind 45 | return SourceKind 46 | } 47 | 48 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 49 | // Initializes a MSSQL source 50 | db, err := initMssqlConnection(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database) 51 | if err != nil { 52 | return nil, fmt.Errorf("unable to create db connection: %w", err) 53 | } 54 | 55 | // Verify db connection 56 | err = db.PingContext(ctx) 57 | if err != nil { 58 | return nil, fmt.Errorf("unable to connect successfully: %w", err) 59 | } 60 | 61 | s := &Source{ 62 | Name: r.Name, 63 | Kind: SourceKind, 64 | Db: db, 65 | } 66 | return s, nil 67 | } 68 | 69 | var _ sources.Source = &Source{} 70 | 71 | type Source struct { 72 | // Cloud SQL MSSQL struct with connection pool 73 | Name string `yaml:"name"` 74 | Kind string `yaml:"kind"` 75 | Db *sql.DB 76 | } 77 | 78 | func (s *Source) SourceKind() string { 79 | // Returns Cloud SQL MSSQL source kind 80 | return SourceKind 81 | } 82 | 83 | func (s *Source) MSSQLDB() *sql.DB { 84 | // Returns a Cloud SQL MSSQL database connection pool 85 | return s.Db 86 | } 87 | 88 | func initMssqlConnection(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*sql.DB, error) { 89 | //nolint:all // Reassigned ctx 90 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 91 | defer span.End() 92 | 93 | // Create dsn 94 | dsn := fmt.Sprintf("sqlserver://%s:%s@%s:%s?database=%s", user, pass, host, port, dbname) 95 | 96 | // Open database connection 97 | db, err := sql.Open("sqlserver", dsn) 98 | if err != nil { 99 | return nil, fmt.Errorf("sql.Open: %w", err) 100 | } 101 | return db, nil 102 | } 103 | -------------------------------------------------------------------------------- /internal/sources/mysql/mysql.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 mysql 16 | 17 | import ( 18 | "context" 19 | "database/sql" 20 | "fmt" 21 | 22 | _ "github.com/go-sql-driver/mysql" 23 | "github.com/googleapis/genai-toolbox/internal/sources" 24 | "go.opentelemetry.io/otel/trace" 25 | ) 26 | 27 | const SourceKind string = "mysql" 28 | 29 | // validate interface 30 | var _ sources.SourceConfig = Config{} 31 | 32 | type Config struct { 33 | Name string `yaml:"name" validate:"required"` 34 | Kind string `yaml:"kind" validate:"required"` 35 | Host string `yaml:"host" validate:"required"` 36 | Port string `yaml:"port" validate:"required"` 37 | User string `yaml:"user" validate:"required"` 38 | Password string `yaml:"password" validate:"required"` 39 | Database string `yaml:"database" validate:"required"` 40 | } 41 | 42 | func (r Config) SourceConfigKind() string { 43 | return SourceKind 44 | } 45 | 46 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 47 | pool, err := initMySQLConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database) 48 | if err != nil { 49 | return nil, fmt.Errorf("unable to create pool: %w", err) 50 | } 51 | 52 | err = pool.PingContext(ctx) 53 | if err != nil { 54 | return nil, fmt.Errorf("unable to connect successfully: %w", err) 55 | } 56 | 57 | s := &Source{ 58 | Name: r.Name, 59 | Kind: SourceKind, 60 | Pool: pool, 61 | } 62 | return s, nil 63 | } 64 | 65 | var _ sources.Source = &Source{} 66 | 67 | type Source struct { 68 | Name string `yaml:"name"` 69 | Kind string `yaml:"kind"` 70 | Pool *sql.DB 71 | } 72 | 73 | func (s *Source) SourceKind() string { 74 | return SourceKind 75 | } 76 | 77 | func (s *Source) MySQLPool() *sql.DB { 78 | return s.Pool 79 | } 80 | 81 | func initMySQLConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*sql.DB, error) { 82 | //nolint:all // Reassigned ctx 83 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 84 | defer span.End() 85 | 86 | // Configure the driver to connect to the database 87 | dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true", user, pass, host, port, dbname) 88 | 89 | // Interact with the driver directly as you normally would 90 | pool, err := sql.Open("mysql", dsn) 91 | if err != nil { 92 | return nil, fmt.Errorf("sql.Open: %w", err) 93 | } 94 | return pool, nil 95 | } 96 | -------------------------------------------------------------------------------- /internal/sources/neo4j/neo4j.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 neo4j 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | 21 | "github.com/googleapis/genai-toolbox/internal/sources" 22 | "github.com/neo4j/neo4j-go-driver/v5/neo4j" 23 | "go.opentelemetry.io/otel/trace" 24 | ) 25 | 26 | const SourceKind string = "neo4j" 27 | 28 | // validate interface 29 | var _ sources.SourceConfig = Config{} 30 | 31 | type Config struct { 32 | Name string `yaml:"name" validate:"required"` 33 | Kind string `yaml:"kind" validate:"required"` 34 | Uri string `yaml:"uri" validate:"required"` 35 | User string `yaml:"user" validate:"required"` 36 | Password string `yaml:"password" validate:"required"` 37 | Database string `yaml:"database" validate:"required"` 38 | } 39 | 40 | func (r Config) SourceConfigKind() string { 41 | return SourceKind 42 | } 43 | 44 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 45 | driver, err := initNeo4jDriver(ctx, tracer, r.Uri, r.User, r.Password, r.Name) 46 | if err != nil { 47 | return nil, fmt.Errorf("unable to create driver: %w", err) 48 | } 49 | 50 | err = driver.VerifyConnectivity(ctx) 51 | if err != nil { 52 | return nil, fmt.Errorf("unable to connect successfully: %w", err) 53 | } 54 | 55 | if r.Database == "" { 56 | r.Database = "neo4j" 57 | } 58 | s := &Source{ 59 | Name: r.Name, 60 | Kind: SourceKind, 61 | Database: r.Database, 62 | Driver: driver, 63 | } 64 | return s, nil 65 | } 66 | 67 | var _ sources.Source = &Source{} 68 | 69 | type Source struct { 70 | Name string `yaml:"name"` 71 | Kind string `yaml:"kind"` 72 | Database string `yaml:"database"` 73 | Driver neo4j.DriverWithContext 74 | } 75 | 76 | func (s *Source) SourceKind() string { 77 | return SourceKind 78 | } 79 | 80 | func (s *Source) Neo4jDriver() neo4j.DriverWithContext { 81 | return s.Driver 82 | } 83 | 84 | func (s *Source) Neo4jDatabase() string { 85 | return s.Database 86 | } 87 | 88 | func initNeo4jDriver(ctx context.Context, tracer trace.Tracer, uri, user, password, name string) (neo4j.DriverWithContext, error) { 89 | //nolint:all // Reassigned ctx 90 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 91 | defer span.End() 92 | 93 | auth := neo4j.BasicAuth(user, password, "") 94 | driver, err := neo4j.NewDriverWithContext(uri, auth) 95 | if err != nil { 96 | return nil, fmt.Errorf("unable to create connection driver: %w", err) 97 | } 98 | return driver, nil 99 | } 100 | -------------------------------------------------------------------------------- /internal/sources/postgres/postgres.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 postgres 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | 21 | "github.com/googleapis/genai-toolbox/internal/sources" 22 | "github.com/jackc/pgx/v5/pgxpool" 23 | "go.opentelemetry.io/otel/trace" 24 | ) 25 | 26 | const SourceKind string = "postgres" 27 | 28 | // validate interface 29 | var _ sources.SourceConfig = Config{} 30 | 31 | type Config struct { 32 | Name string `yaml:"name" validate:"required"` 33 | Kind string `yaml:"kind" validate:"required"` 34 | Host string `yaml:"host" validate:"required"` 35 | Port string `yaml:"port" validate:"required"` 36 | User string `yaml:"user" validate:"required"` 37 | Password string `yaml:"password" validate:"required"` 38 | Database string `yaml:"database" validate:"required"` 39 | } 40 | 41 | func (r Config) SourceConfigKind() string { 42 | return SourceKind 43 | } 44 | 45 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 46 | pool, err := initPostgresConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database) 47 | if err != nil { 48 | return nil, fmt.Errorf("unable to create pool: %w", err) 49 | } 50 | 51 | err = pool.Ping(ctx) 52 | if err != nil { 53 | return nil, fmt.Errorf("unable to connect successfully: %w", err) 54 | } 55 | 56 | s := &Source{ 57 | Name: r.Name, 58 | Kind: SourceKind, 59 | Pool: pool, 60 | } 61 | return s, nil 62 | } 63 | 64 | var _ sources.Source = &Source{} 65 | 66 | type Source struct { 67 | Name string `yaml:"name"` 68 | Kind string `yaml:"kind"` 69 | Pool *pgxpool.Pool 70 | } 71 | 72 | func (s *Source) SourceKind() string { 73 | return SourceKind 74 | } 75 | 76 | func (s *Source) PostgresPool() *pgxpool.Pool { 77 | return s.Pool 78 | } 79 | 80 | func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string) (*pgxpool.Pool, error) { 81 | //nolint:all // Reassigned ctx 82 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 83 | defer span.End() 84 | // urlExample := "postgres:dd//username:password@localhost:5432/database_name" 85 | i := fmt.Sprintf("postgres://%s:%s@%s:%s/%s", user, pass, host, port, dbname) 86 | pool, err := pgxpool.New(ctx, i) 87 | if err != nil { 88 | return nil, fmt.Errorf("unable to create connection pool: %w", err) 89 | } 90 | 91 | return pool, nil 92 | } 93 | -------------------------------------------------------------------------------- /internal/sources/sources.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 sources 16 | 17 | import ( 18 | "context" 19 | 20 | "go.opentelemetry.io/otel/attribute" 21 | "go.opentelemetry.io/otel/trace" 22 | ) 23 | 24 | // SourceConfig is the interface for configuring a source. 25 | type SourceConfig interface { 26 | SourceConfigKind() string 27 | Initialize(ctx context.Context, tracer trace.Tracer) (Source, error) 28 | } 29 | 30 | // Source is the interface for the source itself. 31 | type Source interface { 32 | SourceKind() string 33 | } 34 | 35 | // InitConnectionSpan adds a span for database pool connection initialization 36 | func InitConnectionSpan(ctx context.Context, tracer trace.Tracer, sourceKind, sourceName string) (context.Context, trace.Span) { 37 | ctx, span := tracer.Start( 38 | ctx, 39 | "toolbox/server/source/connect", 40 | trace.WithAttributes(attribute.String("source_kind", sourceKind)), 41 | trace.WithAttributes(attribute.String("source_name", sourceName)), 42 | ) 43 | return ctx, span 44 | } 45 | -------------------------------------------------------------------------------- /internal/sources/sqlite/sqlite.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 sqlite 16 | 17 | import ( 18 | "context" 19 | "database/sql" 20 | "fmt" 21 | 22 | "github.com/googleapis/genai-toolbox/internal/sources" 23 | "go.opentelemetry.io/otel/trace" 24 | _ "modernc.org/sqlite" // Pure Go SQLite driver 25 | ) 26 | 27 | const SourceKind string = "sqlite" 28 | 29 | // validate interface 30 | var _ sources.SourceConfig = Config{} 31 | 32 | type Config struct { 33 | Name string `yaml:"name" validate:"required"` 34 | Kind string `yaml:"kind" validate:"required"` 35 | Database string `yaml:"database" validate:"required"` // Path to SQLite database file 36 | } 37 | 38 | func (r Config) SourceConfigKind() string { 39 | return SourceKind 40 | } 41 | 42 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { 43 | db, err := initSQLiteConnection(ctx, tracer, r.Name, r.Database) 44 | if err != nil { 45 | return nil, fmt.Errorf("unable to create db connection: %w", err) 46 | } 47 | 48 | err = db.PingContext(context.Background()) 49 | if err != nil { 50 | return nil, fmt.Errorf("unable to connect successfully: %w", err) 51 | } 52 | 53 | s := &Source{ 54 | Name: r.Name, 55 | Kind: SourceKind, 56 | Db: db, 57 | } 58 | return s, nil 59 | } 60 | 61 | var _ sources.Source = &Source{} 62 | 63 | type Source struct { 64 | Name string `yaml:"name"` 65 | Kind string `yaml:"kind"` 66 | Db *sql.DB 67 | } 68 | 69 | func (s *Source) SourceKind() string { 70 | return SourceKind 71 | } 72 | 73 | func (s *Source) SQLiteDB() *sql.DB { 74 | return s.Db 75 | } 76 | 77 | func initSQLiteConnection(ctx context.Context, tracer trace.Tracer, name, dbPath string) (*sql.DB, error) { 78 | //nolint:all // Reassigned ctx 79 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) 80 | defer span.End() 81 | 82 | // Open database connection 83 | db, err := sql.Open("sqlite", dbPath) 84 | if err != nil { 85 | return nil, fmt.Errorf("sql.Open: %w", err) 86 | } 87 | 88 | // Set some reasonable defaults for SQLite 89 | db.SetMaxOpenConns(1) // SQLite only supports one writer at a time 90 | db.SetMaxIdleConns(1) 91 | 92 | return db, nil 93 | } 94 | -------------------------------------------------------------------------------- /internal/sources/sqlite/sqlite_test.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 sqlite_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/sources" 24 | "github.com/googleapis/genai-toolbox/internal/sources/sqlite" 25 | "github.com/googleapis/genai-toolbox/internal/testutils" 26 | ) 27 | 28 | func TestParseFromYamlSQLite(t *testing.T) { 29 | tcs := []struct { 30 | desc string 31 | in string 32 | want server.SourceConfigs 33 | }{ 34 | { 35 | desc: "basic example", 36 | in: ` 37 | sources: 38 | my-sqlite-db: 39 | kind: sqlite 40 | database: /path/to/database.db 41 | `, 42 | want: map[string]sources.SourceConfig{ 43 | "my-sqlite-db": sqlite.Config{ 44 | Name: "my-sqlite-db", 45 | Kind: sqlite.SourceKind, 46 | Database: "/path/to/database.db", 47 | }, 48 | }, 49 | }, 50 | } 51 | for _, tc := range tcs { 52 | t.Run(tc.desc, func(t *testing.T) { 53 | got := struct { 54 | Sources server.SourceConfigs `yaml:"sources"` 55 | }{} 56 | // Parse contents 57 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) 58 | if err != nil { 59 | t.Fatalf("unable to unmarshal: %s", err) 60 | } 61 | if !cmp.Equal(tc.want, got.Sources) { 62 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources) 63 | } 64 | }) 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /internal/sources/util.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 sources 16 | 17 | import ( 18 | "context" 19 | "encoding/json" 20 | "fmt" 21 | "io" 22 | "net/http" 23 | "strings" 24 | 25 | "cloud.google.com/go/cloudsqlconn" 26 | "golang.org/x/oauth2/google" 27 | ) 28 | 29 | // GetCloudSQLDialOpts retrieve dial options with the right ip type and user agent for cloud sql 30 | // databases. 31 | func GetCloudSQLOpts(ipType, userAgent string, useIAM bool) ([]cloudsqlconn.Option, error) { 32 | opts := []cloudsqlconn.Option{cloudsqlconn.WithUserAgent(userAgent)} 33 | switch strings.ToLower(ipType) { 34 | case "private": 35 | opts = append(opts, cloudsqlconn.WithDefaultDialOptions(cloudsqlconn.WithPrivateIP())) 36 | case "public": 37 | opts = append(opts, cloudsqlconn.WithDefaultDialOptions(cloudsqlconn.WithPublicIP())) 38 | default: 39 | return nil, fmt.Errorf("invalid ipType %s", ipType) 40 | } 41 | 42 | if useIAM { 43 | opts = append(opts, cloudsqlconn.WithIAMAuthN()) 44 | } 45 | return opts, nil 46 | } 47 | 48 | // GetIAMPrincipalEmailFromADC finds the email associated with ADC 49 | func GetIAMPrincipalEmailFromADC(ctx context.Context) (string, error) { 50 | // Finds ADC and returns an HTTP client associated with it 51 | client, err := google.DefaultClient(ctx, 52 | "https://www.googleapis.com/auth/userinfo.email") 53 | if err != nil { 54 | return "", fmt.Errorf("failed to call userinfo endpoint: %w", err) 55 | } 56 | 57 | // Retrieve the email associated with the token 58 | resp, err := client.Get("https://oauth2.googleapis.com/tokeninfo") 59 | if err != nil { 60 | return "", fmt.Errorf("failed to call tokeninfo endpoint: %w", err) 61 | } 62 | defer resp.Body.Close() 63 | 64 | bodyBytes, err := io.ReadAll(resp.Body) 65 | if err != nil { 66 | return "", fmt.Errorf("error reading response body %d: %s", resp.StatusCode, string(bodyBytes)) 67 | } 68 | if resp.StatusCode != http.StatusOK { 69 | return "", fmt.Errorf("tokeninfo endpoint returned non-OK status %d: %s", resp.StatusCode, string(bodyBytes)) 70 | } 71 | 72 | // Unmarshal response body and get `email` 73 | var responseJSON map[string]any 74 | err = json.Unmarshal(bodyBytes, &responseJSON) 75 | if err != nil { 76 | 77 | return "", fmt.Errorf("error parsing JSON: %v", err) 78 | } 79 | 80 | emailValue, ok := responseJSON["email"] 81 | if !ok { 82 | return "", fmt.Errorf("email not found in response: %v", err) 83 | } 84 | // service account email used for IAM should trim the suffix 85 | email := strings.TrimSuffix(emailValue.(string), ".gserviceaccount.com") 86 | return email, nil 87 | } 88 | -------------------------------------------------------------------------------- /internal/testutils/testutils.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 testutils 16 | 17 | import ( 18 | "context" 19 | "fmt" 20 | "os" 21 | "strings" 22 | 23 | "github.com/googleapis/genai-toolbox/internal/log" 24 | "github.com/googleapis/genai-toolbox/internal/util" 25 | ) 26 | 27 | // formatYaml is a utility function for stripping out tabs in multiline strings 28 | func FormatYaml(in string) []byte { 29 | // removes any leading indentation(tabs) 30 | in = strings.ReplaceAll(in, "\n\t", "\n ") 31 | // converts remaining indentation 32 | in = strings.ReplaceAll(in, "\t", " ") 33 | return []byte(in) 34 | } 35 | 36 | // ContextWithNewLogger create a new context with new logger 37 | func ContextWithNewLogger() (context.Context, error) { 38 | ctx := context.Background() 39 | logger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info") 40 | if err != nil { 41 | return nil, fmt.Errorf("unable to create logger: %s", err) 42 | } 43 | return util.WithLogger(ctx, logger), nil 44 | } 45 | -------------------------------------------------------------------------------- /internal/tools/bigquery/bigquery_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package bigquery_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/bigquery" 26 | ) 27 | 28 | func TestParseFromYamlBigQuery(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: bigquery-sql 44 | source: my-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | parameters: 49 | - name: country 50 | type: string 51 | description: some description 52 | `, 53 | want: server.ToolConfigs{ 54 | "example_tool": bigquery.Config{ 55 | Name: "example_tool", 56 | Kind: bigquery.ToolKind, 57 | Source: "my-instance", 58 | Description: "some description", 59 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 60 | AuthRequired: []string{}, 61 | Parameters: []tools.Parameter{ 62 | tools.NewStringParameter("country", "some description"), 63 | }, 64 | }, 65 | }, 66 | }, 67 | } 68 | for _, tc := range tcs { 69 | t.Run(tc.desc, func(t *testing.T) { 70 | got := struct { 71 | Tools server.ToolConfigs `yaml:"tools"` 72 | }{} 73 | // Parse contents 74 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 75 | if err != nil { 76 | t.Fatalf("unable to unmarshal: %s", err) 77 | } 78 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 79 | t.Fatalf("incorrect parse: diff %v", diff) 80 | } 81 | }) 82 | } 83 | 84 | } 85 | -------------------------------------------------------------------------------- /internal/tools/bigqueryexecutesql/bigqueryexecutesql_test.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 bigqueryexecutesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/bigqueryexecutesql" 25 | ) 26 | 27 | func TestParseFromYamlBigQueryExecuteSql(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: bigquery-execute-sql 43 | source: my-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": bigqueryexecutesql.Config{ 48 | Name: "example_tool", 49 | Kind: bigqueryexecutesql.ToolKind, 50 | Source: "my-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | }, 54 | }, 55 | }, 56 | } 57 | for _, tc := range tcs { 58 | t.Run(tc.desc, func(t *testing.T) { 59 | got := struct { 60 | Tools server.ToolConfigs `yaml:"tools"` 61 | }{} 62 | // Parse contents 63 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 64 | if err != nil { 65 | t.Fatalf("unable to unmarshal: %s", err) 66 | } 67 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 68 | t.Fatalf("incorrect parse: diff %v", diff) 69 | } 70 | }) 71 | } 72 | 73 | } 74 | -------------------------------------------------------------------------------- /internal/tools/bigquerygetdatasetinfo/bigquerygetdatasetinfo_test.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 bigquerygetdatasetinfo_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/bigquerygetdatasetinfo" 25 | ) 26 | 27 | func TestParseFromYamlBigQueryGetDatasetInfo(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: bigquery-get-dataset-info 43 | source: my-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": bigquerygetdatasetinfo.Config{ 48 | Name: "example_tool", 49 | Kind: bigquerygetdatasetinfo.ToolKind, 50 | Source: "my-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | }, 54 | }, 55 | }, 56 | } 57 | for _, tc := range tcs { 58 | t.Run(tc.desc, func(t *testing.T) { 59 | got := struct { 60 | Tools server.ToolConfigs `yaml:"tools"` 61 | }{} 62 | // Parse contents 63 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 64 | if err != nil { 65 | t.Fatalf("unable to unmarshal: %s", err) 66 | } 67 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 68 | t.Fatalf("incorrect parse: diff %v", diff) 69 | } 70 | }) 71 | } 72 | 73 | } 74 | -------------------------------------------------------------------------------- /internal/tools/bigquerygettableinfo/bigquerygettableinfo_test.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 bigquerygettableinfo_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/bigquerygettableinfo" 25 | ) 26 | 27 | func TestParseFromYamlBigQueryGetTableInfo(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: bigquery-get-table-info 43 | source: my-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": bigquerygettableinfo.Config{ 48 | Name: "example_tool", 49 | Kind: bigquerygettableinfo.ToolKind, 50 | Source: "my-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | }, 54 | }, 55 | }, 56 | } 57 | for _, tc := range tcs { 58 | t.Run(tc.desc, func(t *testing.T) { 59 | got := struct { 60 | Tools server.ToolConfigs `yaml:"tools"` 61 | }{} 62 | // Parse contents 63 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 64 | if err != nil { 65 | t.Fatalf("unable to unmarshal: %s", err) 66 | } 67 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 68 | t.Fatalf("incorrect parse: diff %v", diff) 69 | } 70 | }) 71 | } 72 | 73 | } 74 | -------------------------------------------------------------------------------- /internal/tools/bigquerylistdatasetids/bigquerylistdatasetids_test.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 bigquerylistdatasetids_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/bigquerylistdatasetids" 25 | ) 26 | 27 | func TestParseFromYamlBigQueryListDatasetIds(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: bigquery-list-dataset-ids 43 | source: my-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": bigquerylistdatasetids.Config{ 48 | Name: "example_tool", 49 | Kind: bigquerylistdatasetids.ToolKind, 50 | Source: "my-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | }, 54 | }, 55 | }, 56 | } 57 | for _, tc := range tcs { 58 | t.Run(tc.desc, func(t *testing.T) { 59 | got := struct { 60 | Tools server.ToolConfigs `yaml:"tools"` 61 | }{} 62 | // Parse contents 63 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 64 | if err != nil { 65 | t.Fatalf("unable to unmarshal: %s", err) 66 | } 67 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 68 | t.Fatalf("incorrect parse: diff %v", diff) 69 | } 70 | }) 71 | } 72 | 73 | } 74 | -------------------------------------------------------------------------------- /internal/tools/bigquerylisttableids/bigquerylisttableids_test.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 bigquerylisttableids_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/bigquerylisttableids" 25 | ) 26 | 27 | func TestParseFromYamlBigQueryListTableIds(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: bigquery-list-table-ids 43 | source: my-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": bigquerylisttableids.Config{ 48 | Name: "example_tool", 49 | Kind: bigquerylisttableids.ToolKind, 50 | Source: "my-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | }, 54 | }, 55 | }, 56 | } 57 | for _, tc := range tcs { 58 | t.Run(tc.desc, func(t *testing.T) { 59 | got := struct { 60 | Tools server.ToolConfigs `yaml:"tools"` 61 | }{} 62 | // Parse contents 63 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 64 | if err != nil { 65 | t.Fatalf("unable to unmarshal: %s", err) 66 | } 67 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 68 | t.Fatalf("incorrect parse: diff %v", diff) 69 | } 70 | }) 71 | } 72 | 73 | } 74 | -------------------------------------------------------------------------------- /internal/tools/bigtable/bigtable_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package bigtable_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/bigtable" 26 | ) 27 | 28 | func TestParseFromYamlBigtable(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: bigtable-sql 44 | source: my-pg-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | parameters: 49 | - name: country 50 | type: string 51 | description: some description 52 | `, 53 | want: server.ToolConfigs{ 54 | "example_tool": bigtable.Config{ 55 | Name: "example_tool", 56 | Kind: bigtable.ToolKind, 57 | Source: "my-pg-instance", 58 | Description: "some description", 59 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 60 | AuthRequired: []string{}, 61 | Parameters: []tools.Parameter{ 62 | tools.NewStringParameter("country", "some description"), 63 | }, 64 | }, 65 | }, 66 | }, 67 | } 68 | for _, tc := range tcs { 69 | t.Run(tc.desc, func(t *testing.T) { 70 | got := struct { 71 | Tools server.ToolConfigs `yaml:"tools"` 72 | }{} 73 | // Parse contents 74 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 75 | if err != nil { 76 | t.Fatalf("unable to unmarshal: %s", err) 77 | } 78 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 79 | t.Fatalf("incorrect parse: diff %v", diff) 80 | } 81 | }) 82 | } 83 | 84 | } 85 | -------------------------------------------------------------------------------- /internal/tools/common.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 tools 16 | 17 | import ( 18 | "regexp" 19 | ) 20 | 21 | var validName = regexp.MustCompile(`^[a-zA-Z0-9_-]*$`) 22 | 23 | func IsValidName(s string) bool { 24 | return validName.MatchString(s) 25 | } 26 | -------------------------------------------------------------------------------- /internal/tools/couchbase/couchbase_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2024 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | package couchbase_test 16 | 17 | import ( 18 | "testing" 19 | 20 | "github.com/googleapis/genai-toolbox/internal/tools/couchbase" 21 | 22 | yaml "github.com/goccy/go-yaml" 23 | "github.com/google/go-cmp/cmp" 24 | "github.com/googleapis/genai-toolbox/internal/server" 25 | "github.com/googleapis/genai-toolbox/internal/testutils" 26 | "github.com/googleapis/genai-toolbox/internal/tools" 27 | ) 28 | 29 | func TestParseFromYamlCouchbase(t *testing.T) { 30 | tcs := []struct { 31 | desc string 32 | in string 33 | want server.ToolConfigs 34 | }{ 35 | { 36 | desc: "basic example", 37 | in: ` 38 | tools: 39 | example_tool: 40 | kind: couchbase-sql 41 | source: my-couchbase-instance 42 | description: some tool description 43 | statement: | 44 | select * from hotel WHERE name = $hotel; 45 | parameters: 46 | - name: hotel 47 | type: string 48 | description: hotel parameter description 49 | `, 50 | want: server.ToolConfigs{ 51 | "example_tool": couchbase.Config{ 52 | Name: "example_tool", 53 | Kind: couchbase.ToolKind, 54 | AuthRequired: []string{}, 55 | Source: "my-couchbase-instance", 56 | Description: "some tool description", 57 | Statement: "select * from hotel WHERE name = $hotel;\n", 58 | Parameters: []tools.Parameter{ 59 | tools.NewStringParameter("hotel", "hotel parameter description"), 60 | }, 61 | }, 62 | }, 63 | }, 64 | } 65 | for _, tc := range tcs { 66 | t.Run(tc.desc, func(t *testing.T) { 67 | got := struct { 68 | Tools server.ToolConfigs `yaml:"tools"` 69 | }{} 70 | 71 | // Create a context with a logger 72 | ctx, err := testutils.ContextWithNewLogger() 73 | if err != nil { 74 | t.Fatalf("unable to create context with logger: %s", err) 75 | } 76 | 77 | // Parse contents with context 78 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 79 | if err != nil { 80 | t.Fatalf("unable to unmarshal: %s", err) 81 | } 82 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 83 | t.Fatalf("incorrect parse: diff %v", diff) 84 | } 85 | }) 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /internal/tools/dgraph/dgraph_test.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 dgraph_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/dgraph" 25 | ) 26 | 27 | func TestParseFromYamlDgraph(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic query example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: dgraph-dql 43 | source: my-dgraph-instance 44 | description: some tool description 45 | isQuery: true 46 | timeout: 20s 47 | statement: | 48 | query {q(func: eq(email, "example@email.com")) {email}} 49 | `, 50 | want: server.ToolConfigs{ 51 | "example_tool": dgraph.Config{ 52 | Name: "example_tool", 53 | Kind: dgraph.ToolKind, 54 | Source: "my-dgraph-instance", 55 | AuthRequired: []string{}, 56 | Description: "some tool description", 57 | IsQuery: true, 58 | Timeout: "20s", 59 | Statement: "query {q(func: eq(email, \"example@email.com\")) {email}}\n", 60 | }, 61 | }, 62 | }, 63 | { 64 | desc: "basic mutation example", 65 | in: ` 66 | tools: 67 | example_tool: 68 | kind: dgraph-dql 69 | source: my-dgraph-instance 70 | description: some tool description 71 | statement: | 72 | mutation {set { _:a "a@email.com" . _:b "b@email.com" .}} 73 | `, 74 | want: server.ToolConfigs{ 75 | "example_tool": dgraph.Config{ 76 | Name: "example_tool", 77 | Kind: dgraph.ToolKind, 78 | Source: "my-dgraph-instance", 79 | Description: "some tool description", 80 | AuthRequired: []string{}, 81 | Statement: "mutation {set { _:a \"a@email.com\" . _:b \"b@email.com\" .}}\n", 82 | }, 83 | }, 84 | }, 85 | } 86 | for _, tc := range tcs { 87 | t.Run(tc.desc, func(t *testing.T) { 88 | got := struct { 89 | Tools server.ToolConfigs `yaml:"tools"` 90 | }{} 91 | // Parse contents 92 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 93 | if err != nil { 94 | t.Fatalf("unable to unmarshal: %s", err) 95 | } 96 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 97 | t.Fatalf("incorrect parse: diff %v", diff) 98 | } 99 | }) 100 | } 101 | 102 | } 103 | -------------------------------------------------------------------------------- /internal/tools/http_method.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | package tools 15 | 16 | import ( 17 | "context" 18 | "fmt" 19 | "net/http" 20 | "strings" 21 | ) 22 | 23 | // HTTPMethod is a string of a valid HTTP method (e.g "GET") 24 | type HTTPMethod string 25 | 26 | // isValidHTTPMethod checks if the input string matches one of the method constants defined in the net/http package 27 | func isValidHTTPMethod(method string) bool { 28 | 29 | switch method { 30 | case http.MethodGet, http.MethodPost, http.MethodPut, http.MethodDelete, 31 | http.MethodPatch, http.MethodHead, http.MethodOptions, http.MethodTrace, 32 | http.MethodConnect: 33 | return true 34 | } 35 | return false 36 | } 37 | 38 | func (i *HTTPMethod) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error { 39 | var httpMethod string 40 | if err := unmarshal(&httpMethod); err != nil { 41 | return fmt.Errorf(`error unmarshalling HTTP method: %s`, err) 42 | } 43 | httpMethod = strings.ToUpper(httpMethod) 44 | if !isValidHTTPMethod(httpMethod) { 45 | return fmt.Errorf(`%s is not a valid http method`, httpMethod) 46 | } 47 | *i = HTTPMethod(httpMethod) 48 | return nil 49 | } 50 | -------------------------------------------------------------------------------- /internal/tools/mssqlexecutesql/mssqlexecutesql_test.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 mssqlexecutesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/mssqlexecutesql" 25 | ) 26 | 27 | func TestParseFromYamlExecuteSql(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: mssql-execute-sql 43 | source: my-instance 44 | description: some description 45 | authRequired: 46 | - my-google-auth-service 47 | - other-auth-service 48 | `, 49 | want: server.ToolConfigs{ 50 | "example_tool": mssqlexecutesql.Config{ 51 | Name: "example_tool", 52 | Kind: mssqlexecutesql.ToolKind, 53 | Source: "my-instance", 54 | Description: "some description", 55 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 56 | }, 57 | }, 58 | }, 59 | } 60 | for _, tc := range tcs { 61 | t.Run(tc.desc, func(t *testing.T) { 62 | got := struct { 63 | Tools server.ToolConfigs `yaml:"tools"` 64 | }{} 65 | // Parse contents 66 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 67 | if err != nil { 68 | t.Fatalf("unable to unmarshal: %s", err) 69 | } 70 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 71 | t.Fatalf("incorrect parse: diff %v", diff) 72 | } 73 | }) 74 | } 75 | 76 | } 77 | -------------------------------------------------------------------------------- /internal/tools/mssqlsql/mssqlsql_test.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 mssqlsql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" 26 | ) 27 | 28 | func TestParseFromYamlMssql(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: mssql-sql 44 | source: my-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | authRequired: 49 | - my-google-auth-service 50 | - other-auth-service 51 | parameters: 52 | - name: country 53 | type: string 54 | description: some description 55 | authServices: 56 | - name: my-google-auth-service 57 | field: user_id 58 | - name: other-auth-service 59 | field: user_id 60 | `, 61 | want: server.ToolConfigs{ 62 | "example_tool": mssqlsql.Config{ 63 | Name: "example_tool", 64 | Kind: mssqlsql.ToolKind, 65 | Source: "my-instance", 66 | Description: "some description", 67 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 68 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 69 | Parameters: []tools.Parameter{ 70 | tools.NewStringParameterWithAuth("country", "some description", 71 | []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"}, 72 | {Name: "other-auth-service", Field: "user_id"}}), 73 | }, 74 | }, 75 | }, 76 | }, 77 | } 78 | for _, tc := range tcs { 79 | t.Run(tc.desc, func(t *testing.T) { 80 | got := struct { 81 | Tools server.ToolConfigs `yaml:"tools"` 82 | }{} 83 | // Parse contents 84 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 85 | if err != nil { 86 | t.Fatalf("unable to unmarshal: %s", err) 87 | } 88 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 89 | t.Fatalf("incorrect parse: diff %v", diff) 90 | } 91 | }) 92 | } 93 | 94 | } 95 | -------------------------------------------------------------------------------- /internal/tools/mysqlexecutesql/mysqlexecutesql_test.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 mysqlexecutesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/mysqlexecutesql" 25 | ) 26 | 27 | func TestParseFromYamlExecuteSql(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: mysql-execute-sql 43 | source: my-instance 44 | description: some description 45 | authRequired: 46 | - my-google-auth-service 47 | - other-auth-service 48 | `, 49 | want: server.ToolConfigs{ 50 | "example_tool": mysqlexecutesql.Config{ 51 | Name: "example_tool", 52 | Kind: mysqlexecutesql.ToolKind, 53 | Source: "my-instance", 54 | Description: "some description", 55 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 56 | }, 57 | }, 58 | }, 59 | } 60 | for _, tc := range tcs { 61 | t.Run(tc.desc, func(t *testing.T) { 62 | got := struct { 63 | Tools server.ToolConfigs `yaml:"tools"` 64 | }{} 65 | // Parse contents 66 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 67 | if err != nil { 68 | t.Fatalf("unable to unmarshal: %s", err) 69 | } 70 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 71 | t.Fatalf("incorrect parse: diff %v", diff) 72 | } 73 | }) 74 | } 75 | 76 | } 77 | -------------------------------------------------------------------------------- /internal/tools/mysqlsql/mysqlsql_test.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 mysqlsql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/mysqlsql" 26 | ) 27 | 28 | func TestParseFromYamlMySQL(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: mysql-sql 44 | source: my-mysql-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | authRequired: 49 | - my-google-auth-service 50 | - other-auth-service 51 | parameters: 52 | - name: country 53 | type: string 54 | description: some description 55 | authServices: 56 | - name: my-google-auth-service 57 | field: user_id 58 | - name: other-auth-service 59 | field: user_id 60 | `, 61 | want: server.ToolConfigs{ 62 | "example_tool": mysqlsql.Config{ 63 | Name: "example_tool", 64 | Kind: mysqlsql.ToolKind, 65 | Source: "my-mysql-instance", 66 | Description: "some description", 67 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 68 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 69 | Parameters: []tools.Parameter{ 70 | tools.NewStringParameterWithAuth("country", "some description", 71 | []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"}, 72 | {Name: "other-auth-service", Field: "user_id"}}), 73 | }, 74 | }, 75 | }, 76 | }, 77 | } 78 | for _, tc := range tcs { 79 | t.Run(tc.desc, func(t *testing.T) { 80 | got := struct { 81 | Tools server.ToolConfigs `yaml:"tools"` 82 | }{} 83 | // Parse contents 84 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 85 | if err != nil { 86 | t.Fatalf("unable to unmarshal: %s", err) 87 | } 88 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 89 | t.Fatalf("incorrect parse: diff %v", diff) 90 | } 91 | }) 92 | } 93 | 94 | } 95 | -------------------------------------------------------------------------------- /internal/tools/neo4j/neo4j_test.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 neo4j_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/neo4j" 26 | ) 27 | 28 | func TestParseFromYamlNeo4j(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: neo4j-cypher 44 | source: my-neo4j-instance 45 | description: some tool description 46 | authRequired: 47 | - my-google-auth-service 48 | - other-auth-service 49 | statement: | 50 | MATCH (c:Country) WHERE c.name = $country RETURN c.id as id; 51 | parameters: 52 | - name: country 53 | type: string 54 | description: country parameter description 55 | `, 56 | want: server.ToolConfigs{ 57 | "example_tool": neo4j.Config{ 58 | Name: "example_tool", 59 | Kind: neo4j.ToolKind, 60 | Source: "my-neo4j-instance", 61 | Description: "some tool description", 62 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 63 | Statement: "MATCH (c:Country) WHERE c.name = $country RETURN c.id as id;\n", 64 | Parameters: []tools.Parameter{ 65 | tools.NewStringParameter("country", "country parameter description"), 66 | }, 67 | }, 68 | }, 69 | }, 70 | } 71 | for _, tc := range tcs { 72 | t.Run(tc.desc, func(t *testing.T) { 73 | got := struct { 74 | Tools server.ToolConfigs `yaml:"tools"` 75 | }{} 76 | // Parse contents 77 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 78 | if err != nil { 79 | t.Fatalf("unable to unmarshal: %s", err) 80 | } 81 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 82 | t.Fatalf("incorrect parse: diff %v", diff) 83 | } 84 | }) 85 | } 86 | 87 | } 88 | -------------------------------------------------------------------------------- /internal/tools/postgresexecutesql/postgresexecutesql_test.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 postgresexecutesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/postgresexecutesql" 25 | ) 26 | 27 | func TestParseFromYamlExecuteSql(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: postgres-execute-sql 43 | source: my-instance 44 | description: some description 45 | authRequired: 46 | - my-google-auth-service 47 | - other-auth-service 48 | `, 49 | want: server.ToolConfigs{ 50 | "example_tool": postgresexecutesql.Config{ 51 | Name: "example_tool", 52 | Kind: postgresexecutesql.ToolKind, 53 | Source: "my-instance", 54 | Description: "some description", 55 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 56 | }, 57 | }, 58 | }, 59 | } 60 | for _, tc := range tcs { 61 | t.Run(tc.desc, func(t *testing.T) { 62 | got := struct { 63 | Tools server.ToolConfigs `yaml:"tools"` 64 | }{} 65 | // Parse contents 66 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 67 | if err != nil { 68 | t.Fatalf("unable to unmarshal: %s", err) 69 | } 70 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 71 | t.Fatalf("incorrect parse: diff %v", diff) 72 | } 73 | }) 74 | } 75 | 76 | } 77 | -------------------------------------------------------------------------------- /internal/tools/postgressql/postgressql_test.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 postgressql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/postgressql" 26 | ) 27 | 28 | func TestParseFromYamlPostgres(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: postgres-sql 44 | source: my-pg-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | authRequired: 49 | - my-google-auth-service 50 | - other-auth-service 51 | parameters: 52 | - name: country 53 | type: string 54 | description: some description 55 | authServices: 56 | - name: my-google-auth-service 57 | field: user_id 58 | - name: other-auth-service 59 | field: user_id 60 | `, 61 | want: server.ToolConfigs{ 62 | "example_tool": postgressql.Config{ 63 | Name: "example_tool", 64 | Kind: postgressql.ToolKind, 65 | Source: "my-pg-instance", 66 | Description: "some description", 67 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 68 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 69 | Parameters: []tools.Parameter{ 70 | tools.NewStringParameterWithAuth("country", "some description", 71 | []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"}, 72 | {Name: "other-auth-service", Field: "user_id"}}), 73 | }, 74 | }, 75 | }, 76 | }, 77 | } 78 | for _, tc := range tcs { 79 | t.Run(tc.desc, func(t *testing.T) { 80 | got := struct { 81 | Tools server.ToolConfigs `yaml:"tools"` 82 | }{} 83 | // Parse contents 84 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 85 | if err != nil { 86 | t.Fatalf("unable to unmarshal: %s", err) 87 | } 88 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 89 | t.Fatalf("incorrect parse: diff %v", diff) 90 | } 91 | }) 92 | } 93 | 94 | } 95 | -------------------------------------------------------------------------------- /internal/tools/spannerexecutesql/spannerexecutesql_test.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 spannerexecutesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools/spannerexecutesql" 25 | ) 26 | 27 | func TestParseFromYamlExecuteSql(t *testing.T) { 28 | ctx, err := testutils.ContextWithNewLogger() 29 | if err != nil { 30 | t.Fatalf("unexpected error: %s", err) 31 | } 32 | tcs := []struct { 33 | desc string 34 | in string 35 | want server.ToolConfigs 36 | }{ 37 | { 38 | desc: "basic example", 39 | in: ` 40 | tools: 41 | example_tool: 42 | kind: spanner-execute-sql 43 | source: my-spanner-instance 44 | description: some description 45 | `, 46 | want: server.ToolConfigs{ 47 | "example_tool": spannerexecutesql.Config{ 48 | Name: "example_tool", 49 | Kind: spannerexecutesql.ToolKind, 50 | Source: "my-spanner-instance", 51 | Description: "some description", 52 | AuthRequired: []string{}, 53 | ReadOnly: false, 54 | }, 55 | }, 56 | }, 57 | { 58 | desc: "read only set to true", 59 | in: ` 60 | tools: 61 | example_tool: 62 | kind: spanner-execute-sql 63 | source: my-spanner-instance 64 | description: some description 65 | readOnly: true 66 | `, 67 | want: server.ToolConfigs{ 68 | "example_tool": spannerexecutesql.Config{ 69 | Name: "example_tool", 70 | Kind: spannerexecutesql.ToolKind, 71 | Source: "my-spanner-instance", 72 | Description: "some description", 73 | AuthRequired: []string{}, 74 | ReadOnly: true, 75 | }, 76 | }, 77 | }, 78 | } 79 | for _, tc := range tcs { 80 | t.Run(tc.desc, func(t *testing.T) { 81 | got := struct { 82 | Tools server.ToolConfigs `yaml:"tools"` 83 | }{} 84 | // Parse contents 85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 86 | if err != nil { 87 | t.Fatalf("unable to unmarshal: %s", err) 88 | } 89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 90 | t.Fatalf("incorrect parse: diff %v", diff) 91 | } 92 | }) 93 | } 94 | 95 | } 96 | -------------------------------------------------------------------------------- /internal/tools/sqlitesql/sqlitesql_test.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 sqlitesql_test 16 | 17 | import ( 18 | "testing" 19 | 20 | yaml "github.com/goccy/go-yaml" 21 | "github.com/google/go-cmp/cmp" 22 | "github.com/googleapis/genai-toolbox/internal/server" 23 | "github.com/googleapis/genai-toolbox/internal/testutils" 24 | "github.com/googleapis/genai-toolbox/internal/tools" 25 | "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql" 26 | ) 27 | 28 | func TestParseFromYamlSQLite(t *testing.T) { 29 | ctx, err := testutils.ContextWithNewLogger() 30 | if err != nil { 31 | t.Fatalf("unexpected error: %s", err) 32 | } 33 | tcs := []struct { 34 | desc string 35 | in string 36 | want server.ToolConfigs 37 | }{ 38 | { 39 | desc: "basic example", 40 | in: ` 41 | tools: 42 | example_tool: 43 | kind: sqlite-sql 44 | source: my-sqlite-instance 45 | description: some description 46 | statement: | 47 | SELECT * FROM SQL_STATEMENT; 48 | authRequired: 49 | - my-google-auth-service 50 | - other-auth-service 51 | parameters: 52 | - name: country 53 | type: string 54 | description: some description 55 | authServices: 56 | - name: my-google-auth-service 57 | field: user_id 58 | - name: other-auth-service 59 | field: user_id 60 | `, 61 | want: server.ToolConfigs{ 62 | "example_tool": sqlitesql.Config{ 63 | Name: "example_tool", 64 | Kind: sqlitesql.ToolKind, 65 | Source: "my-sqlite-instance", 66 | Description: "some description", 67 | Statement: "SELECT * FROM SQL_STATEMENT;\n", 68 | AuthRequired: []string{"my-google-auth-service", "other-auth-service"}, 69 | Parameters: []tools.Parameter{ 70 | tools.NewStringParameterWithAuth("country", "some description", 71 | []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"}, 72 | {Name: "other-auth-service", Field: "user_id"}}), 73 | }, 74 | }, 75 | }, 76 | }, 77 | } 78 | for _, tc := range tcs { 79 | t.Run(tc.desc, func(t *testing.T) { 80 | got := struct { 81 | Tools server.ToolConfigs `yaml:"tools"` 82 | }{} 83 | // Parse contents 84 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) 85 | if err != nil { 86 | t.Fatalf("unable to unmarshal: %s", err) 87 | } 88 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" { 89 | t.Fatalf("incorrect parse: diff %v", diff) 90 | } 91 | }) 92 | } 93 | 94 | } 95 | -------------------------------------------------------------------------------- /internal/tools/tools.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 tools 16 | 17 | import ( 18 | "context" 19 | "slices" 20 | 21 | "github.com/googleapis/genai-toolbox/internal/sources" 22 | ) 23 | 24 | type ToolConfig interface { 25 | ToolConfigKind() string 26 | Initialize(map[string]sources.Source) (Tool, error) 27 | } 28 | 29 | type Tool interface { 30 | Invoke(context.Context, ParamValues) ([]any, error) 31 | ParseParams(map[string]any, map[string]map[string]any) (ParamValues, error) 32 | Manifest() Manifest 33 | McpManifest() McpManifest 34 | Authorized([]string) bool 35 | } 36 | 37 | // Manifest is the representation of tools sent to Client SDKs. 38 | type Manifest struct { 39 | Description string `json:"description"` 40 | Parameters []ParameterManifest `json:"parameters"` 41 | AuthRequired []string `json:"authRequired"` 42 | } 43 | 44 | // Definition for a tool the MCP client can call. 45 | type McpManifest struct { 46 | // The name of the tool. 47 | Name string `json:"name"` 48 | // A human-readable description of the tool. 49 | Description string `json:"description,omitempty"` 50 | // A JSON Schema object defining the expected parameters for the tool. 51 | InputSchema McpToolsSchema `json:"inputSchema,omitempty"` 52 | } 53 | 54 | // Helper function that returns if a tool invocation request is authorized 55 | func IsAuthorized(authRequiredSources []string, verifiedAuthServices []string) bool { 56 | if len(authRequiredSources) == 0 { 57 | // no authorization requirement 58 | return true 59 | } 60 | for _, a := range authRequiredSources { 61 | if slices.Contains(verifiedAuthServices, a) { 62 | return true 63 | } 64 | } 65 | return false 66 | } 67 | -------------------------------------------------------------------------------- /internal/tools/toolsets.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 tools 16 | 17 | import ( 18 | "fmt" 19 | ) 20 | 21 | type ToolsetConfig struct { 22 | Name string `yaml:"name"` 23 | ToolNames []string `yaml:",inline"` 24 | } 25 | 26 | type Toolset struct { 27 | Name string `yaml:"name"` 28 | Tools []*Tool `yaml:",inline"` 29 | Manifest ToolsetManifest `yaml:",inline"` 30 | McpManifest []McpManifest `yaml:",inline"` 31 | } 32 | 33 | type ToolsetManifest struct { 34 | ServerVersion string `json:"serverVersion"` 35 | ToolsManifest map[string]Manifest `json:"tools"` 36 | } 37 | 38 | func (t ToolsetConfig) Initialize(serverVersion string, toolsMap map[string]Tool) (Toolset, error) { 39 | // finish toolset setup 40 | // Check each declared tool name exists 41 | var toolset Toolset 42 | toolset.Name = t.Name 43 | if !IsValidName(toolset.Name) { 44 | return toolset, fmt.Errorf("invalid toolset name: %s", t) 45 | } 46 | toolset.Tools = make([]*Tool, len(t.ToolNames)) 47 | toolset.Manifest = ToolsetManifest{ 48 | ServerVersion: serverVersion, 49 | ToolsManifest: make(map[string]Manifest), 50 | } 51 | for _, toolName := range t.ToolNames { 52 | tool, ok := toolsMap[toolName] 53 | if !ok { 54 | return toolset, fmt.Errorf("tool does not exist: %s", t) 55 | } 56 | toolset.Tools = append(toolset.Tools, &tool) 57 | toolset.Manifest.ToolsManifest[toolName] = tool.Manifest() 58 | toolset.McpManifest = append(toolset.McpManifest, tool.McpManifest()) 59 | } 60 | 61 | return toolset, nil 62 | } 63 | -------------------------------------------------------------------------------- /internal/util/util.go: -------------------------------------------------------------------------------- 1 | // Copyright 2025 Google LLC 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | package util 15 | 16 | import ( 17 | "bytes" 18 | "context" 19 | "fmt" 20 | 21 | "github.com/go-playground/validator/v10" 22 | yaml "github.com/goccy/go-yaml" 23 | "github.com/googleapis/genai-toolbox/internal/log" 24 | ) 25 | 26 | var _ yaml.InterfaceUnmarshalerContext = &DelayedUnmarshaler{} 27 | 28 | // DelayedUnmarshaler is struct that saves the provided unmarshal function 29 | // passed to UnmarshalYAML so it can be re-used later once the target interface 30 | // is known. 31 | type DelayedUnmarshaler struct { 32 | unmarshal func(interface{}) error 33 | } 34 | 35 | func (d *DelayedUnmarshaler) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error { 36 | d.unmarshal = unmarshal 37 | return nil 38 | } 39 | 40 | func (d *DelayedUnmarshaler) Unmarshal(v interface{}) error { 41 | if d.unmarshal == nil { 42 | return fmt.Errorf("nothing to unmarshal") 43 | } 44 | return d.unmarshal(v) 45 | } 46 | 47 | type contextKey string 48 | 49 | // userAgentKey is the key used to store userAgent within context 50 | const userAgentKey contextKey = "userAgent" 51 | 52 | // WithUserAgent adds a user agent into the context as a value 53 | func WithUserAgent(ctx context.Context, versionString string) context.Context { 54 | userAgent := "genai-toolbox/" + versionString 55 | return context.WithValue(ctx, userAgentKey, userAgent) 56 | } 57 | 58 | // UserAgentFromContext retrieves the user agent or return an error 59 | func UserAgentFromContext(ctx context.Context) (string, error) { 60 | if ua := ctx.Value(userAgentKey); ua != nil { 61 | return ua.(string), nil 62 | } else { 63 | return "", fmt.Errorf("unable to retrieve user agent") 64 | } 65 | } 66 | 67 | func NewStrictDecoder(v interface{}) (*yaml.Decoder, error) { 68 | b, err := yaml.Marshal(v) 69 | if err != nil { 70 | return nil, fmt.Errorf("fail to marshal %q: %w", v, err) 71 | } 72 | 73 | dec := yaml.NewDecoder( 74 | bytes.NewReader(b), 75 | yaml.Strict(), 76 | yaml.Validator(validator.New()), 77 | ) 78 | return dec, nil 79 | } 80 | 81 | // loggerKey is the key used to store logger within context 82 | const loggerKey contextKey = "logger" 83 | 84 | // WithLogger adds a logger into the context as a value 85 | func WithLogger(ctx context.Context, logger log.Logger) context.Context { 86 | return context.WithValue(ctx, loggerKey, logger) 87 | } 88 | 89 | // LoggerFromContext retreives the logger or return an error 90 | func LoggerFromContext(ctx context.Context) (log.Logger, error) { 91 | if logger, ok := ctx.Value(loggerKey).(log.Logger); ok { 92 | return logger, nil 93 | } 94 | return nil, fmt.Errorf("unable to retrieve logger") 95 | } 96 | -------------------------------------------------------------------------------- /logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/googleapis/genai-toolbox/4700dd363c1559193d4e59dfd5897e11e445e31d/logo.png -------------------------------------------------------------------------------- /main.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 main 16 | 17 | import "github.com/googleapis/genai-toolbox/cmd" 18 | 19 | func main() { 20 | cmd.Execute() 21 | } 22 | -------------------------------------------------------------------------------- /tests/auth.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 tests 16 | 17 | import ( 18 | "context" 19 | "os" 20 | "os/exec" 21 | "strings" 22 | 23 | "google.golang.org/api/idtoken" 24 | ) 25 | 26 | var SERVICE_ACCOUNT_EMAIL = os.Getenv("SERVICE_ACCOUNT_EMAIL") 27 | var ClientId = os.Getenv("CLIENT_ID") 28 | 29 | // GetGoogleIdToken retrieve and return the Google ID token 30 | func GetGoogleIdToken(audience string) (string, error) { 31 | // For local testing - use gcloud command to print personal ID token 32 | cmd := exec.Command("gcloud", "auth", "print-identity-token") 33 | output, err := cmd.Output() 34 | if err == nil { 35 | return strings.TrimSpace(string(output)), nil 36 | } 37 | // For Cloud Build testing - retrieve ID token from GCE metadata server 38 | ts, err := idtoken.NewTokenSource(context.Background(), ClientId) 39 | if err != nil { 40 | return "", err 41 | } 42 | token, err := ts.Token() 43 | if err != nil { 44 | return "", err 45 | } 46 | return token.AccessToken, nil 47 | } 48 | -------------------------------------------------------------------------------- /tests/source.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 tests contains end to end tests meant to verify the Toolbox Server 16 | // works as expected when executed as a binary. 17 | 18 | package tests 19 | 20 | import ( 21 | "context" 22 | "fmt" 23 | "regexp" 24 | "strings" 25 | "testing" 26 | "time" 27 | 28 | "cloud.google.com/go/cloudsqlconn" 29 | ) 30 | 31 | // RunSourceConnection test for source connection 32 | func RunSourceConnectionTest(t *testing.T, sourceConfig map[string]any, toolKind string) error { 33 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute) 34 | defer cancel() 35 | 36 | var args []string 37 | 38 | // Write config into a file and pass it to command 39 | toolsFile := map[string]any{ 40 | "sources": map[string]any{ 41 | "my-instance": sourceConfig, 42 | }, 43 | "tools": map[string]any{ 44 | "my-simple-tool": map[string]any{ 45 | "kind": toolKind, 46 | "source": "my-instance", 47 | "description": "Simple tool to test end to end functionality.", 48 | "statement": "SELECT 1;", 49 | }, 50 | }, 51 | } 52 | cmd, cleanup, err := StartCmd(ctx, toolsFile, args...) 53 | if err != nil { 54 | return fmt.Errorf("command initialization returned an error: %s", err) 55 | } 56 | defer cleanup() 57 | 58 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) 59 | defer cancel() 60 | out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`)) 61 | if err != nil { 62 | t.Logf("toolbox command logs: \n%s", out) 63 | return fmt.Errorf("toolbox didn't start successfully: %s", err) 64 | } 65 | return nil 66 | } 67 | 68 | // GetCloudSQLDialOpts returns cloud sql connector's dial option for ip type. 69 | func GetCloudSQLDialOpts(ipType string) ([]cloudsqlconn.DialOption, error) { 70 | switch strings.ToLower(ipType) { 71 | case "private": 72 | return []cloudsqlconn.DialOption{cloudsqlconn.WithPrivateIP()}, nil 73 | case "public": 74 | return []cloudsqlconn.DialOption{cloudsqlconn.WithPublicIP()}, nil 75 | default: 76 | return nil, fmt.Errorf("invalid ipType %s", ipType) 77 | } 78 | } 79 | --------------------------------------------------------------------------------