@framers/agentos-ext-content-policy-rewriter 0.1.0

package/test/ContentPolicyRewriter.spec.ts

@@ -0,0 +1,79 @@
+import { describe, it, expect, vi } from 'vitest';
+import { ContentPolicyRewriterGuardrail } from '../src/ContentPolicyRewriterGuardrail.js';
+
+describe('ContentPolicyRewriterGuardrail', () => {
+  it('has correct guardrail config', () => {
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: vi.fn(),
+      categories: { profanity: { enabled: true } },
+    });
+    expect(g.config.canSanitize).toBe(true);
+    expect(g.config.evaluateStreamingChunks).toBe(true);
+  });
+
+  it('allows clean streaming chunk', async () => {
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: vi.fn(),
+      categories: { profanity: { enabled: true } },
+    });
+    const result = await g.evaluateOutput({
+      chunk: { type: 'TEXT_DELTA', text: 'Hello world' },
+    });
+    expect(result).toBeNull();
+  });
+
+  it('blocks streaming chunk with keyword match', async () => {
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: vi.fn(),
+      categories: { profanity: { enabled: true, action: 'block' } },
+    });
+    const result = await g.evaluateOutput({
+      chunk: { type: 'TEXT_DELTA', text: 'What the fuck' },
+    });
+    expect(result).not.toBeNull();
+    expect(result!.action).toBe('block');
+  });
+
+  it('sanitizes final response via LLM', async () => {
+    const mockLlm = vi.fn();
+    mockLlm.mockResolvedValueOnce('{ "violations": [{ "category": "profanity", "severity": "medium", "spans": ["damn"] }] }');
+    mockLlm.mockResolvedValueOnce('The person was very upset.');
+
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: mockLlm,
+      categories: { profanity: { enabled: true, action: 'sanitize' } },
+    });
+    const result = await g.evaluateOutput({
+      chunk: { type: 'FINAL_RESPONSE', finalResponseText: 'The person was damn upset.' },
+    });
+    expect(result).not.toBeNull();
+    expect(result!.action).toBe('sanitize');
+    expect(result!.modifiedText).toBe('The person was very upset.');
+  });
+
+  it('allows clean final response (no LLM rewrite call)', async () => {
+    const mockLlm = vi.fn();
+    mockLlm.mockResolvedValueOnce('{ "violations": [] }');
+
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: mockLlm,
+      categories: { profanity: { enabled: true } },
+    });
+    const result = await g.evaluateOutput({
+      chunk: { type: 'FINAL_RESPONSE', finalResponseText: 'Everything is fine.' },
+    });
+    expect(result).toBeNull();
+    expect(mockLlm).toHaveBeenCalledTimes(1);
+  });
+
+  it('returns null when no categories enabled', async () => {
+    const g = new ContentPolicyRewriterGuardrail({
+      llmInvoker: vi.fn(),
+      categories: {},
+    });
+    const result = await g.evaluateOutput({
+      chunk: { type: 'FINAL_RESPONSE', finalResponseText: 'Anything goes.' },
+    });
+    expect(result).toBeNull();
+  });
+});

package/test/KeywordPreFilter.spec.ts

@@ -0,0 +1,46 @@
+import { describe, it, expect } from 'vitest';
+import { KeywordPreFilter } from '../src/KeywordPreFilter.js';
+
+describe('KeywordPreFilter', () => {
+  const filter = new KeywordPreFilter();
+
+  it('returns null for clean text', () => {
+    const result = filter.scan('Hello, how are you?', {
+      profanity: { enabled: true, action: 'sanitize' },
+    });
+    expect(result).toBeNull();
+  });
+
+  it('detects profanity keyword', () => {
+    const result = filter.scan('What the fuck is this?', {
+      profanity: { enabled: true, action: 'block' },
+    });
+    expect(result).not.toBeNull();
+    expect(result!.category).toBe('profanity');
+  });
+
+  it('skips disabled categories', () => {
+    const result = filter.scan('What the fuck is this?', {
+      profanity: { enabled: false, action: 'block' },
+    });
+    expect(result).toBeNull();
+  });
+
+  it('uses custom keyword list when provided', () => {
+    const customFilter = new KeywordPreFilter({
+      profanity: ['dingus', 'bozo'],
+    });
+    const result = customFilter.scan('You absolute bozo', {
+      profanity: { enabled: true, action: 'sanitize' },
+    });
+    expect(result).not.toBeNull();
+    expect(result!.category).toBe('profanity');
+  });
+
+  it('is case-insensitive', () => {
+    const result = filter.scan('What the FUCK', {
+      profanity: { enabled: true, action: 'block' },
+    });
+    expect(result).not.toBeNull();
+  });
+});

package/test/LlmPolicyJudge.spec.ts

