fix(rm IS_VERCEL):

This commit is contained in:
Alexander Myasoedov committed 2024-11-30 12:53:34 +02:00
1 parent 7b086242a3
commit 12bd95b74d
3 files changed
+2 -31

No files matched your search

-14
View File
@@ -1,5 +1,4 @@
import asyncio import asyncio
import os
import random import random
from collections.abc import AsyncGenerator 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_actor.refusal import refusal_heuristic
from agentic_security.probe_data.data import prepare_prompts 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]: async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
if isinstance(prompts, list): if isinstance(prompts, list):
@@ -34,11 +31,6 @@ async def perform_scan(
stop_event: asyncio.Event = None, stop_event: asyncio.Event = None,
) -> AsyncGenerator[str, None]: ) -> AsyncGenerator[str, None]:
"""Perform a standard security scan.""" """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: try:
yield ScanResult.status_msg("Loading datasets...") yield ScanResult.status_msg("Loading datasets...")
@@ -152,12 +144,6 @@ async def perform_multi_step_scan(
probe_frequency: float = 0.2, probe_frequency: float = 0.2,
) -> AsyncGenerator[str, None]: ) -> AsyncGenerator[str, None]:
"""Perform a multi-step security scan with probe injection.""" """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: try:
# Load main and probe datasets # Load main and probe datasets
yield ScanResult.status_msg("Loading datasets...") yield ScanResult.status_msg("Loading datasets...")
+1 -15
View File
@@ -6,6 +6,7 @@ from functools import lru_cache
import httpx import httpx
import pandas as pd import pandas as pd
from cache_to_disk import cache_to_disk
from loguru import logger from loguru import logger
from agentic_security.probe_data import stenography_fn from agentic_security.probe_data import stenography_fn
@@ -15,21 +16,6 @@ from agentic_security.probe_data.modules import (
inspect_ai_tool, 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 @dataclass
class ProbeDataset: class ProbeDataset:
+1 -2
View File
@@ -1,5 +1,4 @@
from pathlib import Path from pathlib import Path
from typing import Optional
from fastapi import APIRouter, HTTPException, Request from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse 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. Helper function to serve static files with proper error handling and caching.