mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-10 07:58:12 -05:00
feat(tools/dataplex-search-aspect-types): Add support for dataplex-search-aspect-types tool (#1061)
Added support for search aspect types tool in Dataplex. Fixes #1056 --------- Co-authored-by: Averi Kitsch <akitsch@google.com>
This commit is contained in:
committed by
GitHub
parent
31ed87861d
commit
d940187c85
@@ -52,6 +52,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/bigtable"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
|
||||
_ "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/firestore/firestoredeletedocuments"
|
||||
|
||||
@@ -1250,7 +1250,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"dataplex-tools": tools.ToolsetConfig{
|
||||
Name: "dataplex-tools",
|
||||
ToolNames: []string{"dataplex_search_entries", "dataplex_lookup_entry"},
|
||||
ToolNames: []string{"dataplex_search_entries", "dataplex_lookup_entry", "dataplex_search_aspect_types"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -218,22 +218,45 @@ Search syntax supports the following qualifiers:
|
||||
- "label.foo" - Matches BigQuery resources that have a label whose key equals foo as a string.
|
||||
- "type=TYPE" - Matches resources of a specific entry type or its type alias.
|
||||
- "projectid:bar" - Matches resources within Google Cloud projects that match bar as a substring in the ID.
|
||||
- "parent:x" - Matches x as a substring of the hierarchical path of a resource. The parent path is a fully_qualified_name of the parent resource.
|
||||
- "parent:x" - Matches x as a substring of the hierarchical path of a resource. It supports same syntax as `name` predicate.
|
||||
- "orgid=number" - Matches resources within a Google Cloud organization with the exact ID value of the number.
|
||||
- "system=SYSTEM" - Matches resources from a specified system. For example, system=bigquery matches BigQuery resources.
|
||||
- "location=LOCATION" - Matches resources in a specified location with an exact name. For example, location=us-central1 matches assets hosted in Iowa. BigQuery Omni assets support this qualifier by using the BigQuery Omni location name. For example, location=aws-us-east-1 matches BigQuery Omni assets in Northern Virginia.
|
||||
- "createtime" -
|
||||
Finds resources that were created within, before, or after a given date or time. For example "createtime:2019-01-01" matches resources created on 2019-01-01.
|
||||
- "updatetime" - Finds resources that were updated within, before, or after a given date or time. For example "updatetime>2019-01-01" matches resources updated after 2019-01-01.
|
||||
- "fully_qualified_name:x" - Matches x as a substring of fully_qualified_name.
|
||||
- "fully_qualified_name=x" - Matches x as fully_qualified_name.
|
||||
|
||||
### Aspect Search
|
||||
To search for entries based on their attached aspects, use the following query syntax.
|
||||
|
||||
aspect:x Matches x as a substring of the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
|
||||
aspect=x Matches x as the full path to the aspect type of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID
|
||||
aspect:xOPERATORvalue
|
||||
Searches for aspect field values. Matches x as a substring of the full path to the aspect type and field name of an aspect that is attached to the entry, in the format projectid.location.ASPECT_TYPE_ID.FIELD_NAME
|
||||
|
||||
The list of supported {OPERATOR}s depends on the type of field in the aspect, as follows:
|
||||
- String: = (exact match) and : (substring)
|
||||
- All number types: =, :, <, >, <=, >=, =>, =<
|
||||
- Enum: =
|
||||
- Datetime: same as for numbers, but the values to compare are treated as datetimes instead of numbers
|
||||
- Boolean: =
|
||||
|
||||
Only top-level fields of the aspect are searchable. For example, all of the following queries match entries where the value of the is-enrolled field in the employee-info aspect type is true. Other entries that match on the substring are also returned.
|
||||
- aspect:example-project.us-central1.employee-info.is-enrolled=true
|
||||
- aspect:example-project.us-central1.employee=true
|
||||
- aspect:employee=true
|
||||
|
||||
Example:-
|
||||
You can use following filters
|
||||
- dataplex-types.global.bigquery-table.type={BIGLAKE_TABLE, BIGLAKE_OBJECT_TABLE, EXTERNAL_TABLE, TABLE}
|
||||
- dataplex-types.global.storage.type={STRUCTURED, UNSTRUCTURED}
|
||||
|
||||
### Logical operators
|
||||
A query can consist of several predicates with logical operators. If you don't specify an operator, logical AND is implied. For example, foo bar returns resources that match both predicate foo and predicate bar.
|
||||
Logical AND and logical OR are supported. For example, foo OR bar.
|
||||
|
||||
You can negate a predicate with a - (hyphen) or NOT prefix. For example, -name:foo returns resources with names that don't match the predicate foo.
|
||||
Logical operators aren't case-sensitive. For example, both or and OR are acceptable.
|
||||
Logical operators are case-sensitive. `OR` and `AND` are acceptable whereas `or` and `and` are not.
|
||||
|
||||
### Request
|
||||
1. Always try to rewrite the prompt using search syntax.
|
||||
@@ -287,7 +310,7 @@ Logical operators aren't case-sensitive. For example, both or and OR are accepta
|
||||
|
||||
## Tool: dataplex_lookup_entry
|
||||
### Request
|
||||
1. Always try to limit the size of the response by specifying `aspect_types` parameter. Make sure to include to select view=CUSTOM when using aspect_types parameter.
|
||||
1. Always try to limit the size of the response by specifying `aspect_types` parameter. Make sure to include to select view=CUSTOM when using aspect_types parameter. If you do not know the name of the aspect type, use the `dataplex_search_aspect_types` tool.
|
||||
2. If you do not know the name of the entry, use `dataplex_search_entries` tool
|
||||
### Response
|
||||
1. Unless asked for a specific aspect, respond with all aspects attached to the entry.
|
||||
@@ -298,4 +321,4 @@ Logical operators aren't case-sensitive. For example, both or and OR are accepta
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-----------|:--------:|:------------:|----------------------------------------------------------------------------------|
|
||||
| kind | string | true | Must be "dataplex". |
|
||||
| project | string | true | ID of the GCP project used for quota and billing purposes (e.g. "my-project-id").|
|
||||
| project | string | true | ID of the GCP project used for quota and billing purposes (e.g. "my-project-id").|
|
||||
@@ -0,0 +1,62 @@
|
||||
---
|
||||
title: "dataplex-search-aspect-types"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
A "dataplex-search-aspect-types" tool allows to to find aspect types relevant to the query.
|
||||
aliases:
|
||||
- /resources/tools/dataplex-search-aspect-types
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
A `dataplex-search-aspect-types` tool allows to fetch the metadata template of aspect types based on search query.
|
||||
It's compatible with the following sources:
|
||||
|
||||
- [dataplex](../../sources/dataplex.md)
|
||||
|
||||
`dataplex-search-aspect-types` accepts following parameters optionally:
|
||||
|
||||
- `query` - Narrows down the search of aspect types to value of this parameter. If not provided, it fetches all aspect types available to the user.
|
||||
- `pageSize` - Number of returned aspect types in the search page. Defaults to `5`.
|
||||
- `orderBy` - Specifies the ordering of results. Supported values are: relevance (default), last_modified_timestamp, last_modified_timestamp asc.
|
||||
|
||||
## Requirements
|
||||
|
||||
### IAM Permissions
|
||||
|
||||
Dataplex uses [Identity and Access Management (IAM)][iam-overview] to control
|
||||
user and group access to Dataplex resources. Toolbox will use your
|
||||
[Application Default Credentials (ADC)][adc] to authorize and authenticate when
|
||||
interacting with [Dataplex][dataplex-docs].
|
||||
|
||||
In addition to [setting the ADC for your server][set-adc], you need to ensure
|
||||
the IAM identity has been given the correct IAM permissions for the tasks you
|
||||
intend to perform. See [Dataplex Universal Catalog IAM permissions][iam-permissions]
|
||||
and [Dataplex Universal Catalog IAM roles][iam-roles] for more information on
|
||||
applying IAM permissions and roles to an identity.
|
||||
|
||||
[iam-overview]: https://cloud.google.com/dataplex/docs/iam-and-access-control
|
||||
[adc]: https://cloud.google.com/docs/authentication#adc
|
||||
[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc
|
||||
[iam-permissions]: https://cloud.google.com/dataplex/docs/iam-permissions
|
||||
[iam-roles]: https://cloud.google.com/dataplex/docs/iam-roles
|
||||
[dataplex-docs]: https://cloud.google.com/dataplex
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
dataplex-search-aspect-types:
|
||||
kind: dataplex-search-aspect-types
|
||||
source: my-dataplex-source
|
||||
description: Use this tool to find aspect types relevant to the query.
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
|
||||
| kind | string | true | Must be "dataplex-search-aspect-types". |
|
||||
| source | string | true | Name of the source the tool should execute on. |
|
||||
| description | string | true | Description of the tool that is passed to the LLM. |
|
||||
@@ -17,22 +17,11 @@ It's compatible with the following sources:
|
||||
- [dataplex](../../sources/dataplex.md)
|
||||
|
||||
`dataplex-search-entries` takes a required `query` parameter based on which
|
||||
entries are filtered and returned to the user and a required `name` parameter
|
||||
which is constructed using source's project if user does not provide it
|
||||
explicitly and has the following format: projects/{project}/locations/global. It
|
||||
also optionally accepts following parameters:
|
||||
entries are filtered and returned to the user. It also optionally accepts following parameters:
|
||||
|
||||
- `pageSize` - Number of results in the search page. Defaults to `5`.
|
||||
- `pageToken` - Page token received from a previous locations.searchEntries
|
||||
call.
|
||||
- `orderBy` - Specifies the ordering of results. Supported values are: relevance
|
||||
(default), last_modified_timestamp, last_modified_timestamp asc
|
||||
- `semanticSearch` - Specifies whether the search should understand the meaning
|
||||
and intent behind the query, rather than just matching keywords. Defaults to
|
||||
`true`.
|
||||
- `scope` - The scope under which the search should be operating. Since this
|
||||
parameter is not exposed to the toolbox user, it defaults to the organization
|
||||
where the project provided in name is located.
|
||||
(default), last_modified_timestamp, last_modified_timestamp asc.
|
||||
|
||||
## Requirements
|
||||
|
||||
|
||||
2
go.mod
2
go.mod
@@ -14,6 +14,7 @@ require (
|
||||
cloud.google.com/go/spanner v1.84.1
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0
|
||||
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.29.0
|
||||
github.com/cenkalti/backoff/v5 v5.0.2
|
||||
github.com/couchbase/gocb/v2 v2.10.1
|
||||
github.com/couchbase/tools-common/http v1.0.9
|
||||
github.com/fsnotify/fsnotify v1.9.0
|
||||
@@ -73,7 +74,6 @@ require (
|
||||
github.com/PuerkitoBio/goquery v1.10.3 // indirect
|
||||
github.com/ajg/form v1.5.1 // indirect
|
||||
github.com/apache/arrow/go/v15 v15.0.2 // indirect
|
||||
github.com/cenkalti/backoff/v5 v5.0.2 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect
|
||||
github.com/couchbase/gocbcore/v10 v10.7.1 // indirect
|
||||
|
||||
@@ -12,8 +12,13 @@ tools:
|
||||
kind: dataplex-lookup-entry
|
||||
source: dataplex-source
|
||||
description: Use this tool to retrieve a specific entry from Dataplex Catalog.
|
||||
dataplex_search_aspect_types:
|
||||
kind: dataplex-search-aspect-types
|
||||
source: dataplex-source
|
||||
description: Use this tool to find aspect types relevant to the query.
|
||||
|
||||
toolsets:
|
||||
dataplex-tools:
|
||||
- dataplex_search_entries
|
||||
- dataplex_lookup_entry
|
||||
- dataplex_lookup_entry
|
||||
- dataplex_search_aspect_types
|
||||
@@ -106,7 +106,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
InputSchema: parameters.McpManifest(),
|
||||
}
|
||||
|
||||
t := &Tool{
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: parameters,
|
||||
@@ -132,11 +132,11 @@ type Tool struct {
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t *Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
viewMap := map[int]dataplexpb.EntryView{
|
||||
1: dataplexpb.EntryView_BASIC,
|
||||
@@ -167,17 +167,17 @@ func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (t *Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
// Parse parameters from the provided data
|
||||
return tools.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t *Tool) Manifest() tools.Manifest {
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
// Returns the tool manifest
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t *Tool) McpManifest() tools.McpManifest {
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
// Returns the tool MCP manifest
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
// Copyright 2025 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 dataplexsearchaspecttypes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
dataplexapi "cloud.google.com/go/dataplex/apiv1"
|
||||
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
|
||||
"github.com/cenkalti/backoff/v5"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
dataplexds "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
)
|
||||
|
||||
const kind string = "dataplex-search-aspect-types"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
CatalogClient() *dataplexapi.CatalogClient
|
||||
ProjectID() string
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &dataplexds.Source{}
|
||||
|
||||
var compatibleSources = [...]string{dataplexds.SourceKind}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// Initialize the search configuration with the provided sources
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
query := tools.NewStringParameter("query", "The query against which aspect type should be matched.")
|
||||
pageSize := tools.NewIntParameterWithDefault("pageSize", 5, "Number of returned aspect types in the search page.")
|
||||
orderBy := tools.NewStringParameterWithDefault("orderBy", "relevance", "Specifies the ordering of results. Supported values are: relevance, last_modified_timestamp, last_modified_timestamp asc")
|
||||
parameters := tools.Parameters{query, pageSize, orderBy}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
Description: cfg.Description,
|
||||
InputSchema: parameters.McpManifest(),
|
||||
}
|
||||
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: parameters,
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
CatalogClient: s.CatalogClient(),
|
||||
ProjectID: s.ProjectID(),
|
||||
manifest: tools.Manifest{
|
||||
Description: cfg.Description,
|
||||
Parameters: parameters.Manifest(),
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
},
|
||||
mcpManifest: mcpManifest,
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
type Tool struct {
|
||||
Name string
|
||||
Kind string
|
||||
Parameters tools.Parameters
|
||||
AuthRequired []string
|
||||
CatalogClient *dataplexapi.CatalogClient
|
||||
ProjectID string
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
|
||||
// Invoke the tool with the provided parameters
|
||||
paramsMap := params.AsMap()
|
||||
query, _ := paramsMap["query"].(string)
|
||||
pageSize := int32(paramsMap["pageSize"].(int))
|
||||
orderBy, _ := paramsMap["orderBy"].(string)
|
||||
|
||||
// Create SearchEntriesRequest with the provided parameters
|
||||
req := &dataplexpb.SearchEntriesRequest{
|
||||
Query: query + " type=projects/dataplex-types/locations/global/entryTypes/aspecttype",
|
||||
Name: fmt.Sprintf("projects/%s/locations/global", t.ProjectID),
|
||||
PageSize: pageSize,
|
||||
OrderBy: orderBy,
|
||||
SemanticSearch: true,
|
||||
}
|
||||
|
||||
// Perform the search using the CatalogClient - this will return an iterator
|
||||
it := t.CatalogClient.SearchEntries(ctx, req)
|
||||
if it == nil {
|
||||
return nil, fmt.Errorf("failed to create search entries iterator for project %q", t.ProjectID)
|
||||
}
|
||||
|
||||
// Create an instance of exponential backoff with default values for retrying GetAspectType calls
|
||||
// InitialInterval, RandomizationFactor, Multiplier, MaxInterval = 500 ms, 0.5, 1.5, 60 s
|
||||
getAspectBackOff := backoff.NewExponentialBackOff()
|
||||
|
||||
// Iterate through the search results and call GetAspectType for each result using the resource name
|
||||
var results []*dataplexpb.AspectType
|
||||
for {
|
||||
entry, err := it.Next()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
resourceName := entry.DataplexEntry.GetEntrySource().Resource
|
||||
getAspectTypeReq := &dataplexpb.GetAspectTypeRequest{
|
||||
Name: resourceName,
|
||||
}
|
||||
|
||||
operation := func() (*dataplexpb.AspectType, error) {
|
||||
aspectType, err := t.CatalogClient.GetAspectType(ctx, getAspectTypeReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get aspect type for entry %q: %w", resourceName, err)
|
||||
}
|
||||
return aspectType, nil
|
||||
}
|
||||
|
||||
// Retry the GetAspectType operation with exponential backoff
|
||||
aspectType, err := backoff.Retry(ctx, operation, backoff.WithBackOff(getAspectBackOff))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get aspect type after retries for entry %q: %w", resourceName, err)
|
||||
}
|
||||
|
||||
results = append(results, aspectType)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
// Parse parameters from the provided data
|
||||
return tools.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
// Returns the tool manifest
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
// Returns the tool MCP manifest
|
||||
return t.mcpManifest
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
// Copyright 2025 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 dataplexsearchaspecttypes_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
|
||||
)
|
||||
|
||||
func TestParseFromYamlDataplexSearchAspectTypes(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
example_tool:
|
||||
kind: dataplex-search-aspect-types
|
||||
source: my-instance
|
||||
description: some description
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": dataplexsearchaspecttypes.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "dataplex-search-aspect-types",
|
||||
Source: "my-instance",
|
||||
Description: "some description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -81,10 +81,8 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
|
||||
query := tools.NewStringParameter("query", "The query against which entries in scope should be matched.")
|
||||
pageSize := tools.NewIntParameterWithDefault("pageSize", 5, "Number of results in the search page.")
|
||||
pageToken := tools.NewStringParameterWithDefault("pageToken", "", "Page token received from a previous locations.searchEntries call. Provide this to retrieve the subsequent page.")
|
||||
orderBy := tools.NewStringParameterWithDefault("orderBy", "relevance", "Specifies the ordering of results. Supported values are: relevance, last_modified_timestamp, last_modified_timestamp asc")
|
||||
semanticSearch := tools.NewBooleanParameterWithDefault("semanticSearch", true, "Whether to use semantic search for the query. If true, the query will be processed using semantic search capabilities.")
|
||||
parameters := tools.Parameters{query, pageSize, pageToken, orderBy, semanticSearch}
|
||||
parameters := tools.Parameters{query, pageSize, orderBy}
|
||||
|
||||
mcpManifest := tools.McpManifest{
|
||||
Name: cfg.Name,
|
||||
@@ -92,7 +90,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
|
||||
InputSchema: parameters.McpManifest(),
|
||||
}
|
||||
|
||||
t := &Tool{
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: parameters,
|
||||
@@ -120,25 +118,22 @@ type Tool struct {
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t *Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
query, _ := paramsMap["query"].(string)
|
||||
pageSize, _ := paramsMap["pageSize"].(int32)
|
||||
pageToken, _ := paramsMap["pageToken"].(string)
|
||||
pageSize := int32(paramsMap["pageSize"].(int))
|
||||
orderBy, _ := paramsMap["orderBy"].(string)
|
||||
semanticSearch, _ := paramsMap["semanticSearch"].(bool)
|
||||
|
||||
req := &dataplexpb.SearchEntriesRequest{
|
||||
Query: query,
|
||||
Name: fmt.Sprintf("projects/%s/locations/global", t.ProjectID),
|
||||
PageSize: pageSize,
|
||||
PageToken: pageToken,
|
||||
OrderBy: orderBy,
|
||||
SemanticSearch: semanticSearch,
|
||||
SemanticSearch: true,
|
||||
}
|
||||
|
||||
it := t.CatalogClient.SearchEntries(ctx, req)
|
||||
@@ -157,17 +152,17 @@ func (t *Tool) Invoke(ctx context.Context, params tools.ParamValues) (any, error
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (t *Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
// Parse parameters from the provided data
|
||||
return tools.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
func (t *Tool) Manifest() tools.Manifest {
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
// Returns the tool manifest
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t *Tool) McpManifest() tools.McpManifest {
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
// Returns the tool MCP manifest
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
@@ -28,6 +28,8 @@ import (
|
||||
"time"
|
||||
|
||||
bigqueryapi "cloud.google.com/go/bigquery"
|
||||
dataplex "cloud.google.com/go/dataplex/apiv1"
|
||||
dataplexpb "cloud.google.com/go/dataplex/apiv1/dataplexpb"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
@@ -38,10 +40,11 @@ import (
|
||||
)
|
||||
|
||||
var (
|
||||
DataplexSourceKind = "dataplex"
|
||||
DataplexSearchEntriesToolKind = "dataplex-search-entries"
|
||||
DataplexLookupEntryToolKind = "dataplex-lookup-entry"
|
||||
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
|
||||
DataplexSourceKind = "dataplex"
|
||||
DataplexSearchEntriesToolKind = "dataplex-search-entries"
|
||||
DataplexLookupEntryToolKind = "dataplex-lookup-entry"
|
||||
DataplexSearchAspectTypesToolKind = "dataplex-search-aspect-types"
|
||||
DataplexProject = os.Getenv("DATAPLEX_PROJECT")
|
||||
)
|
||||
|
||||
func getDataplexVars(t *testing.T) map[string]any {
|
||||
@@ -69,6 +72,19 @@ func initBigQueryConnection(ctx context.Context, project string) (*bigqueryapi.C
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func initDataplexConnection(ctx context.Context) (*dataplex.CatalogClient, error) {
|
||||
cred, err := google.FindDefaultCredentials(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to find default Google Cloud credentials: %w", err)
|
||||
}
|
||||
|
||||
client, err := dataplex.NewCatalogClient(ctx, option.WithCredentials(cred))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create Dataplex client %w", err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func TestDataplexToolEndpoints(t *testing.T) {
|
||||
sourceConfig := getDataplexVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
|
||||
@@ -81,12 +97,21 @@ func TestDataplexToolEndpoints(t *testing.T) {
|
||||
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
|
||||
}
|
||||
|
||||
// create table name with UUID
|
||||
dataplexClient, err := initDataplexConnection(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create Dataplex connection: %s", err)
|
||||
}
|
||||
|
||||
// create resources with UUID
|
||||
datasetName := fmt.Sprintf("temp_toolbox_test_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
aspectTypeId := fmt.Sprintf("param-aspect-type-%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
|
||||
teardownTable1 := setupBigQueryTable(t, ctx, bigqueryClient, datasetName, tableName)
|
||||
teardownAspectType1 := setupDataplexThirdPartyAspectType(t, ctx, dataplexClient, aspectTypeId)
|
||||
time.Sleep(2 * time.Minute) // wait for table and aspect type to be ingested
|
||||
defer teardownTable1(t)
|
||||
defer teardownAspectType1(t)
|
||||
|
||||
toolsFile := getDataplexToolsConfig(sourceConfig)
|
||||
|
||||
@@ -107,6 +132,7 @@ func TestDataplexToolEndpoints(t *testing.T) {
|
||||
runDataplexToolGetTest(t)
|
||||
runDataplexSearchEntriesToolInvokeTest(t, tableName, datasetName)
|
||||
runDataplexLookupEntryToolInvokeTest(t, tableName, datasetName)
|
||||
runDataplexSearchAspectTypesToolInvokeTest(t, aspectTypeId)
|
||||
}
|
||||
|
||||
func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.Client, datasetName string, tableName string) func(*testing.T) {
|
||||
@@ -132,8 +158,6 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C
|
||||
t.Fatalf("Create table job for %s failed: %v", tableName, err)
|
||||
}
|
||||
|
||||
time.Sleep(2 * time.Minute) // wait for table to be ingested
|
||||
|
||||
return func(t *testing.T) {
|
||||
// tear down table
|
||||
dropSQL := fmt.Sprintf("drop table %s.%s", datasetName, tableName)
|
||||
@@ -166,6 +190,35 @@ func setupBigQueryTable(t *testing.T, ctx context.Context, client *bigqueryapi.C
|
||||
}
|
||||
}
|
||||
|
||||
func setupDataplexThirdPartyAspectType(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, aspectTypeId string) func(*testing.T) {
|
||||
parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject)
|
||||
createAspectTypeReq := &dataplexpb.CreateAspectTypeRequest{
|
||||
Parent: parent,
|
||||
AspectTypeId: aspectTypeId,
|
||||
AspectType: &dataplexpb.AspectType{
|
||||
Name: fmt.Sprintf("%s/aspectTypes/%s", parent, aspectTypeId),
|
||||
MetadataTemplate: &dataplexpb.AspectType_MetadataTemplate{
|
||||
Name: "UserSchema",
|
||||
Type: "record",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := client.CreateAspectType(ctx, createAspectTypeReq)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create aspect type %s: %v", aspectTypeId, err)
|
||||
}
|
||||
|
||||
return func(t *testing.T) {
|
||||
// tear down aspect type
|
||||
deleteAspectTypeReq := &dataplexpb.DeleteAspectTypeRequest{
|
||||
Name: fmt.Sprintf("%s/aspectTypes/%s", parent, aspectTypeId),
|
||||
}
|
||||
if _, err := client.DeleteAspectType(ctx, deleteAspectTypeReq); err != nil {
|
||||
t.Errorf("Failed to delete aspect type %s: %v", aspectTypeId, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := map[string]any{
|
||||
@@ -182,12 +235,12 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
|
||||
"my-dataplex-search-entries-tool": map[string]any{
|
||||
"kind": DataplexSearchEntriesToolKind,
|
||||
"source": "my-dataplex-instance",
|
||||
"description": "Simple tool to test end to end functionality.",
|
||||
"description": "Simple dataplex search entries tool to test end to end functionality.",
|
||||
},
|
||||
"my-auth-dataplex-search-entries-tool": map[string]any{
|
||||
"kind": DataplexSearchEntriesToolKind,
|
||||
"source": "my-dataplex-instance",
|
||||
"description": "Simple tool to test end to end functionality.",
|
||||
"description": "Simple dataplex search entries tool to test end to end functionality.",
|
||||
"authRequired": []string{"my-google-auth"},
|
||||
},
|
||||
"my-dataplex-lookup-entry-tool": map[string]any{
|
||||
@@ -201,6 +254,17 @@ func getDataplexToolsConfig(sourceConfig map[string]any) map[string]any {
|
||||
"description": "Simple dataplex lookup entry tool to test end to end functionality.",
|
||||
"authRequired": []string{"my-google-auth"},
|
||||
},
|
||||
"my-dataplex-search-aspect-types-tool": map[string]any{
|
||||
"kind": DataplexSearchAspectTypesToolKind,
|
||||
"source": "my-dataplex-instance",
|
||||
"description": "Simple dataplex search aspect types tool to test end to end functionality.",
|
||||
},
|
||||
"my-auth-dataplex-search-aspect-types-tool": map[string]any{
|
||||
"kind": DataplexSearchAspectTypesToolKind,
|
||||
"source": "my-dataplex-instance",
|
||||
"description": "Simple dataplex search aspect types tool to test end to end functionality.",
|
||||
"authRequired": []string{"my-google-auth"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -216,13 +280,18 @@ func runDataplexToolGetTest(t *testing.T) {
|
||||
{
|
||||
name: "get my-dataplex-search-entries-tool",
|
||||
toolName: "my-dataplex-search-entries-tool",
|
||||
expectedParams: []string{"pageSize", "pageToken", "query", "orderBy", "semanticSearch"},
|
||||
expectedParams: []string{"pageSize", "query", "orderBy"},
|
||||
},
|
||||
{
|
||||
name: "get my-dataplex-lookup-entry-tool",
|
||||
toolName: "my-dataplex-lookup-entry-tool",
|
||||
expectedParams: []string{"name", "view", "aspectTypes", "entry"},
|
||||
},
|
||||
{
|
||||
name: "get my-dataplex-search-aspect-types-tool",
|
||||
toolName: "my-dataplex-search-aspect-types-tool",
|
||||
expectedParams: []string{"pageSize", "query", "orderBy"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
@@ -559,3 +628,119 @@ func runDataplexLookupEntryToolInvokeTest(t *testing.T, tableName string, datase
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runDataplexSearchAspectTypesToolInvokeTest(t *testing.T, aspectTypeId string) {
|
||||
idToken, err := tests.GetGoogleIdToken(tests.ClientId)
|
||||
if err != nil {
|
||||
t.Fatalf("error getting Google ID token: %s", err)
|
||||
}
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
api string
|
||||
requestHeader map[string]string
|
||||
requestBody io.Reader
|
||||
wantStatusCode int
|
||||
expectResult bool
|
||||
wantContentKey string
|
||||
}{
|
||||
{
|
||||
name: "Success - Aspect Type Found",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-aspect-types-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
|
||||
wantStatusCode: 200,
|
||||
expectResult: true,
|
||||
wantContentKey: "metadata_template",
|
||||
},
|
||||
{
|
||||
name: "Success - Aspect Type Found with Authorization",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
|
||||
wantStatusCode: 200,
|
||||
expectResult: true,
|
||||
wantContentKey: "metadata_template",
|
||||
},
|
||||
{
|
||||
name: "Failure - Aspect Type Not Found",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-dataplex-search-aspect-types-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(`"{\"query\":\"name=_aspectType\"}"`)),
|
||||
wantStatusCode: 400,
|
||||
expectResult: false,
|
||||
},
|
||||
{
|
||||
name: "Failure - Invalid Authorization Token",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
|
||||
requestHeader: map[string]string{"my-google-auth_token": "invalid_token"},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
|
||||
wantStatusCode: 401,
|
||||
expectResult: false,
|
||||
},
|
||||
{
|
||||
name: "Failure - No Authorization Token",
|
||||
api: "http://127.0.0.1:5000/api/tool/my-auth-dataplex-search-aspect-types-tool/invoke",
|
||||
requestHeader: map[string]string{},
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf("{\"query\":\"name=%s_aspectType\"}", aspectTypeId))),
|
||||
wantStatusCode: 401,
|
||||
expectResult: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create request: %s", err)
|
||||
}
|
||||
req.Header.Add("Content-type", "application/json")
|
||||
for k, v := range tc.requestHeader {
|
||||
req.Header.Add(k, v)
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to send request: %s", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != tc.wantStatusCode {
|
||||
t.Fatalf("response status code is not %d. It is %d", tc.wantStatusCode, resp.StatusCode)
|
||||
}
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
t.Fatalf("error parsing response body: %s", err)
|
||||
}
|
||||
resultStr, ok := result["result"].(string)
|
||||
if !ok {
|
||||
if result["result"] == nil && !tc.expectResult {
|
||||
return
|
||||
}
|
||||
t.Fatalf("expected 'result' field to be a string, got %T", result["result"])
|
||||
}
|
||||
if !tc.expectResult && (resultStr == "" || resultStr == "[]") {
|
||||
return
|
||||
}
|
||||
var entries []interface{}
|
||||
if err := json.Unmarshal([]byte(resultStr), &entries); err != nil {
|
||||
t.Fatalf("error unmarshalling result string: %v", err)
|
||||
}
|
||||
|
||||
if tc.expectResult {
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("expected exactly one entry, but got %d", len(entries))
|
||||
}
|
||||
entry, ok := entries[0].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected entry to be a map, got %T", entries[0])
|
||||
}
|
||||
if _, ok := entry[tc.wantContentKey]; !ok {
|
||||
t.Fatalf("expected entry to have key '%s', but it was not found in %v", tc.wantContentKey, entry)
|
||||
}
|
||||
} else {
|
||||
if len(entries) != 0 {
|
||||
t.Fatalf("expected 0 entries, but got %d", len(entries))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user