Maintenance: Refactor AI filter code

ai/openai_omni_moderation.go

@@ -0,0 +1,97 @@
+package ai
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"log"
+
+	"github.com/matrix-org/policyserv/config"
+	"github.com/matrix-org/policyserv/event"
+	"github.com/matrix-org/policyserv/filter/classification"
+	"github.com/openai/openai-go/v3"
+	"github.com/openai/openai-go/v3/option"
+)
+
+type OpenAIOmniModerationConfig struct {
+	FailSecure bool
+}
+
+type OpenAIOmniModeration struct {
+	// Implements Provider[*OpenAIOmniModerationConfig]
+
+	client openai.Client
+}
+
+func NewOpenAIOmniModeration(cnf *config.InstanceConfig, additionalClientOptions ...option.RequestOption) (Provider[*OpenAIOmniModerationConfig], error) {
+	apiKey := cnf.OpenAIApiKey
+	if len(apiKey) == 0 {
+		return nil, errors.New("api key not set")
+	}
+	options := append([]option.RequestOption{option.WithAPIKey(apiKey)}, additionalClientOptions...)
+	client := openai.NewClient(options...)
+	return &OpenAIOmniModeration{
+		client: client,
+	}, nil
+}
+
+func (m *OpenAIOmniModeration) CheckEvent(ctx context.Context, cnf *OpenAIOmniModerationConfig, input *Input) ([]classification.Classification, error) {
+	messages, err := event.RenderToText(input.Event)
+	if err != nil {
+		return nil, err
+	}
+	for _, message := range messages {
+		// Note: we don't want to log message contents in production
+		log.Printf("[%s | %s] Message sent by %s", input.Event.EventID(), input.Event.RoomID(), input.Event.SenderID())
+		res, err := m.client.Moderations.New(ctx, openai.ModerationNewParams{
+			Model: openai.ModerationModelOmniModerationLatest,
+			Input: openai.ModerationNewParamsInputUnion{
+				OfString: openai.String(message),
+			},
+		})
+		if err != nil {
+			log.Printf("[%s | %s] Error checking message: %s", input.Event.EventID(), input.Event.RoomID(), err)
+			if cnf.FailSecure {
+				log.Printf("[%s | %s] Returning spam response to block events and discourage retries", input.Event.EventID(), input.Event.RoomID())
+				return []classification.Classification{classification.Spam, classification.Frequency}, nil
+			} else {
+				log.Printf("[%s | %s] Returning neutral response despite error, per config", input.Event.EventID(), input.Event.RoomID())
+				return nil, nil
+			}
+		}
+		for _, r := range res.Results {
+			// Note: we compress JSON here because the OpenAI library tends to return *a lot* of redundant detail, including JSON with newlines in it.
+			log.Printf("[%s | %s] Result for sender %s: Flagged=%t Flags=%s Scores=%s", input.Event.EventID(), input.Event.RoomID(), input.Event.SenderID(), r.Flagged, compressJsonResponse(r.Categories), compressJsonResponse(r.CategoryScores))
+			if r.Flagged {
+				flags := []classification.Classification{classification.Spam}
+				if r.Categories.SexualMinors {
+					flags = append(flags, classification.CSAM)
+				}
+				return flags, nil
+			}
+		}
+	}
+	return nil, nil
+}
+
+type compressible interface {
+	RawJSON() string // same definition that's shared with the OpenAI response parts
+}
+
+func compressJsonResponse(target compressible) string {
+	raw := target.RawJSON()
+
+	val := make(map[string]any)
+	err := json.Unmarshal([]byte(raw), &val)
+	if err != nil {
+		log.Printf("Non-fatal error compressing JSON. Using uncompressed instead. %s", err)
+		return raw
+	}
+	b, err := json.Marshal(val)
+	if err != nil {
+		log.Printf("Non-fatal error compressing JSON. Using uncompressed instead. %s", err)
+		return raw
+	}
+
+	return string(b)
+}

ai/openai_omni_moderation_test.go

