Video check (#566)

- remove "inline" mode
- add gcs delete prefix to allow deleting the `msg_XYZ` directory in GCS
(so that we dont rely on message in the DB)
- add cleanup in the handle_retry_exhausted
- change log error before retry_exhausted to warning
- propagate trace through workers

~~Have not verified traces on google~~ Look good there
~~Still looking into why message_id isn't there~~ I think messages
erroring (saw some from inference error) preventing creation of message
(no message_Id)

Author: bronzle

.env.test

@@ -20,6 +20,6 @@ AI2_MODEL_HUB_API_KEY=fake
 RECAPTCHA_ENABLED=false
 RECAPTCHA_KEY=fake
 
-VIDEO_SAFETY_CHECK_WORKER_STRATEGY=fake
+SAFETY_QUEUE_ENABLED=false
 SAFETY_QUEUE_URL=fake
 SAFTEY_GCS_UPLOAD_BUCKET=fake

.vscode/launch.json

@@ -53,6 +53,25 @@
             "jinja": true,
             "justMyCode": false
         },
+        {
+            "name": "Python Debugger: API Safety Worker",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "dramatiq",
+            "env": {
+                "PYTHONPATH": "${workspaceFolder}/apps/api/src",
+                "ENV": "development",
+                "GOOGLE_APPLICATION_CREDENTIALS": "${workspaceFolder}/service_account.json"
+            },
+            "args": [
+                "--processes", "1",
+                "--threads", "1",
+                "api.safety_queue:setup_safety_queue",
+                "api.thread.chat.safety.safety_checkers.google_video_safety_checker"
+            ],
+            "jinja": true,
+            "justMyCode": false
+        },
         {
             "name": "Alembic Auto-generate Debug Migration",
             "type": "debugpy",

apps/api/src/api/config.py

@@ -72,9 +72,9 @@ class Settings(BaseSettings):
     RECAPTCHA_KEY: str = Field(init=False)
     RECAPTCHA_MIN_SCORE_REQUIREMENT: float = 0.3
 
+    SAFETY_QUEUE_ENABLED: bool = True
     SAFETY_QUEUE_URL: str = Field(init=False)
     SAFTEY_GCS_UPLOAD_BUCKET: str = Field(init=False)
-    VIDEO_SAFETY_CHECK_WORKER_STRATEGY: str = "deferred"
 
     model_config = SettingsConfigDict(
         extra="ignore",

apps/api/src/api/safety_queue.py

@@ -1,33 +1,83 @@
 import dramatiq
+import structlog
 from dramatiq.brokers.redis import RedisBroker
 from dramatiq.brokers.stub import StubBroker
 from dramatiq.middleware.asyncio import AsyncIO
 from dramatiq.middleware.prometheus import Prometheus
+from opentelemetry import context, propagate, trace
 from typing_extensions import override
 
 from api.config import settings
+from api.logging.fastapi_logger import FastAPIStructLogger
 from api.logging.setup import setup_logging
 from api.otel.setup import setup_otel
 
 # Dramatiq requires a broker to be set before actors are declared.
 # This sets a stub so actor modules can be imported safely before setup_safety_queue() is called.
 # https://github.com/Bogdanp/dramatiq/pull/762
 dramatiq.set_broker(StubBroker())
+logger = FastAPIStructLogger()
+
+SAFETY_QUEUE_NAMESPACE = "playground_api_safety_queue"
 
 
 class OtelMiddleware(dramatiq.Middleware):
     @override
     def after_worker_boot(self, broker: dramatiq.Broker, worker: dramatiq.Worker) -> None:
         setup_otel()
 
+    @override
+    def before_enqueue(self, broker: dramatiq.Broker, message: dramatiq.Message, delay: int) -> None:
+        if "otel_context" not in message.options:
+            carrier: dict = {}
+            propagate.inject(carrier)
+            message.options["otel_context"] = carrier
+
+    @override
+    def before_process_message(self, broker: dramatiq.Broker, message: dramatiq.MessageProxy) -> None:
+        # Context is in message.options for a regular dramatiq actor
+        # or in the message.kwargs.failed_message.options for retry_exhausted actor
+        carrier = message.options.get("otel_context") or (
+            message.kwargs.get("failed_message", {}).get("options", {}).get("otel_context", {})
+        )
+
+        ctx = propagate.extract(carrier)
+        token = context.attach(ctx)
+        message.options["_otel_token"] = token
+
+        structlog.contextvars.clear_contextvars()
+        span_ctx = trace.get_current_span().get_span_context()
+        if span_ctx.is_valid:
+            structlog.contextvars.bind_contextvars(
+                trace_id=format(span_ctx.trace_id, "032x"),
+                span_id=format(span_ctx.span_id, "016x"),
+                trace_flags=span_ctx.trace_flags,
+            )
+
+    @override
+    def after_process_message(
+        self, broker: dramatiq.Broker, message: dramatiq.MessageProxy, *, result=None, exception=None
+    ) -> None:
+        token = message.options.pop("_otel_token", None)
+        if token is not None:
+            context.detach(token)
+        structlog.contextvars.clear_contextvars()
+
+    @override
+    def after_skip_message(self, broker: dramatiq.Broker, message: dramatiq.MessageProxy) -> None:
+        token = message.options.pop("_otel_token", None)
+        if token is not None:
+            context.detach(token)
+        structlog.contextvars.clear_contextvars()
+
 
 def setup_safety_queue() -> None:
     setup_logging(json_logs=settings.LOG_JSON_FORMAT, log_level=settings.LOG_LEVEL)
 
-    if settings.VIDEO_SAFETY_CHECK_WORKER_STRATEGY != "deferred":
+    if not settings.SAFETY_QUEUE_ENABLED:
         return
 
-    redis_broker = RedisBroker(url=settings.SAFETY_QUEUE_URL, namespace="playground_safety_queue")
+    redis_broker = RedisBroker(url=settings.SAFETY_QUEUE_URL, namespace=SAFETY_QUEUE_NAMESPACE)
 
     old_broker = dramatiq.get_broker()
     for existing_actor_name in old_broker.get_declared_actors():
@@ -36,7 +86,7 @@ def setup_safety_queue() -> None:
         redis_broker.declare_actor(actor)
 
     redis_broker.add_middleware(Prometheus())
-    redis_broker.add_middleware(AsyncIO())
     redis_broker.add_middleware(OtelMiddleware())
+    redis_broker.add_middleware(AsyncIO())
 
     dramatiq.set_broker(redis_broker)

apps/api/src/api/thread/chat/safety/safety_checkers/google_video_safety_checker.py

@@ -26,7 +26,6 @@
     SafetyChecker,
     SafetyCheckRequest,
     SafetyCheckResponse,
-    SafetyCheckUnsafeError,
     SkippedSafetyCheckResponse,
 )
 
@@ -56,9 +55,6 @@ def _make_worker_sessionmaker():
 class VideoIntelligenceOperationNotAvailableError(Exception): ...
 
 
-class VideoIntelligenceOperationNotFinishedError(Exception): ...
-
-
 class VideoIntelligenceOperationMessageNotFoundError(Exception): ...
 
 
@@ -108,46 +104,65 @@ async def check_request(self, request: SafetyCheckRequest, *, throw: bool = Fals
 
         span.set_attribute("operation_name", operation_name)
 
-        if settings.VIDEO_SAFETY_CHECK_WORKER_STRATEGY == "deferred":
-            handle_video_safety_check.send(
-                operation_name=operation_name, file_url=request.content, message_id=request.message_id
-            )
-
-        if settings.VIDEO_SAFETY_CHECK_WORKER_STRATEGY == "inline":
-            response = await _handle_video_safety_check_async(
-                operation_name=operation_name, file_url=request.content, message_id=request.message_id
-            )
-            if not response.is_safe() and throw:
-                raise SafetyCheckUnsafeError
-
-            return response
+        handle_video_safety_check.send(
+            operation_name=operation_name, file_url=request.content, message_id=request.message_id
+        )
 
         return SkippedSafetyCheckResponse()
 
 
 @dramatiq.actor
-def handle_retry_exhausted(*args, **kwargs) -> None:
-    # Logs retry limits so we can alert of off them
+@tracer.start_as_current_span("handle_retry_exhausted")
+async def handle_retry_exhausted(failed_message: dict, retry_info: dict):
+    kwargs: dict = failed_message.get("kwargs", {})
+    operation_name: str = kwargs.get("operation_name", "unknown")
+    file_url: str = kwargs.get("file_url", "unknown")
+    message_id: str = kwargs.get("message_id", "unknown")
+
+    span = trace.get_current_span()
+    span.set_attributes({
+        "operation_name": operation_name,
+        "file_url": file_url,
+        "message_id": message_id,
+        "retries": retry_info.get("retries", 0),
+    })
+    span.set_status(trace.StatusCode.ERROR, "video safety check retries exhausted")
+
     logger.error(
         "video_safety_worker.retry_exhausted",
-        job_args=str(args),
-        job_kwargs=str(kwargs),
+        operation_name=operation_name,
+        file_url=file_url,
+        message_id=message_id,
+        retries=retry_info.get("retries"),
+        max_retries=retry_info.get("max_retries"),
+        traceback=failed_message.get("options", {}).get("traceback"),
     )
 
+    storage_client = get_google_cloud_storage()
+    safety_file_name = Path(file_url).parts[-1]
+
+    await storage_client.delete_file(filename=safety_file_name, bucket_name=settings.SAFTEY_GCS_UPLOAD_BUCKET)
+
+    await storage_client.delete_prefix(prefix=message_id, bucket_name=settings.USER_CONTENT_BUCKET)
+
+    Session = _make_worker_sessionmaker()  # noqa: N806
+    async with Session() as session:
+        message_repository = AsyncMessageRepository(session)
+        message = await message_repository.get_message_by_id(message_id)
+        if message is not None:
+            message.harmful = True
+            message.file_urls = None
+            await message_repository.update(message)
+            await session.commit()
+
 
-@dramatiq.actor(  # type: ignore
+@dramatiq.actor(
     queue_name=SAFETY_QUEUE_NAME,
     max_retries=5,
     on_retry_exhausted=handle_retry_exhausted.actor_name,
 )
-async def handle_video_safety_check(operation_name: str, file_url: str, message_id: str) -> None:
-    await _handle_video_safety_check_async(operation_name=operation_name, file_url=file_url, message_id=message_id)
-
-
 @tracer.start_as_current_span("handle_video_safety_check")
-async def _handle_video_safety_check_async(
-    operation_name: str, file_url: str, message_id: str
-) -> GoogleVideoIntelligenceResponse:
+async def handle_video_safety_check(operation_name: str, file_url: str, message_id: str):
     span = trace.get_current_span()
     span.set_attributes({
         "operation_name": operation_name,
@@ -168,6 +183,11 @@ async def _handle_video_safety_check_async(
                 trace.StatusCode.ERROR,
                 f"Operation {operation_name} not found. The Operation endpoint may not have the operation yet.",
             )
+            logger.warning(
+                "video_safety.operation_not_available",
+                operation_name=operation_name,
+                message_id=message_id,
+            )
             raise VideoIntelligenceOperationNotAvailableError
 
         try:
@@ -182,35 +202,30 @@ async def _handle_video_safety_check_async(
                 trace.StatusCode.ERROR,
                 f"Operation {operation_name} not found or not done. The Operation endpoint may not have the operation yet.",
             )
+            logger.warning(
+                "video_safety.operation_not_available",
+                operation_name=operation_name,
+                message_id=message_id,
+            )
             raise VideoIntelligenceOperationNotAvailableError from e
 
         result = await operation.result()
 
         if not isinstance(result, AnnotateVideoResponse):
             msg = "Unexpected result from google video checker"
+            span.set_status(trace.StatusCode.ERROR, msg)
             raise TypeError(msg)
 
         mapped_response = GoogleVideoIntelligenceResponse(result)
         span.set_attribute("is_safe", mapped_response.is_safe())
 
-        # if we are blocking mode -- the message doesn't exist yet and we will return
-        # a response to indicate its not safe back to the main event loop
-        if settings.VIDEO_SAFETY_CHECK_WORKER_STRATEGY == "inline":
-            logger.info(
-                "video_safety.blocking_mode.complete",
-                operation=operation_name,
-                is_safe=mapped_response.is_safe(),
-                message_id=message_id,
-            )
-
-            return mapped_response
-
         message_repository = AsyncMessageRepository(session)
         message = await message_repository.get_message_by_id(message_id)
 
         if message is None:
             not_found_message = f"Message {message_id} not found when evaluating a video safety check"
-            logger.error(
+            span.set_status(trace.StatusCode.ERROR, not_found_message)
+            logger.warning(
                 "video_safety.message_not_found",
                 operation=operation_name,
                 message_id=message_id,
@@ -228,29 +243,22 @@ async def _handle_video_safety_check_async(
             )
 
             storage_client = get_google_cloud_storage()
-
             safety_file_name = Path(file_url).parts[-1]
 
             await storage_client.delete_file(
                 filename=safety_file_name,
                 bucket_name=settings.SAFTEY_GCS_UPLOAD_BUCKET,
             )
 
-            if message.file_urls:
-                await storage_client.delete_multiple_files_by_url(
-                    file_urls=message.file_urls,
-                    bucket_name=settings.USER_CONTENT_BUCKET,
-                )
-                message.file_urls = None
+            await storage_client.delete_prefix(prefix=message_id, bucket_name=settings.USER_CONTENT_BUCKET)
+            message.file_urls = None
 
         await message_repository.update(message)
         await session.commit()
 
         logger.info(
-            "video_safety.queue_mode.complete",
-            operation=operation_name,
+            "video_safety.complete",
+            operation_name=operation_name,
             is_safe=mapped_response.is_safe(),
             message_id=message_id,
         )
-
-        return mapped_response