From 12bd95b74d51da8987c12516005ff390b68edceb Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 30 Nov 2024 12:53:34 +0200
Subject: [PATCH] fix(rm IS_VERCEL):

---
 agentic_security/probe_actor/fuzzer.py | 14 --------------
 agentic_security/probe_data/data.py    | 16 +---------------
 agentic_security/routes/static.py      |  3 +--
 3 files changed, 2 insertions(+), 31 deletions(-)

diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py
index 15924e9..9b44fab 100644
--- a/agentic_security/probe_actor/fuzzer.py
+++ b/agentic_security/probe_actor/fuzzer.py
@@ -1,5 +1,4 @@
 import asyncio
-import os
 import random
 from collections.abc import AsyncGenerator
 
@@ -13,8 +12,6 @@ from agentic_security.models.schemas import ScanResult
 from agentic_security.probe_actor.refusal import refusal_heuristic
 from agentic_security.probe_data.data import prepare_prompts
 
-IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
-
 
 async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
     if isinstance(prompts, list):
@@ -34,11 +31,6 @@ async def perform_scan(
     stop_event: asyncio.Event = None,
 ) -> AsyncGenerator[str, None]:
     """Perform a standard security scan."""
-    if IS_VERCEL:
-        yield ScanResult.status_msg(
-            "Vercel deployment detected. Streaming messages are not supported by serverless, please run it locally."
-        )
-        return
 
     try:
         yield ScanResult.status_msg("Loading datasets...")
@@ -152,12 +144,6 @@ async def perform_multi_step_scan(
     probe_frequency: float = 0.2,
 ) -> AsyncGenerator[str, None]:
     """Perform a multi-step security scan with probe injection."""
-    if IS_VERCEL:
-        yield ScanResult.status_msg(
-            "Vercel deployment detected. Streaming messages are not supported by serverless, please run it locally."
-        )
-        return
-
     try:
         # Load main and probe datasets
         yield ScanResult.status_msg("Loading datasets...")
diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py
index 8d48d8b..2078af8 100644
--- a/agentic_security/probe_data/data.py
+++ b/agentic_security/probe_data/data.py
@@ -6,6 +6,7 @@ from functools import lru_cache
 
 import httpx
 import pandas as pd
+from cache_to_disk import cache_to_disk
 from loguru import logger
 
 from agentic_security.probe_data import stenography_fn
@@ -15,21 +16,6 @@ from agentic_security.probe_data.modules import (
     inspect_ai_tool,
 )
 
-IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
-
-if not IS_VERCEL:
-    from cache_to_disk import cache_to_disk
-else:
-    # Read only fs in vercel, just mock no-op decorator
-    def cache_to_disk(*_):
-        def decorator(fn):
-            def wrapper(*args, **kwargs):
-                return fn(*args, **kwargs)
-
-            return wrapper
-
-        return decorator
-
 
 @dataclass
 class ProbeDataset:
diff --git a/agentic_security/routes/static.py b/agentic_security/routes/static.py
index 0fea2ce..d36ce31 100644
--- a/agentic_security/routes/static.py
+++ b/agentic_security/routes/static.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-from typing import Optional
 
 from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import FileResponse, HTMLResponse
@@ -32,7 +31,7 @@ CONTENT_TYPES = {
 }
 
 
-def get_static_file(filepath: Path, content_type: Optional[str] = None) -> FileResponse:
+def get_static_file(filepath: Path, content_type: str | None = None) -> FileResponse:
     """
     Helper function to serve static files with proper error handling and caching.