Custom Detectors

Rdakt AI's detection pipeline is extensible. You can write custom detectors that implement the Detector protocol and add them as pipeline stages.

The Detector Protocol

A detector is any object with a detect(text: str) -> list[Entity] method:

from rdakt_ai import Entity

class Detector:
    def detect(self, text: str) -> list[Entity]: ...

Example: Custom Regex Detector

Detect employee IDs and internal project codes:

import re
from rdakt_ai import Entity

class InternalDataDetector:
    """Detect company-specific PII patterns."""

    PATTERNS = {
        "EMPLOYEE_ID": r"EMP-\d{6}",
        "PROJECT_CODE": r"PRJ-[A-Z]{2}-\d{4}",
        "INTERNAL_EMAIL": r"[a-zA-Z.]+@mycompany\.com",
    }

    def detect(self, text: str) -> list[Entity]:
        entities = []
        for entity_type, pattern in self.PATTERNS.items():
            for match in re.finditer(pattern, text):
                entities.append(Entity(
                    value=match.group(),
                    type=entity_type,
                    start=match.start(),
                    end=match.end(),
                ))
        return entities

Adding to the Pipeline

Wrap your detector in a DetectorStage and build the pipeline:

from rdakt_ai import (
    RdaktMiddleware,
    RdaktPipeline,
    DetectorStage,
    RegexDetector,
)

pipeline = RdaktPipeline(stages=[
    DetectorStage(RegexDetector()),          # Built-in patterns
    DetectorStage(InternalDataDetector()),    # Your custom detector
])

To use the custom pipeline with the middleware, you'll need to build the middleware state manually or configure it through the config's entity patterns:

from rdakt_ai import RdaktMiddleware, RdaktConfig, EntityConfig

config = RdaktConfig(
    entities={
        "EMPLOYEE_ID": EntityConfig(pattern=r"EMP-\d{6}", strategy="token"),
        "PROJECT_CODE": EntityConfig(pattern=r"PRJ-[A-Z]{2}-\d{4}", strategy="token"),
        "INTERNAL_EMAIL": EntityConfig(pattern=r"[a-zA-Z.]+@mycompany\.com", strategy="token"),
    },
)

middleware = RdaktMiddleware(config=config)

Custom Pipeline Stage

For more control, implement a full RdaktStage:

from rdakt_ai import RdaktStage, RdaktContext

class FilterStage(RdaktStage):
    """Remove entities below a confidence threshold or in an allowlist."""

    def __init__(self, allowlist: set[str]):
        self.allowlist = allowlist

    async def process(self, text, context, call_next):
        # Let downstream detectors run first
        await call_next(text, context)
        # Filter out allowed values
        context.entities = [
            e for e in context.entities
            if e.value.lower() not in self.allowlist
        ]

    def process_sync(self, text, context, call_next_sync):
        call_next_sync(text, context)
        context.entities = [
            e for e in context.entities
            if e.value.lower() not in self.allowlist
        ]

Both process (async) and process_sync must be implemented if you want your stage to work with both RdaktMiddleware (async) and RdaktSyncMiddleware (sync).

Testing Your Detector

def test_internal_data_detector():
    detector = InternalDataDetector()
    entities = detector.detect("Contact EMP-123456 about PRJ-AB-1234")

    assert len(entities) == 2
    assert entities[0].type == "EMPLOYEE_ID"
    assert entities[0].value == "EMP-123456"
    assert entities[1].type == "PROJECT_CODE"
    assert entities[1].value == "PRJ-AB-1234"