mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-02-13 16:45:01 -05:00
Compare commits
10 Commits
processing
...
feat/add-g
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9a2092ffd6 | ||
|
|
5fcf3df882 | ||
|
|
195767bdcd | ||
|
|
c5524d32f5 | ||
|
|
e1739abd81 | ||
|
|
478a0bdb59 | ||
|
|
2d341acaa6 | ||
|
|
f032389a07 | ||
|
|
32610d71a3 | ||
|
|
32cb4db712 |
3
.gitignore
vendored
3
.gitignore
vendored
@@ -18,6 +18,9 @@ node_modules
|
||||
# coverage
|
||||
.coverage
|
||||
|
||||
# python
|
||||
__pycache__/
|
||||
|
||||
# executable
|
||||
genai-toolbox
|
||||
toolbox
|
||||
@@ -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/"
|
||||
|
||||
26
CHANGELOG.md
26
CHANGELOG.md
@@ -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)
|
||||
|
||||
|
||||
|
||||
108
GEMINI.md
Normal file
108
GEMINI.md
Normal file
@@ -0,0 +1,108 @@
|
||||
# MCP Toolbox Context
|
||||
|
||||
This file (symlinked as `CLAUDE.md` and `AGENTS.md`) provides context and guidelines for AI agents working on the MCP Toolbox for Databases project. It summarizes key information from `CONTRIBUTING.md` and `DEVELOPER.md`.
|
||||
|
||||
## Project Overview
|
||||
|
||||
**MCP Toolbox for Databases** is a Go-based project designed to provide Model Context Protocol (MCP) tools for various data sources and services. It allows Large Language Models (LLMs) to interact with databases and other tools safely and efficiently.
|
||||
|
||||
## Tech Stack
|
||||
|
||||
- **Language:** Go (1.23+)
|
||||
- **Documentation:** Hugo (Extended Edition v0.146.0+)
|
||||
- **Containerization:** Docker
|
||||
- **CI/CD:** GitHub Actions, Google Cloud Build
|
||||
- **Linting:** `golangci-lint`
|
||||
|
||||
## Key Directories
|
||||
|
||||
- `cmd/`: Application entry points.
|
||||
- `internal/sources/`: Implementations of database sources (e.g., Postgres, BigQuery).
|
||||
- `internal/tools/`: Implementations of specific tools for each source.
|
||||
- `tests/`: Integration tests.
|
||||
- `docs/`: Project documentation (Hugo site).
|
||||
|
||||
## Development Workflow
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Go 1.23 or later.
|
||||
- Docker (for building container images and running some tests).
|
||||
- Access to necessary Google Cloud resources for integration testing (if applicable).
|
||||
|
||||
### Building and Running
|
||||
|
||||
1. **Build Binary:** `go build -o toolbox`
|
||||
2. **Run Server:** `go run .` (Listens on port 5000 by default)
|
||||
3. **Run with Help:** `go run . --help`
|
||||
4. **Test Endpoint:** `curl http://127.0.0.1:5000`
|
||||
|
||||
### Testing
|
||||
|
||||
- **Unit Tests:** `go test -race -v ./cmd/... ./internal/...`
|
||||
- **Integration Tests:**
|
||||
- Run specific source tests: `go test -race -v ./tests/<source_dir>`
|
||||
- Example: `go test -race -v ./tests/alloydbpg`
|
||||
- Add new sources to `.ci/integration.cloudbuild.yaml`
|
||||
- **Linting:** `golangci-lint run --fix`
|
||||
|
||||
## Developing Documentation
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Hugo (Extended Edition v0.146.0+)
|
||||
- Node.js (for `npm ci`)
|
||||
|
||||
### Running Local Server
|
||||
|
||||
1. Navigate to `.hugo` directory: `cd .hugo`
|
||||
2. Install dependencies: `npm ci`
|
||||
3. Start server: `hugo server`
|
||||
|
||||
### Versioning Workflows
|
||||
|
||||
1. **Deploy In-development docs**: Merges to main -> `/dev/`.
|
||||
2. **Deploy Versioned Docs**: New Release -> `/<version>/` and root.
|
||||
3. **Deploy Previous Version Docs**: Manual workflow for older versions.
|
||||
|
||||
## Coding Conventions
|
||||
|
||||
### Tool Naming
|
||||
|
||||
- **Tool Name:** `snake_case` (e.g., `list_collections`, `run_query`).
|
||||
- Do *not* include the product name (e.g., avoid `firestore_list_collections`).
|
||||
- **Tool Type:** `kebab-case` (e.g., `firestore-list-collections`).
|
||||
- *Must* include the product name.
|
||||
|
||||
### Branching and Commits
|
||||
|
||||
- **Branch Naming:** `feat/`, `fix/`, `docs/`, `chore/` (e.g., `feat/add-gemini-md`).
|
||||
- **Commit Messages:** [Conventional Commits](https://www.conventionalcommits.org/) format.
|
||||
- Format: `<type>(<scope>): <description>`
|
||||
- Example: `feat(source/postgres): add new connection option`
|
||||
- Types: `feat`, `fix`, `docs`, `chore`, `test`, `ci`, `refactor`, `revert`, `style`.
|
||||
|
||||
## Adding New Features
|
||||
|
||||
### Adding a New Data Source
|
||||
|
||||
1. Create a new directory: `internal/sources/<newdb>`.
|
||||
2. Define `Config` and `Source` structs in `internal/sources/<newdb>/<newdb>.go`.
|
||||
3. Implement `SourceConfig` interface (`SourceConfigType`, `Initialize`).
|
||||
4. Implement `Source` interface (`SourceType`).
|
||||
5. Implement `init()` to register the source.
|
||||
6. Add unit tests in `internal/sources/<newdb>/<newdb>_test.go`.
|
||||
|
||||
### Adding a New Tool
|
||||
|
||||
1. Create a new directory: `internal/tools/<newdb>/<toolname>`.
|
||||
2. Define `Config` and `Tool` structs.
|
||||
3. Implement `ToolConfig` interface (`ToolConfigType`, `Initialize`).
|
||||
4. Implement `Tool` interface (`Invoke`, `ParseParams`, `Manifest`, `McpManifest`, `Authorized`).
|
||||
5. Implement `init()` to register the tool.
|
||||
6. Add unit tests.
|
||||
|
||||
### Adding Documentation
|
||||
|
||||
- Add source documentation to `docs/en/resources/sources/`.
|
||||
- Add tool documentation to `docs/en/resources/tools/`.
|
||||
14
README.md
14
README.md
@@ -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
257
cmd/internal/imports.go
Normal 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"
|
||||
)
|
||||
@@ -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), ¶ms); 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
|
||||
}
|
||||
@@ -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
251
cmd/internal/options.go
Normal 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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
46
cmd/internal/persistent_flags.go
Normal file
46
cmd/internal/persistent_flags.go
Normal 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.")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
349
cmd/internal/tools_file.go
Normal 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)
|
||||
}
|
||||
2141
cmd/internal/tools_file_test.go
Normal file
2141
cmd/internal/tools_file_test.go
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
}
|
||||
}
|
||||
891
cmd/root.go
891
cmd/root.go
File diff suppressed because it is too large
Load Diff
2204
cmd/root_test.go
2204
cmd/root_test.go
File diff suppressed because it is too large
Load Diff
@@ -1 +1 @@
|
||||
0.26.0
|
||||
0.27.0
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 %}}
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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} -->
|
||||
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -16,11 +16,13 @@ package cloudloggingadminlistlognames
|
||||
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"
|
||||
)
|
||||
|
||||
@@ -89,10 +91,10 @@ type Tool struct {
|
||||
Parameters parameters.Parameters `yaml:"parameters"`
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
limit := defaultLimit
|
||||
@@ -100,18 +102,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
if val, ok := paramsMap["limit"].(int); ok && val > 0 {
|
||||
limit = val
|
||||
} else if ok && val < 0 {
|
||||
return nil, fmt.Errorf("limit must be greater than or equal to 1")
|
||||
return nil, util.NewAgentError("limit must be greater than or equal to 1", nil)
|
||||
}
|
||||
|
||||
tokenString := ""
|
||||
if source.UseClientAuthorization() {
|
||||
tokenString, err = accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse access token: %w", err)
|
||||
return nil, util.NewClientServerError("failed to parse access token", http.StatusUnauthorized, err)
|
||||
}
|
||||
}
|
||||
|
||||
return source.ListLogNames(ctx, limit, tokenString)
|
||||
resp, err := source.ListLogNames(ctx, limit, tokenString)
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claimsMap map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
|
||||
@@ -16,11 +16,13 @@ package cloudloggingadminlistresourcetypes
|
||||
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"
|
||||
)
|
||||
|
||||
@@ -84,21 +86,25 @@ type Tool struct {
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
tokenString := ""
|
||||
if source.UseClientAuthorization() {
|
||||
tokenString, err = accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse access token: %w", err)
|
||||
return nil, util.NewClientServerError("failed to parse access token", http.StatusUnauthorized, err)
|
||||
}
|
||||
}
|
||||
|
||||
return source.ListResourceTypes(ctx, tokenString)
|
||||
resp, err := source.ListResourceTypes(ctx, tokenString)
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claimsMap map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
|
||||
@@ -16,6 +16,7 @@ package cloudloggingadminquerylogs
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-yaml"
|
||||
@@ -23,6 +24,7 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
cla "github.com/googleapis/genai-toolbox/internal/sources/cloudloggingadmin"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
)
|
||||
|
||||
@@ -104,10 +106,10 @@ type Tool struct {
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
// Parse parameters
|
||||
@@ -119,7 +121,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
if val, ok := paramsMap["limit"].(int); ok && val > 0 {
|
||||
limit = val
|
||||
} else if ok && val < 0 {
|
||||
return nil, fmt.Errorf("limit must be greater than or equal to 1")
|
||||
return nil, util.NewAgentError("limit must be greater than or equal to 1", nil)
|
||||
}
|
||||
|
||||
// Check for verbosity of output
|
||||
@@ -129,7 +131,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
var filter string
|
||||
if f, ok := paramsMap["filter"].(string); ok {
|
||||
if len(f) == 0 {
|
||||
return nil, fmt.Errorf("filter cannot be empty if provided")
|
||||
return nil, util.NewAgentError("filter cannot be empty if provided", nil)
|
||||
}
|
||||
filter = f
|
||||
}
|
||||
@@ -138,7 +140,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
var startTime string
|
||||
if val, ok := paramsMap["startTime"].(string); ok && val != "" {
|
||||
if _, err := time.Parse(time.RFC3339, val); err != nil {
|
||||
return nil, fmt.Errorf("startTime must be in RFC3339 format (e.g., 2025-12-09T00:00:00Z): %w", err)
|
||||
return nil, util.NewAgentError(fmt.Sprintf("startTime must be in RFC3339 format (e.g., 2025-12-09T00:00:00Z): %v", err), err)
|
||||
}
|
||||
startTime = val
|
||||
} else {
|
||||
@@ -149,7 +151,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
var endTime string
|
||||
if val, ok := paramsMap["endTime"].(string); ok && val != "" {
|
||||
if _, err := time.Parse(time.RFC3339, val); err != nil {
|
||||
return nil, fmt.Errorf("endTime must be in RFC3339 format (e.g., 2025-12-09T23:59:59Z): %w", err)
|
||||
return nil, util.NewAgentError(fmt.Sprintf("endTime must be in RFC3339 format (e.g., 2025-12-09T23:59:59Z): %v", err), err)
|
||||
}
|
||||
endTime = val
|
||||
}
|
||||
@@ -158,7 +160,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
if source.UseClientAuthorization() {
|
||||
tokenString, err = accessToken.ParseBearerToken()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse access token: %w", err)
|
||||
return nil, util.NewClientServerError("failed to parse access token", http.StatusUnauthorized, err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,7 +173,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
||||
Limit: limit,
|
||||
}
|
||||
|
||||
return source.QueryLogs(ctx, queryParams, tokenString)
|
||||
resp, err := source.QueryLogs(ctx, queryParams, tokenString)
|
||||
if err != nil {
|
||||
return nil, util.ProcessGcpError(err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claimsMap map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
|
||||
@@ -23,6 +23,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"
|
||||
)
|
||||
|
||||
@@ -93,22 +94,26 @@ 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()
|
||||
projectID, ok := paramsMap["projectId"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("projectId parameter not found or not a string")
|
||||
return nil, util.NewAgentError("projectId parameter not found or not a string", nil)
|
||||
}
|
||||
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)
|
||||
}
|
||||
return source.RunQuery(projectID, query)
|
||||
resp, err := source.RunQuery(projectID, query)
|
||||
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
Reference in New Issue
Block a user