fix(rm IS_VERCEL):

This commit is contained in:
Alexander Myasoedov
2024-11-30 12:53:34 +02:00
parent 7b086242a3
commit 12bd95b74d
3 changed files with 2 additions and 31 deletions
-14
View File
@@ -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...")
+1 -15
View File
@@ -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:
+1 -2
View File
@@ -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.