refactor detector management to store multiple detector types

Author: saichandrapandraju

detectors/built_in/app.py

@@ -1,5 +1,5 @@
 from fastapi import HTTPException
-
+from contextlib import asynccontextmanager
 from base_detector_registry import BaseDetectorRegistry
 from regex_detectors import RegexDetectorRegistry
 from file_type_detectors import FileTypeDetectorRegistry
@@ -8,21 +8,31 @@
 from detectors.common.scheme import ContentAnalysisHttpRequest,  ContentsAnalysisResponse
 from detectors.common.app import DetectorBaseAPI as FastAPI
 
-app = FastAPI()
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.set_detector(RegexDetectorRegistry(), "regex")
+    app.set_detector(FileTypeDetectorRegistry(), "file_type")
+    yield
+    
+    app.cleanup_detector()
+
+
+app = FastAPI(lifespan=lifespan)
 Instrumentator().instrument(app).expose(app)
 
 
-registry : dict[str, BaseDetectorRegistry] = {
-    "regex": RegexDetectorRegistry(),
-    "file_type": FileTypeDetectorRegistry(),
-}
+# registry : dict[str, BaseDetectorRegistry] = {
+#     "regex": RegexDetectorRegistry(),
+#     "file_type": FileTypeDetectorRegistry(),
+# }
 
 @app.post("/api/v1/text/contents", response_model=ContentsAnalysisResponse)
 def detect_content(request: ContentAnalysisHttpRequest):
     detections = []
     for content in request.contents:
         message_detections = []
-        for detector_kind, detector_registry in registry.items():
+        for detector_kind, detector_registry in app.get_all_detectors().items():
+            assert isinstance(detector_registry, BaseDetectorRegistry), f"Detector {detector_kind} is not a valid BaseDetectorRegistry"
             if detector_kind in request.detector_params:
                 try:
                     message_detections += detector_registry.handle_request(content, request.detector_params)
@@ -37,7 +47,8 @@ def detect_content(request: ContentAnalysisHttpRequest):
 @app.get("/registry")
 def get_registry():
     result = {}
-    for detector_type, detector_registry in registry.items():
+    for detector_type, detector_registry in app.get_all_detectors().items():
+        assert isinstance(detector_registry, BaseDetectorRegistry), f"Detector {detector_type} is not a valid BaseDetectorRegistry"
         result[detector_type] = {}
         for detector_name, detector_fn in detector_registry.get_registry().items():
             result[detector_type][detector_name] = detector_fn.__doc__

detectors/common/app.py

@@ -34,6 +34,7 @@
 class DetectorBaseAPI(FastAPI):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.state.detectors = {}
         self.add_exception_handler(
             RequestValidationError, self.validation_exception_handler
         )
@@ -94,17 +95,21 @@ async def http_exception_handler(self, request, exc):
             content={"code": exc.status_code, "message": exc.detail},
         )
 
-    def set_detector(self, detector) -> None:
+    def set_detector(self, detector, detector_name="default") -> None:
         """Store detector in app.state"""
-        self.state.detector = detector
+        self.state.detectors[detector_name] = detector
 
-    def get_detector(self):
+    def get_detector(self, detector_name="default"):
         """Retrieve detector from app.state"""
-        return getattr(self.state, 'detector', None)
+        return self.state.detectors.get(detector_name)
+
+    def get_all_detectors(self) -> dict:
+        """Retrieve all detectors from app.state"""
+        return self.state.detectors
 
     def cleanup_detector(self) -> None:
         """Clean up detector resources"""
-        self.state.detector = None
+        self.state.detectors.clear()
 
 async def health():
     return "ok"

detectors/llm_judge/app.py

@@ -1,7 +1,7 @@
 from contextlib import asynccontextmanager
 from typing import Annotated, Dict
 
-from fastapi import Header
+from fastapi import Header, HTTPException
 from prometheus_fastapi_instrumentator import Instrumentator
 
 from detectors.common.app import DetectorBaseAPI as FastAPI
@@ -48,6 +48,8 @@ async def detector_unary_handler(
 ):
     """Analyze content using LLM-as-Judge evaluation."""
     detector: LLMJudgeDetector = app.get_detector()
+    if not detector:
+        raise HTTPException(status_code=404, detail="Detector not found")
     return ContentsAnalysisResponse(root=await detector.run(request))
 
 
@@ -63,7 +65,7 @@ async def list_metrics():
     """List all available evaluation metrics."""
     detector: LLMJudgeDetector = app.get_detector()
     if not detector:
-        return {"metrics": [], "total": 0}
+        raise HTTPException(status_code=404, detail="Detector not found")
 
     metrics = detector.list_available_metrics()
     return MetricsListResponse(metrics=metrics, total=len(metrics))

tests/detectors/builtIn/test_filetype.py

@@ -7,6 +7,10 @@ class TestFileTypeDetectors:
     @pytest.fixture
     def client(self):
         from detectors.built_in.app import app
+        from detectors.built_in.file_type_detectors import FileTypeDetectorRegistry
+
+        app.set_detector(FileTypeDetectorRegistry(), "file_type")
+
         return TestClient(app)
 
     @pytest.fixture

tests/detectors/builtIn/test_regex.py

@@ -5,6 +5,10 @@ class TestRegexDetectors:
     @pytest.fixture
     def client(self):
         from detectors.built_in.app import app
+        from detectors.built_in.regex_detectors import RegexDetectorRegistry
+
+        app.set_detector(RegexDetectorRegistry(), "regex")
+
         return TestClient(app)
 
     @pytest.mark.parametrize(

tox.ini

@@ -10,6 +10,7 @@ deps =
     -r{toxinidir}/detectors/huggingface/requirements.txt
     -r{toxinidir}/detectors/common/requirements.txt
     -r{toxinidir}/detectors/llm_judge/requirements.txt
+    -r{toxinidir}/detectors/built_in/requirements.txt
 setenv =
     PYTHONPATH = {toxinidir}/detectors/huggingface:{toxinidir}/detectors/llm_judge:{toxinidir}/detectors:{toxinidir}
 commands =