diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go index 4b5822e17..0968f2743 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go @@ -18,6 +18,8 @@ package agentanalytics import ( "context" "fmt" + "reflect" + "strings" "time" bq "cloud.google.com/go/bigquery" @@ -102,6 +104,10 @@ func NewBigQueryAgentAnalyticsPluginWithClients( } err = tableRef.Create(ctx, &bq.TableMetadata{ Schema: EventsSchema(), + TimePartitioning: &bq.TimePartitioning{ + Field: "timestamp", + Type: bq.DayPartitioningType, + }, Clustering: &bq.Clustering{ Fields: config.ClusteringFields, }, @@ -268,18 +274,25 @@ func NewBigQueryAgentAnalyticsPluginWithClients( }, BeforeToolCallback: func(ctx agent.Context, t tool.Tool, args map[string]any) (map[string]any, error) { - attrs := map[string]any{"tool_name": t.Name()} + attrs := map[string]any{ + "tool_name": t.Name(), + "tool_origin": getToolOrigin(t), + } logEvent(ctx, "TOOL_START", args, attrs) return nil, nil }, AfterToolCallback: func(ctx agent.Context, t tool.Tool, args, res map[string]any, err error) (map[string]any, error) { - attrs := map[string]any{"tool_name": t.Name()} + attrs := map[string]any{ + "tool_name": t.Name(), + "tool_origin": getToolOrigin(t), + } logEvent(ctx, "TOOL_END", res, attrs) return nil, nil }, OnToolErrorCallback: func(ctx agent.Context, t tool.Tool, args map[string]any, err error) (map[string]any, error) { attrs := map[string]any{ "tool_name": t.Name(), + "tool_origin": getToolOrigin(t), "error_message": err.Error(), } logEvent(ctx, "TOOL_ERROR", nil, attrs) @@ -297,3 +310,33 @@ func NewBigQueryAgentAnalyticsPluginWithClients( return baseplugin.New(cfg) } + +func getToolOrigin(t tool.Tool) string { + tType := reflect.TypeOf(t).String() + if strings.Contains(tType, "mcptoolset.mcpTool") { + return "MCP" + } + if strings.Contains(tType, "agenttool.agentTool") { + val := reflect.ValueOf(t) + if val.Kind() == reflect.Ptr { + val = val.Elem() + } + if val.Kind() == reflect.Struct { + agentField := val.FieldByName("agent") + if agentField.IsValid() { + agentType := agentField.Type().String() + if strings.Contains(agentType, "remoteagent.a2aAgent") || strings.Contains(agentType, "remoteagent/v2.a2aAgent") { + return "A2A" + } + } + } + return "SUB_AGENT" + } + if strings.Contains(tType, "functiontool.functionTool") { + return "LOCAL" + } + if strings.Contains(tType, "TransferToAgent") { + return "TRANSFER_AGENT" + } + return "UNKNOWN" +} diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go index 2a59eefcf..85a402317 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go @@ -15,6 +15,7 @@ package agentanalytics import ( + "bytes" "context" "errors" "io" @@ -37,6 +38,8 @@ import ( "google.golang.org/adk/v2/model" baseplugin "google.golang.org/adk/v2/plugin" "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/tool/agenttool" + "google.golang.org/adk/v2/tool/functiontool" ) type mockTransport struct { @@ -342,3 +345,112 @@ func TestLogEvent_ExtractsTraceInfo(t *testing.T) { t.Error("Timed out waiting for request") } } + +func TestNewBigQueryAgentAnalyticsPlugin_CreateTable_WithPartitioning(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + + createCalled := false + var requestBody string + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + // Table metadata request: returns 404 Not Found to trigger creation + if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") { + return &http.Response{ + StatusCode: http.StatusNotFound, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":404,"message":"Not found"}}`)), + }, nil + } + // Table creation request + if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") { + createCalled = true + bodyBytes, _ := io.ReadAll(r.Body) + requestBody = string(bodyBytes) + r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("{}")), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("{}")), + }, nil + }, + } + httpClient := &http.Client{Transport: mockTransport} + bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient)) + if err != nil { + t.Fatalf("Failed to create bigquery client: %v", err) + } + + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + gSrv := grpc.NewServer() + storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{}) + go func() { _ = gSrv.Serve(lis) }() + t.Cleanup(gSrv.Stop) + + conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to dial test server: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn)) + if err != nil { + t.Fatalf("Failed to create BigQuery write client: %v", err) + } + + _, err = NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err != nil { + t.Fatalf("Plugin initialization error: %v", err) + } + + if !createCalled { + t.Error("Expected table creation to be called") + } + + if !strings.Contains(requestBody, "timePartitioning") { + t.Errorf("Expected request body to contain 'timePartitioning', got: %s", requestBody) + } + if !strings.Contains(requestBody, "DAY") { + t.Errorf("Expected partitioning type to be 'DAY', got request body: %s", requestBody) + } + if !strings.Contains(requestBody, "timestamp") { + t.Errorf("Expected partitioning field to be 'timestamp', got request body: %s", requestBody) + } +} + +func TestToolProvenance(t *testing.T) { + // 1. Test LOCAL tool origin + localTool, err := functiontool.New[map[string]any, map[string]any](functiontool.Config{ + Name: "local_test_tool", + Description: "Local tool description", + }, func(ctx agent.Context, args map[string]any) (map[string]any, error) { + return map[string]any{"status": "ok"}, nil + }) + if err != nil { + t.Fatalf("Failed to create functiontool: %v", err) + } + + origin := getToolOrigin(localTool) + if origin != "LOCAL" { + t.Errorf("Expected LOCAL tool origin, got: %s", origin) + } + + // 2. Test SUB_AGENT tool origin + subAgentTool := agenttool.New(nil, nil) + origin = getToolOrigin(subAgentTool) + if origin != "SUB_AGENT" { + t.Errorf("Expected SUB_AGENT tool origin, got: %s", origin) + } +}