@@ -0,0 +1,72 @@
+package ai
+
+import (
+	"context"
+	"testing"
+
+	"github.com/matrix-org/policyserv/config"
+	"github.com/matrix-org/policyserv/filter/classification"
+	"github.com/matrix-org/policyserv/test"
+	"github.com/openai/openai-go/v3/option"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestOpenAIOmniModeration(t *testing.T) {
+	t.Parallel()
+
+	// Caution: this is a long test. It mocks the OpenAI API to drive our Omni provider through 4 test cases:
+	//  1. Regular spammy message (according to omni)
+	//  2. Spammy message flagged for CSAM by omni
+	//  3. Neutral message (according to omni)
+	//  4. Simulates a server error to test `FailSecure` in both operation modes
+
+	// We create our own HTTP client to intercept and act as the OpenAI API
+	apiKey := "not_a_real_key"
+	mockApi := test.MakeOpenAIModerationServer(t, apiKey)
+	defer mockApi.Close()
+	client := mockApi.Client() // get a client instance that trusts the mockApi certificate
+
+	// Create the provider
+	provider, err := NewOpenAIOmniModeration(
+		&config.InstanceConfig{OpenAIApiKey: apiKey},
+		option.WithHTTPClient(client),
+		option.WithBaseURL(mockApi.URL),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, provider)
+
+	spammyEvent1 := test.MustMakeKeywordEvent(test.KeywordSpammy)
+	ret, err := provider.CheckEvent(context.Background(), &OpenAIOmniModerationConfig{FailSecure: true}, &Input{Event: spammyEvent1})
+	assert.NoError(t, err)
+	assert.Equal(t, []classification.Classification{
+		classification.Spam,
+	}, ret)
+
+	spammyEvent2 := test.MustMakeKeywordEvent(test.KeywordSpammyCSAM)
+	ret, err = provider.CheckEvent(context.Background(), &OpenAIOmniModerationConfig{FailSecure: true}, &Input{Event: spammyEvent2})
+	assert.NoError(t, err)
+	assert.Equal(t, []classification.Classification{
+		classification.Spam,
+		classification.CSAM, // should have been detected
+	}, ret)
+
+	neutralEvent := test.MustMakeKeywordEvent(test.KeywordNeutral)
+	ret, err = provider.CheckEvent(context.Background(), &OpenAIOmniModerationConfig{FailSecure: true}, &Input{Event: neutralEvent})
+	assert.NoError(t, err)
+	assert.Nil(t, ret) // no classifications
+
+	failEvent := test.MustMakeKeywordEvent(test.KeywordIntentionalFail)
+
+	// First test that when FailSecure: true we return a spam classification
+	ret, err = provider.CheckEvent(context.Background(), &OpenAIOmniModerationConfig{FailSecure: true}, &Input{Event: failEvent})
+	assert.NoError(t, err)
+	assert.Equal(t, []classification.Classification{
+		classification.Spam,
+		classification.Frequency, // also added by the provider
+	}, ret)
+
+	// Now when FailSecure: false, we should get no classifications (but also no errors)
+	ret, err = provider.CheckEvent(context.Background(), &OpenAIOmniModerationConfig{FailSecure: false}, &Input{Event: failEvent})
+	assert.NoError(t, err)
+	assert.Nil(t, ret)
+}

ai/types.go

@@ -0,0 +1,18 @@
+package ai
+
+import (
+	"context"
+
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/policyserv/filter/classification"
+	"github.com/matrix-org/policyserv/media"
+)
+
+type Input struct {
+	Event  gomatrixserverlib.PDU
+	Medias []*media.Item
+}
+
+type Provider[ConfigT any] interface {
+	CheckEvent(ctx context.Context, cnf ConfigT, input *Input) ([]classification.Classification, error)
+}

community/manager.go

@@ -137,7 +137,7 @@ func (m *Manager) getCommunityFilterSet(ctx context.Context, communityId string)
 	}
 	if m.instanceConfig.OpenAIApiKey != "" {
 		// Access to this filter is gated by further instance config (namely, the room IDs allowed to use it)
-		filters = append(filters, filter.OpenAIFilterName)
+		filters = append(filters, filter.OpenAIOmniFilterName)
 	}
 	if !internal.Dereference(communityConfig.StickyEventsFilterAllowStickyEvents) {
 		filters = append(filters, filter.StickyEventsFilterName)

filter/render_event.go → event/renderer.go

@@ -1,4 +1,4 @@
-package filter
+package event
 
 import (
 	"encoding/json"
@@ -22,7 +22,7 @@ type reactionEventContent struct {
 	} `json:"m.relates_to"`
 }
 
-func renderEventToText(event gomatrixserverlib.PDU) ([]string, error) {
+func RenderToText(event gomatrixserverlib.PDU) ([]string, error) {
 	switch event.Type() {
 	case "m.room.message":
 		content := &messageEventContent{}

filter/render_event_test.go → event/renderer_test.go

@@ -1,4 +1,4 @@
-package filter
+package event
 
 import (
 	"testing"
@@ -20,7 +20,7 @@ func TestRenderEventMText(t *testing.T) {
 			"msgtype": "m.text",
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: hello world"}, render)
 
@@ -36,7 +36,7 @@ func TestRenderEventMText(t *testing.T) {
 			"formatted_body": "<b>hello world</b>",
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: hello world", "@alice:example.org says: <b>hello world</b>"}, render)
 }
@@ -54,7 +54,7 @@ func TestRenderEventMNotice(t *testing.T) {
 			"msgtype": "m.notice",
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: hello world"}, render)
 
@@ -70,7 +70,7 @@ func TestRenderEventMNotice(t *testing.T) {
 			"formatted_body": "<b>hello world</b>",
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: hello world", "@alice:example.org says: <b>hello world</b>"}, render)
 }
@@ -88,7 +88,7 @@ func TestRenderEventMEmote(t *testing.T) {
 			"msgtype": "m.emote",
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: /me waves hello"}, render)
 
@@ -104,7 +104,7 @@ func TestRenderEventMEmote(t *testing.T) {
 			"formatted_body": "<b>waves hello</b>",
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org says: /me waves hello", "@alice:example.org says: /me <b>waves hello</b>"}, render)
 }
@@ -121,7 +121,7 @@ func TestRenderEventUnrenderableMessages(t *testing.T) {
 			"msgtype": "m.image", // not a supported message type
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 
@@ -134,7 +134,7 @@ func TestRenderEventUnrenderableMessages(t *testing.T) {
 			"msgtype": "m.video", // not a supported message type
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 
@@ -147,7 +147,7 @@ func TestRenderEventUnrenderableMessages(t *testing.T) {
 			"msgtype": "m.audio", // not a supported message type
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 
@@ -160,7 +160,7 @@ func TestRenderEventUnrenderableMessages(t *testing.T) {
 			"msgtype": "m.file", // not a supported message type
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 
@@ -173,7 +173,7 @@ func TestRenderEventUnrenderableMessages(t *testing.T) {
 			"msgtype": "org.example.custom", // not a supported message type
 		},
 	})
-	render, err = renderEventToText(event)
+	render, err = RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 }
@@ -194,7 +194,7 @@ func TestRenderEventReaction(t *testing.T) {
 			},
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"@alice:example.org reacted with 💖"}, render)
 }
@@ -211,7 +211,7 @@ func TestRenderEventUnknownType(t *testing.T) {
 			"this": "does not matter",
 		},
 	})
-	render, err := renderEventToText(event)
+	render, err := RenderToText(event)
 	assert.NoError(t, err)
 	assert.Equal(t, []string(nil), render)
 }

filter/filter_ai_executor.go

@@ -0,0 +1,41 @@
+package filter
+
+import (
+	"context"
+	"slices"
+
+	"github.com/matrix-org/policyserv/ai"
+	"github.com/matrix-org/policyserv/filter/classification"
+)
+
+type InstancedAIExecutorFilter[ConfigT any] struct {
+	name       string
+	set        *Set
+	config     ConfigT
+	aiProvider ai.Provider[ConfigT]
+	inRoomIds  []string
+}
+
+func NewInstancedAIExecutorFilter[ConfigT any](name string, set *Set, config ConfigT, aiProvider ai.Provider[ConfigT], inRoomIds []string) *InstancedAIExecutorFilter[ConfigT] {
+	return &InstancedAIExecutorFilter[ConfigT]{
+		name:       name,
+		set:        set,
+		config:     config,
+		aiProvider: aiProvider,
+		inRoomIds:  inRoomIds,
+	}
+}
+
+func (f *InstancedAIExecutorFilter[ConfigT]) Name() string {
+	return f.name
+}
+
+func (f *InstancedAIExecutorFilter[ConfigT]) CheckEvent(ctx context.Context, input *Input) ([]classification.Classification, error) {
+	if !slices.Contains(f.inRoomIds, input.Event.RoomID().String()) {
+		return nil, nil // this filter isn't allowed to execute in here
+	}
+	return f.aiProvider.CheckEvent(ctx, f.config, &ai.Input{
+		Event:  input.Event,
+		Medias: input.Medias,
+	})
+}