Compare commits

...

14 Commits

Author SHA1 Message Date
rahulpinto19
4cfcdcba33 docs: optimize link checker 2026-02-16 13:34:17 +00:00
Anubhav Dhawan
41afeafaae refactor: remove explicit Protocol import and argument from ToolboxClient initialization example in deployment documentation (#2480)
We now use `Protocol.MCP` as default and would be deprecating
`Protocol.TOOLBOX` going forward. This PR makes the guide simpler and
more future-proof by removing passing `Protocol.TOOLBOX`.
2026-02-16 16:49:11 +05:30
Twisha Bansal
598b56f478 refactor: remove golden file from quickstart tests (#2456)
## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [x] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-02-16 13:15:19 +05:30
Averi Kitsch
a678886ee3 docs: update prebuilt tools ref (#2459)
## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>
2026-02-13 18:57:55 +00:00
Averi Kitsch
6602abd059 docs: update diagram (#2461)
## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>
2026-02-13 10:25:34 -08:00
Binh Tran
62b830987d fix: Deflake alloydb omni (#2431)
## Description

Improve logic to check that database is up.

**IMPORTANT** DO NOT MERGE until I have reverted
f7d7d9e708

## PR Checklist

- [x] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [x] https://github.com/googleapis/genai-toolbox/issues/2422
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #2422

Co-authored-by: Averi Kitsch <akitsch@google.com>
2026-02-13 09:55:53 -08:00
Yuan Teoh
195767bdcd refactor: refactor subcommands and move tests to its own package (#2439)
This PR refactors the command structure to decouple subcommands from the
root command, improving modularity and testability.

**Key Changes:**
- Move `internal/cli` to `cmd/internal`. Being in a `internal` folder,
other packages outside of `cmd` will not be able to import them.
- Encapsulated I/O: Introduced a new IOStreams struct to standardize in,
out, and errOut handling.
- Shared Dependencies: Extracted shared fields (including IOStreams,
Logger, ServerConfig, and various Tools paths) from the root `Command`
into a new `ToolboxOptions` struct. This also includes moving
`cmd/options.go` to be part of `ToolboxOptions`.
- Logic Migration: Moved setup logic, such as `Setup()` and
`LoadConfig()`, into `ToolboxOptions`. Removing the need to import
`rootCmd` to subcommands.
- Package Reorganization:
- Relocated PersistentFlag and ToolsFiles to the cli package to remove
base command dependencies. This removes dependencies on the base
command, allowing subcommands to consume these utilities independently.
- Moved all side-effect registration to the `cmd/internal` package,
enabling other packages to import it safely for unit tests.

**Testing Improvements:**
- Subcommand packages can now be tested in isolation without relying on
the base command package.
- Added `TestSubcommandWiring()` to the base command tests to verify
proper subcommand registration.
2026-02-13 02:28:58 +00:00
release-please[bot]
c5524d32f5 chore(main): release 0.27.0 (#2363)
🤖 I have created a release *beep* *boop*
---


##
[0.27.0](https://github.com/googleapis/genai-toolbox/compare/v0.26.0...v0.27.0)
(2026-02-12)


### ⚠ BREAKING CHANGES

* Update configuration file v2
([#2369](293c1d6889))
* Update/add detailed telemetry for mcp endpoint compliant with OTEL
semantic convention
([#1987](https://github.com/googleapis/genai-toolbox/issues/1987))
([478a0bd](478a0bdb59))

### Features

* **cli/invoke:** Add support for direct tool invocation from CLI
([#2353](https://github.com/googleapis/genai-toolbox/issues/2353))
([6e49ba4](6e49ba436e))
* **cli/skills:** Add support for generating agent skills from toolset
([#2392](https://github.com/googleapis/genai-toolbox/issues/2392))
([80ef346](80ef346214))
* **cloud-logging-admin:** Add source, tools, integration test and docs
([#2137](https://github.com/googleapis/genai-toolbox/issues/2137))
([252fc30](252fc3091a))
* **cockroachdb:** Add CockroachDB integration with cockroach-go
([#2006](https://github.com/googleapis/genai-toolbox/issues/2006))
([1fdd99a](1fdd99a9b6))
* **prebuiltconfigs/alloydb-omni:** Implement Alloydb omni dataplane
tools ([#2340](https://github.com/googleapis/genai-toolbox/issues/2340))
([e995349](e995349ea0))
* **server:** Add Tool call error categories
([#2387](https://github.com/googleapis/genai-toolbox/issues/2387))
([32cb4db](32cb4db712))
* **tools/looker:** support `looker-validate-project` tool
([#2430](https://github.com/googleapis/genai-toolbox/issues/2430))
([a15a128](a15a12873f))



### Bug Fixes

* **dataplex:** Capture GCP HTTP errors in MCP Toolbox
([#2347](https://github.com/googleapis/genai-toolbox/issues/2347))
([1d7c498](1d7c498116))
* **sources/cockroachdb:** Update kind to type
([#2465](https://github.com/googleapis/genai-toolbox/issues/2465))
([2d341ac](2d341acaa6))
* Surface Dataplex API errors in MCP results
([#2347](1d7c498116))


---
This PR was generated with [Release
Please](https://github.com/googleapis/release-please). See
[documentation](https://github.com/googleapis/release-please#release-please).

---------

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>
Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
2026-02-12 18:03:05 -08:00
Yuan Teoh
e1739abd81 chore: release 0.27.0 (#2467)
Release-As: 0.27.0
2026-02-13 01:16:23 +00:00
Parth Ajmera
478a0bdb59 feat: update/add detailed telemetry for stdio and http mcp transports (#1987)
## Description

This PR adds consistent and actionable telemetry for MCP sessions across
HTTP and STDIO transports, enabling quick visibility into toolset
discovery and tool invocation activity with minimal setup.

## PR Checklist

- [x] Make sure you reviewed
[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [x] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

---------

Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
2026-02-13 00:29:53 +00:00
Yuan Teoh
2d341acaa6 fix(sources/cockroachdb): update kind to type (#2465)
Fix failing integration test, clean up source code from `kind` to
`type`.
2026-02-12 23:59:55 +00:00
Huan Chen
f032389a07 chore(tools/bigquery&looker-conversational-analytics): add X-Goog-API-Client header (#2462)
## Description

Add X-Goog-API-Client header.

Change entry point to v1beta.

## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>
2026-02-12 12:55:12 -08:00
Wenxin Du
32610d71a3 refactor(server): standardize tool error handling and status code mapping (#2402)
- Detect errors and return error codes accordingly in the tool call
handler functions.
- Replace the old `util.ErrUnauthorized` with the new Toolbox error
type.
2026-02-12 11:42:27 -05:00
Wenxin Du
32cb4db712 feat(server): add Tool call error categories (#2387)
Create Agent vs Server error types to distinguish between the two types.

---------

Co-authored-by: Averi Kitsch <akitsch@google.com>
2026-02-12 13:34:28 +00:00
294 changed files with 7784 additions and 6050 deletions

View File

@@ -15,6 +15,9 @@ name: Link Checker
on:
pull_request:
permissions:
contents: read
pull-requests: write
jobs:
@@ -23,8 +26,33 @@ jobs:
steps:
- name: Checkout Repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
with:
fetch-depth: 0
- name: Identify Changed Files
id: changed-files
shell: bash
run: |
git fetch origin main
CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRT origin/main...HEAD)
if [ -z "$CHANGED_FILES" ]; then
echo "No markdown files changed. Skipping checks."
echo "HAS_CHANGES=false" >> $GITHUB_ENV
else
echo "--- Changed Files to Scan ---"
echo "$CHANGED_FILES"
echo "-----------------------------"
# Flatten newlines to spaces for the args list
FILES_FLAT=$(echo "$CHANGED_FILES" | tr '\n' ' ')
echo "CHECK_FILES=$FILES_FLAT" >> $GITHUB_ENV
echo "HAS_CHANGES=true" >> $GITHUB_ENV
fi
- name: Restore lychee cache
if: env.HAS_CHANGES == 'true'
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5
with:
path: .lycheecache
@@ -33,6 +61,7 @@ jobs:
- name: Link Checker
id: lychee-check
if: env.HAS_CHANGES == 'true'
uses: lycheeverse/lychee-action@a8c4c7cb88f0c7386610c35eb25108e448569cb0 # v2
continue-on-error: true
with:
@@ -42,8 +71,7 @@ jobs:
--cache
--max-cache-age 1d
--exclude '^neo4j\+.*' --exclude '^bolt://.*'
README.md
docs/
${{ env.CHECK_FILES }}
output: lychee-report.md
format: markdown
fail: true
@@ -51,18 +79,44 @@ jobs:
debug: false
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Find comment
if: env.HAS_CHANGES == 'true' && steps.lychee-check.outcome == 'failure'
uses: peter-evans/find-comment@v3
id: find-comment
with:
issue-number: ${{ github.event.pull_request.number }}
comment-author: 'github-actions[bot]'
body-includes: "## Link Resolution Note"
- name: Prepare Report
if: env.HAS_CHANGES == 'true' && steps.lychee-check.outcome == 'failure'
run: |
# Create a new file 'full-report.md'
echo "## Link Resolution Note" > full-report.md
echo "Local links and directory changes work differently on GitHub than on the docsite." >> full-report.md
echo "You must ensure fixes pass the **GitHub check** and also work with **\`hugo server\`**." >> full-report.md
echo "See [Link Checking and Fixing with Lychee](https://github.com/googleapis/genai-toolbox/blob/main/DEVELOPER.md#link-checking-and-fixing-with-lychee) for more details." >> full-report.md
echo "---" >> full-report.md
echo "### Broken Links Found" >> full-report.md
# Clean the report (remove redirects) AND append it to the file
sed -E '/(Redirect|Redirects per input)/d' lychee-report.md >> full-report.md
- name: Create PR Comment
if: env.HAS_CHANGES == 'true' && steps.lychee-check.outcome == 'failure'
uses: peter-evans/create-or-update-comment@v4
with:
comment-id: ${{ steps.find-comment.outputs.comment-id }}
issue-number: ${{ github.event.pull_request.number }}
body-path: full-report.md
edit-mode: replace
- name: Display Failure Report
# Run this ONLY if the link checker failed
if: steps.lychee-check.outcome == 'failure'
run: |
echo "## Link Resolution Note" >> $GITHUB_STEP_SUMMARY
echo "Local links and directory changes work differently on GitHub than on the docsite." >> $GITHUB_STEP_SUMMARY
echo "You must ensure fixes pass the **GitHub check** and also work with **\`hugo server\`**." >> $GITHUB_STEP_SUMMARY
echo "See [Link Checking and Fixing with Lychee](https://github.com/googleapis/genai-toolbox/blob/main/DEVELOPER.md#link-checking-and-fixing-with-lychee) for more details." >> $GITHUB_STEP_SUMMARY
echo "---" >> $GITHUB_STEP_SUMMARY
echo "### Broken Links Found" >> $GITHUB_STEP_SUMMARY
cat ./lychee-report.md >> $GITHUB_STEP_SUMMARY
# We can now simply output the prepared file to the job summary
cat full-report.md >> $GITHUB_STEP_SUMMARY
# Fail the job
exit 1

View File

@@ -51,6 +51,10 @@ ignoreFiles = ["quickstart/shared", "quickstart/python", "quickstart/js", "quick
# Add a new version block here before every release
# The order of versions in this file is mirrored into the dropdown
[[params.versions]]
version = "v0.27.0"
url = "https://googleapis.github.io/genai-toolbox/v0.27.0/"
[[params.versions]]
version = "v0.26.0"
url = "https://googleapis.github.io/genai-toolbox/v0.26.0/"

View File

@@ -1,5 +1,31 @@
# Changelog
## [0.27.0](https://github.com/googleapis/genai-toolbox/compare/v0.26.0...v0.27.0) (2026-02-12)
### ⚠ BREAKING CHANGES
* Update configuration file v2 ([#2369](https://github.com/googleapis/genai-toolbox/issues/2369))([293c1d6](https://github.com/googleapis/genai-toolbox/commit/293c1d6889c39807855ba5e01d4c13ba2a4c50ce))
* Update/add detailed telemetry for mcp endpoint compliant with OTEL semantic convention ([#1987](https://github.com/googleapis/genai-toolbox/issues/1987)) ([478a0bd](https://github.com/googleapis/genai-toolbox/commit/478a0bdb59288c1213f83862f95a698b4c2c0aab))
### Features
* **cli/invoke:** Add support for direct tool invocation from CLI ([#2353](https://github.com/googleapis/genai-toolbox/issues/2353)) ([6e49ba4](https://github.com/googleapis/genai-toolbox/commit/6e49ba436ef2390c13feaf902b29f5907acffb57))
* **cli/skills:** Add support for generating agent skills from toolset ([#2392](https://github.com/googleapis/genai-toolbox/issues/2392)) ([80ef346](https://github.com/googleapis/genai-toolbox/commit/80ef34621453b77bdf6a6016c354f102a17ada04))
* **cloud-logging-admin:** Add source, tools, integration test and docs ([#2137](https://github.com/googleapis/genai-toolbox/issues/2137)) ([252fc30](https://github.com/googleapis/genai-toolbox/commit/252fc3091af10d25d8d7af7e047b5ac87a5dd041))
* **cockroachdb:** Add CockroachDB integration with cockroach-go ([#2006](https://github.com/googleapis/genai-toolbox/issues/2006)) ([1fdd99a](https://github.com/googleapis/genai-toolbox/commit/1fdd99a9b609a5e906acce414226ff44d75d5975))
* **prebuiltconfigs/alloydb-omni:** Implement Alloydb omni dataplane tools ([#2340](https://github.com/googleapis/genai-toolbox/issues/2340)) ([e995349](https://github.com/googleapis/genai-toolbox/commit/e995349ea0756c700d188b8f04e9459121219f0c))
* **server:** Add Tool call error categories ([#2387](https://github.com/googleapis/genai-toolbox/issues/2387)) ([32cb4db](https://github.com/googleapis/genai-toolbox/commit/32cb4db712d27579c1bf29e61cbd0bed02286c28))
* **tools/looker:** support `looker-validate-project` tool ([#2430](https://github.com/googleapis/genai-toolbox/issues/2430)) ([a15a128](https://github.com/googleapis/genai-toolbox/commit/a15a12873f936b0102aeb9500cc3bcd71bb38c34))
### Bug Fixes
* **dataplex:** Capture GCP HTTP errors in MCP Toolbox ([#2347](https://github.com/googleapis/genai-toolbox/issues/2347)) ([1d7c498](https://github.com/googleapis/genai-toolbox/commit/1d7c4981164c34b4d7bc8edecfd449f57ad11e15))
* **sources/cockroachdb:** Update kind to type ([#2465](https://github.com/googleapis/genai-toolbox/issues/2465)) ([2d341ac](https://github.com/googleapis/genai-toolbox/commit/2d341acaa61c3c1fe908fceee8afbd90fb646d3a))
* Surface Dataplex API errors in MCP results ([#2347](https://github.com/googleapis/genai-toolbox/pull/2347))([1d7c498](https://github.com/googleapis/genai-toolbox/commit/1d7c4981164c34b4d7bc8edecfd449f57ad11e15))
## [0.26.0](https://github.com/googleapis/genai-toolbox/compare/v0.25.0...v0.26.0) (2026-01-22)

View File

@@ -142,7 +142,7 @@ To install Toolbox as a binary:
>
> ```sh
> # see releases page for other versions
> export VERSION=0.26.0
> export VERSION=0.27.0
> curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/linux/amd64/toolbox
> chmod +x toolbox
> ```
@@ -155,7 +155,7 @@ To install Toolbox as a binary:
>
> ```sh
> # see releases page for other versions
> export VERSION=0.26.0
> export VERSION=0.27.0
> curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/darwin/arm64/toolbox
> chmod +x toolbox
> ```
@@ -168,7 +168,7 @@ To install Toolbox as a binary:
>
> ```sh
> # see releases page for other versions
> export VERSION=0.26.0
> export VERSION=0.27.0
> curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/darwin/amd64/toolbox
> chmod +x toolbox
> ```
@@ -181,7 +181,7 @@ To install Toolbox as a binary:
>
> ```cmd
> :: see releases page for other versions
> set VERSION=0.26.0
> set VERSION=0.27.0
> curl -o toolbox.exe "https://storage.googleapis.com/genai-toolbox/v%VERSION%/windows/amd64/toolbox.exe"
> ```
>
@@ -193,7 +193,7 @@ To install Toolbox as a binary:
>
> ```powershell
> # see releases page for other versions
> $VERSION = "0.26.0"
> $VERSION = "0.27.0"
> curl.exe -o toolbox.exe "https://storage.googleapis.com/genai-toolbox/v$VERSION/windows/amd64/toolbox.exe"
> ```
>
@@ -206,7 +206,7 @@ You can also install Toolbox as a container:
```sh
# see releases page for other versions
export VERSION=0.26.0
export VERSION=0.27.0
docker pull us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:$VERSION
```
@@ -230,7 +230,7 @@ To install from source, ensure you have the latest version of
[Go installed](https://go.dev/doc/install), and then run the following command:
```sh
go install github.com/googleapis/genai-toolbox@v0.26.0
go install github.com/googleapis/genai-toolbox@v0.27.0
```
<!-- {x-release-please-end} -->

257
cmd/internal/imports.go Normal file
View File

@@ -0,0 +1,257 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package internal
import (
// Import prompt packages for side effect of registration
_ "github.com/googleapis/genai-toolbox/internal/prompts/custom"
// Import tool packages for side effect of registration
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbcreatecluster"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbcreateinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbcreateuser"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetcluster"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetuser"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistclusters"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistinstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistusers"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbwaitforoperation"
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryanalyzecontribution"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryconversationalanalytics"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigqueryforecast"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerygetdatasetinfo"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerygettableinfo"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerylistdatasetids"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerylisttableids"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysearchcatalog"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigtable"
_ "github.com/googleapis/genai-toolbox/internal/tools/cassandra/cassandracql"
_ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouseexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselistdatabases"
_ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhouselisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/clickhouse/clickhousesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudgda"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcarefhirfetchpage"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcarefhirpatienteverything"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcarefhirpatientsearch"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetdataset"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetdicomstore"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetdicomstoremetrics"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetfhirresource"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetfhirstore"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaregetfhirstoremetrics"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcarelistdicomstores"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcarelistfhirstores"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcareretrieverendereddicominstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaresearchdicominstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaresearchdicomseries"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/cloudhealthcaresearchdicomstudies"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudloggingadmin/cloudloggingadminlistlognames"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudloggingadmin/cloudloggingadminlistresourcetypes"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudloggingadmin/cloudloggingadminquerylogs"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudmonitoring"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcloneinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreatebackup"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreatedatabase"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlcreateusers"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlgetinstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistdatabases"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistinstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlrestorebackup"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlwaitforoperation"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdbexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdblistschemas"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdblisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdbsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataform/dataformcompilelocal"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/tools/elasticsearch/elasticsearchesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreadddocuments"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoredeletedocuments"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetdocuments"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetrules"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorelistcollections"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorequery"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorequerycollection"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreupdatedocument"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorevalidaterules"
_ "github.com/googleapis/genai-toolbox/internal/tools/http"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookeradddashboardelement"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookeradddashboardfilter"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerconversationalanalytics"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercreateprojectfile"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerdeleteprojectfile"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerdevmode"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergenerateembedurl"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiondatabases"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnections"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectionschemas"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontablecolumns"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetconnectiontables"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdashboards"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetdimensions"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetexplores"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetfilters"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetlooks"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetmeasures"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetmodels"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetparameters"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetprojectfile"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetprojectfiles"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookergetprojects"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerhealthanalyze"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerhealthpulse"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerhealthvacuum"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookermakedashboard"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookermakelook"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerquery"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerquerysql"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerqueryurl"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerrundashboard"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerrunlook"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookerupdateprojectfile"
_ "github.com/googleapis/genai-toolbox/internal/tools/looker/lookervalidateproject"
_ "github.com/googleapis/genai-toolbox/internal/tools/mindsdb/mindsdbexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/mindsdb/mindsdbsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbaggregate"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbdeletemany"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbdeleteone"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbfind"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbfindone"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbinsertmany"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbinsertone"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdatemany"
_ "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdateone"
_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqlexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqllisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqlsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlgetqueryplan"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllistactivequeries"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllisttablefragmentation"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllisttablesmissinguniqueindexes"
_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jcypher"
_ "github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jexecutecypher"
_ "github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jschema"
_ "github.com/googleapis/genai-toolbox/internal/tools/oceanbase/oceanbaseexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/oceanbase/oceanbasesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/oracle/oracleexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/oracle/oraclesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresdatabaseoverview"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresgetcolumncardinality"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistactivequeries"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistavailableextensions"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistdatabasestats"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistindexes"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistinstalledextensions"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistlocks"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistpgsettings"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistpublicationtables"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistquerystats"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistroles"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistsequences"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresliststoredprocedure"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttablespaces"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttablestats"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttriggers"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistviews"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslongrunningtransactions"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresreplicationstats"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
_ "github.com/googleapis/genai-toolbox/internal/tools/redis"
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcancelbatch"
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcreatepysparkbatch"
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcreatesparkbatch"
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkgetbatch"
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparklistbatches"
_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoreexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoresql"
_ "github.com/googleapis/genai-toolbox/internal/tools/snowflake/snowflakeexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/snowflake/snowflakesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlistgraphs"
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/trino/trinoexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/trino/trinosql"
_ "github.com/googleapis/genai-toolbox/internal/tools/utility/wait"
_ "github.com/googleapis/genai-toolbox/internal/tools/valkey"
_ "github.com/googleapis/genai-toolbox/internal/tools/yugabytedbsql"
_ "github.com/googleapis/genai-toolbox/internal/sources/alloydbadmin"
_ "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
_ "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
_ "github.com/googleapis/genai-toolbox/internal/sources/bigtable"
_ "github.com/googleapis/genai-toolbox/internal/sources/cassandra"
_ "github.com/googleapis/genai-toolbox/internal/sources/clickhouse"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudgda"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudhealthcare"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudloggingadmin"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudmonitoring"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqladmin"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
_ "github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/sources/elasticsearch"
_ "github.com/googleapis/genai-toolbox/internal/sources/firebird"
_ "github.com/googleapis/genai-toolbox/internal/sources/firestore"
_ "github.com/googleapis/genai-toolbox/internal/sources/http"
_ "github.com/googleapis/genai-toolbox/internal/sources/looker"
_ "github.com/googleapis/genai-toolbox/internal/sources/mindsdb"
_ "github.com/googleapis/genai-toolbox/internal/sources/mongodb"
_ "github.com/googleapis/genai-toolbox/internal/sources/mssql"
_ "github.com/googleapis/genai-toolbox/internal/sources/mysql"
_ "github.com/googleapis/genai-toolbox/internal/sources/neo4j"
_ "github.com/googleapis/genai-toolbox/internal/sources/oceanbase"
_ "github.com/googleapis/genai-toolbox/internal/sources/oracle"
_ "github.com/googleapis/genai-toolbox/internal/sources/postgres"
_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
_ "github.com/googleapis/genai-toolbox/internal/sources/serverlessspark"
_ "github.com/googleapis/genai-toolbox/internal/sources/singlestore"
_ "github.com/googleapis/genai-toolbox/internal/sources/snowflake"
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
_ "github.com/googleapis/genai-toolbox/internal/sources/trino"
_ "github.com/googleapis/genai-toolbox/internal/sources/valkey"
_ "github.com/googleapis/genai-toolbox/internal/sources/yugabytedb"
)

View File

@@ -18,37 +18,15 @@ import (
"context"
"encoding/json"
"fmt"
"io"
"github.com/googleapis/genai-toolbox/internal/log"
"github.com/googleapis/genai-toolbox/cmd/internal"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/server/resources"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/spf13/cobra"
)
// RootCommand defines the interface for required by invoke subcommand.
// This allows subcommands to access shared resources and functionality without
// direct coupling to the root command's implementation.
type RootCommand interface {
// Config returns a copy of the current server configuration.
Config() server.ServerConfig
// Out returns the writer used for standard output.
Out() io.Writer
// LoadConfig loads and merges the configuration from files, folders, and prebuilts.
LoadConfig(ctx context.Context) error
// Setup initializes the runtime environment, including logging and telemetry.
// It returns the updated context and a shutdown function to be called when finished.
Setup(ctx context.Context) (context.Context, func(context.Context) error, error)
// Logger returns the logger instance.
Logger() log.Logger
}
func NewCommand(rootCmd RootCommand) *cobra.Command {
func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
cmd := &cobra.Command{
Use: "invoke <tool-name> [params]",
Short: "Execute a tool directly",
@@ -58,17 +36,17 @@ Example:
toolbox invoke my-tool '{"param1": "value1"}'`,
Args: cobra.MinimumNArgs(1),
RunE: func(c *cobra.Command, args []string) error {
return runInvoke(c, args, rootCmd)
return runInvoke(c, args, opts)
},
}
return cmd
}
func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
func runInvoke(cmd *cobra.Command, args []string, opts *internal.ToolboxOptions) error {
ctx, cancel := context.WithCancel(cmd.Context())
defer cancel()
ctx, shutdown, err := rootCmd.Setup(ctx)
ctx, shutdown, err := opts.Setup(ctx)
if err != nil {
return err
}
@@ -76,16 +54,16 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
_ = shutdown(ctx)
}()
// Load and merge tool configurations
if err := rootCmd.LoadConfig(ctx); err != nil {
_, err = opts.LoadConfig(ctx)
if err != nil {
return err
}
// Initialize Resources
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, rootCmd.Config())
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, opts.Cfg)
if err != nil {
errMsg := fmt.Errorf("failed to initialize resources: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -96,7 +74,7 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
errMsg := fmt.Errorf("tool %q not found", toolName)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -109,7 +87,7 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
if paramsInput != "" {
if err := json.Unmarshal([]byte(paramsInput), &params); err != nil {
errMsg := fmt.Errorf("params must be a valid JSON string: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
}
@@ -117,14 +95,14 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
parsedParams, err := parameters.ParseParams(tool.GetParameters(), params, nil)
if err != nil {
errMsg := fmt.Errorf("invalid parameters: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
parsedParams, err = tool.EmbedParams(ctx, parsedParams, resourceMgr.GetEmbeddingModelMap())
if err != nil {
errMsg := fmt.Errorf("error embedding parameters: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -132,19 +110,19 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
requiresAuth, err := tool.RequiresClientAuthorization(resourceMgr)
if err != nil {
errMsg := fmt.Errorf("failed to check auth requirements: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
if requiresAuth {
errMsg := fmt.Errorf("client authorization is not supported")
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
result, err := tool.Invoke(ctx, resourceMgr, parsedParams, "")
if err != nil {
errMsg := fmt.Errorf("tool execution failed: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -152,10 +130,10 @@ func runInvoke(cmd *cobra.Command, args []string, rootCmd RootCommand) error {
output, err := json.MarshalIndent(result, "", " ")
if err != nil {
errMsg := fmt.Errorf("failed to marshal result: %w", err)
rootCmd.Logger().ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
fmt.Fprintln(rootCmd.Out(), string(output))
fmt.Fprintln(opts.IOStreams.Out, string(output))
return nil
}

View File

@@ -12,16 +12,38 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package cmd
package invoke
import (
"context"
"bytes"
"os"
"path/filepath"
"strings"
"testing"
"github.com/googleapis/genai-toolbox/cmd/internal"
_ "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
_ "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
"github.com/spf13/cobra"
)
func invokeCommand(args []string) (string, error) {
parentCmd := &cobra.Command{Use: "toolbox"}
buf := new(bytes.Buffer)
opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
internal.PersistentFlags(parentCmd, opts)
cmd := NewCommand(opts)
parentCmd.AddCommand(cmd)
parentCmd.SetArgs(args)
err := parentCmd.Execute()
return buf.String(), err
}
func TestInvokeTool(t *testing.T) {
// Create a temporary tools file
tmpDir := t.TempDir()
@@ -86,7 +108,7 @@ tools:
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
_, got, err := invokeCommandWithContext(context.Background(), tc.args)
got, err := invokeCommand(tc.args)
if (err != nil) != tc.wantErr {
t.Fatalf("got error %v, wantErr %v", err, tc.wantErr)
}
@@ -121,7 +143,7 @@ tools:
}
args := []string{"invoke", "bq-tool", "--tools-file", toolsFilePath}
_, _, err := invokeCommandWithContext(context.Background(), args)
_, err := invokeCommand(args)
if err == nil {
t.Fatal("expected error for tool requiring client auth, but got nil")
}

251
cmd/internal/options.go Normal file
View File

@@ -0,0 +1,251 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package internal
import (
"context"
"fmt"
"io"
"os"
"slices"
"strings"
"github.com/googleapis/genai-toolbox/internal/log"
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/telemetry"
"github.com/googleapis/genai-toolbox/internal/util"
)
type IOStreams struct {
In io.Reader
Out io.Writer
ErrOut io.Writer
}
// ToolboxOptions holds dependencies shared by all commands.
type ToolboxOptions struct {
IOStreams IOStreams
Logger log.Logger
Cfg server.ServerConfig
ToolsFile string
ToolsFiles []string
ToolsFolder string
PrebuiltConfigs []string
}
// Option defines a function that modifies the ToolboxOptions struct.
type Option func(*ToolboxOptions)
// NewToolboxOptions creates a new instance with defaults, then applies any
// provided options.
func NewToolboxOptions(opts ...Option) *ToolboxOptions {
o := &ToolboxOptions{
IOStreams: IOStreams{
In: os.Stdin,
Out: os.Stdout,
ErrOut: os.Stderr,
},
}
for _, opt := range opts {
opt(o)
}
return o
}
// Apply allows you to update an EXISTING ToolboxOptions instance.
// This is useful for "late binding".
func (o *ToolboxOptions) Apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithIOStreams updates the IO streams.
func WithIOStreams(out, err io.Writer) Option {
return func(o *ToolboxOptions) {
o.IOStreams.Out = out
o.IOStreams.ErrOut = err
}
}
// Setup create logger and telemetry instrumentations.
func (opts *ToolboxOptions) Setup(ctx context.Context) (context.Context, func(context.Context) error, error) {
// If stdio, set logger's out stream (usually DEBUG and INFO logs) to
// errStream
loggerOut := opts.IOStreams.Out
if opts.Cfg.Stdio {
loggerOut = opts.IOStreams.ErrOut
}
// Handle logger separately from config
logger, err := log.NewLogger(opts.Cfg.LoggingFormat.String(), opts.Cfg.LogLevel.String(), loggerOut, opts.IOStreams.ErrOut)
if err != nil {
return ctx, nil, fmt.Errorf("unable to initialize logger: %w", err)
}
ctx = util.WithLogger(ctx, logger)
opts.Logger = logger
// Set up OpenTelemetry
otelShutdown, err := telemetry.SetupOTel(ctx, opts.Cfg.Version, opts.Cfg.TelemetryOTLP, opts.Cfg.TelemetryGCP, opts.Cfg.TelemetryServiceName)
if err != nil {
errMsg := fmt.Errorf("error setting up OpenTelemetry: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
return ctx, nil, errMsg
}
shutdownFunc := func(ctx context.Context) error {
err := otelShutdown(ctx)
if err != nil {
errMsg := fmt.Errorf("error shutting down OpenTelemetry: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
return err
}
return nil
}
instrumentation, err := telemetry.CreateTelemetryInstrumentation(opts.Cfg.Version)
if err != nil {
errMsg := fmt.Errorf("unable to create telemetry instrumentation: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
return ctx, shutdownFunc, errMsg
}
ctx = util.WithInstrumentation(ctx, instrumentation)
return ctx, shutdownFunc, nil
}
// LoadConfig checks and merge files that should be loaded into the server
func (opts *ToolboxOptions) LoadConfig(ctx context.Context) (bool, error) {
// Determine if Custom Files should be loaded
// Check for explicit custom flags
isCustomConfigured := opts.ToolsFile != "" || len(opts.ToolsFiles) > 0 || opts.ToolsFolder != ""
// Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags)
useDefaultToolsFile := len(opts.PrebuiltConfigs) == 0 && !isCustomConfigured
if useDefaultToolsFile {
opts.ToolsFile = "tools.yaml"
isCustomConfigured = true
}
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return isCustomConfigured, err
}
var allToolsFiles []ToolsFile
// Load Prebuilt Configuration
if len(opts.PrebuiltConfigs) > 0 {
slices.Sort(opts.PrebuiltConfigs)
sourcesList := strings.Join(opts.PrebuiltConfigs, ", ")
logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList)
logger.InfoContext(ctx, logMsg)
for _, configName := range opts.PrebuiltConfigs {
buf, err := prebuiltconfigs.Get(configName)
if err != nil {
logger.ErrorContext(ctx, err.Error())
return isCustomConfigured, err
}
// Parse into ToolsFile struct
parsed, err := parseToolsFile(ctx, buf)
if err != nil {
errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err)
logger.ErrorContext(ctx, errMsg.Error())
return isCustomConfigured, errMsg
}
allToolsFiles = append(allToolsFiles, parsed)
}
}
// Load Custom Configurations
if isCustomConfigured {
// Enforce exclusivity among custom flags (tools-file vs tools-files vs tools-folder)
if (opts.ToolsFile != "" && len(opts.ToolsFiles) > 0) ||
(opts.ToolsFile != "" && opts.ToolsFolder != "") ||
(len(opts.ToolsFiles) > 0 && opts.ToolsFolder != "") {
errMsg := fmt.Errorf("--tools-file, --tools-files, and --tools-folder flags cannot be used simultaneously")
logger.ErrorContext(ctx, errMsg.Error())
return isCustomConfigured, errMsg
}
var customTools ToolsFile
var err error
if len(opts.ToolsFiles) > 0 {
// Use tools-files
logger.InfoContext(ctx, fmt.Sprintf("Loading and merging %d tool configuration files", len(opts.ToolsFiles)))
customTools, err = LoadAndMergeToolsFiles(ctx, opts.ToolsFiles)
} else if opts.ToolsFolder != "" {
// Use tools-folder
logger.InfoContext(ctx, fmt.Sprintf("Loading and merging all YAML files from directory: %s", opts.ToolsFolder))
customTools, err = LoadAndMergeToolsFolder(ctx, opts.ToolsFolder)
} else {
// Use single file (tools-file or default `tools.yaml`)
buf, readFileErr := os.ReadFile(opts.ToolsFile)
if readFileErr != nil {
errMsg := fmt.Errorf("unable to read tool file at %q: %w", opts.ToolsFile, readFileErr)
logger.ErrorContext(ctx, errMsg.Error())
return isCustomConfigured, errMsg
}
customTools, err = parseToolsFile(ctx, buf)
if err != nil {
err = fmt.Errorf("unable to parse tool file at %q: %w", opts.ToolsFile, err)
}
}
if err != nil {
logger.ErrorContext(ctx, err.Error())
return isCustomConfigured, err
}
allToolsFiles = append(allToolsFiles, customTools)
}
// Modify version string based on loaded configurations
if len(opts.PrebuiltConfigs) > 0 {
tag := "prebuilt"
if isCustomConfigured {
tag = "custom"
}
// prebuiltConfigs is already sorted above
for _, configName := range opts.PrebuiltConfigs {
opts.Cfg.Version += fmt.Sprintf("+%s.%s", tag, configName)
}
}
// Merge Everything
// This will error if custom tools collide with prebuilt tools
finalToolsFile, err := mergeToolsFiles(allToolsFiles...)
if err != nil {
logger.ErrorContext(ctx, err.Error())
return isCustomConfigured, err
}
opts.Cfg.SourceConfigs = finalToolsFile.Sources
opts.Cfg.AuthServiceConfigs = finalToolsFile.AuthServices
opts.Cfg.EmbeddingModelConfigs = finalToolsFile.EmbeddingModels
opts.Cfg.ToolConfigs = finalToolsFile.Tools
opts.Cfg.ToolsetConfigs = finalToolsFile.Toolsets
opts.Cfg.PromptConfigs = finalToolsFile.Prompts
return isCustomConfigured, nil
}

View File

@@ -12,57 +12,38 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package cmd
package internal
import (
"errors"
"io"
"testing"
"github.com/spf13/cobra"
)
func TestCommandOptions(t *testing.T) {
func TestToolboxOptions(t *testing.T) {
w := io.Discard
tcs := []struct {
desc string
isValid func(*Command) error
isValid func(*ToolboxOptions) error
option Option
}{
{
desc: "with logger",
isValid: func(c *Command) error {
if c.outStream != w || c.errStream != w {
isValid: func(o *ToolboxOptions) error {
if o.IOStreams.Out != w || o.IOStreams.ErrOut != w {
return errors.New("loggers do not match")
}
return nil
},
option: WithStreams(w, w),
option: WithIOStreams(w, w),
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got, err := invokeProxyWithOption(tc.option)
if err != nil {
t.Fatal(err)
}
got := NewToolboxOptions(tc.option)
if err := tc.isValid(got); err != nil {
t.Errorf("option did not initialize command correctly: %v", err)
}
})
}
}
func invokeProxyWithOption(o Option) (*Command, error) {
c := NewCommand(o)
// Keep the test output quiet
c.SilenceUsage = true
c.SilenceErrors = true
// Disable execute behavior
c.RunE = func(*cobra.Command, []string) error {
return nil
}
err := c.Execute()
return c, err
}

View File

@@ -0,0 +1,46 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package internal
import (
"fmt"
"strings"
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
"github.com/spf13/cobra"
)
// PersistentFlags sets up flags that are available for all commands and
// subcommands
// It is also used to set up persistent flags during subcommand unit tests
func PersistentFlags(parentCmd *cobra.Command, opts *ToolboxOptions) {
persistentFlags := parentCmd.PersistentFlags()
persistentFlags.StringVar(&opts.ToolsFile, "tools-file", "", "File path specifying the tool configuration. Cannot be used with --tools-files, or --tools-folder.")
persistentFlags.StringSliceVar(&opts.ToolsFiles, "tools-files", []string{}, "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --tools-file, or --tools-folder.")
persistentFlags.StringVar(&opts.ToolsFolder, "tools-folder", "", "Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with --tools-file, or --tools-files.")
persistentFlags.Var(&opts.Cfg.LogLevel, "log-level", "Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.")
persistentFlags.Var(&opts.Cfg.LoggingFormat, "logging-format", "Specify logging format to use. Allowed: 'standard' or 'JSON'.")
persistentFlags.BoolVar(&opts.Cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
persistentFlags.StringVar(&opts.Cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')")
persistentFlags.StringVar(&opts.Cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
// Fetch prebuilt tools sources to customize the help description
prebuiltHelp := fmt.Sprintf(
"Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.",
strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
)
persistentFlags.StringSliceVar(&opts.PrebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
}

View File

@@ -22,7 +22,7 @@ import (
"path/filepath"
"sort"
"github.com/googleapis/genai-toolbox/internal/log"
"github.com/googleapis/genai-toolbox/cmd/internal"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/server/resources"
"github.com/googleapis/genai-toolbox/internal/tools"
@@ -30,28 +30,9 @@ import (
"github.com/spf13/cobra"
)
// RootCommand defines the interface for required by skills-generate subcommand.
// This allows subcommands to access shared resources and functionality without
// direct coupling to the root command's implementation.
type RootCommand interface {
// Config returns a copy of the current server configuration.
Config() server.ServerConfig
// LoadConfig loads and merges the configuration from files, folders, and prebuilts.
LoadConfig(ctx context.Context) error
// Setup initializes the runtime environment, including logging and telemetry.
// It returns the updated context and a shutdown function to be called when finished.
Setup(ctx context.Context) (context.Context, func(context.Context) error, error)
// Logger returns the logger instance.
Logger() log.Logger
}
// Command is the command for generating skills.
type Command struct {
// skillsCmd is the command for generating skills.
type skillsCmd struct {
*cobra.Command
rootCmd RootCommand
name string
description string
toolset string
@@ -59,15 +40,13 @@ type Command struct {
}
// NewCommand creates a new Command.
func NewCommand(rootCmd RootCommand) *cobra.Command {
cmd := &Command{
rootCmd: rootCmd,
}
func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
cmd := &skillsCmd{}
cmd.Command = &cobra.Command{
Use: "skills-generate",
Short: "Generate skills from tool configurations",
RunE: func(c *cobra.Command, args []string) error {
return cmd.run(c)
return run(cmd, opts)
},
}
@@ -81,11 +60,11 @@ func NewCommand(rootCmd RootCommand) *cobra.Command {
return cmd.Command
}
func (c *Command) run(cmd *cobra.Command) error {
func run(cmd *skillsCmd, opts *internal.ToolboxOptions) error {
ctx, cancel := context.WithCancel(cmd.Context())
defer cancel()
ctx, shutdown, err := c.rootCmd.Setup(ctx)
ctx, shutdown, err := opts.Setup(ctx)
if err != nil {
return err
}
@@ -93,39 +72,37 @@ func (c *Command) run(cmd *cobra.Command) error {
_ = shutdown(ctx)
}()
logger := c.rootCmd.Logger()
// Load and merge tool configurations
if err := c.rootCmd.LoadConfig(ctx); err != nil {
_, err = opts.LoadConfig(ctx)
if err != nil {
return err
}
if err := os.MkdirAll(c.outputDir, 0755); err != nil {
if err := os.MkdirAll(cmd.outputDir, 0755); err != nil {
errMsg := fmt.Errorf("error creating output directory: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
logger.InfoContext(ctx, fmt.Sprintf("Generating skill '%s'...", c.name))
opts.Logger.InfoContext(ctx, fmt.Sprintf("Generating skill '%s'...", cmd.name))
// Initialize toolbox and collect tools
allTools, err := c.collectTools(ctx)
allTools, err := cmd.collectTools(ctx, opts)
if err != nil {
errMsg := fmt.Errorf("error collecting tools: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
if len(allTools) == 0 {
logger.InfoContext(ctx, "No tools found to generate.")
opts.Logger.InfoContext(ctx, "No tools found to generate.")
return nil
}
// Generate the combined skill directory
skillPath := filepath.Join(c.outputDir, c.name)
skillPath := filepath.Join(cmd.outputDir, cmd.name)
if err := os.MkdirAll(skillPath, 0755); err != nil {
errMsg := fmt.Errorf("error creating skill directory: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -133,7 +110,7 @@ func (c *Command) run(cmd *cobra.Command) error {
assetsPath := filepath.Join(skillPath, "assets")
if err := os.MkdirAll(assetsPath, 0755); err != nil {
errMsg := fmt.Errorf("error creating assets dir: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -141,7 +118,7 @@ func (c *Command) run(cmd *cobra.Command) error {
scriptsPath := filepath.Join(skillPath, "scripts")
if err := os.MkdirAll(scriptsPath, 0755); err != nil {
errMsg := fmt.Errorf("error creating scripts dir: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -154,10 +131,10 @@ func (c *Command) run(cmd *cobra.Command) error {
for _, toolName := range toolNames {
// Generate YAML config in asset directory
minimizedContent, err := generateToolConfigYAML(c.rootCmd.Config(), toolName)
minimizedContent, err := generateToolConfigYAML(opts.Cfg, toolName)
if err != nil {
errMsg := fmt.Errorf("error generating filtered config for %s: %w", toolName, err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
@@ -166,7 +143,7 @@ func (c *Command) run(cmd *cobra.Command) error {
destPath := filepath.Join(assetsPath, specificToolsFileName)
if err := os.WriteFile(destPath, minimizedContent, 0644); err != nil {
errMsg := fmt.Errorf("error writing filtered config for %s: %w", toolName, err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
}
@@ -175,40 +152,40 @@ func (c *Command) run(cmd *cobra.Command) error {
scriptContent, err := generateScriptContent(toolName, specificToolsFileName)
if err != nil {
errMsg := fmt.Errorf("error generating script content for %s: %w", toolName, err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
scriptFilename := filepath.Join(scriptsPath, fmt.Sprintf("%s.js", toolName))
if err := os.WriteFile(scriptFilename, []byte(scriptContent), 0755); err != nil {
errMsg := fmt.Errorf("error writing script %s: %w", scriptFilename, err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
}
// Generate SKILL.md
skillContent, err := generateSkillMarkdown(c.name, c.description, allTools)
skillContent, err := generateSkillMarkdown(cmd.name, cmd.description, allTools)
if err != nil {
errMsg := fmt.Errorf("error generating SKILL.md content: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
skillMdPath := filepath.Join(skillPath, "SKILL.md")
if err := os.WriteFile(skillMdPath, []byte(skillContent), 0644); err != nil {
errMsg := fmt.Errorf("error writing SKILL.md: %w", err)
logger.ErrorContext(ctx, errMsg.Error())
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
logger.InfoContext(ctx, fmt.Sprintf("Successfully generated skill '%s' with %d tools.", c.name, len(allTools)))
opts.Logger.InfoContext(ctx, fmt.Sprintf("Successfully generated skill '%s' with %d tools.", cmd.name, len(allTools)))
return nil
}
func (c *Command) collectTools(ctx context.Context) (map[string]tools.Tool, error) {
func (c *skillsCmd) collectTools(ctx context.Context, opts *internal.ToolboxOptions) (map[string]tools.Tool, error) {
// Initialize Resources
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, c.rootCmd.Config())
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, opts.Cfg)
if err != nil {
return nil, fmt.Errorf("failed to initialize resources: %w", err)
}

View File

@@ -12,17 +12,36 @@
// See the License for the specific language governing permissions and
// limitations under the License.
package cmd
package skills
import (
"context"
"bytes"
"os"
"path/filepath"
"strings"
"testing"
"time"
"github.com/googleapis/genai-toolbox/cmd/internal"
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
"github.com/spf13/cobra"
)
func invokeCommand(args []string) (string, error) {
parentCmd := &cobra.Command{Use: "toolbox"}
buf := new(bytes.Buffer)
opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
internal.PersistentFlags(parentCmd, opts)
cmd := NewCommand(opts)
parentCmd.AddCommand(cmd)
parentCmd.SetArgs(args)
err := parentCmd.Execute()
return buf.String(), err
}
func TestGenerateSkill(t *testing.T) {
// Create a temporary directory for tests
tmpDir := t.TempDir()
@@ -55,10 +74,7 @@ tools:
"--description", "hello tool",
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
_, got, err := invokeCommandWithContext(ctx, args)
got, err := invokeCommand(args)
if err != nil {
t.Fatalf("command failed: %v\nOutput: %s", err, got)
}
@@ -136,7 +152,7 @@ func TestGenerateSkill_NoConfig(t *testing.T) {
"--description", "test",
}
_, _, err := invokeCommandWithContext(context.Background(), args)
_, err := invokeCommand(args)
if err == nil {
t.Fatal("expected command to fail when no configuration is provided and tools.yaml is missing")
}
@@ -170,7 +186,7 @@ func TestGenerateSkill_MissingArguments(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, got, err := invokeCommandWithContext(context.Background(), tt.args)
got, err := invokeCommand(tt.args)
if err == nil {
t.Fatalf("expected command to fail due to missing arguments, but it succeeded\nOutput: %s", got)
}

349
cmd/internal/tools_file.go Normal file
View File

@@ -0,0 +1,349 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package internal
import (
"bytes"
"context"
"fmt"
"io"
"os"
"path/filepath"
"regexp"
"slices"
"strings"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/server"
)
type ToolsFile struct {
Sources server.SourceConfigs `yaml:"sources"`
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
EmbeddingModels server.EmbeddingModelConfigs `yaml:"embeddingModels"`
Tools server.ToolConfigs `yaml:"tools"`
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
Prompts server.PromptConfigs `yaml:"prompts"`
}
// parseEnv replaces environment variables ${ENV_NAME} with their values.
// also support ${ENV_NAME:default_value}.
func parseEnv(input string) (string, error) {
re := regexp.MustCompile(`\$\{(\w+)(:([^}]*))?\}`)
var err error
output := re.ReplaceAllStringFunc(input, func(match string) string {
parts := re.FindStringSubmatch(match)
// extract the variable name
variableName := parts[1]
if value, found := os.LookupEnv(variableName); found {
return value
}
if len(parts) >= 4 && parts[2] != "" {
return parts[3]
}
err = fmt.Errorf("environment variable not found: %q", variableName)
return ""
})
return output, err
}
// parseToolsFile parses the provided yaml into appropriate configs.
func parseToolsFile(ctx context.Context, raw []byte) (ToolsFile, error) {
var toolsFile ToolsFile
// Replace environment variables if found
output, err := parseEnv(string(raw))
if err != nil {
return toolsFile, fmt.Errorf("error parsing environment variables: %s", err)
}
raw = []byte(output)
raw, err = convertToolsFile(raw)
if err != nil {
return toolsFile, fmt.Errorf("error converting tools file: %s", err)
}
// Parse contents
toolsFile.Sources, toolsFile.AuthServices, toolsFile.EmbeddingModels, toolsFile.Tools, toolsFile.Toolsets, toolsFile.Prompts, err = server.UnmarshalResourceConfig(ctx, raw)
if err != nil {
return toolsFile, err
}
return toolsFile, nil
}
func convertToolsFile(raw []byte) ([]byte, error) {
var input yaml.MapSlice
decoder := yaml.NewDecoder(bytes.NewReader(raw), yaml.UseOrderedMap())
// convert to tools file v2
var buf bytes.Buffer
encoder := yaml.NewEncoder(&buf)
v1keys := []string{"sources", "authSources", "authServices", "embeddingModels", "tools", "toolsets", "prompts"}
for {
if err := decoder.Decode(&input); err != nil {
if err == io.EOF {
break
}
return nil, err
}
for _, item := range input {
key, ok := item.Key.(string)
if !ok {
return nil, fmt.Errorf("unexpected non-string key in input: %v", item.Key)
}
// check if the key is config file v1's key
if slices.Contains(v1keys, key) {
// check if value conversion to yaml.MapSlice successfully
// fields such as "tools" in toolsets might pass the first check but
// fail to convert to MapSlice
if slice, ok := item.Value.(yaml.MapSlice); ok {
// Deprecated: convert authSources to authServices
if key == "authSources" {
key = "authServices"
}
transformed, err := transformDocs(key, slice)
if err != nil {
return nil, err
}
// encode per-doc
for _, doc := range transformed {
if err := encoder.Encode(doc); err != nil {
return nil, err
}
}
} else {
// invalid input will be ignored
// we don't want to throw error here since the config could
// be valid but with a different order such as:
// ---
// tools:
// - tool_a
// kind: toolsets
// ---
continue
}
} else {
// this doc is already v2, encode to buf
if err := encoder.Encode(input); err != nil {
return nil, err
}
break
}
}
}
return buf.Bytes(), nil
}
// transformDocs transforms the configuration file from v1 format to v2
// yaml.MapSlice will preserve the order in a map
func transformDocs(kind string, input yaml.MapSlice) ([]yaml.MapSlice, error) {
var transformed []yaml.MapSlice
for _, entry := range input {
entryName, ok := entry.Key.(string)
if !ok {
return nil, fmt.Errorf("unexpected non-string key for entry in '%s': %v", kind, entry.Key)
}
entryBody := ProcessValue(entry.Value, kind == "toolsets")
currentTransformed := yaml.MapSlice{
{Key: "kind", Value: kind},
{Key: "name", Value: entryName},
}
// Merge the transformed body into our result
if bodySlice, ok := entryBody.(yaml.MapSlice); ok {
currentTransformed = append(currentTransformed, bodySlice...)
} else {
return nil, fmt.Errorf("unable to convert entryBody to MapSlice")
}
transformed = append(transformed, currentTransformed)
}
return transformed, nil
}
// ProcessValue recursively looks for MapSlices to rename 'kind' -> 'type'
func ProcessValue(v any, isToolset bool) any {
switch val := v.(type) {
case yaml.MapSlice:
// creating a new MapSlice is safer for recursive transformation
newVal := make(yaml.MapSlice, len(val))
for i, item := range val {
// Perform renaming
if item.Key == "kind" {
item.Key = "type"
}
// Recursive call for nested values (e.g., nested objects or lists)
item.Value = ProcessValue(item.Value, false)
newVal[i] = item
}
return newVal
case []any:
// Process lists: If it's a toolset top-level list, wrap it.
if isToolset {
return yaml.MapSlice{{Key: "tools", Value: val}}
}
// Otherwise, recurse into list items (to catch nested objects)
newVal := make([]any, len(val))
for i := range val {
newVal[i] = ProcessValue(val[i], false)
}
return newVal
default:
return val
}
}
// mergeToolsFiles merges multiple ToolsFile structs into one.
// Detects and raises errors for resource conflicts in sources, authServices, tools, and toolsets.
// All resource names (sources, authServices, tools, toolsets) must be unique across all files.
func mergeToolsFiles(files ...ToolsFile) (ToolsFile, error) {
merged := ToolsFile{
Sources: make(server.SourceConfigs),
AuthServices: make(server.AuthServiceConfigs),
EmbeddingModels: make(server.EmbeddingModelConfigs),
Tools: make(server.ToolConfigs),
Toolsets: make(server.ToolsetConfigs),
Prompts: make(server.PromptConfigs),
}
var conflicts []string
for fileIndex, file := range files {
// Check for conflicts and merge sources
for name, source := range file.Sources {
if _, exists := merged.Sources[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("source '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Sources[name] = source
}
}
// Check for conflicts and merge authServices
for name, authService := range file.AuthServices {
if _, exists := merged.AuthServices[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("authService '%s' (file #%d)", name, fileIndex+1))
} else {
merged.AuthServices[name] = authService
}
}
// Check for conflicts and merge embeddingModels
for name, em := range file.EmbeddingModels {
if _, exists := merged.EmbeddingModels[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("embedding model '%s' (file #%d)", name, fileIndex+1))
} else {
merged.EmbeddingModels[name] = em
}
}
// Check for conflicts and merge tools
for name, tool := range file.Tools {
if _, exists := merged.Tools[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("tool '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Tools[name] = tool
}
}
// Check for conflicts and merge toolsets
for name, toolset := range file.Toolsets {
if _, exists := merged.Toolsets[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("toolset '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Toolsets[name] = toolset
}
}
// Check for conflicts and merge prompts
for name, prompt := range file.Prompts {
if _, exists := merged.Prompts[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("prompt '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Prompts[name] = prompt
}
}
}
// If conflicts were detected, return an error
if len(conflicts) > 0 {
return ToolsFile{}, fmt.Errorf("resource conflicts detected:\n - %s\n\nPlease ensure each source, authService, tool, toolset and prompt has a unique name across all files", strings.Join(conflicts, "\n - "))
}
return merged, nil
}
// LoadAndMergeToolsFiles loads multiple YAML files and merges them
func LoadAndMergeToolsFiles(ctx context.Context, filePaths []string) (ToolsFile, error) {
var toolsFiles []ToolsFile
for _, filePath := range filePaths {
buf, err := os.ReadFile(filePath)
if err != nil {
return ToolsFile{}, fmt.Errorf("unable to read tool file at %q: %w", filePath, err)
}
toolsFile, err := parseToolsFile(ctx, buf)
if err != nil {
return ToolsFile{}, fmt.Errorf("unable to parse tool file at %q: %w", filePath, err)
}
toolsFiles = append(toolsFiles, toolsFile)
}
mergedFile, err := mergeToolsFiles(toolsFiles...)
if err != nil {
return ToolsFile{}, fmt.Errorf("unable to merge tools files: %w", err)
}
return mergedFile, nil
}
// LoadAndMergeToolsFolder loads all YAML files from a directory and merges them
func LoadAndMergeToolsFolder(ctx context.Context, folderPath string) (ToolsFile, error) {
// Check if directory exists
info, err := os.Stat(folderPath)
if err != nil {
return ToolsFile{}, fmt.Errorf("unable to access tools folder at %q: %w", folderPath, err)
}
if !info.IsDir() {
return ToolsFile{}, fmt.Errorf("path %q is not a directory", folderPath)
}
// Find all YAML files in the directory
pattern := filepath.Join(folderPath, "*.yaml")
yamlFiles, err := filepath.Glob(pattern)
if err != nil {
return ToolsFile{}, fmt.Errorf("error finding YAML files in %q: %w", folderPath, err)
}
// Also find .yml files
ymlPattern := filepath.Join(folderPath, "*.yml")
ymlFiles, err := filepath.Glob(ymlPattern)
if err != nil {
return ToolsFile{}, fmt.Errorf("error finding YML files in %q: %w", folderPath, err)
}
// Combine both file lists
allFiles := append(yamlFiles, ymlFiles...)
if len(allFiles) == 0 {
return ToolsFile{}, fmt.Errorf("no YAML files found in directory %q", folderPath)
}
// Use existing LoadAndMergeToolsFiles function
return LoadAndMergeToolsFiles(ctx, allFiles)
}

File diff suppressed because it is too large Load Diff

View File

@@ -1,30 +0,0 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package cmd
import (
"io"
)
// Option is a function that configures a Command.
type Option func(*Command)
// WithStreams overrides the default writer.
func WithStreams(out, err io.Writer) Option {
return func(c *Command) {
c.outStream = out
c.errStream = err
}
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -1 +1 @@
0.26.0
0.27.0

View File

@@ -234,7 +234,7 @@
},
"outputs": [],
"source": [
"version = \"0.26.0\" # x-release-please-version\n",
"version = \"0.27.0\" # x-release-please-version\n",
"! curl -O https://storage.googleapis.com/genai-toolbox/v{version}/linux/amd64/toolbox\n",
"\n",
"# Make the binary executable\n",

View File

@@ -109,7 +109,7 @@ To install Toolbox as a binary on Linux (AMD64):
```sh
# see releases page for other versions
export VERSION=0.26.0
export VERSION=0.27.0
curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/linux/amd64/toolbox
chmod +x toolbox
```
@@ -120,7 +120,7 @@ To install Toolbox as a binary on macOS (Apple Silicon):
```sh
# see releases page for other versions
export VERSION=0.26.0
export VERSION=0.27.0
curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/darwin/arm64/toolbox
chmod +x toolbox
```
@@ -131,7 +131,7 @@ To install Toolbox as a binary on macOS (Intel):
```sh
# see releases page for other versions
export VERSION=0.26.0
export VERSION=0.27.0
curl -L -o toolbox https://storage.googleapis.com/genai-toolbox/v$VERSION/darwin/amd64/toolbox
chmod +x toolbox
```
@@ -142,7 +142,7 @@ To install Toolbox as a binary on Windows (Command Prompt):
```cmd
:: see releases page for other versions
set VERSION=0.26.0
set VERSION=0.27.0
curl -o toolbox.exe "https://storage.googleapis.com/genai-toolbox/v%VERSION%/windows/amd64/toolbox.exe"
```
@@ -152,7 +152,7 @@ To install Toolbox as a binary on Windows (PowerShell):
```powershell
# see releases page for other versions
$VERSION = "0.26.0"
$VERSION = "0.27.0"
curl.exe -o toolbox.exe "https://storage.googleapis.com/genai-toolbox/v$VERSION/windows/amd64/toolbox.exe"
```
@@ -164,7 +164,7 @@ You can also install Toolbox as a container:
```sh
# see releases page for other versions
export VERSION=0.26.0
export VERSION=0.27.0
docker pull us-central1-docker.pkg.dev/database-toolbox/toolbox/toolbox:$VERSION
```
@@ -183,7 +183,7 @@ To install from source, ensure you have the latest version of
[Go installed](https://go.dev/doc/install), and then run the following command:
```sh
go install github.com/googleapis/genai-toolbox@v0.26.0
go install github.com/googleapis/genai-toolbox@v0.27.0
```
{{% /tab %}}

Binary file not shown.

Before

Width:  |  Height:  |  Size: 241 KiB

After

Width:  |  Height:  |  Size: 271 KiB

View File

@@ -105,7 +105,7 @@ In this section, we will download Toolbox, configure our tools in a
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -20,6 +20,8 @@ import (
"testing"
)
var goldenKeywords = []string{"Hilton Basel", "Hyatt Regency", "book"}
func TestQuickstartSample(t *testing.T) {
framework := os.Getenv("ORCH_NAME")
if framework == "" {
@@ -59,16 +61,10 @@ func TestQuickstartSample(t *testing.T) {
t.Fatal("Script ran successfully but produced no output.")
}
goldenFile, err := os.ReadFile("../golden.txt")
if err != nil {
t.Fatalf("Could not read golden.txt to check for keywords: %v", err)
}
keywords := strings.Split(string(goldenFile), "\n")
var missingKeywords []string
outputLower := strings.ToLower(actualOutput)
for _, keyword := range keywords {
for _, keyword := range goldenKeywords {
kw := strings.TrimSpace(keyword)
if kw != "" && !strings.Contains(outputLower, strings.ToLower(kw)) {
missingKeywords = append(missingKeywords, kw)

View File

@@ -1,3 +0,0 @@
Hilton Basel
Hyatt Regency
book

View File

@@ -25,7 +25,7 @@ const quickstartPath = path.join(orchDir, "quickstart.js");
const { main: runAgent } = await import(quickstartPath);
const GOLDEN_FILE_PATH = path.resolve(__dirname, "../golden.txt");
const GOLDEN_KEYWORDS = ["Hilton Basel", "Hyatt Regency", "book"];
describe(`${ORCH_NAME} Quickstart Agent`, () => {
let capturedOutput = [];
@@ -52,11 +52,8 @@ describe(`${ORCH_NAME} Quickstart Agent`, () => {
"Assertion Failed: Script ran successfully but produced no output."
);
const goldenFile = fs.readFileSync(GOLDEN_FILE_PATH, "utf8");
const keywords = goldenFile.split("\n").filter((kw) => kw.trim() !== "");
const missingKeywords = [];
for (const keyword of keywords) {
for (const keyword of GOLDEN_KEYWORDS) {
if (!actualOutput.toLowerCase().includes(keyword.toLowerCase())) {
missingKeywords.push(keyword);
}

View File

@@ -24,18 +24,7 @@ module_path = f"python.{ORCH_NAME}.quickstart"
quickstart = importlib.import_module(module_path)
@pytest.fixture(scope="module")
def golden_keywords():
"""Loads expected keywords from the golden.txt file."""
golden_file_path = Path("../golden.txt")
if not golden_file_path.exists():
pytest.fail(f"Golden file not found: {golden_file_path}")
try:
with open(golden_file_path, 'r') as f:
return [line.strip() for line in f.readlines() if line.strip()]
except Exception as e:
pytest.fail(f"Could not read golden.txt: {e}")
GOLDEN_KEYWORDS = ["Hilton Basel", "Hyatt Regency", "book"]
# --- Execution Tests ---
class TestExecution:
@@ -62,8 +51,8 @@ class TestExecution:
"""Test that the script runs and produces no stderr."""
assert script_output.err == "", f"Script produced stderr: {script_output.err}"
def test_keywords_in_output(self, script_output, golden_keywords):
def test_keywords_in_output(self, script_output):
"""Test that expected keywords are present in the script's output."""
output = script_output.out
missing_keywords = [kw for kw in golden_keywords if kw not in output]
missing_keywords = [kw for kw in GOLDEN_KEYWORDS if kw.lower() not in output.lower()]
assert not missing_keywords, f"Missing keywords in output: {missing_keywords}"

View File

@@ -13,7 +13,7 @@ In this section, we will download Toolbox, configure our tools in a
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -100,19 +100,19 @@ After you install Looker in the MCP Store, resources and tools from the server a
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -45,19 +45,19 @@ instance:
<!-- {x-release-please-start-version} -->
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -43,19 +43,19 @@ expose your developer assistant tools to a MySQL instance:
<!-- {x-release-please-start-version} -->
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -44,19 +44,19 @@ expose your developer assistant tools to a Neo4j instance:
<!-- {x-release-please-start-version} -->
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -56,19 +56,19 @@ Omni](https://cloud.google.com/alloydb/omni/docs/overview).
<!-- {x-release-please-start-version} -->
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -43,19 +43,19 @@ to expose your developer assistant tools to a SQLite instance:
<!-- {x-release-please-start-version} -->
{{< tabpane persist=header >}}
{{< tab header="linux/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/linux/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/linux/amd64/toolbox
{{< /tab >}}
{{< tab header="darwin/arm64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/arm64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/arm64/toolbox
{{< /tab >}}
{{< tab header="darwin/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/darwin/amd64/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/darwin/amd64/toolbox
{{< /tab >}}
{{< tab header="windows/amd64" lang="bash" >}}
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/windows/amd64/toolbox.exe
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/windows/amd64/toolbox.exe
{{< /tab >}}
{{< /tabpane >}}
<!-- {x-release-please-end} -->

View File

@@ -207,7 +207,6 @@ You can connect to Toolbox Cloud Run instances directly through the SDK.
{{< tab header="Python" lang="python" >}}
import asyncio
from toolbox_core import ToolboxClient, auth_methods
from toolbox_core.protocol import Protocol
# Replace with the Cloud Run service URL generated in the previous step
URL = "https://cloud-run-url.app"
@@ -218,7 +217,6 @@ async def main():
async with ToolboxClient(
URL,
client_headers={"Authorization": auth_token_provider},
protocol=Protocol.TOOLBOX,
) as toolbox:
toolset = await toolbox.load_toolset()
# ...

View File

@@ -44,6 +44,12 @@ See [Usage Examples](../reference/cli.md#examples).
* **Tools:**
* `execute_sql`: Executes a SQL query.
* `list_tables`: Lists tables in the database.
* `list_active_queries`: Lists ongoing queries.
* `list_available_extensions`: Discover all PostgreSQL extensions available for installation.
* `list_installed_extensions`: List all installed PostgreSQL extensions.
* `long_running_transactions`: Identifies and lists database transactions that exceed a specified time limit.
* `list_locks`: Identifies all locks held by active processes.
* `replication_stats`: Lists each replica's process ID and sync state.
* `list_autovacuum_configurations`: Lists autovacuum configurations in the
database.
* `list_memory_configurations`: Lists memory-related configurations in the
@@ -59,12 +65,16 @@ See [Usage Examples](../reference/cli.md#examples).
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
* `list_query_stats`: Lists query statistics.
* `get_column_cardinality`: Gets column cardinality.
* `list_table_stats`: Lists table statistics.
* `list_publication_tables`: List publication tables in a PostgreSQL database.
* `list_tablespaces`: Lists tablespaces in the database.
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
* `list_database_stats`: Lists the key performance and activity statistics for
each database in the AlloyDB instance.
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
* `list_stored_procedure`: Lists stored procedures.
## AlloyDB Postgres Admin
@@ -113,6 +123,12 @@ See [Usage Examples](../reference/cli.md#examples).
* **Tools:**
* `execute_sql`: Executes a SQL query.
* `list_tables`: Lists tables in the database.
* `list_active_queries`: Lists ongoing queries.
* `list_available_extensions`: Discover all PostgreSQL extensions available for installation.
* `list_installed_extensions`: List all installed PostgreSQL extensions.
* `long_running_transactions`: Identifies and lists database transactions that exceed a specified time limit.
* `list_locks`: Identifies all locks held by active processes.
* `replication_stats`: Lists each replica's process ID and sync state.
* `list_autovacuum_configurations`: Lists autovacuum configurations in the
database.
* `list_columnar_configurations`: List AlloyDB Omni columnar-related configurations.
@@ -130,12 +146,16 @@ See [Usage Examples](../reference/cli.md#examples).
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
* `list_query_stats`: Lists query statistics.
* `get_column_cardinality`: Gets column cardinality.
* `list_table_stats`: Lists table statistics.
* `list_publication_tables`: List publication tables in a PostgreSQL database.
* `list_tablespaces`: Lists tablespaces in the database.
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
* `list_database_stats`: Lists the key performance and activity statistics for
each database in the AlloyDB instance.
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
* `list_stored_procedure`: Lists stored procedures.
## BigQuery
@@ -173,6 +193,21 @@ See [Usage Examples](../reference/cli.md#examples).
* `list_table_ids`: Lists tables.
* `search_catalog`: Search for entries based on the provided query.
## ClickHouse
* `--prebuilt` value: `clickhouse`
* **Environment Variables:**
* `CLICKHOUSE_HOST`: The hostname or IP address of the ClickHouse server.
* `CLICKHOUSE_PORT`: The port number of the ClickHouse server.
* `CLICKHOUSE_USER`: The database username.
* `CLICKHOUSE_PASSWORD`: The password for the database user.
* `CLICKHOUSE_DATABASE`: The name of the database to connect to.
* `CLICKHOUSE_PROTOCOL`: The protocol to use (e.g., http).
* **Tools:**
* `execute_sql`: Use this tool to execute SQL.
* `list_databases`: Use this tool to list all databases in ClickHouse.
* `list_tables`: Use this tool to list all tables in a specific ClickHouse database.
## Cloud SQL for MySQL
* `--prebuilt` value: `cloud-sql-mysql`
@@ -270,6 +305,12 @@ See [Usage Examples](../reference/cli.md#examples).
* **Tools:**
* `execute_sql`: Executes a SQL query.
* `list_tables`: Lists tables in the database.
* `list_active_queries`: Lists ongoing queries.
* `list_available_extensions`: Discover all PostgreSQL extensions available for installation.
* `list_installed_extensions`: List all installed PostgreSQL extensions.
* `long_running_transactions`: Identifies and lists database transactions that exceed a specified time limit.
* `list_locks`: Identifies all locks held by active processes.
* `replication_stats`: Lists each replica's process ID and sync state.
* `list_autovacuum_configurations`: Lists autovacuum configurations in the
database.
* `list_memory_configurations`: Lists memory-related configurations in the
@@ -285,12 +326,16 @@ See [Usage Examples](../reference/cli.md#examples).
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
* `list_query_stats`: Lists query statistics.
* `get_column_cardinality`: Gets column cardinality.
* `list_table_stats`: Lists table statistics.
* `list_publication_tables`: List publication tables in a PostgreSQL database.
* `list_tablespaces`: Lists tablespaces in the database.
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
* `list_database_stats`: Lists the key performance and activity statistics for
each database in the postgreSQL instance.
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
* `list_stored_procedure`: Lists stored procedures.
## Cloud SQL for PostgreSQL Observability
@@ -336,6 +381,7 @@ See [Usage Examples](../reference/cli.md#examples).
* `create_user`: Creates a new user in a Cloud SQL instance.
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
* `clone_instance`: Creates a clone for an existing Cloud SQL for PostgreSQL instance.
* `postgres_upgrade_precheck`: Performs a precheck for a major version upgrade of a Cloud SQL for PostgreSQL instance.
* `create_backup`: Creates a backup on a Cloud SQL instance.
* `restore_backup`: Restores a backup of a Cloud SQL instance.
@@ -420,6 +466,15 @@ See [Usage Examples](../reference/cli.md#examples).
* `search_aspect_types`: Finds aspect types relevant to the
query.
## Elasticsearch
* `--prebuilt` value: `elasticsearch`
* **Environment Variables:**
* `ELASTICSEARCH_HOST`: The hostname or IP address of the Elasticsearch server.
* `ELASTICSEARCH_APIKEY`: The API key for authentication.
* **Tools:**
* `execute_esql_query`: Use this tool to execute ES|QL queries.
## Firestore
* `--prebuilt` value: `firestore`
@@ -537,6 +592,19 @@ See [Usage Examples](../reference/cli.md#examples).
* `execute_sql`: Executes a SQL query.
* `list_tables`: Lists tables in the database.
## MindsDB
* `--prebuilt` value: `mindsdb`
* **Environment Variables:**
* `MINDSDB_HOST`: The hostname or IP address of the MindsDB server.
* `MINDSDB_PORT`: The port number of the MindsDB server.
* `MINDSDB_DATABASE`: The name of the database to connect to.
* `MINDSDB_USER`: The database username.
* `MINDSDB_PASS`: The password for the database user.
* **Tools:**
* `mindsdb-execute-sql`: Execute SQL queries directly on MindsDB database.
* `mindsdb-sql`: Execute parameterized SQL queries on MindsDB database.
## MySQL
* `--prebuilt` value: `mysql`
@@ -592,6 +660,12 @@ See [Usage Examples](../reference/cli.md#examples).
* **Tools:**
* `execute_sql`: Executes a SQL query.
* `list_tables`: Lists tables in the database.
* `list_active_queries`: Lists ongoing queries.
* `list_available_extensions`: Discover all PostgreSQL extensions available for installation.
* `list_installed_extensions`: List all installed PostgreSQL extensions.
* `long_running_transactions`: Identifies and lists database transactions that exceed a specified time limit.
* `list_locks`: Identifies all locks held by active processes.
* `replication_stats`: Lists each replica's process ID and sync state.
* `list_autovacuum_configurations`: Lists autovacuum configurations in the
database.
* `list_memory_configurations`: Lists memory-related configurations in the
@@ -607,12 +681,16 @@ See [Usage Examples](../reference/cli.md#examples).
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
* `list_query_stats`: Lists query statistics.
* `get_column_cardinality`: Gets column cardinality.
* `list_table_stats`: Lists table statistics.
* `list_publication_tables`: List publication tables in a PostgreSQL database.
* `list_tablespaces`: Lists tablespaces in the database.
* `list_pg_settings`: List configuration parameters for the PostgreSQL server.
* `list_database_stats`: Lists the key performance and activity statistics for
each database in the PostgreSQL server.
* `list_roles`: Lists all the user-created roles in PostgreSQL database.
* `list_stored_procedure`: Lists stored procedures.
## Google Cloud Serverless for Apache Spark
@@ -627,6 +705,38 @@ See [Usage Examples](../reference/cli.md#examples).
view serverless batches.
* **Tools:**
* `list_batches`: Lists Spark batches.
* `get_batch`: Gets information about a Spark batch.
* `cancel_batch`: Cancels a Spark batch.
* `create_pyspark_batch`: Creates a PySpark batch.
* `create_spark_batch`: Creates a Spark batch.
## SingleStore
* `--prebuilt` value: `singlestore`
* **Environment Variables:**
* `SINGLESTORE_HOST`: The hostname or IP address of the SingleStore server.
* `SINGLESTORE_PORT`: The port number of the SingleStore server.
* `SINGLESTORE_DATABASE`: The name of the database to connect to.
* `SINGLESTORE_USER`: The database username.
* `SINGLESTORE_PASSWORD`: The password for the database user.
* **Tools:**
* `execute_sql`: Use this tool to execute SQL.
* `list_tables`: Lists detailed schema information for user-created tables.
## Snowflake
* `--prebuilt` value: `snowflake`
* **Environment Variables:**
* `SNOWFLAKE_ACCOUNT`: The Snowflake account.
* `SNOWFLAKE_USER`: The database username.
* `SNOWFLAKE_PASSWORD`: The password for the database user.
* `SNOWFLAKE_DATABASE`: The name of the database to connect to.
* `SNOWFLAKE_SCHEMA`: The schema name.
* `SNOWFLAKE_WAREHOUSE`: The warehouse name.
* `SNOWFLAKE_ROLE`: The role name.
* **Tools:**
* `execute_sql`: Use this tool to execute SQL.
* `list_tables`: Lists detailed schema information for user-created tables.
## Spanner (GoogleSQL dialect)

View File

@@ -771,7 +771,7 @@
},
"outputs": [],
"source": [
"version = \"0.26.0\" # x-release-please-version\n",
"version = \"0.27.0\" # x-release-please-version\n",
"! curl -L -o /content/toolbox https://storage.googleapis.com/genai-toolbox/v{version}/linux/amd64/toolbox\n",
"\n",
"# Make the binary executable\n",

View File

@@ -123,7 +123,7 @@ In this section, we will download and install the Toolbox binary.
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
export VERSION="0.26.0"
export VERSION="0.27.0"
curl -O https://storage.googleapis.com/genai-toolbox/v$VERSION/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -220,7 +220,7 @@
},
"outputs": [],
"source": [
"version = \"0.26.0\" # x-release-please-version\n",
"version = \"0.27.0\" # x-release-please-version\n",
"! curl -O https://storage.googleapis.com/genai-toolbox/v{version}/linux/amd64/toolbox\n",
"\n",
"# Make the binary executable\n",

View File

@@ -179,7 +179,7 @@ to use BigQuery, and then run the Toolbox server.
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -98,7 +98,7 @@ In this section, we will download Toolbox, configure our tools in a
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -34,7 +34,7 @@ In this section, we will download Toolbox and run the Toolbox server.
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -48,7 +48,7 @@ In this section, we will download Toolbox and run the Toolbox server.
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -34,7 +34,7 @@ In this section, we will download Toolbox and run the Toolbox server.
<!-- {x-release-please-start-version} -->
```bash
export OS="linux/amd64" # one of linux/amd64, darwin/arm64, darwin/amd64, or windows/amd64
curl -O https://storage.googleapis.com/genai-toolbox/v0.26.0/$OS/toolbox
curl -O https://storage.googleapis.com/genai-toolbox/v0.27.0/$OS/toolbox
```
<!-- {x-release-please-end} -->

View File

@@ -1,6 +1,6 @@
{
"name": "mcp-toolbox-for-databases",
"version": "0.26.0",
"version": "0.27.0",
"description": "MCP Toolbox for Databases is an open-source MCP server for more than 30 different datasources.",
"contextFileName": "MCP-TOOLBOX-EXTENSION.md"
}

View File

@@ -19,7 +19,6 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
@@ -216,7 +215,7 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("tool invocation not authorized. Please make sure your specify correct auth headers")
err = fmt.Errorf("tool invocation not authorized. Please make sure you specify correct auth headers")
s.logger.DebugContext(ctx, err.Error())
_ = render.Render(w, r, newErrResponse(err, http.StatusUnauthorized))
return
@@ -234,15 +233,28 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
params, err := parameters.ParseParams(tool.GetParameters(), data, claimsFromAuth)
if err != nil {
// If auth error, return 401
if errors.Is(err, util.ErrUnauthorized) {
s.logger.DebugContext(ctx, fmt.Sprintf("error parsing authenticated parameters from ID token: %s", err))
var clientServerErr *util.ClientServerError
// Return 401 Authentication errors
if errors.As(err, &clientServerErr) && clientServerErr.Code == http.StatusUnauthorized {
s.logger.DebugContext(ctx, fmt.Sprintf("auth error: %v", err))
_ = render.Render(w, r, newErrResponse(err, http.StatusUnauthorized))
return
}
err = fmt.Errorf("provided parameters were invalid: %w", err)
s.logger.DebugContext(ctx, err.Error())
_ = render.Render(w, r, newErrResponse(err, http.StatusBadRequest))
var agentErr *util.AgentError
if errors.As(err, &agentErr) {
s.logger.DebugContext(ctx, fmt.Sprintf("agent validation error: %v", err))
errMap := map[string]string{"error": err.Error()}
errMarshal, _ := json.Marshal(errMap)
_ = render.Render(w, r, &resultResponse{Result: string(errMarshal)})
return
}
// Return 500 if it's a specific ClientServerError that isn't a 401, or any other unexpected error
s.logger.ErrorContext(ctx, fmt.Sprintf("internal server error: %v", err))
_ = render.Render(w, r, newErrResponse(err, http.StatusInternalServerError))
return
}
s.logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
@@ -259,34 +271,50 @@ func toolInvokeHandler(s *Server, w http.ResponseWriter, r *http.Request) {
// Determine what error to return to the users.
if err != nil {
errStr := err.Error()
var statusCode int
var tbErr util.ToolboxError
// Upstream API auth error propagation
switch {
case strings.Contains(errStr, "Error 401"):
statusCode = http.StatusUnauthorized
case strings.Contains(errStr, "Error 403"):
statusCode = http.StatusForbidden
}
if errors.As(err, &tbErr) {
switch tbErr.Category() {
case util.CategoryAgent:
// Agent Errors -> 200 OK
s.logger.DebugContext(ctx, fmt.Sprintf("Tool invocation agent error: %v", err))
res = map[string]string{
"error": err.Error(),
}
if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden {
if clientAuth {
// Propagate the original 401/403 error.
s.logger.DebugContext(ctx, fmt.Sprintf("error invoking tool. Client credentials lack authorization to the source: %v", err))
case util.CategoryServer:
// Server Errors -> Check the specific code inside
var clientServerErr *util.ClientServerError
statusCode := http.StatusInternalServerError // Default to 500
if errors.As(err, &clientServerErr) {
if clientServerErr.Code != 0 {
statusCode = clientServerErr.Code
}
}
// Process auth error
if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden {
if clientAuth {
// Token error, pass through 401/403
s.logger.DebugContext(ctx, fmt.Sprintf("Client credentials lack authorization: %v", err))
_ = render.Render(w, r, newErrResponse(err, statusCode))
return
}
// ADC/Config error, return 500
statusCode = http.StatusInternalServerError
}
s.logger.ErrorContext(ctx, fmt.Sprintf("Tool invocation server error: %v", err))
_ = render.Render(w, r, newErrResponse(err, statusCode))
return
}
// ADC lacking permission or credentials configuration error.
internalErr := fmt.Errorf("unexpected auth error occured during Tool invocation: %w", err)
s.logger.ErrorContext(ctx, internalErr.Error())
_ = render.Render(w, r, newErrResponse(internalErr, http.StatusInternalServerError))
} else {
// Unknown error -> 500
s.logger.ErrorContext(ctx, fmt.Sprintf("Tool invocation unknown error: %v", err))
_ = render.Render(w, r, newErrResponse(err, http.StatusInternalServerError))
return
}
err = fmt.Errorf("error while invoking tool: %w", err)
s.logger.DebugContext(ctx, err.Error())
_ = render.Render(w, r, newErrResponse(err, http.StatusBadRequest))
return
}
resMarshal, err := json.Marshal(res)

View File

@@ -23,7 +23,6 @@ import (
"fmt"
"io"
"net/http"
"strings"
"sync"
"time"
@@ -37,9 +36,11 @@ import (
v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
"github.com/googleapis/genai-toolbox/internal/util"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/metric"
"go.opentelemetry.io/otel/trace"
)
type sseSession struct {
@@ -117,6 +118,55 @@ type stdioSession struct {
writer io.Writer
}
// traceContextCarrier implements propagation.TextMapCarrier for extracting trace context from _meta
type traceContextCarrier map[string]string
func (c traceContextCarrier) Get(key string) string {
return c[key]
}
func (c traceContextCarrier) Set(key, value string) {
c[key] = value
}
func (c traceContextCarrier) Keys() []string {
keys := make([]string, 0, len(c))
for k := range c {
keys = append(keys, k)
}
return keys
}
// extractTraceContext extracts W3C Trace Context from params._meta
func extractTraceContext(ctx context.Context, body []byte) context.Context {
// Try to parse the request to extract _meta
var req struct {
Params struct {
Meta struct {
Traceparent string `json:"traceparent,omitempty"`
Tracestate string `json:"tracestate,omitempty"`
} `json:"_meta,omitempty"`
} `json:"params,omitempty"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ctx
}
// If traceparent is present, extract the context
if req.Params.Meta.Traceparent != "" {
carrier := traceContextCarrier{
"traceparent": req.Params.Meta.Traceparent,
}
if req.Params.Meta.Tracestate != "" {
carrier["tracestate"] = req.Params.Meta.Tracestate
}
return otel.GetTextMapPropagator().Extract(ctx, carrier)
}
return ctx
}
func NewStdioSession(s *Server, stdin io.Reader, stdout io.Writer) *stdioSession {
stdioSession := &stdioSession{
server: s,
@@ -143,18 +193,29 @@ func (s *stdioSession) readInputStream(ctx context.Context) error {
}
return err
}
v, res, err := processMcpMessage(ctx, []byte(line), s.server, s.protocol, "", "", nil)
// This ensures the transport span becomes a child of the client span
msgCtx := extractTraceContext(ctx, []byte(line))
// Create span for STDIO transport
msgCtx, span := s.server.instrumentation.Tracer.Start(msgCtx, "toolbox/server/mcp/stdio",
trace.WithSpanKind(trace.SpanKindServer),
)
defer span.End()
v, res, err := processMcpMessage(msgCtx, []byte(line), s.server, s.protocol, "", "", nil, "")
if err != nil {
// errors during the processing of message will generate a valid MCP Error response.
// server can continue to run.
s.server.logger.ErrorContext(ctx, err.Error())
s.server.logger.ErrorContext(msgCtx, err.Error())
span.SetStatus(codes.Error, err.Error())
}
if v != "" {
s.protocol = v
}
// no responses for notifications
if res != nil {
if err = s.write(ctx, res); err != nil {
if err = s.write(msgCtx, res); err != nil {
return err
}
}
@@ -240,7 +301,9 @@ func mcpRouter(s *Server) (chi.Router, error) {
// sseHandler handles sse initialization and message.
func sseHandler(s *Server, w http.ResponseWriter, r *http.Request) {
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp/sse")
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp/sse",
trace.WithSpanKind(trace.SpanKindServer),
)
r = r.WithContext(ctx)
sessionId := uuid.New().String()
@@ -336,9 +399,27 @@ func methodNotAllowed(s *Server, w http.ResponseWriter, r *http.Request) {
func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp")
ctx := r.Context()
ctx = util.WithLogger(ctx, s.logger)
// Read body first so we can extract trace context
body, err := io.ReadAll(r.Body)
if err != nil {
// Generate a new uuid if unable to decode
id := uuid.New().String()
s.logger.DebugContext(ctx, err.Error())
render.JSON(w, r, jsonrpc.NewError(id, jsonrpc.PARSE_ERROR, err.Error(), nil))
return
}
// This ensures the transport span becomes a child of the client span
ctx = extractTraceContext(ctx, body)
// Create span for HTTP transport
ctx, span := s.instrumentation.Tracer.Start(ctx, "toolbox/server/mcp/http",
trace.WithSpanKind(trace.SpanKindServer),
)
r = r.WithContext(ctx)
ctx = util.WithLogger(r.Context(), s.logger)
var sessionId, protocolVersion string
var session *sseSession
@@ -380,7 +461,6 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", toolsetName))
span.SetAttributes(attribute.String("toolset_name", toolsetName))
var err error
defer func() {
if err != nil {
span.SetStatus(codes.Error, err.Error())
@@ -399,17 +479,9 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
)
}()
// Read and returns a body from io.Reader
body, err := io.ReadAll(r.Body)
if err != nil {
// Generate a new uuid if unable to decode
id := uuid.New().String()
s.logger.DebugContext(ctx, err.Error())
render.JSON(w, r, jsonrpc.NewError(id, jsonrpc.PARSE_ERROR, err.Error(), nil))
return
}
networkProtocolVersion := fmt.Sprintf("%d.%d", r.ProtoMajor, r.ProtoMinor)
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, toolsetName, promptsetName, r.Header)
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, toolsetName, promptsetName, r.Header, networkProtocolVersion)
if err != nil {
s.logger.DebugContext(ctx, fmt.Errorf("error processing message: %w", err).Error())
}
@@ -444,15 +516,12 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
code := rpcResponse.Error.Code
switch code {
case jsonrpc.INTERNAL_ERROR:
// Map Internal RPC Error (-32603) to HTTP 500
w.WriteHeader(http.StatusInternalServerError)
case jsonrpc.INVALID_REQUEST:
errStr := err.Error()
if errors.Is(err, util.ErrUnauthorized) {
w.WriteHeader(http.StatusUnauthorized)
} else if strings.Contains(errStr, "Error 401") {
w.WriteHeader(http.StatusUnauthorized)
} else if strings.Contains(errStr, "Error 403") {
w.WriteHeader(http.StatusForbidden)
var clientServerErr *util.ClientServerError
if errors.As(err, &clientServerErr) {
w.WriteHeader(clientServerErr.Code)
}
}
}
@@ -462,7 +531,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
}
// processMcpMessage process the messages received from clients
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, toolsetName string, promptsetName string, header http.Header) (string, any, error) {
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, toolsetName string, promptsetName string, header http.Header, networkProtocolVersion string) (string, any, error) {
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return "", jsonrpc.NewError("", jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -498,31 +567,95 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Create method-specific span with semantic conventions
// Note: Trace context is already extracted and set in ctx by the caller
ctx, span := s.instrumentation.Tracer.Start(ctx, baseMessage.Method,
trace.WithSpanKind(trace.SpanKindServer),
)
defer span.End()
// Determine network transport and protocol based on header presence
networkTransport := "pipe" // default for stdio
networkProtocolName := "stdio"
if header != nil {
networkTransport = "tcp" // HTTP/SSE transport
networkProtocolName = "http"
}
// Set required semantic attributes for span according to OTEL MCP semcov
// ref: https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#server
span.SetAttributes(
attribute.String("mcp.method.name", baseMessage.Method),
attribute.String("network.transport", networkTransport),
attribute.String("network.protocol.name", networkProtocolName),
)
// Set network protocol version if available
if networkProtocolVersion != "" {
span.SetAttributes(attribute.String("network.protocol.version", networkProtocolVersion))
}
// Set MCP protocol version if available
if protocolVersion != "" {
span.SetAttributes(attribute.String("mcp.protocol.version", protocolVersion))
}
// Set request ID
if baseMessage.Id != nil {
span.SetAttributes(attribute.String("jsonrpc.request.id", fmt.Sprintf("%v", baseMessage.Id)))
}
// Set toolset name
span.SetAttributes(attribute.String("toolset.name", toolsetName))
// Check if message is a notification
if baseMessage.Id == nil {
err := mcp.NotificationHandler(ctx, body)
if err != nil {
span.SetStatus(codes.Error, err.Error())
}
return "", nil, err
}
// Process the method
switch baseMessage.Method {
case mcputil.INITIALIZE:
res, v, err := mcp.InitializeResponse(ctx, baseMessage.Id, body, s.version)
result, version, err := mcp.InitializeResponse(ctx, baseMessage.Id, body, s.version)
if err != nil {
return "", res, err
span.SetStatus(codes.Error, err.Error())
if rpcErr, ok := result.(jsonrpc.JSONRPCError); ok {
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
}
return "", result, err
}
return v, res, err
span.SetAttributes(attribute.String("mcp.protocol.version", version))
return version, result, err
default:
toolset, ok := s.ResourceMgr.GetToolset(toolsetName)
if !ok {
err = fmt.Errorf("toolset does not exist")
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
err := fmt.Errorf("toolset does not exist")
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
span.SetStatus(codes.Error, err.Error())
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
return "", rpcErr, err
}
promptset, ok := s.ResourceMgr.GetPromptset(promptsetName)
if !ok {
err = fmt.Errorf("promptset does not exist")
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
err := fmt.Errorf("promptset does not exist")
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
span.SetStatus(codes.Error, err.Error())
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
return "", rpcErr, err
}
res, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, promptset, s.ResourceMgr, body, header)
return "", res, err
result, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, promptset, s.ResourceMgr, body, header)
if err != nil {
span.SetStatus(codes.Error, err.Error())
// Set error.type based on JSON-RPC error code
if rpcErr, ok := result.(jsonrpc.JSONRPCError); ok {
span.SetAttributes(attribute.Int("jsonrpc.error.code", rpcErr.Error.Code))
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
}
}
return "", result, err
}
}

View File

@@ -45,6 +45,9 @@ type Request struct {
// notifications. The receiver is not obligated to provide these
// notifications.
ProgressToken ProgressToken `json:"progressToken,omitempty"`
// W3C Trace Context fields for distributed tracing
Traceparent string `json:"traceparent,omitempty"`
Tracestate string `json:"tracestate,omitempty"`
} `json:"_meta,omitempty"`
} `json:"params,omitempty"`
}
@@ -97,6 +100,24 @@ type Error struct {
Data interface{} `json:"data,omitempty"`
}
// String returns the error type as a string based on the error code.
func (e Error) String() string {
switch e.Code {
case METHOD_NOT_FOUND:
return "method_not_found"
case INVALID_PARAMS:
return "invalid_params"
case INTERNAL_ERROR:
return "internal_error"
case PARSE_ERROR:
return "parse_error"
case INVALID_REQUEST:
return "invalid_request"
default:
return "jsonrpc_error"
}
}
// JSONRPCError represents a non-successful (error) response to a request.
type JSONRPCError struct {
Jsonrpc string `json:"jsonrpc"`

View File

@@ -21,7 +21,6 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
@@ -29,6 +28,8 @@ import (
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)
// ProcessMethod returns a response for the request.
@@ -102,6 +103,14 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
toolName := req.Params.Name
toolArgument := req.Params.Arguments
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
span.SetAttributes(
attribute.String("gen_ai.tool.name", toolName),
attribute.String("gen_ai.operation.name", "execute_tool"),
)
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -124,7 +133,12 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
if clientAuth {
if accessToken == "" {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
err := util.NewClientServerError(
"missing access token in the 'Authorization' header",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
}
@@ -172,7 +186,11 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
err = util.NewClientServerError(
"unauthorized Tool call: Please make sure you specify correct auth headers",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
logger.DebugContext(ctx, "tool invocation authorized")
@@ -194,30 +212,44 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {
errStr := err.Error()
// Missing authService tokens.
if errors.Is(err, util.ErrUnauthorized) {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Upstream auth error
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
if clientAuth {
// Error with client credentials should pass down to the client
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
var tbErr util.ToolboxError
if errors.As(err, &tbErr) {
switch tbErr.Category() {
case util.CategoryAgent:
// MCP - Tool execution error
// Return SUCCESS but with IsError: true
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
case util.CategoryServer:
// MCP Spec - Protocol error
// Return JSON-RPC ERROR
var clientServerErr *util.ClientServerError
rpcCode := jsonrpc.INTERNAL_ERROR // Default to Internal Error (-32603)
if errors.As(err, &clientServerErr) {
if clientServerErr.Code == http.StatusUnauthorized || clientServerErr.Code == http.StatusForbidden {
if clientAuth {
rpcCode = jsonrpc.INVALID_REQUEST
} else {
rpcCode = jsonrpc.INTERNAL_ERROR
}
}
}
return jsonrpc.NewError(id, rpcCode, err.Error(), nil), err
}
// Auth error with ADC should raise internal 500 error
} else {
// Unknown error -> 500
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
}
content := make([]TextContent, 0)
@@ -288,6 +320,11 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
promptName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
prompt, ok := resourceMgr.GetPrompt(promptName)
if !ok {
err := fmt.Errorf("prompt with name %q does not exist", promptName)

View File

@@ -21,7 +21,6 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
@@ -29,6 +28,8 @@ import (
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)
// ProcessMethod returns a response for the request.
@@ -102,6 +103,15 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
toolName := req.Params.Name
toolArgument := req.Params.Arguments
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
span.SetAttributes(
attribute.String("gen_ai.tool.name", toolName),
attribute.String("gen_ai.operation.name", "execute_tool"),
)
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -124,7 +134,12 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
if clientAuth {
if accessToken == "" {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
err := util.NewClientServerError(
"missing access token in the 'Authorization' header",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
}
@@ -172,7 +187,11 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
err = util.NewClientServerError(
"unauthorized Tool call: Please make sure you specify correct auth headers",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
logger.DebugContext(ctx, "tool invocation authorized")
@@ -194,31 +213,45 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {
errStr := err.Error()
// Missing authService tokens.
if errors.Is(err, util.ErrUnauthorized) {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Upstream auth error
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
if clientAuth {
// Error with client credentials should pass down to the client
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
var tbErr util.ToolboxError
if errors.As(err, &tbErr) {
switch tbErr.Category() {
case util.CategoryAgent:
// MCP - Tool execution error
// Return SUCCESS but with IsError: true
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
case util.CategoryServer:
// MCP Spec - Protocol error
// Return JSON-RPC ERROR
var clientServerErr *util.ClientServerError
rpcCode := jsonrpc.INTERNAL_ERROR // Default to Internal Error (-32603)
if errors.As(err, &clientServerErr) {
if clientServerErr.Code == http.StatusUnauthorized || clientServerErr.Code == http.StatusForbidden {
if clientAuth {
rpcCode = jsonrpc.INVALID_REQUEST
} else {
rpcCode = jsonrpc.INTERNAL_ERROR
}
}
}
return jsonrpc.NewError(id, rpcCode, err.Error(), nil), err
}
// Auth error with ADC should raise internal 500 error
} else {
// Unknown error -> 500
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
}
content := make([]TextContent, 0)
sliceRes, ok := results.([]any)
@@ -287,6 +320,12 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
promptName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
prompt, ok := resourceMgr.GetPrompt(promptName)
if !ok {
err := fmt.Errorf("prompt with name %q does not exist", promptName)

View File

@@ -21,7 +21,6 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
@@ -29,6 +28,8 @@ import (
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)
// ProcessMethod returns a response for the request.
@@ -95,6 +96,15 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
toolName := req.Params.Name
toolArgument := req.Params.Arguments
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
span.SetAttributes(
attribute.String("gen_ai.tool.name", toolName),
attribute.String("gen_ai.operation.name", "execute_tool"),
)
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -117,7 +127,12 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
if clientAuth {
if accessToken == "" {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
err := util.NewClientServerError(
"missing access token in the 'Authorization' header",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
}
@@ -165,7 +180,11 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
err = util.NewClientServerError(
"unauthorized Tool call: Please make sure you specify correct auth headers",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
logger.DebugContext(ctx, "tool invocation authorized")
@@ -187,29 +206,44 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {
errStr := err.Error()
// Missing authService tokens.
if errors.Is(err, util.ErrUnauthorized) {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Upstream auth error
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
if clientAuth {
// Error with client credentials should pass down to the client
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
var tbErr util.ToolboxError
if errors.As(err, &tbErr) {
switch tbErr.Category() {
case util.CategoryAgent:
// MCP - Tool execution error
// Return SUCCESS but with IsError: true
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
case util.CategoryServer:
// MCP Spec - Protocol error
// Return JSON-RPC ERROR
var clientServerErr *util.ClientServerError
rpcCode := jsonrpc.INTERNAL_ERROR // Default to Internal Error (-32603)
if errors.As(err, &clientServerErr) {
if clientServerErr.Code == http.StatusUnauthorized || clientServerErr.Code == http.StatusForbidden {
if clientAuth {
rpcCode = jsonrpc.INVALID_REQUEST
} else {
rpcCode = jsonrpc.INTERNAL_ERROR
}
}
}
return jsonrpc.NewError(id, rpcCode, err.Error(), nil), err
}
// Auth error with ADC should raise internal 500 error
} else {
// Unknown error -> 500
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
}
content := make([]TextContent, 0)
@@ -280,6 +314,12 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
promptName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
prompt, ok := resourceMgr.GetPrompt(promptName)
if !ok {
err := fmt.Errorf("prompt with name %q does not exist", promptName)

View File

@@ -21,7 +21,6 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
@@ -29,6 +28,8 @@ import (
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)
// ProcessMethod returns a response for the request.
@@ -95,6 +96,15 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
toolName := req.Params.Name
toolArgument := req.Params.Arguments
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
span.SetAttributes(
attribute.String("gen_ai.tool.name", toolName),
attribute.String("gen_ai.operation.name", "execute_tool"),
)
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -117,7 +127,12 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
}
if clientAuth {
if accessToken == "" {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
err := util.NewClientServerError(
"missing access token in the 'Authorization' header",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
}
@@ -165,7 +180,11 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
err = util.NewClientServerError(
"unauthorized Tool call: Please make sure you specify correct auth headers",
http.StatusUnauthorized,
nil,
)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
logger.DebugContext(ctx, "tool invocation authorized")
@@ -187,29 +206,44 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
// run tool invocation and generate response.
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
if err != nil {
errStr := err.Error()
// Missing authService tokens.
if errors.Is(err, util.ErrUnauthorized) {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Upstream auth error
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
if clientAuth {
// Error with client credentials should pass down to the client
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
var tbErr util.ToolboxError
if errors.As(err, &tbErr) {
switch tbErr.Category() {
case util.CategoryAgent:
// MCP - Tool execution error
// Return SUCCESS but with IsError: true
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
case util.CategoryServer:
// MCP Spec - Protocol error
// Return JSON-RPC ERROR
var clientServerErr *util.ClientServerError
rpcCode := jsonrpc.INTERNAL_ERROR // Default to Internal Error (-32603)
if errors.As(err, &clientServerErr) {
if clientServerErr.Code == http.StatusUnauthorized || clientServerErr.Code == http.StatusForbidden {
if clientAuth {
rpcCode = jsonrpc.INVALID_REQUEST
} else {
rpcCode = jsonrpc.INTERNAL_ERROR
}
}
}
return jsonrpc.NewError(id, rpcCode, err.Error(), nil), err
}
// Auth error with ADC should raise internal 500 error
} else {
// Unknown error -> 500
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
}
content := make([]TextContent, 0)
@@ -280,6 +314,12 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
promptName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
prompt, ok := resourceMgr.GetPrompt(promptName)
if !ok {
err := fmt.Errorf("prompt with name %q does not exist", promptName)

View File

@@ -231,7 +231,7 @@ func TestMcpEndpointWithoutInitialized(t *testing.T) {
"id": "tools-call-tool4",
"error": map[string]any{
"code": -32600.0,
"message": "unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized",
"message": "unauthorized Tool call: Please make sure you specify correct auth headers",
},
},
},
@@ -320,7 +320,7 @@ func TestMcpEndpointWithoutInitialized(t *testing.T) {
Params: map[string]any{
"name": "prompt2",
"arguments": map[string]any{
"arg1": 42, // prompt2 expects a string, we send a number
"arg1": 42,
},
},
},
@@ -834,7 +834,7 @@ func TestMcpEndpoint(t *testing.T) {
"id": "tools-call-tool4",
"error": map[string]any{
"code": -32600.0,
"message": "unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized",
"message": "unauthorized Tool call: Please make sure you specify correct auth headers",
},
},
},

View File

@@ -21,6 +21,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -34,7 +35,7 @@ type MockTool struct {
requiresClientAuthrorization bool
}
func (t MockTool) Invoke(context.Context, tools.SourceProvider, parameters.ParamValues, tools.AccessToken) (any, error) {
func (t MockTool) Invoke(context.Context, tools.SourceProvider, parameters.ParamValues, tools.AccessToken) (any, util.ToolboxError) {
mock := []any{t.Name}
return mock, nil
}

View File

@@ -361,7 +361,11 @@ func (s *Source) GetOperations(ctx context.Context, project, location, operation
}
}
return string(opBytes), nil
var result any
if err := json.Unmarshal(opBytes, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal operation bytes: %w", err)
}
return result, nil
}
logger.DebugContext(ctx, fmt.Sprintf("Operation not complete, retrying in %v\n", delay))
}

View File

@@ -34,14 +34,13 @@ import (
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "cockroachdb"
const SourceType string = "cockroachdb"
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
if !sources.Register(SourceType, newConfig) {
panic(fmt.Sprintf("source type %q already registered", SourceType))
}
}
@@ -94,10 +93,6 @@ type Config struct {
ClusterID string `yaml:"clusterID"` // Optional cluster identifier for telemetry
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) SourceConfigType() string {
return SourceType
}
@@ -127,10 +122,6 @@ type Source struct {
Pool *pgxpool.Pool
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) SourceType() string {
return SourceType
}
@@ -379,7 +370,7 @@ func (s *Source) EmitTelemetry(ctx context.Context, event TelemetryEvent) {
func initCockroachDBConnectionPoolWithRetry(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, maxRetries int, baseDelay time.Duration) (*pgxpool.Pool, error) {
//nolint:all
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceType, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)

View File

@@ -17,11 +17,13 @@ package alloydbcreatecluster
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -122,44 +124,49 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok || project == "" {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a non-empty string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
return nil, util.NewAgentError("invalid 'location' parameter; expected a string", nil)
}
clusterID, ok := paramsMap["cluster"].(string)
if !ok || clusterID == "" {
return nil, fmt.Errorf("invalid or missing 'cluster' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a non-empty string", nil)
}
password, ok := paramsMap["password"].(string)
if !ok || password == "" {
return nil, fmt.Errorf("invalid or missing 'password' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'password' parameter; expected a non-empty string", nil)
}
network, ok := paramsMap["network"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'network' parameter; expected a string")
return nil, util.NewAgentError("invalid 'network' parameter; expected a string", nil)
}
user, ok := paramsMap["user"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'user' parameter; expected a string")
return nil, util.NewAgentError("invalid 'user' parameter; expected a string", nil)
}
resp, err := source.CreateCluster(ctx, project, location, network, user, password, clusterID, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return source.CreateCluster(ctx, project, location, network, user, password, clusterID, string(accessToken))
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydbcreateinstance
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -123,36 +125,36 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok || project == "" {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a non-empty string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok || location == "" {
return nil, fmt.Errorf("invalid or missing 'location' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'location' parameter; expected a non-empty string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok || cluster == "" {
return nil, fmt.Errorf("invalid or missing 'cluster' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a non-empty string", nil)
}
instanceID, ok := paramsMap["instance"].(string)
if !ok || instanceID == "" {
return nil, fmt.Errorf("invalid or missing 'instance' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'instance' parameter; expected a non-empty string", nil)
}
instanceType, ok := paramsMap["instanceType"].(string)
if !ok || (instanceType != "READ_POOL" && instanceType != "PRIMARY") {
return nil, fmt.Errorf("invalid 'instanceType' parameter; expected 'PRIMARY' or 'READ_POOL'")
return nil, util.NewAgentError("invalid 'instanceType' parameter; expected 'PRIMARY' or 'READ_POOL'", nil)
}
displayName, _ := paramsMap["displayName"].(string)
@@ -161,11 +163,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if instanceType == "READ_POOL" {
nodeCount, ok = paramsMap["nodeCount"].(int)
if !ok {
return nil, fmt.Errorf("invalid 'nodeCount' parameter; expected an integer for READ_POOL")
return nil, util.NewAgentError("invalid 'nodeCount' parameter; expected an integer for READ_POOL", nil)
}
}
return source.CreateInstance(ctx, project, location, cluster, instanceID, instanceType, displayName, nodeCount, string(accessToken))
resp, err := source.CreateInstance(ctx, project, location, cluster, instanceID, instanceType, displayName, nodeCount, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydbcreateuser
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -122,43 +124,43 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok || project == "" {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a non-empty string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok || location == "" {
return nil, fmt.Errorf("invalid or missing'location' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing'location' parameter; expected a non-empty string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok || cluster == "" {
return nil, fmt.Errorf("invalid or missing 'cluster' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a non-empty string", nil)
}
userID, ok := paramsMap["user"].(string)
if !ok || userID == "" {
return nil, fmt.Errorf("invalid or missing 'user' parameter; expected a non-empty string")
return nil, util.NewAgentError("invalid or missing 'user' parameter; expected a non-empty string", nil)
}
userType, ok := paramsMap["userType"].(string)
if !ok || (userType != "ALLOYDB_BUILT_IN" && userType != "ALLOYDB_IAM_USER") {
return nil, fmt.Errorf("invalid or missing 'userType' parameter; expected 'ALLOYDB_BUILT_IN' or 'ALLOYDB_IAM_USER'")
return nil, util.NewAgentError("invalid or missing 'userType' parameter; expected 'ALLOYDB_BUILT_IN' or 'ALLOYDB_IAM_USER'", nil)
}
var password string
if userType == "ALLOYDB_BUILT_IN" {
password, ok = paramsMap["password"].(string)
if !ok || password == "" {
return nil, fmt.Errorf("password is required when userType is ALLOYDB_BUILT_IN")
return nil, util.NewAgentError("password is required when userType is ALLOYDB_BUILT_IN", nil)
}
}
@@ -170,7 +172,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
}
}
return source.CreateUser(ctx, userType, password, roles, string(accessToken), project, location, cluster, userID)
resp, err := source.CreateUser(ctx, userType, password, roles, string(accessToken), project, location, cluster, userID)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydbgetcluster
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -120,28 +122,32 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
if !ok || location == "" {
return nil, util.NewAgentError("invalid or missing 'location' parameter; expected a string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
if !ok || cluster == "" {
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a string", nil)
}
return source.GetCluster(ctx, project, location, cluster, string(accessToken))
resp, err := source.GetCluster(ctx, project, location, cluster, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydbgetinstance
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -120,32 +122,36 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
if !ok || location == "" {
return nil, util.NewAgentError("invalid or missing 'location' parameter; expected a string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
if !ok || cluster == "" {
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a string", nil)
}
instance, ok := paramsMap["instance"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'instance' parameter; expected a string")
if !ok || instance == "" {
return nil, util.NewAgentError("invalid or missing 'instance' parameter; expected a string", nil)
}
return source.GetInstance(ctx, project, location, cluster, instance, string(accessToken))
resp, err := source.GetInstance(ctx, project, location, cluster, instance, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydbgetuser
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -120,32 +122,36 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
if !ok || location == "" {
return nil, util.NewAgentError("invalid or missing 'location' parameter; expected a string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
if !ok || cluster == "" {
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a string", nil)
}
user, ok := paramsMap["user"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'user' parameter; expected a string")
if !ok || user == "" {
return nil, util.NewAgentError("invalid or missing 'user' parameter; expected a string", nil)
}
return source.GetUsers(ctx, project, location, cluster, user, string(accessToken))
resp, err := source.GetUsers(ctx, project, location, cluster, user, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydblistclusters
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -118,24 +120,28 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
return nil, util.NewAgentError("invalid 'location' parameter; expected a string", nil)
}
return source.ListCluster(ctx, project, location, string(accessToken))
resp, err := source.ListCluster(ctx, project, location, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydblistinstances
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -119,28 +121,32 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
return nil, util.NewAgentError("invalid 'location' parameter; expected a string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
return nil, util.NewAgentError("invalid 'cluster' parameter; expected a string", nil)
}
return source.ListInstance(ctx, project, location, cluster, string(accessToken))
resp, err := source.ListInstance(ctx, project, location, cluster, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package alloydblistusers
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -119,28 +121,32 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
if !ok || project == "" {
return nil, util.NewAgentError("invalid or missing 'project' parameter; expected a string", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
if !ok || location == "" {
return nil, util.NewAgentError("invalid or missing 'location' parameter; expected a string", nil)
}
cluster, ok := paramsMap["cluster"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
if !ok || cluster == "" {
return nil, util.NewAgentError("invalid or missing 'cluster' parameter; expected a string", nil)
}
return source.ListUsers(ctx, project, location, cluster, string(accessToken))
resp, err := source.ListUsers(ctx, project, location, cluster, string(accessToken))
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -24,6 +24,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -213,25 +214,25 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke executes the tool's logic.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
project, ok := paramsMap["project"].(string)
if !ok {
return nil, fmt.Errorf("missing 'project' parameter")
return nil, util.NewAgentError("missing 'project' parameter", nil)
}
location, ok := paramsMap["location"].(string)
if !ok {
return nil, fmt.Errorf("missing 'location' parameter")
return nil, util.NewAgentError("missing 'location' parameter", nil)
}
operation, ok := paramsMap["operation"].(string)
if !ok {
return nil, fmt.Errorf("missing 'operation' parameter")
return nil, util.NewAgentError("missing 'operation' parameter", nil)
}
ctx, cancel := context.WithTimeout(ctx, 30*time.Minute)
@@ -246,14 +247,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
for retries < maxRetries {
select {
case <-ctx.Done():
return nil, fmt.Errorf("timed out waiting for operation: %w", ctx.Err())
return nil, util.NewAgentError("timed out waiting for operation", ctx.Err())
default:
}
op, err := source.GetOperations(ctx, project, location, operation, alloyDBConnectionMessageTemplate, delay, string(accessToken))
if err != nil {
return nil, err
} else if op != nil {
return nil, util.ProcessGeneralError(err)
}
if op != nil {
return op, nil
}
@@ -264,7 +266,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
retries++
}
return nil, fmt.Errorf("exceeded max retries waiting for operation")
return nil, util.NewAgentError("exceeded max retries waiting for operation", nil)
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package alloydbainl
import (
"context"
"fmt"
"net/http"
"strings"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/jackc/pgx/v5/pgxpool"
)
@@ -127,10 +129,10 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
sliceParams := params.AsSlice()
@@ -143,7 +145,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
resp, err := source.RunSQL(ctx, t.Statement, allParamValues)
if err != nil {
return nil, fmt.Errorf("%w. Query: %v , Values: %v. Toolbox v0.19.0+ is only compatible with AlloyDB AI NL v1.0.3+. Please ensure that you are using the latest AlloyDB AI NL extension", err, t.Statement, allParamValues)
return nil, util.NewClientServerError(fmt.Sprintf("error running SQL query: %v. Query: %v , Values: %v. Toolbox v0.19.0+ is only compatible with AlloyDB AI NL v1.0.3+. Please ensure that you are using the latest AlloyDB AI NL extension", err, t.Statement, allParamValues), http.StatusBadRequest, err)
}
return resp, nil
}

View File

@@ -17,6 +17,7 @@ package bigqueryanalyzecontribution
import (
"context"
"fmt"
"net/http"
"strings"
bigqueryapi "cloud.google.com/go/bigquery"
@@ -27,6 +28,7 @@ import (
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
"github.com/googleapis/genai-toolbox/internal/tools"
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
)
@@ -154,21 +156,21 @@ func (t Tool) ToConfig() tools.ToolConfig {
}
// Invoke runs the contribution analysis.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
inputData, ok := paramsMap["input_data"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast input_data parameter %s", paramsMap["input_data"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast input_data parameter %s", paramsMap["input_data"]), nil)
}
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
modelID := fmt.Sprintf("contribution_analysis_model_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
@@ -186,7 +188,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
options = append(options, fmt.Sprintf("DIMENSION_ID_COLS = [%s]", strings.Join(strCols, ", ")))
} else {
return nil, fmt.Errorf("unable to cast dimension_id_cols parameter %s", paramsMap["dimension_id_cols"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast dimension_id_cols parameter %s", paramsMap["dimension_id_cols"]), nil)
}
}
if val, ok := paramsMap["top_k_insights_by_apriori_support"]; ok {
@@ -195,7 +197,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if val, ok := paramsMap["pruning_method"].(string); ok {
upperVal := strings.ToUpper(val)
if upperVal != "NO_PRUNING" && upperVal != "PRUNE_REDUNDANT_INSIGHTS" {
return nil, fmt.Errorf("invalid pruning_method: %s", val)
return nil, util.NewAgentError(fmt.Sprintf("invalid pruning_method: %s", val), nil)
}
options = append(options, fmt.Sprintf("PRUNING_METHOD = '%s'", upperVal))
}
@@ -207,7 +209,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
var connProps []*bigqueryapi.ConnectionProperty
session, err := source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
}
if session != nil {
connProps = []*bigqueryapi.ConnectionProperty{
@@ -216,22 +218,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
dryRunJob, err := bqutil.DryRunQuery(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, inputData, nil, connProps)
if err != nil {
return nil, fmt.Errorf("query validation failed: %w", err)
return nil, util.ProcessGcpError(err)
}
statementType := dryRunJob.Statistics.Query.StatementType
if statementType != "SELECT" {
return nil, fmt.Errorf("the 'input_data' parameter only supports a table ID or a SELECT query. The provided query has statement type '%s'", statementType)
return nil, util.NewAgentError(fmt.Sprintf("the 'input_data' parameter only supports a table ID or a SELECT query. The provided query has statement type '%s'", statementType), nil)
}
queryStats := dryRunJob.Statistics.Query
if queryStats != nil {
for _, tableRef := range queryStats.ReferencedTables {
if !source.IsDatasetAllowed(tableRef.ProjectId, tableRef.DatasetId) {
return nil, fmt.Errorf("query in input_data accesses dataset '%s.%s', which is not in the allowed list", tableRef.ProjectId, tableRef.DatasetId)
return nil, util.NewAgentError(fmt.Sprintf("query in input_data accesses dataset '%s.%s', which is not in the allowed list", tableRef.ProjectId, tableRef.DatasetId), nil)
}
}
} else {
return nil, fmt.Errorf("could not analyze query in input_data to validate against allowed datasets")
return nil, util.NewAgentError("could not analyze query in input_data to validate against allowed datasets", nil)
}
}
inputDataSource = fmt.Sprintf("(%s)", inputData)
@@ -245,10 +247,10 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
case 2: // dataset.table
projectID, datasetID = source.BigQueryClient().Project(), parts[0]
default:
return nil, fmt.Errorf("invalid table ID format for 'input_data': %q. Expected 'dataset.table' or 'project.dataset.table'", inputData)
return nil, util.NewAgentError(fmt.Sprintf("invalid table ID format for 'input_data': %q. Expected 'dataset.table' or 'project.dataset.table'", inputData), nil)
}
if !source.IsDatasetAllowed(projectID, datasetID) {
return nil, fmt.Errorf("access to dataset '%s.%s' (from table '%s') is not allowed", projectID, datasetID, inputData)
return nil, util.NewAgentError(fmt.Sprintf("access to dataset '%s.%s' (from table '%s') is not allowed", projectID, datasetID, inputData), nil)
}
}
inputDataSource = fmt.Sprintf("SELECT * FROM `%s`", inputData)
@@ -268,7 +270,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
// Otherwise, a new session will be created by the first query.
session, err := source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
}
if session != nil {
@@ -281,15 +283,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
createModelJob, err := createModelQuery.Run(ctx)
if err != nil {
return nil, fmt.Errorf("failed to start create model job: %w", err)
return nil, util.ProcessGcpError(err)
}
status, err := createModelJob.Wait(ctx)
if err != nil {
return nil, fmt.Errorf("failed to wait for create model job: %w", err)
return nil, util.ProcessGcpError(err)
}
if err := status.Err(); err != nil {
return nil, fmt.Errorf("create model job failed: %w", err)
return nil, util.ProcessGcpError(err)
}
// Determine the session ID to use for subsequent queries.
@@ -300,12 +302,17 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
} else if status.Statistics != nil && status.Statistics.SessionInfo != nil {
sessionID = status.Statistics.SessionInfo.SessionID
} else {
return nil, fmt.Errorf("failed to get or create a BigQuery session ID")
return nil, util.NewClientServerError("failed to get or create a BigQuery session ID", http.StatusInternalServerError, nil)
}
getInsightsSQL := fmt.Sprintf("SELECT * FROM ML.GET_INSIGHTS(MODEL %s)", modelID)
connProps := []*bigqueryapi.ConnectionProperty{{Key: "session_id", Value: sessionID}}
return source.RunSQL(ctx, bqClient, getInsightsSQL, "SELECT", nil, connProps)
resp, err := source.RunSQL(ctx, bqClient, getInsightsSQL, "SELECT", nil, connProps)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -35,6 +35,8 @@ import (
const resourceType string = "bigquery-conversational-analytics"
const gdaURLFormat = "https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s:chat"
const instructions = `**INSTRUCTIONS - FOLLOW THESE RULES:**
1. **CONTENT:** Your answer should present the supporting data and then provide a conclusion based on that data.
2. **OUTPUT FORMAT:** Your entire response MUST be in plain text format ONLY.
@@ -172,10 +174,10 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
var tokenStr string
@@ -184,26 +186,26 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if source.UseClientAuthorization() {
// Use client-side access token
if accessToken == "" {
return nil, fmt.Errorf("tool is configured for client OAuth but no token was provided in the request header: %w", util.ErrUnauthorized)
return nil, util.NewClientServerError("tool is configured for client OAuth but no token was provided in the request header", http.StatusUnauthorized, nil)
}
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
} else {
// Get a token source for the Gemini Data Analytics API.
tokenSource, err := source.BigQueryTokenSourceWithScope(ctx, nil)
if err != nil {
return nil, fmt.Errorf("failed to get token source: %w", err)
return nil, util.NewClientServerError("failed to get token source", http.StatusInternalServerError, err)
}
// Use cloud-platform token source for Gemini Data Analytics API
if tokenSource == nil {
return nil, fmt.Errorf("cloud-platform token source is missing")
return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil)
}
token, err := tokenSource.Token()
if err != nil {
return nil, fmt.Errorf("failed to get token from cloud-platform token source: %w", err)
return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err)
}
tokenStr = token.AccessToken
}
@@ -218,14 +220,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
var tableRefs []BQTableReference
if tableRefsJSON != "" {
if err := json.Unmarshal([]byte(tableRefsJSON), &tableRefs); err != nil {
return nil, fmt.Errorf("failed to parse 'table_references' JSON string: %w", err)
return nil, util.NewAgentError("failed to parse 'table_references' JSON string", err)
}
}
if len(source.BigQueryAllowedDatasets()) > 0 {
for _, tableRef := range tableRefs {
if !source.IsDatasetAllowed(tableRef.ProjectID, tableRef.DatasetID) {
return nil, fmt.Errorf("access to dataset '%s.%s' (from table '%s') is not allowed", tableRef.ProjectID, tableRef.DatasetID, tableRef.TableID)
return nil, util.NewAgentError(fmt.Sprintf("access to dataset '%s.%s' (from table '%s') is not allowed", tableRef.ProjectID, tableRef.DatasetID, tableRef.TableID), nil)
}
}
}
@@ -236,11 +238,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if location == "" {
location = "us"
}
caURL := fmt.Sprintf("https://geminidataanalytics.googleapis.com/v1alpha/projects/%s/locations/%s:chat", projectID, location)
caURL := fmt.Sprintf(gdaURLFormat, projectID, location)
headers := map[string]string{
"Authorization": fmt.Sprintf("Bearer %s", tokenStr),
"Content-Type": "application/json",
"Authorization": fmt.Sprintf("Bearer %s", tokenStr),
"Content-Type": "application/json",
"X-Goog-API-Client": util.GDAClientID,
}
payload := CAPayload{
@@ -252,13 +255,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
},
Options: Options{Chart: ChartOptions{Image: ImageOptions{NoImage: map[string]any{}}}},
},
ClientIdEnum: "GENAI_TOOLBOX",
ClientIdEnum: util.GDAClientID,
}
// Call the streaming API
response, err := getStream(caURL, payload, headers, source.GetMaxQueryResultRows())
if err != nil {
return nil, fmt.Errorf("failed to get response from conversational analytics API: %w", err)
// getStream wraps network errors or non-200 responses
return nil, util.NewClientServerError("failed to get response from conversational analytics API", http.StatusInternalServerError, err)
}
return response, nil

View File

@@ -18,6 +18,7 @@ import (
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
bigqueryapi "cloud.google.com/go/bigquery"
@@ -152,25 +153,25 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
sql, ok := paramsMap["sql"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast sql parameter %s", paramsMap["sql"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast sql parameter %s", paramsMap["sql"]), nil)
}
dryRun, ok := paramsMap["dry_run"].(bool)
if !ok {
return nil, fmt.Errorf("unable to cast dry_run parameter %s", paramsMap["dry_run"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast dry_run parameter %s", paramsMap["dry_run"]), nil)
}
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
var connProps []*bigqueryapi.ConnectionProperty
@@ -178,7 +179,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if source.BigQueryWriteMode() == bigqueryds.WriteModeProtected {
session, err = source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session for protected mode: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session for protected mode", http.StatusInternalServerError, err)
}
connProps = []*bigqueryapi.ConnectionProperty{
{Key: "session_id", Value: session.ID},
@@ -187,7 +188,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
dryRunJob, err := bqutil.DryRunQuery(ctx, restService, bqClient.Project(), bqClient.Location, sql, nil, connProps)
if err != nil {
return nil, fmt.Errorf("query validation failed: %w", err)
return nil, util.NewClientServerError("query validation failed", http.StatusInternalServerError, err)
}
statementType := dryRunJob.Statistics.Query.StatementType
@@ -195,13 +196,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
switch source.BigQueryWriteMode() {
case bigqueryds.WriteModeBlocked:
if statementType != "SELECT" {
return nil, fmt.Errorf("write mode is 'blocked', only SELECT statements are allowed")
return nil, util.NewAgentError("write mode is 'blocked', only SELECT statements are allowed", nil)
}
case bigqueryds.WriteModeProtected:
if dryRunJob.Configuration != nil && dryRunJob.Configuration.Query != nil {
if dest := dryRunJob.Configuration.Query.DestinationTable; dest != nil && dest.DatasetId != session.DatasetID {
return nil, fmt.Errorf("protected write mode only supports SELECT statements, or write operations in the anonymous "+
"dataset of a BigQuery session, but destination was %q", dest.DatasetId)
return nil, util.NewAgentError(fmt.Sprintf("protected write mode only supports SELECT statements, or write operations in the anonymous "+
"dataset of a BigQuery session, but destination was %q", dest.DatasetId), nil)
}
}
}
@@ -209,11 +210,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if len(source.BigQueryAllowedDatasets()) > 0 {
switch statementType {
case "CREATE_SCHEMA", "DROP_SCHEMA", "ALTER_SCHEMA":
return nil, fmt.Errorf("dataset-level operations like '%s' are not allowed when dataset restrictions are in place", statementType)
return nil, util.NewAgentError(fmt.Sprintf("dataset-level operations like '%s' are not allowed when dataset restrictions are in place", statementType), nil)
case "CREATE_FUNCTION", "CREATE_TABLE_FUNCTION", "CREATE_PROCEDURE":
return nil, fmt.Errorf("creating stored routines ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType)
return nil, util.NewAgentError(fmt.Sprintf("creating stored routines ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType), nil)
case "CALL":
return nil, fmt.Errorf("calling stored procedures ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType)
return nil, util.NewAgentError(fmt.Sprintf("calling stored procedures ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType), nil)
}
// Use a map to avoid duplicate table names.
@@ -244,7 +245,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
parsedTables, parseErr := bqutil.TableParser(sql, source.BigQueryClient().Project())
if parseErr != nil {
// If parsing fails (e.g., EXECUTE IMMEDIATE), we cannot guarantee safety, so we must fail.
return nil, fmt.Errorf("could not parse tables from query to validate against allowed datasets: %w", parseErr)
return nil, util.NewAgentError("could not parse tables from query to validate against allowed datasets", parseErr)
}
tableNames = parsedTables
}
@@ -254,7 +255,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if len(parts) == 3 {
projectID, datasetID := parts[0], parts[1]
if !source.IsDatasetAllowed(projectID, datasetID) {
return nil, fmt.Errorf("query accesses dataset '%s.%s', which is not in the allowed list", projectID, datasetID)
return nil, util.NewAgentError(fmt.Sprintf("query accesses dataset '%s.%s', which is not in the allowed list", projectID, datasetID), nil)
}
}
}
@@ -264,7 +265,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if dryRunJob != nil {
jobJSON, err := json.MarshalIndent(dryRunJob, "", " ")
if err != nil {
return nil, fmt.Errorf("failed to marshal dry run job to JSON: %w", err)
return nil, util.NewClientServerError("failed to marshal dry run job to JSON", http.StatusInternalServerError, err)
}
return string(jobJSON), nil
}
@@ -275,10 +276,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
return nil, util.NewClientServerError("error getting logger", http.StatusInternalServerError, err)
}
logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", resourceType, sql))
return source.RunSQL(ctx, bqClient, sql, statementType, nil, connProps)
resp, err := source.RunSQL(ctx, bqClient, sql, statementType, nil, connProps)
if err != nil {
return nil, util.NewClientServerError("error running sql", http.StatusInternalServerError, err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,6 +17,7 @@ package bigqueryforecast
import (
"context"
"fmt"
"net/http"
"strings"
bigqueryapi "cloud.google.com/go/bigquery"
@@ -133,34 +134,34 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
historyData, ok := paramsMap["history_data"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast history_data parameter %v", paramsMap["history_data"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast history_data parameter %v", paramsMap["history_data"]), nil)
}
timestampCol, ok := paramsMap["timestamp_col"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast timestamp_col parameter %v", paramsMap["timestamp_col"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast timestamp_col parameter %v", paramsMap["timestamp_col"]), nil)
}
dataCol, ok := paramsMap["data_col"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast data_col parameter %v", paramsMap["data_col"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast data_col parameter %v", paramsMap["data_col"]), nil)
}
idColsRaw, ok := paramsMap["id_cols"].([]any)
if !ok {
return nil, fmt.Errorf("unable to cast id_cols parameter %v", paramsMap["id_cols"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast id_cols parameter %v", paramsMap["id_cols"]), nil)
}
var idCols []string
for _, v := range idColsRaw {
s, ok := v.(string)
if !ok {
return nil, fmt.Errorf("id_cols contains non-string value: %v", v)
return nil, util.NewAgentError(fmt.Sprintf("id_cols contains non-string value: %v", v), nil)
}
idCols = append(idCols, s)
}
@@ -169,13 +170,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if h, ok := paramsMap["horizon"].(float64); ok {
horizon = int(h)
} else {
return nil, fmt.Errorf("unable to cast horizon parameter %v", paramsMap["horizon"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast horizon parameter %v", paramsMap["horizon"]), nil)
}
}
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
var historyDataSource string
@@ -185,7 +186,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
var connProps []*bigqueryapi.ConnectionProperty
session, err := source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
}
if session != nil {
connProps = []*bigqueryapi.ConnectionProperty{
@@ -194,22 +195,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
dryRunJob, err := bqutil.DryRunQuery(ctx, restService, source.BigQueryClient().Project(), source.BigQueryClient().Location, historyData, nil, connProps)
if err != nil {
return nil, fmt.Errorf("query validation failed: %w", err)
return nil, util.ProcessGcpError(err)
}
statementType := dryRunJob.Statistics.Query.StatementType
if statementType != "SELECT" {
return nil, fmt.Errorf("the 'history_data' parameter only supports a table ID or a SELECT query. The provided query has statement type '%s'", statementType)
return nil, util.NewAgentError(fmt.Sprintf("the 'history_data' parameter only supports a table ID or a SELECT query. The provided query has statement type '%s'", statementType), nil)
}
queryStats := dryRunJob.Statistics.Query
if queryStats != nil {
for _, tableRef := range queryStats.ReferencedTables {
if !source.IsDatasetAllowed(tableRef.ProjectId, tableRef.DatasetId) {
return nil, fmt.Errorf("query in history_data accesses dataset '%s.%s', which is not in the allowed list", tableRef.ProjectId, tableRef.DatasetId)
return nil, util.NewAgentError(fmt.Sprintf("query in history_data accesses dataset '%s.%s', which is not in the allowed list", tableRef.ProjectId, tableRef.DatasetId), nil)
}
}
} else {
return nil, fmt.Errorf("could not analyze query in history_data to validate against allowed datasets")
return nil, util.NewAgentError("could not analyze query in history_data to validate against allowed datasets", nil)
}
}
historyDataSource = fmt.Sprintf("(%s)", historyData)
@@ -226,11 +227,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
projectID = source.BigQueryClient().Project()
datasetID = parts[0]
default:
return nil, fmt.Errorf("invalid table ID format for 'history_data': %q. Expected 'dataset.table' or 'project.dataset.table'", historyData)
return nil, util.NewAgentError(fmt.Sprintf("invalid table ID format for 'history_data': %q. Expected 'dataset.table' or 'project.dataset.table'", historyData), nil)
}
if !source.IsDatasetAllowed(projectID, datasetID) {
return nil, fmt.Errorf("access to dataset '%s.%s' (from table '%s') is not allowed", projectID, datasetID, historyData)
return nil, util.NewAgentError(fmt.Sprintf("access to dataset '%s.%s' (from table '%s') is not allowed", projectID, datasetID, historyData), nil)
}
}
historyDataSource = fmt.Sprintf("TABLE `%s`", historyData)
@@ -243,15 +244,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
sql := fmt.Sprintf(`SELECT *
FROM AI.FORECAST(
%s,
data_col => '%s',
timestamp_col => '%s',
horizon => %d%s)`,
%s,
data_col => '%s',
timestamp_col => '%s',
horizon => %d%s)`,
historyDataSource, dataCol, timestampCol, horizon, idColsArg)
session, err := source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
}
var connProps []*bigqueryapi.ConnectionProperty
if session != nil {
@@ -264,11 +265,15 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
// Log the query executed for debugging.
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
return nil, util.NewClientServerError("error getting logger", http.StatusInternalServerError, err)
}
logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", resourceType, sql))
return source.RunSQL(ctx, bqClient, sql, "SELECT", nil, connProps)
resp, err := source.RunSQL(ctx, bqClient, sql, "SELECT", nil, connProps)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,6 +17,7 @@ package bigquerygetdatasetinfo
import (
"context"
"fmt"
"net/http"
bigqueryapi "cloud.google.com/go/bigquery"
yaml "github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
)
@@ -120,38 +122,38 @@ type Tool struct {
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
mapParams := params.AsMap()
projectId, ok := mapParams[projectKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
// Updated: Use fmt.Sprintf for formatting, pass nil as cause
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil)
}
datasetId, ok := mapParams[datasetKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil)
}
bqClient, _, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
if !source.IsDatasetAllowed(projectId, datasetId) {
return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId)
return nil, util.NewAgentError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), nil)
}
dsHandle := bqClient.DatasetInProject(projectId, datasetId)
metadata, err := dsHandle.Metadata(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get metadata for dataset %s (in project %s): %w", datasetId, projectId, err)
return nil, util.ProcessGcpError(err)
}
return metadata, nil

View File

@@ -17,6 +17,7 @@ package bigquerygettableinfo
import (
"context"
"fmt"
"net/http"
bigqueryapi "cloud.google.com/go/bigquery"
yaml "github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
)
@@ -125,35 +127,35 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
mapParams := params.AsMap()
projectId, ok := mapParams[projectKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil)
}
datasetId, ok := mapParams[datasetKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil)
}
tableId, ok := mapParams[tableKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", tableKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", tableKey), nil)
}
if !source.IsDatasetAllowed(projectId, datasetId) {
return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId)
return nil, util.NewAgentError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), nil)
}
bqClient, _, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
dsHandle := bqClient.DatasetInProject(projectId, datasetId)
@@ -161,7 +163,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
metadata, err := tableHandle.Metadata(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get metadata for table %s.%s.%s: %w", projectId, datasetId, tableId, err)
return nil, util.ProcessGcpError(err)
}
return metadata, nil

View File

@@ -17,12 +17,14 @@ package bigquerylistdatasetids
import (
"context"
"fmt"
"net/http"
bigqueryapi "cloud.google.com/go/bigquery"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
"google.golang.org/api/iterator"
@@ -120,10 +122,10 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
if len(source.BigQueryAllowedDatasets()) > 0 {
@@ -132,12 +134,12 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
mapParams := params.AsMap()
projectId, ok := mapParams[projectKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil)
}
bqClient, _, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
datasetIterator := bqClient.Datasets(ctx)
datasetIterator.ProjectID = projectId
@@ -149,7 +151,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
break
}
if err != nil {
return nil, fmt.Errorf("unable to iterate through datasets: %w", err)
return nil, util.ProcessGcpError(err)
}
// Remove leading and trailing quotes

View File

@@ -17,6 +17,7 @@ package bigquerylisttableids
import (
"context"
"fmt"
"net/http"
bigqueryapi "cloud.google.com/go/bigquery"
yaml "github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
"google.golang.org/api/iterator"
@@ -123,31 +125,30 @@ type Tool struct {
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
mapParams := params.AsMap()
projectId, ok := mapParams[projectKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", projectKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", projectKey), nil)
}
datasetId, ok := mapParams[datasetKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", datasetKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", datasetKey), nil)
}
if !source.IsDatasetAllowed(projectId, datasetId) {
return nil, fmt.Errorf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId)
return nil, util.NewAgentError(fmt.Sprintf("access denied to dataset '%s' because it is not in the configured list of allowed datasets for project '%s'", datasetId, projectId), nil)
}
bqClient, _, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
dsHandle := bqClient.DatasetInProject(projectId, datasetId)
@@ -160,7 +161,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
break
}
if err != nil {
return nil, fmt.Errorf("failed to iterate through tables in dataset %s.%s: %w", projectId, datasetId, err)
return nil, util.ProcessGcpError(err)
}
// Remove leading and trailing quotes

View File

@@ -17,6 +17,7 @@ package bigquerysearchcatalog
import (
"context"
"fmt"
"net/http"
"strings"
dataplexapi "cloud.google.com/go/dataplex/apiv1"
@@ -26,6 +27,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/iterator"
)
@@ -186,28 +188,31 @@ func ExtractType(resourceString string) string {
return typeMap[resourceString[lastIndex+1:]]
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
pageSize := int32(paramsMap["pageSize"].(int))
prompt, _ := paramsMap["prompt"].(string)
projectIdSlice, err := parameters.ConvertAnySliceToTyped(paramsMap["projectIds"].([]any), "string")
if err != nil {
return nil, fmt.Errorf("can't convert projectIds to array of strings: %s", err)
return nil, util.NewAgentError(fmt.Sprintf("can't convert projectIds to array of strings: %s", err), err)
}
projectIds := projectIdSlice.([]string)
datasetIdSlice, err := parameters.ConvertAnySliceToTyped(paramsMap["datasetIds"].([]any), "string")
if err != nil {
return nil, fmt.Errorf("can't convert datasetIds to array of strings: %s", err)
return nil, util.NewAgentError(fmt.Sprintf("can't convert datasetIds to array of strings: %s", err), err)
}
datasetIds := datasetIdSlice.([]string)
typesSlice, err := parameters.ConvertAnySliceToTyped(paramsMap["types"].([]any), "string")
if err != nil {
return nil, fmt.Errorf("can't convert types to array of strings: %s", err)
return nil, util.NewAgentError(fmt.Sprintf("can't convert types to array of strings: %s", err), err)
}
types := typesSlice.([]string)
@@ -223,17 +228,17 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if source.UseClientAuthorization() {
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
catalogClient, err = dataplexClientCreator(tokenStr)
if err != nil {
return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
return nil, util.NewClientServerError("error creating client from OAuth access token", http.StatusInternalServerError, err)
}
}
it := catalogClient.SearchEntries(ctx, req)
if it == nil {
return nil, fmt.Errorf("failed to create search entries iterator for project %q", source.BigQueryProject())
return nil, util.NewClientServerError(fmt.Sprintf("failed to create search entries iterator for project %q", source.BigQueryProject()), http.StatusInternalServerError, nil)
}
var results []Response
@@ -243,7 +248,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
break
}
if err != nil {
break
return nil, util.ProcessGcpError(err)
}
entrySource := entry.DataplexEntry.GetEntrySource()
resp := Response{

View File

@@ -17,6 +17,7 @@ package bigquerysql
import (
"context"
"fmt"
"net/http"
"reflect"
"strings"
@@ -27,6 +28,7 @@ import (
bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
"github.com/googleapis/genai-toolbox/internal/tools"
bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
)
@@ -103,11 +105,10 @@ type Tool struct {
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
highLevelParams := make([]bigqueryapi.QueryParameter, 0, len(t.Parameters))
@@ -116,7 +117,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
paramsMap := params.AsMap()
newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract template params %w", err)
return nil, util.NewAgentError("unable to extract template params", err)
}
for _, p := range t.Parameters {
@@ -127,13 +128,13 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if arrayParam, ok := p.(*parameters.ArrayParameter); ok {
arrayParamValue, ok := value.([]any)
if !ok {
return nil, fmt.Errorf("unable to convert parameter `%s` to []any", name)
return nil, util.NewAgentError(fmt.Sprintf("unable to convert parameter `%s` to []any", name), nil)
}
itemType := arrayParam.GetItems().GetType()
var err error
value, err = parameters.ConvertAnySliceToTyped(arrayParamValue, itemType)
if err != nil {
return nil, fmt.Errorf("unable to convert parameter `%s` from []any to typed slice: %w", name, err)
return nil, util.NewAgentError(fmt.Sprintf("unable to convert parameter `%s` from []any to typed slice", name), err)
}
}
@@ -161,7 +162,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
lowLevelParam.ParameterType.Type = "ARRAY"
itemType, err := bqutil.BQTypeStringFromToolType(arrayParam.GetItems().GetType())
if err != nil {
return nil, err
return nil, util.NewAgentError("unable to get BigQuery type from tool parameter type", err)
}
lowLevelParam.ParameterType.ArrayType = &bigqueryrestapi.QueryParameterType{Type: itemType}
@@ -178,7 +179,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
// Handle scalar types based on their defined type.
bqType, err := bqutil.BQTypeStringFromToolType(p.GetType())
if err != nil {
return nil, err
return nil, util.NewAgentError("unable to get BigQuery type from tool parameter type", err)
}
lowLevelParam.ParameterType.Type = bqType
lowLevelParam.ParameterValue.Value = fmt.Sprintf("%v", value)
@@ -190,7 +191,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if source.BigQuerySession() != nil {
session, err := source.BigQuerySession()(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get BigQuery session: %w", err)
return nil, util.NewClientServerError("failed to get BigQuery session", http.StatusInternalServerError, err)
}
if session != nil {
// Add session ID to the connection properties for subsequent calls.
@@ -200,17 +201,20 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
bqClient, restService, err := source.RetrieveClientAndService(accessToken)
if err != nil {
return nil, err
return nil, util.NewClientServerError("failed to retrieve BigQuery client", http.StatusInternalServerError, err)
}
dryRunJob, err := bqutil.DryRunQuery(ctx, restService, bqClient.Project(), bqClient.Location, newStatement, lowLevelParams, connProps)
if err != nil {
return nil, fmt.Errorf("query validation failed: %w", err)
return nil, util.ProcessGcpError(err)
}
statementType := dryRunJob.Statistics.Query.StatementType
return source.RunSQL(ctx, bqClient, newStatement, statementType, highLevelParams, connProps)
resp, err := source.RunSQL(ctx, bqClient, newStatement, statementType, highLevelParams, connProps)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package bigtable
import (
"context"
"fmt"
"net/http"
"cloud.google.com/go/bigtable"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -96,24 +98,28 @@ type Tool struct {
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract template params %w", err)
return nil, util.NewAgentError("unable to extract template params", err)
}
newParams, err := parameters.GetParams(t.Parameters, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract standard params %w", err)
return nil, util.NewAgentError("unable to extract standard params", err)
}
return source.RunSQL(ctx, newStatement, t.Parameters, newParams)
resp, err := source.RunSQL(ctx, newStatement, t.Parameters, newParams)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -4,7 +4,7 @@
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
@@ -17,12 +17,14 @@ package cassandracql
import (
"context"
"fmt"
"net/http"
gocql "github.com/apache/cassandra-gocql-driver/v2"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -107,23 +109,27 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool {
}
// Invoke implements tools.Tool.
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract template params %w", err)
return nil, util.NewAgentError("unable to extract template params", err)
}
newParams, err := parameters.GetParams(t.Parameters, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract standard params %w", err)
return nil, util.NewAgentError("unable to extract standard params", err)
}
return source.RunSQL(ctx, newStatement, newParams)
resp, err := source.RunSQL(ctx, newStatement, newParams)
if err != nil {
return nil, util.ProcessGeneralError(err)
}
return resp, nil
}
// Manifest implements tools.Tool.

View File

@@ -17,11 +17,13 @@ package clickhouse
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -87,18 +89,22 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
sql, ok := paramsMap["sql"].(string)
if !ok {
return nil, fmt.Errorf("unable to cast sql parameter %s", paramsMap["sql"])
return nil, util.NewAgentError(fmt.Sprintf("unable to cast sql parameter %s", paramsMap["sql"]), nil)
}
return source.RunSQL(ctx, sql, nil)
resp, err := source.RunSQL(ctx, sql, nil)
if err != nil {
return nil, util.ProcessGeneralError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package clickhouse
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -86,10 +88,10 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
// Query to list all databases
@@ -97,7 +99,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
out, err := source.RunSQL(ctx, query, nil)
if err != nil {
return nil, err
return nil, util.ProcessGeneralError(err)
}
return out, nil

View File

@@ -17,11 +17,13 @@ package clickhouse
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -90,34 +92,37 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
mapParams := params.AsMap()
database, ok := mapParams[databaseKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", databaseKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", databaseKey), nil)
}
// Query to list all tables in the specified database
// Note: formatting identifier directly is risky if input is untrusted, but standard for this tool structure.
query := fmt.Sprintf("SHOW TABLES FROM %s", database)
out, err := source.RunSQL(ctx, query, nil)
if err != nil {
return nil, err
return nil, util.ProcessGeneralError(err)
}
res, ok := out.([]any)
if !ok {
return nil, fmt.Errorf("unable to convert result to list")
return nil, util.NewClientServerError("unable to convert result to list", http.StatusInternalServerError, nil)
}
var tables []map[string]any
for _, item := range res {
tableMap, ok := item.(map[string]any)
if !ok {
return nil, fmt.Errorf("unexpected type in result: got %T, want map[string]any", item)
return nil, util.NewClientServerError(fmt.Sprintf("unexpected type in result: got %T, want map[string]any", item), http.StatusInternalServerError, nil)
}
tableMap["database"] = database
tables = append(tables, tableMap)

View File

@@ -17,11 +17,13 @@ package clickhouse
import (
"context"
"fmt"
"net/http"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -88,24 +90,28 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, token tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract template params: %w", err)
return nil, util.NewAgentError("unable to extract template params", err)
}
newParams, err := parameters.GetParams(t.Parameters, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract standard params: %w", err)
return nil, util.NewAgentError("unable to extract standard params", err)
}
return source.RunSQL(ctx, newStatement, newParams)
resp, err := source.RunSQL(ctx, newStatement, newParams)
if err != nil {
return nil, util.ProcessGeneralError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -18,11 +18,13 @@ import (
"context"
"encoding/json"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -119,17 +121,16 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
// Invoke executes the tool logic
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
paramsMap := params.AsMap()
query, ok := paramsMap["query"].(string)
if !ok {
return nil, fmt.Errorf("query parameter not found or not a string")
return nil, util.NewAgentError("query parameter not found or not a string", nil)
}
// Parse the access token if provided
@@ -138,7 +139,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
var err error
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
@@ -154,9 +155,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
bodyBytes, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request payload: %w", err)
return nil, util.NewClientServerError("failed to marshal request payload", http.StatusInternalServerError, err)
}
return source.RunQuery(ctx, tokenStr, bodyBytes)
resp, err := source.RunQuery(ctx, tokenStr, bodyBytes)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package fhirfetchpage
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -93,24 +95,31 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
url, ok := params.AsMap()[pageURLKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", pageURLKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", pageURLKey), nil)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.FHIRFetchPage(ctx, url, tokenStr)
resp, err := source.FHIRFetchPage(ctx, url, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,6 +17,7 @@ package fhirpatienteverything
import (
"context"
"fmt"
"net/http"
"strings"
"github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/googleapi"
)
@@ -116,26 +118,27 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedFHIRStores())
if err != nil {
return nil, err
// ValidateAndFetchStoreID usually returns input validation errors
return nil, util.NewAgentError("failed to validate store ID", err)
}
patientID, ok := params.AsMap()[patientIDKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", patientIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", patientIDKey), nil)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
@@ -143,11 +146,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if val, ok := params.AsMap()[typeFilterKey]; ok {
types, ok := val.([]any)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string array", typeFilterKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string array", typeFilterKey), nil)
}
typeFilterSlice, err := parameters.ConvertAnySliceToTyped(types, "string")
if err != nil {
return nil, fmt.Errorf("can't convert '%s' to array of strings: %s", typeFilterKey, err)
return nil, util.NewAgentError(fmt.Sprintf("can't convert '%s' to array of strings: %s", typeFilterKey, err), err)
}
if len(typeFilterSlice.([]string)) != 0 {
opts = append(opts, googleapi.QueryParameter("_type", strings.Join(typeFilterSlice.([]string), ",")))
@@ -156,13 +159,18 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if since, ok := params.AsMap()[sinceFilterKey]; ok {
sinceStr, ok := since.(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", sinceFilterKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", sinceFilterKey), nil)
}
if sinceStr != "" {
opts = append(opts, googleapi.QueryParameter("_since", sinceStr))
}
}
return source.FHIRPatientEverything(storeID, patientID, tokenStr, opts)
resp, err := source.FHIRPatientEverything(storeID, patientID, tokenStr, opts)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,6 +17,7 @@ package fhirpatientsearch
import (
"context"
"fmt"
"net/http"
"strings"
"github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/googleapi"
)
@@ -150,22 +152,22 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedFHIRStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
@@ -179,14 +181,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
var ok bool
summary, ok = v.(bool)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a boolean", summaryKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a boolean", summaryKey), nil)
}
continue
}
val, ok := v.(string)
if !ok {
return nil, fmt.Errorf("invalid parameter '%s'; expected a string", k)
return nil, util.NewAgentError(fmt.Sprintf("invalid parameter '%s'; expected a string", k), nil)
}
if val == "" {
continue
@@ -205,7 +207,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
parts := strings.Split(val, "/")
if len(parts) != 2 {
return nil, fmt.Errorf("invalid '%s' format; expected YYYY-MM-DD/YYYY-MM-DD", k)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' format; expected YYYY-MM-DD/YYYY-MM-DD", k), nil)
}
var values []string
if parts[0] != "" {
@@ -229,13 +231,17 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
case familyNameKey:
opts = append(opts, googleapi.QueryParameter("family", val))
default:
return nil, fmt.Errorf("unexpected parameter key %q", k)
return nil, util.NewAgentError(fmt.Sprintf("unexpected parameter key %q", k), nil)
}
}
if summary {
opts = append(opts, googleapi.QueryParameter("_summary", "text"))
}
return source.FHIRPatientSearch(storeID, tokenStr, opts)
resp, err := source.FHIRPatientSearch(storeID, tokenStr, opts)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package gethealthcaredataset
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -90,19 +92,23 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetDataset(tokenStr)
resp, err := source.GetDataset(tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package getdicomstore
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -107,23 +109,27 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetDICOMStore(storeID, tokenStr)
resp, err := source.GetDICOMStore(storeID, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package getdicomstoremetrics
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -107,23 +109,27 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetDICOMStoreMetrics(storeID, tokenStr)
resp, err := source.GetDICOMStoreMetrics(storeID, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package getfhirresource
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -112,32 +114,36 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedFHIRStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
resType, ok := params.AsMap()[typeKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", typeKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", typeKey), nil)
}
resID, ok := params.AsMap()[idKey].(string)
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", idKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid or missing '%s' parameter; expected a string", idKey), nil)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetFHIRResource(storeID, resType, resID, tokenStr)
resp, err := source.GetFHIRResource(storeID, resType, resID, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package getfhirstore
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -107,23 +109,27 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedFHIRStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetFHIRStore(storeID, tokenStr)
resp, err := source.GetFHIRStore(storeID, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package getfhirstoremetrics
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -107,23 +109,27 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedFHIRStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.GetFHIRStoreMetrics(storeID, tokenStr)
resp, err := source.GetFHIRStoreMetrics(storeID, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package listdicomstores
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -90,19 +92,23 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.ListDICOMStores(tokenStr)
resp, err := source.ListDICOMStores(tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,11 +17,13 @@ package listfhirstores
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/healthcare/v1"
)
@@ -90,19 +92,23 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
return source.ListFHIRStores(tokenStr)
resp, err := source.ListFHIRStores(tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package retrieverendereddicominstance
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
@@ -117,40 +119,44 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
study, ok := params.AsMap()[studyInstanceUIDKey].(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", studyInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", studyInstanceUIDKey), nil)
}
series, ok := params.AsMap()[seriesInstanceUIDKey].(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", seriesInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", seriesInstanceUIDKey), nil)
}
sop, ok := params.AsMap()[sopInstanceUIDKey].(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", sopInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", sopInstanceUIDKey), nil)
}
frame, ok := params.AsMap()[frameNumberKey].(int)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected an integer", frameNumberKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected an integer", frameNumberKey), nil)
}
return source.RetrieveRenderedDICOMInstance(storeID, study, series, sop, frame, tokenStr)
resp, err := source.RetrieveRenderedDICOMInstance(storeID, study, series, sop, frame, tokenStr)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,6 +17,7 @@ package searchdicominstances
import (
"context"
"fmt"
"net/http"
"strings"
"github.com/goccy/go-yaml"
@@ -24,6 +25,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/googleapi"
)
@@ -131,33 +133,33 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
opts, err := common.ParseDICOMSearchParameters(params, []string{sopInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to parse DICOM search parameters", err)
}
paramsMap := params.AsMap()
dicomWebPath := "instances"
if studyInstanceUID, ok := paramsMap[studyInstanceUIDKey]; ok {
id, ok := studyInstanceUID.(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", studyInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", studyInstanceUIDKey), nil)
}
if id != "" {
dicomWebPath = fmt.Sprintf("studies/%s/instances", id)
@@ -166,7 +168,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if seriesInstanceUID, ok := paramsMap[seriesInstanceUIDKey]; ok {
id, ok := seriesInstanceUID.(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", seriesInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", seriesInstanceUIDKey), nil)
}
if id != "" {
if dicomWebPath != "instances" {
@@ -176,7 +178,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
}
}
}
return source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
resp, err := source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package searchdicomseries
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/googleapi"
)
@@ -128,40 +130,44 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
opts, err := common.ParseDICOMSearchParameters(params, []string{seriesInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to parse DICOM search parameters", err)
}
paramsMap := params.AsMap()
dicomWebPath := "series"
if studyInstanceUID, ok := paramsMap[studyInstanceUIDKey]; ok {
id, ok := studyInstanceUID.(string)
if !ok {
return nil, fmt.Errorf("invalid '%s' parameter; expected a string", studyInstanceUIDKey)
return nil, util.NewAgentError(fmt.Sprintf("invalid '%s' parameter; expected a string", studyInstanceUIDKey), nil)
}
if id != "" {
dicomWebPath = fmt.Sprintf("studies/%s/series", id)
}
}
return source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
resp, err := source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

View File

@@ -17,12 +17,14 @@ package searchdicomstudies
import (
"context"
"fmt"
"net/http"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/cloudhealthcare/common"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"google.golang.org/api/googleapi"
)
@@ -124,28 +126,32 @@ func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
}
storeID, err := common.ValidateAndFetchStoreID(params, source.AllowedDICOMStores())
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to validate store ID", err)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
}
}
opts, err := common.ParseDICOMSearchParameters(params, []string{studyInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey})
if err != nil {
return nil, err
return nil, util.NewAgentError("failed to parse DICOM search parameters", err)
}
dicomWebPath := "studies"
return source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
resp, err := source.SearchDICOM(t.Type, storeID, dicomWebPath, tokenStr, opts)
if err != nil {
return nil, util.ProcessGcpError(err)
}
return resp, nil
}
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {

Some files were not shown because too many files have changed in this diff Show More