feat: Add configuration options to filter logged events via allowlist and denylist

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

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
+	}
+}

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,
 	}
 }