diff --git a/ai/openai_omni_moderation.go b/ai/openai_omni_moderation.go
new file mode 100644
index 0000000..732caa7
--- /dev/null
+++ b/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)
+}
diff --git a/ai/openai_omni_moderation_test.go b/ai/openai_omni_moderation_test.go
new file mode 100644
index 0000000..7a5ea3a
--- /dev/null
+++ 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)
+}
diff --git a/ai/types.go b/ai/types.go
new file mode 100644
index 0000000..f063c91
--- /dev/null
+++ b/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)
+}
diff --git a/community/manager.go b/community/manager.go
index 8a127c4..240306b 100644
--- a/community/manager.go
+++ b/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)
diff --git a/filter/render_event.go b/event/renderer.go
similarity index 94%
rename from filter/render_event.go
rename to event/renderer.go
index 9bf3a01..1a5b5a4 100644
--- a/filter/render_event.go
+++ b/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{}
diff --git a/filter/render_event_test.go b/event/renderer_test.go
similarity index 90%
rename from filter/render_event_test.go
rename to event/renderer_test.go
index 38a2bf6..b8ba151 100644
--- a/filter/render_event_test.go
+++ b/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": "hello world",
},
})
- 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: hello world"}, 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": "hello world",
},
})
- 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: hello world"}, 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": "waves hello",
},
})
- 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 waves hello"}, 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)
}
diff --git a/filter/filter_ai_executor.go b/filter/filter_ai_executor.go
new file mode 100644
index 0000000..924723b
--- /dev/null
+++ b/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,
+ })
+}
diff --git a/filter/filter_openai_test.go b/filter/filter_ai_executor_test.go
similarity index 55%
rename from filter/filter_openai_test.go
rename to filter/filter_ai_executor_test.go
index 8c1870a..e12a7d3 100644
--- a/filter/filter_openai_test.go
+++ b/filter/filter_ai_executor_test.go
@@ -5,54 +5,59 @@ import (
"errors"
"testing"
+ "github.com/matrix-org/policyserv/ai"
"github.com/matrix-org/policyserv/config"
"github.com/matrix-org/policyserv/filter/classification"
- "github.com/matrix-org/policyserv/internal"
"github.com/matrix-org/policyserv/test"
"github.com/stretchr/testify/assert"
)
-type TestAIProvider struct {
- T *testing.T
- Called bool
- Return []classification.Classification
- ReturnErr error
+type arbitraryConfig struct {
+ SomeVal bool
}
-func (p *TestAIProvider) CheckEvent(ctx context.Context, cnf *aiFilterConfig, input *Input) ([]classification.Classification, error) {
+type TestAIProvider[ConfigT any] struct {
+ // Implements ai.Provider[ConfigT]
+
+ T *testing.T
+ Called bool
+ ExpectedConfig ConfigT
+ Return []classification.Classification
+ ReturnErr error
+}
+
+func (p *TestAIProvider[ConfigT]) CheckEvent(ctx context.Context, cnf ConfigT, input *ai.Input) ([]classification.Classification, error) {
assert.NotNil(p.T, ctx, "context is required")
assert.NotNil(p.T, cnf, "config is required")
assert.NotNil(p.T, input, "input is required")
- assert.Equal(p.T, true, cnf.FailSecure)
+ assert.Equal(p.T, p.ExpectedConfig, cnf)
p.Called = true
return p.Return, p.ReturnErr
}
-func TestOpenAIFilter(t *testing.T) {
+func TestInstancedAIExecutorFilter(t *testing.T) {
t.Parallel()
- ctx := context.Background()
- // Note: typically filter tests go as far as creating a whole filter set to ensure they are
- // created appropriately, however because we're unable to realistically test against OpenAI
- // directly, we just test that the filter is calling a function properly. This means we create
- // the instanced filter structure manually instead of letting the filter set do it (as letting
- // the set do it would require a real OpenAI API key).
+ // Here we're aiming to test that the instanced filter properly calls the provider with what it's supposed to. Other
+ // tests cover whether it gets created properly in-context (ie: whether the OpenAI Omni filter creates the right
+ // provider & instanced filter).
+ ctx := context.Background()
+ name := "TestAIExecutor"
allowedRoomId := "!allowed:example.org"
- aiProvider := &TestAIProvider{
- T: t,
+ // the actual config type doesn't really matter - we just want to make sure it gets passed down appropriately
+ provider := &TestAIProvider[*arbitraryConfig]{
+ T: t,
+ ExpectedConfig: &arbitraryConfig{SomeVal: true},
}
- instance := &InstancedOpenAIFilter{
- set: &Set{ // the filter doesn't use most of this
- communityConfig: &config.CommunityConfig{
- OpenAIFilterFailSecure: internal.Pointer(true), // asserted in TestAIProvider
- },
- },
- aiProvider: aiProvider,
- inRoomIds: []string{allowedRoomId},
+ set := &Set{
+ communityConfig: &config.CommunityConfig{},
}
+ instance := NewInstancedAIExecutorFilter(name, set, provider.ExpectedConfig, provider, []string{allowedRoomId})
+ assert.NotNil(t, instance)
+ assert.Equal(t, name, instance.Name()) // it should carry the name we gave it
// Ensure the AI Provider isn't called for rooms it's not supposed to
event := test.MustMakePDU(&test.BaseClientEvent{
@@ -68,7 +73,7 @@ func TestOpenAIFilter(t *testing.T) {
vecs, err := instance.CheckEvent(ctx, &Input{Event: event})
assert.NoError(t, err)
assert.Equal(t, 0, len(vecs))
- assert.False(t, aiProvider.Called)
+ assert.False(t, provider.Called)
// Ensure the AI Provider is called for rooms it's supposed to
event = test.MustMakePDU(&test.BaseClientEvent{
@@ -84,18 +89,18 @@ func TestOpenAIFilter(t *testing.T) {
vecs, err = instance.CheckEvent(ctx, &Input{Event: event})
assert.NoError(t, err)
assert.Equal(t, 0, len(vecs))
- assert.True(t, aiProvider.Called)
+ assert.True(t, provider.Called)
// Ensure errors are passed through
retErr := errors.New("test error")
- aiProvider.ReturnErr = retErr
+ provider.ReturnErr = retErr
_, err = instance.CheckEvent(ctx, &Input{Event: event})
assert.Equal(t, retErr, err)
// Ensure classifications are passed through
ret := []classification.Classification{classification.Spam, classification.Volumetric}
- aiProvider.Return = ret
- aiProvider.ReturnErr = nil
+ provider.Return = ret
+ provider.ReturnErr = nil
vecs, err = instance.CheckEvent(ctx, &Input{Event: event})
assert.NoError(t, err)
assert.Equal(t, ret, vecs)
diff --git a/filter/filter_hellban_test.go b/filter/filter_hellban_test.go
index ea5f202..403fc71 100644
--- a/filter/filter_hellban_test.go
+++ b/filter/filter_hellban_test.go
@@ -10,6 +10,7 @@ import (
"github.com/matrix-org/policyserv/filter/classification"
"github.com/matrix-org/policyserv/filter/confidence"
"github.com/matrix-org/policyserv/internal"
+ "github.com/matrix-org/policyserv/media"
"github.com/matrix-org/policyserv/pubsub"
"github.com/matrix-org/policyserv/test"
"github.com/stretchr/testify/assert"
@@ -218,7 +219,7 @@ func TestHellbanPostfilter(t *testing.T) {
assertSpamVector := func(event gomatrixserverlib.PDU, isSpam bool) {
fixedFilter.Expect = &Input{
Event: event,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5},
}
if isSpam {
@@ -323,7 +324,7 @@ func TestHellbanFiltersCombined(t *testing.T) {
// Step 2 prep
fixedFilter.Expect = &Input{
Event: spammyEvent1,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5}, // just the starting value
}
// We set a Mentions classification so we can detect that the filter ran
@@ -347,7 +348,7 @@ func TestHellbanFiltersCombined(t *testing.T) {
// Step 6 prep
fixedFilter.Expect = &Input{
Event: spammyEvent1,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{
classification.Spam: 1.0,
classification.Frequency: 1.0,
diff --git a/filter/filter_media_scanning.go b/filter/filter_media_scanning.go
index 2f7a45b..475f406 100644
--- a/filter/filter_media_scanning.go
+++ b/filter/filter_media_scanning.go
@@ -8,9 +8,10 @@ import (
"strings"
"time"
+ "github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/policyserv/content"
"github.com/matrix-org/policyserv/filter/classification"
- "github.com/matrix-org/gomatrixserverlib"
+ "github.com/matrix-org/policyserv/media"
)
const MediaScanningFilterName = "MediaScanningFilter"
@@ -85,7 +86,7 @@ func (f *InstancedMediaScanningFilter) CheckEvent(ctx context.Context, input *In
return classifications, nil
}
-func (f *InstancedMediaScanningFilter) scanMedia(ctx context.Context, event gomatrixserverlib.PDU, media *Media, ch chan<- []classification.Classification) {
+func (f *InstancedMediaScanningFilter) scanMedia(ctx context.Context, event gomatrixserverlib.PDU, media *media.Item, ch chan<- []classification.Classification) {
log.Printf("[%s | %s] Downloading media %s", event.EventID(), event.RoomID().String(), media)
b, err := media.Download()
if err != nil {
diff --git a/filter/filter_openai.go b/filter/filter_openai.go
deleted file mode 100644
index 0c0bc3b..0000000
--- a/filter/filter_openai.go
+++ /dev/null
@@ -1,103 +0,0 @@
-package filter
-
-import (
- "context"
- "log"
- "slices"
-
- "github.com/matrix-org/policyserv/config"
- "github.com/matrix-org/policyserv/filter/classification"
- "github.com/matrix-org/policyserv/internal"
- "github.com/sashabaranov/go-openai"
-)
-
-const OpenAIFilterName = "OpenAIFilter"
-
-func init() {
- mustRegister(OpenAIFilterName, &OpenAIFilter{})
-}
-
-type aiFilterConfig struct {
- FailSecure bool
-}
-
-type OpenAIFilter struct {
-}
-
-func (o *OpenAIFilter) MakeFor(set *Set) (Instanced, error) {
- return &InstancedOpenAIFilter{
- set: set,
- aiProvider: NewOpenAIModerationModelAIProvider(set.instanceConfig),
- inRoomIds: set.instanceConfig.OpenAIAllowedRoomIds,
- }, nil
-}
-
-type InstancedOpenAIFilter struct {
- set *Set
- aiProvider AIProvider // we use indirection because we can't (realistically) test that OpenAI is working
- inRoomIds []string
-}
-
-func (f *InstancedOpenAIFilter) Name() string {
- return OpenAIFilterName
-}
-
-func (f *InstancedOpenAIFilter) CheckEvent(ctx context.Context, input *Input) ([]classification.Classification, error) {
- if !slices.Contains(f.inRoomIds, input.Event.RoomID().String()) {
- return nil, nil
- }
- return f.aiProvider.CheckEvent(ctx, &aiFilterConfig{
- FailSecure: internal.Dereference(f.set.communityConfig.OpenAIFilterFailSecure),
- }, input)
-}
-
-type AIProvider interface {
- CheckEvent(ctx context.Context, cnf *aiFilterConfig, input *Input) ([]classification.Classification, error)
-}
-
-type OpenAIModerationModelAIProvider struct {
- client *openai.Client
-}
-
-func NewOpenAIModerationModelAIProvider(instanceConfig *config.InstanceConfig) *OpenAIModerationModelAIProvider {
- client := openai.NewClient(instanceConfig.OpenAIApiKey)
- return &OpenAIModerationModelAIProvider{
- client: client,
- }
-}
-
-func (p *OpenAIModerationModelAIProvider) CheckEvent(ctx context.Context, cnf *aiFilterConfig, input *Input) ([]classification.Classification, error) {
- messages, err := renderEventToText(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 := p.client.Moderations(ctx, openai.ModerationRequest{
- Input: message,
- Model: openai.ModerationOmniLatest,
- })
- 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 {
- log.Printf("[%s | %s] Result for sender %s: %+v", input.Event.EventID(), input.Event.RoomID(), input.Event.SenderID(), r)
- if r.Flagged {
- flags := []classification.Classification{classification.Spam}
- if r.Categories.SexualMinors {
- flags = append(flags, classification.CSAM)
- }
- return flags, nil
- }
- }
- }
- return nil, nil
-}
diff --git a/filter/filter_openai_omni.go b/filter/filter_openai_omni.go
new file mode 100644
index 0000000..67be5e7
--- /dev/null
+++ b/filter/filter_openai_omni.go
@@ -0,0 +1,27 @@
+package filter
+
+import (
+ "github.com/matrix-org/policyserv/ai"
+ "github.com/matrix-org/policyserv/internal"
+)
+
+const OpenAIOmniFilterName = "OpenAIOmniFilter"
+
+func init() {
+ mustRegister(OpenAIOmniFilterName, &OpenAIOmniFilter{})
+}
+
+type OpenAIOmniFilter struct {
+}
+
+func (o *OpenAIOmniFilter) MakeFor(set *Set) (Instanced, error) {
+ provider, err := ai.NewOpenAIOmniModeration(set.instanceConfig)
+ if err != nil {
+ return nil, err
+ }
+ providerConfig := &ai.OpenAIOmniModerationConfig{
+ FailSecure: internal.Dereference(set.communityConfig.OpenAIFilterFailSecure),
+ }
+ instanced := NewInstancedAIExecutorFilter(OpenAIOmniFilterName, set, providerConfig, provider, set.instanceConfig.OpenAIAllowedRoomIds)
+ return instanced, nil
+}
diff --git a/filter/filter_openai_omni_test.go b/filter/filter_openai_omni_test.go
new file mode 100644
index 0000000..a6ccba3
--- /dev/null
+++ b/filter/filter_openai_omni_test.go
@@ -0,0 +1,65 @@
+package filter
+
+import (
+ "testing"
+
+ "github.com/matrix-org/policyserv/ai"
+ "github.com/matrix-org/policyserv/config"
+ "github.com/matrix-org/policyserv/internal"
+ "github.com/matrix-org/policyserv/test"
+ "github.com/stretchr/testify/assert"
+)
+
+func TestOpenAIOmniFilter(t *testing.T) {
+ t.Parallel()
+
+ // Note: typically filter tests go as far as creating a whole filter set to ensure they are
+ // created appropriately, however because we're unable to realistically test against OpenAI
+ // directly, we just test that the filter is creating the appropriate provider and instanced
+ // executor filter, without calling it.
+ //
+ // This test ultimately gives us confidence that the filter will use the executor and therefore
+ // won't randomly be turned on in rooms.
+
+ cnf := &SetConfig{
+ InstanceConfig: &config.InstanceConfig{
+ OpenAIAllowedRoomIds: []string{"!allowed:example.org"},
+ OpenAIApiKey: "not a real key",
+ },
+ CommunityConfig: &config.CommunityConfig{
+ OpenAIFilterFailSecure: internal.Pointer(true),
+ },
+ Groups: []*SetGroupConfig{{
+ EnabledNames: []string{OpenAIOmniFilterName},
+ MinimumSpamVectorValue: 0.0,
+ MaximumSpamVectorValue: 1.0,
+ }},
+ }
+ memStorage := test.NewMemoryStorage(t)
+ defer memStorage.Close()
+ ps := test.NewMemoryPubsub(t)
+ defer ps.Close()
+ set, err := NewSet(cnf, memStorage, ps, test.MustMakeAuditQueue(5), nil)
+ assert.NoError(t, err)
+ assert.NotNil(t, set)
+
+ // Pull the filter out of the set to inspect its type & config.
+ // If this for some reason isn't "ok", then the filter created the very wrong thing.
+ instanced, ok := set.groups[0].filters[0].(*InstancedAIExecutorFilter[*ai.OpenAIOmniModerationConfig])
+ assert.True(t, ok)
+ assert.NotNil(t, instanced)
+
+ // Verify that the config/setup of the executor are carried through correctly
+ assert.Equal(t, set, instanced.set) // should have been set during filter creation
+ assert.Equal(t, OpenAIOmniFilterName, instanced.Name()) // should have been set during filter creation
+ assert.Equal(t, &ai.OpenAIOmniModerationConfig{
+ FailSecure: true, // should have been set by pulling in the community config above
+ }, instanced.config)
+ assert.Equal(t, []string{"!allowed:example.org"}, instanced.inRoomIds) // should have been pulled from instance config
+
+ // Verify that the provider is correct. Note that the provider verifies it got a non-empty API key, so this also
+ // tests that we got *something* resembling an API key as far as the provider. We can't really test that the exact
+ // API key from the instance config made it there, but we can be pretty sure. Live testing would reveal whether the
+ // correct API key is making it.
+ assert.IsType(t, &ai.OpenAIOmniModeration{}, instanced.aiProvider)
+}
diff --git a/filter/set.go b/filter/set.go
index 03bdd05..3b4784b 100644
--- a/filter/set.go
+++ b/filter/set.go
@@ -91,7 +91,7 @@ func (s *Set) CheckEvent(ctx context.Context, event gomatrixserverlib.PDU, media
Event: event,
IncrementalConfidenceVectors: vecs,
auditContext: auditCtx,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
}
if mediaDownloader != nil {
@@ -107,7 +107,7 @@ func (s *Set) CheckEvent(ctx context.Context, event gomatrixserverlib.PDU, media
if len(url) == 0 {
return // the field we're trying to parse probably wasn't present on the event
}
- m, err := NewMedia(url, mediaDownloader)
+ m, err := media.NewItem(url, mediaDownloader)
if err != nil {
log.Printf("[%s | %s] Non-fatal error creating new media object for '%s': %s", event.EventID(), event.RoomID().String(), url, err)
}
diff --git a/filter/set_test.go b/filter/set_test.go
index 4dc76ca..bfc1d8c 100644
--- a/filter/set_test.go
+++ b/filter/set_test.go
@@ -15,6 +15,7 @@ import (
"github.com/matrix-org/policyserv/filter/classification"
"github.com/matrix-org/policyserv/filter/confidence"
"github.com/matrix-org/policyserv/internal"
+ "github.com/matrix-org/policyserv/media"
"github.com/matrix-org/policyserv/test"
"github.com/stretchr/testify/assert"
)
@@ -119,14 +120,14 @@ func TestSetCheckEvent(t *testing.T) {
Type: "m.room.message",
Content: make(map[string]any),
})
- f1.Expect = &Input{Event: event, Medias: make([]*Media, 0), IncrementalConfidenceVectors: confidence.Vectors{
+ f1.Expect = &Input{Event: event, Medias: make([]*media.Item, 0), IncrementalConfidenceVectors: confidence.Vectors{
classification.Spam: 0.5,
}}
f1.ReturnClasses = []classification.Classification{
classification.Spam,
classification.Mentions,
}
- f2.Expect = &Input{Event: event, Medias: make([]*Media, 0), IncrementalConfidenceVectors: confidence.Vectors{
+ f2.Expect = &Input{Event: event, Medias: make([]*media.Item, 0), IncrementalConfidenceVectors: confidence.Vectors{
classification.Spam: 1.0,
classification.Mentions: 1.0,
}}
@@ -179,13 +180,13 @@ func TestCheckEventWithErrorInGroup(t *testing.T) {
inputs := []*Input{
{
Event: event,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{
classification.Spam: 0.5,
},
}, {
Event: event,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{
classification.Spam: 1.0, // group 0 result
},
@@ -317,7 +318,7 @@ func TestCallsWebhook(t *testing.T) {
fixedFilter.Set = set
fixedFilter.Expect = &Input{
Event: event,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5},
}
fixedFilter.ReturnClasses = []classification.Classification{}
@@ -393,7 +394,7 @@ func TestCallsWebhookErrorNonFatal(t *testing.T) {
fixedFilter.Set = set
fixedFilter.Expect = &Input{
Event: event,
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5},
}
fixedFilter.ReturnClasses = []classification.Classification{classification.Spam}
@@ -457,7 +458,7 @@ func TestExtractsMedia(t *testing.T) {
fixedFilter.Expect = &Input{
Event: event,
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5},
- Medias: []*Media{
+ Medias: []*media.Item{
{
Origin: origin1,
MediaId: id1,
@@ -478,7 +479,7 @@ func TestExtractsMedia(t *testing.T) {
fixedFilter.Expect = &Input{
Event: event,
IncrementalConfidenceVectors: confidence.Vectors{classification.Spam: 0.5},
- Medias: make([]*Media, 0),
+ Medias: make([]*media.Item, 0),
}
res, err = set.CheckEvent(context.Background(), event, nil)
assert.NoError(t, err)
diff --git a/filter/types.go b/filter/types.go
index 2aef173..45d0661 100644
--- a/filter/types.go
+++ b/filter/types.go
@@ -2,55 +2,13 @@ package filter
import (
"context"
- "errors"
- "fmt"
- "net/url"
- "time"
+ "github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/policyserv/filter/classification"
"github.com/matrix-org/policyserv/filter/confidence"
"github.com/matrix-org/policyserv/media"
- "github.com/matrix-org/gomatrixserverlib"
)
-var InvalidMediaUrlError = errors.New("invalid media url")
-
-// Media - A piece of media that was extracted from an event.
-type Media struct {
- Origin string
- MediaId string
-
- downloader media.Downloader
-}
-
-// NewMedia - Creates a new Media instance from the given URL. Returns an InvalidMediaUrlError if the URL is invalid
-// or not an MXC URL.
-func NewMedia(mediaUrl string, downloader media.Downloader) (*Media, error) {
- parsed, err := url.Parse(mediaUrl)
- if err != nil {
- return nil, errors.Join(InvalidMediaUrlError, err)
- }
- if parsed.Scheme != "mxc" {
- return nil, errors.Join(InvalidMediaUrlError, errors.New("not an mxc uri"))
- }
-
- return &Media{
- Origin: parsed.Host,
- MediaId: parsed.Path[1:], // strip leading slash,
- downloader: downloader,
- }, nil
-}
-
-func (m *Media) Download() ([]byte, error) {
- ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) // 30s is somewhat arbitrary - we just don't want to wait forever
- defer cancel()
- return m.downloader.DownloadMedia(ctx, m.Origin, m.MediaId)
-}
-
-func (m *Media) String() string {
- return fmt.Sprintf("mxc://%s/%s", m.Origin, m.MediaId)
-}
-
// Input - A filter input.
type Input struct {
// The event to process/check.
@@ -60,7 +18,7 @@ type Input struct {
IncrementalConfidenceVectors confidence.Vectors
// Extracted media items from the event.
- Medias []*Media
+ Medias []*media.Item
// The context used for auditing the performance of policyserv's filters.
auditContext *auditContext
diff --git a/go.mod b/go.mod
index 3918ecf..8bc361f 100644
--- a/go.mod
+++ b/go.mod
@@ -10,10 +10,10 @@ require (
github.com/kelseyhightower/envconfig v1.4.0
github.com/lib/pq v1.10.9
github.com/matrix-org/gomatrixserverlib v0.0.0-20251103190139-eb14008fe6d1
+ github.com/openai/openai-go/v3 v3.15.0
github.com/panjf2000/ants/v2 v2.11.3
github.com/prometheus/client_golang v1.23.2
github.com/ryanuber/go-glob v1.0.0
- github.com/sashabaranov/go-openai v1.41.2
github.com/segmentio/ksuid v1.0.4
github.com/stretchr/testify v1.11.1
github.com/t2bot/go-matrix-signing-key v0.0.0-20250823034611-5de5f6ca0042
diff --git a/go.sum b/go.sum
index 5514d36..ede21d4 100644
--- a/go.sum
+++ b/go.sum
@@ -51,6 +51,8 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
github.com/oleiade/lane/v2 v2.0.0 h1:XW/ex/Inr+bPkLd3O240xrFOhUkTd4Wy176+Gv0E3Qw=
github.com/oleiade/lane/v2 v2.0.0/go.mod h1:i5FBPFAYSWCgLh58UkUGCChjcCzef/MI7PlQm2TKCeg=
+github.com/openai/openai-go/v3 v3.15.0 h1:hk99rM7YPz+M99/5B/zOQcVwFRLLMdprVGx1vaZ8XMo=
+github.com/openai/openai-go/v3 v3.15.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
github.com/panjf2000/ants/v2 v2.11.3 h1:AfI0ngBoXJmYOpDh9m516vjqoUu2sLrIVgppI9TZVpg=
github.com/panjf2000/ants/v2 v2.11.3/go.mod h1:8u92CYMUc6gyvTIw8Ru7Mt7+/ESnJahz5EVtqfrilek=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -69,8 +71,6 @@ github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR
github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
github.com/ryanuber/go-glob v1.0.0 h1:iQh3xXAumdQ+4Ufa5b25cRpC5TYKlno6hsv6Cb3pkBk=
github.com/ryanuber/go-glob v1.0.0/go.mod h1:807d1WSdnB0XRJzKNil9Om6lcp/3a0v4qIHxIXzX/Yc=
-github.com/sashabaranov/go-openai v1.41.2 h1:vfPRBZNMpnqu8ELsclWcAvF19lDNgh1t6TVfFFOPiSM=
-github.com/sashabaranov/go-openai v1.41.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
github.com/segmentio/ksuid v1.0.4 h1:sBo2BdShXjmcugAMwjugoGUdUV0pcxY5mW4xKRn3v4c=
github.com/segmentio/ksuid v1.0.4/go.mod h1:/XUiZBD3kVx5SmUOl55voK5yeAbBNNIed+2O73XgrPE=
github.com/shoenig/test v1.12.1 h1:mLHfnMv7gmhhP44WrvT+nKSxKkPDiNkIuHGdIGI9RLU=
diff --git a/media/item.go b/media/item.go
new file mode 100644
index 0000000..cd65cae
--- /dev/null
+++ b/media/item.go
@@ -0,0 +1,47 @@
+package media
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "net/url"
+ "time"
+)
+
+var InvalidMediaUrlError = errors.New("invalid media url")
+
+// Item - A piece of media that was extracted from an event.
+type Item struct {
+ Origin string
+ MediaId string
+
+ downloader Downloader
+}
+
+// NewItem - Creates a new Item instance from the given URL. Returns an InvalidMediaUrlError if the URL is invalid
+// or not an MXC URL.
+func NewItem(mediaUrl string, downloader Downloader) (*Item, error) {
+ parsed, err := url.Parse(mediaUrl)
+ if err != nil {
+ return nil, errors.Join(InvalidMediaUrlError, err)
+ }
+ if parsed.Scheme != "mxc" {
+ return nil, errors.Join(InvalidMediaUrlError, errors.New("not an mxc uri"))
+ }
+
+ return &Item{
+ Origin: parsed.Host,
+ MediaId: parsed.Path[1:], // strip leading slash,
+ downloader: downloader,
+ }, nil
+}
+
+func (m *Item) Download() ([]byte, error) {
+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) // 30s is somewhat arbitrary - we just don't want to wait forever
+ defer cancel()
+ return m.downloader.DownloadMedia(ctx, m.Origin, m.MediaId)
+}
+
+func (m *Item) String() string {
+ return fmt.Sprintf("mxc://%s/%s", m.Origin, m.MediaId)
+}
diff --git a/test/openai_moderation_server.go b/test/openai_moderation_server.go
new file mode 100644
index 0000000..050e9c7
--- /dev/null
+++ b/test/openai_moderation_server.go
@@ -0,0 +1,176 @@
+package test
+
+import (
+ "encoding/json"
+ "fmt"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+
+ "github.com/matrix-org/gomatrixserverlib"
+ "github.com/openai/openai-go/v3"
+ "github.com/stretchr/testify/assert"
+)
+
+// Dev note: Usually we'd write a dedicated test for utilities like this, however the entire functionality is covered by
+// other tests using it, so it should be fine.
+
+// KeywordSpammy - Used by tests to always flag a message as generic spam.
+const KeywordSpammy = "MX_SPAMMY"
+
+// KeywordSpammyCSAM - Used by tests to always flag a message as "containing CSAM".
+const KeywordSpammyCSAM = "MX_SPAMMY_CSAM"
+
+// KeywordNeutral - Used by tests to always flag a message as neutral ("not spammy").
+const KeywordNeutral = "MX_NEUTRAL"
+
+// KeywordIntentionalFail - Used by tests to always cause a 500 Internal Server Error response.
+const KeywordIntentionalFail = "MX_INTENTIONAL_FAIL"
+
+// MustMakeKeywordEvent - Creates a consistent event using the specified keyword
+func MustMakeKeywordEvent(keyword string) gomatrixserverlib.PDU {
+ return MustMakePDU(&BaseClientEvent{
+ EventId: "$openai." + keyword,
+ RoomId: "!foo:example.org",
+ Type: "m.room.message",
+ Sender: "@user:example.org",
+ Content: map[string]any{
+ "msgtype": "m.text",
+ "body": keyword + " | This is a message.",
+ "format": "org.matrix.custom.html",
+ "formatted_body": "" + keyword + " | This is a message.",
+ },
+ })
+}
+
+// MakeOpenAIModerationServer - Creates a mock OpenAI Moderation API server for use in tests.
+func MakeOpenAIModerationServer(t *testing.T, apiKey string) *httptest.Server {
+ return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ assert.Equal(t, r.Header.Get("Authorization"), "Bearer "+apiKey)
+
+ // Dev note: this HTTP handler is sensitive to changes in the OpenAI library. If it makes additional
+ // calls ahead of the moderation test or changes what it supplies as a request body, then this test
+ // will suddenly start failing. It's recommended to make changes to the Omni provider code separate
+ // from OpenAI library upgrades to detect these failures more easily.
+
+ assert.Equal(t, r.URL.Path, "/moderations") // we only handle Moderations API stuff here
+
+ b, err := io.ReadAll(r.Body)
+ if err != nil {
+ t.Fatal(err) // "should never happen"
+ }
+ req := string(b)
+
+ if strings.Contains(req, KeywordSpammyCSAM) {
+ modServerHandleKeywordSpammyCSAM(t, w, req)
+ } else if strings.Contains(req, KeywordSpammy) {
+ modServerHandleKeywordSpammy(t, w, req)
+ } else if strings.Contains(req, KeywordNeutral) {
+ modServerHandleKeywordNeutral(t, w, req)
+ } else if strings.Contains(req, KeywordIntentionalFail) {
+ modServerHandleKeywordIntentionalFail(t, w, req)
+ } else {
+ t.Fatalf("Unexpected request: %s", req)
+ }
+ }))
+}
+
+func assertInputMatchesKeyword(t *testing.T, keyword string, body string) {
+ ev := MustMakeKeywordEvent(keyword)
+ content := struct {
+ Body string `json:"body"`
+ FormattedBody string `json:"formatted_body"`
+ }{}
+ err := json.Unmarshal(ev.Content(), &content)
+ assert.NoError(t, err)
+
+ inputs := make([]string, 2)
+ for i, val := range []string{content.Body, content.FormattedBody} {
+ inputs[i] = fmt.Sprintf(`{"input":"%s says: %s","model":"omni-moderation-latest"}`, ev.SenderID().ToUserID().String(), val)
+ }
+
+ assert.Contains(t, inputs, body)
+}
+
+func modServerHandleKeywordSpammyCSAM(t *testing.T, w http.ResponseWriter, body string) {
+ assertInputMatchesKeyword(t, KeywordSpammyCSAM, body)
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ res := openai.ModerationNewResponse{
+ ID: "1",
+ Model: openai.ModerationModelOmniModerationLatest,
+ Results: []openai.Moderation{{
+ Flagged: true, // flagged, and...
+ Categories: openai.ModerationCategories{
+ SexualMinors: true, // ... is detected as CSAM
+ },
+ CategoryScores: openai.ModerationCategoryScores{
+ SexualMinors: 1.0,
+ },
+ CategoryAppliedInputTypes: openai.ModerationCategoryAppliedInputTypes{
+ SexualMinors: []string{"text"},
+ },
+ }},
+ }
+ b, err := json.Marshal(res)
+ assert.NoError(t, err)
+ _, _ = w.Write(b)
+}
+
+func modServerHandleKeywordSpammy(t *testing.T, w http.ResponseWriter, body string) {
+ assertInputMatchesKeyword(t, KeywordSpammy, body)
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ res := openai.ModerationNewResponse{
+ ID: "1",
+ Model: openai.ModerationModelOmniModerationLatest,
+ Results: []openai.Moderation{{
+ Flagged: true, // flagged, but...
+ Categories: openai.ModerationCategories{
+ SexualMinors: false, // ... not CSAM to avoid accidentally causing that code to activate
+ },
+ CategoryScores: openai.ModerationCategoryScores{
+ SexualMinors: 0.0,
+ },
+ CategoryAppliedInputTypes: openai.ModerationCategoryAppliedInputTypes{
+ SexualMinors: []string{"text"},
+ },
+ }},
+ }
+ b, err := json.Marshal(res)
+ assert.NoError(t, err)
+ _, _ = w.Write(b)
+}
+
+func modServerHandleKeywordNeutral(t *testing.T, w http.ResponseWriter, body string) {
+ assertInputMatchesKeyword(t, KeywordNeutral, body)
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ res := openai.ModerationNewResponse{
+ ID: "1",
+ Model: openai.ModerationModelOmniModerationLatest,
+ Results: []openai.Moderation{{
+ Flagged: false, // "not flagged"
+ Categories: openai.ModerationCategories{},
+ CategoryScores: openai.ModerationCategoryScores{},
+ CategoryAppliedInputTypes: openai.ModerationCategoryAppliedInputTypes{},
+ }},
+ }
+ b, err := json.Marshal(res)
+ assert.NoError(t, err)
+ _, _ = w.Write(b)
+}
+
+func modServerHandleKeywordIntentionalFail(t *testing.T, w http.ResponseWriter, body string) {
+ assertInputMatchesKeyword(t, KeywordIntentionalFail, body)
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusInternalServerError) // should prevent automatic retry from happening
+ // This is a mock OpenAI API error
+ _, _ = w.Write([]byte(`{"error":{"code": "X-ERROR","message":"Intentional fail","param":"x","type":"x"}}`))
+}