@@ -0,0 +1,51 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { LlmPolicyJudge } from '../src/LlmPolicyJudge.js';
+import type { LlmInvoker } from '../src/types.js';
+
+describe('LlmPolicyJudge', () => {
+  let mockLlm: LlmInvoker;
+
+  beforeEach(() => {
+    mockLlm = vi.fn();
+  });
+
+  it('returns empty violations for clean content', async () => {
+    (mockLlm as any).mockResolvedValueOnce('{ "violations": [] }');
+    const judge = new LlmPolicyJudge(mockLlm);
+    const result = await judge.classify('Hello world', {
+      profanity: { enabled: true, action: 'sanitize' },
+    });
+    expect(result.violations).toEqual([]);
+  });
+
+  it('parses violation response correctly', async () => {
+    (mockLlm as any).mockResolvedValueOnce(JSON.stringify({
+      violations: [{ category: 'profanity', severity: 'medium', spans: ['bad word'] }],
+    }));
+    const judge = new LlmPolicyJudge(mockLlm);
+    const result = await judge.classify('text with bad word', {
+      profanity: { enabled: true, action: 'sanitize' },
+    });
+    expect(result.violations).toHaveLength(1);
+    expect(result.violations[0].category).toBe('profanity');
+  });
+
+  it('handles malformed LLM response gracefully', async () => {
+    (mockLlm as any).mockResolvedValueOnce('not json at all');
+    const judge = new LlmPolicyJudge(mockLlm);
+    const result = await judge.classify('some text', {
+      profanity: { enabled: true, action: 'sanitize' },
+    });
+    expect(result.violations).toEqual([]);
+  });
+
+  it('includes custom rules in prompt', async () => {
+    (mockLlm as any).mockResolvedValueOnce('{ "violations": [] }');
+    const judge = new LlmPolicyJudge(mockLlm);
+    await judge.classify('text', {
+      custom: { enabled: true, action: 'sanitize' },
+    }, 'Never mention competitor X');
+    const call = (mockLlm as any).mock.calls[0];
+    expect(call[1]).toContain('Never mention competitor X');
+  });
+});

package/test/LlmRewriter.spec.ts

@@ -0,0 +1,34 @@
+import { describe, it, expect, vi } from 'vitest';
+import { LlmRewriter } from '../src/LlmRewriter.js';
+import type { LlmInvoker, PolicyViolation } from '../src/types.js';
+
+describe('LlmRewriter', () => {
+  const mockLlm: LlmInvoker = vi.fn();
+
+  it('returns rewritten text from LLM', async () => {
+    (mockLlm as any).mockResolvedValueOnce('The person expressed frustration.');
+    const rewriter = new LlmRewriter(mockLlm);
+    const violations: PolicyViolation[] = [
+      { category: 'profanity', severity: 'medium', spans: ['bad word'] },
+    ];
+    const result = await rewriter.rewrite('The person said bad word in anger.', violations);
+    expect(result).toBe('The person expressed frustration.');
+  });
+
+  it('returns original text on LLM failure', async () => {
+    (mockLlm as any).mockRejectedValueOnce(new Error('LLM unavailable'));
+    const rewriter = new LlmRewriter(mockLlm);
+    const violations = [{ category: 'profanity' as const, severity: 'low' as const, spans: ['x'] }];
+    const result = await rewriter.rewrite('original text', violations);
+    expect(result).toBe('original text');
+  });
+
+  it('strips markdown fences from LLM response', async () => {
+    (mockLlm as any).mockResolvedValueOnce('```\nClean version of the text.\n```');
+    const rewriter = new LlmRewriter(mockLlm);
+    const result = await rewriter.rewrite('dirty text', [
+      { category: 'profanity', severity: 'low', spans: ['dirty'] },
+    ]);
+    expect(result).toBe('Clean version of the text.');
+  });
+});

package/test/presets.spec.ts

@@ -0,0 +1,30 @@
+import { describe, it, expect } from 'vitest';
+import { resolvePreset } from '../src/presets.js';
+
+describe('presets', () => {
+  it('uncensored disables all categories', () => {
+    const cfg = resolvePreset('uncensored');
+    for (const cat of Object.values(cfg.categories ?? {})) {
+      expect(cat?.enabled).toBe(false);
+    }
+  });
+
+  it('uncensored-safe enables only illegal_harmful', () => {
+    const cfg = resolvePreset('uncensored-safe');
+    expect(cfg.categories?.illegal_harmful?.enabled).toBe(true);
+    expect(cfg.categories?.adult?.enabled).toBeUndefined();
+  });
+
+  it('family-friendly enables all categories', () => {
+    const cfg = resolvePreset('family-friendly');
+    expect(cfg.categories?.adult?.enabled).toBe(true);
+    expect(cfg.categories?.profanity?.enabled).toBe(true);
+    expect(cfg.categories?.violence?.enabled).toBe(true);
+  });
+
+  it('passes through config objects unchanged', () => {
+    const input = { categories: { adult: { enabled: true } } };
+    const cfg = resolvePreset(input);
+    expect(cfg).toEqual(input);
+  });
+});

package/tsconfig.json

@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "declaration": true,
+    "outDir": "dist",
+    "rootDir": "src",
+    "strict": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true
+  },
+  "include": ["src"]
+}

package/vitest.config.ts

@@ -0,0 +1,8 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    include: ['test/**/*.spec.ts'],
+    testTimeout: 30_000,
+  },
+});