diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go index 4b5822e17..786a06317 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go +++ b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go @@ -102,6 +102,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, }, @@ -129,6 +133,26 @@ func NewBigQueryAgentAnalyticsPluginWithClients( // Helper closure func to construct log events logEvent := func(ctx context.Context, eventType string, content any, extraAttrs map[string]any) { + if len(config.EventDenylist) > 0 { + for _, d := range config.EventDenylist { + if d == eventType { + return + } + } + } + if len(config.EventAllowlist) > 0 { + allowed := false + for _, a := range config.EventAllowlist { + if a == eventType { + allowed = true + break + } + } + if !allowed { + return + } + } + row := make(map[string]any) row["timestamp"] = time.Now() row["event_type"] = eventType diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go index 2a59eefcf..972562146 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" @@ -342,3 +343,208 @@ 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 TestEventFiltering(t *testing.T) { + ctx := context.Background() + config := DefaultConfig() + config.Enabled = true + config.ProjectID = "test-project" + config.DatasetID = "test-dataset" + config.TableName = "test-table" + config.EventAllowlist = []string{"USER_MESSAGE"} + config.EventDenylist = []string{"MODEL_REQUEST"} + + mockTransport := &mockTransport{ + roundTrip: func(r *http.Request) (*http.Response, error) { + 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() + requestsChan := make(chan *storagepb.AppendRowsRequest, 10) + storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{requests: requestsChan}) + 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) + } + + p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient) + if err != nil { + t.Fatalf("Expected no error, got %v", err) + } + + mCtx := &mockInvocationContext{ + ctx: ctx, + agentName: "test-agent", + sessionID: "sess-id", + invocationID: "inv-id", + userID: "user-id", + } + + // Test USER_MESSAGE (in allowlist) + userMsgCb := p.OnUserMessageCallback() + if userMsgCb == nil { + t.Fatal("OnUserMessageCallback is nil") + } + _, err = userMsgCb(mCtx, &genai.Content{Parts: []*genai.Part{{Text: "hello"}}}) + if err != nil { + t.Fatalf("OnUserMessageCallback error: %v", err) + } + + p.AfterRunCallback()(mCtx) + + select { + case req := <-requestsChan: + if req == nil { + t.Error("Received nil request") + } + case <-time.After(1 * time.Second): + t.Error("Timed out waiting for USER_MESSAGE event, expected it to be logged") + } + + // Test MODEL_REQUEST (in denylist) + modelReqCb := p.BeforeModelCallback() + if modelReqCb == nil { + t.Fatal("BeforeModelCallback is nil") + } + // Clear requestsChan + for len(requestsChan) > 0 { + <-requestsChan + } + + _, err = modelReqCb(mCtx, &model.LLMRequest{}) + if err != nil { + t.Fatalf("BeforeModelCallback error: %v", err) + } + p.AfterRunCallback()(mCtx) + + select { + case req := <-requestsChan: + t.Errorf("Expected no event to be logged for MODEL_REQUEST due to denylist, but got: %v", req) + case <-time.After(100 * time.Millisecond): + // Success, no event logged + } + + // Test MODEL_RESPONSE (not in allowlist) + modelRespCb := p.AfterModelCallback() + if modelRespCb == nil { + t.Fatal("AfterModelCallback is nil") + } + _, err = modelRespCb(mCtx, &model.LLMResponse{}, nil) + if err != nil { + t.Fatalf("AfterModelCallback error: %v", err) + } + p.AfterRunCallback()(mCtx) + + select { + case req := <-requestsChan: + t.Errorf("Expected no event to be logged for MODEL_RESPONSE due to allowlist, but got: %v", req) + case <-time.After(100 * time.Millisecond): + // Success, no event logged + } +} diff --git a/plugin/agentanalytics/config.go b/plugin/agentanalytics/config.go index 40e3b5595..2eff0b1dc 100644 --- a/plugin/agentanalytics/config.go +++ b/plugin/agentanalytics/config.go @@ -63,6 +63,10 @@ type Config struct { // Retry configuration for appending rows. RetryConfig RetryConfig + + // Configuration filters for events + EventAllowlist []string + EventDenylist []string } // DefaultConfig returns the default configuration for the agent analytics plugin. @@ -86,5 +90,7 @@ func DefaultConfig() Config { MaxDelay: 10 * time.Second, Multiplier: 2.0, }, + EventAllowlist: nil, + EventDenylist: nil, } }