diff --git a/agentic_security/lib.py b/agentic_security/lib.py index 0c300bf..67ddd32 100644 --- a/agentic_security/lib.py +++ b/agentic_security/lib.py @@ -10,7 +10,7 @@ from rich.table import Table from tabulate import tabulate from agentic_security.config import SettingsMixin # Importing the configuration mixin -from agentic_security.models.schemas import Scan +from agentic_security.primitives import Scan from agentic_security.probe_data import REGISTRY from agentic_security.routes.scan import streaming_response_generator diff --git a/agentic_security/primitives/__init__.py b/agentic_security/primitives/__init__.py new file mode 100644 index 0000000..913ef1d --- /dev/null +++ b/agentic_security/primitives/__init__.py @@ -0,0 +1,11 @@ +from agentic_security.primitives.models import ( + Settings, + LLMInfo, + Scan, + ScanResult, + Probe, + Message, + CompletionRequest, + FileProbeResponse, + Table, +) # noqa diff --git a/agentic_security/models/schemas.py b/agentic_security/primitives/models.py similarity index 100% rename from agentic_security/models/schemas.py rename to agentic_security/primitives/models.py diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index 773ad47..76e4842 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -11,7 +11,7 @@ from skopt import Optimizer from skopt.space import Real from agentic_security.http_spec import Modality -from agentic_security.models.schemas import Scan, ScanResult +from agentic_security.primitives import Scan, ScanResult from agentic_security.probe_actor.cost_module import calculate_cost from agentic_security.probe_actor.refusal import refusal_heuristic from agentic_security.probe_data import audio_generator, image_generator, msj_data diff --git a/agentic_security/probe_actor/test_fuzzer.py b/agentic_security/probe_actor/test_fuzzer.py index 830d606..325c715 100644 --- a/agentic_security/probe_actor/test_fuzzer.py +++ b/agentic_security/probe_actor/test_fuzzer.py @@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest -from agentic_security.models.schemas import Scan +from agentic_security.primitives import Scan from agentic_security.probe_actor.fuzzer import ( generate_prompts, perform_many_shot_scan, diff --git a/agentic_security/routes/probe.py b/agentic_security/routes/probe.py index 2c8d498..523b3d0 100644 --- a/agentic_security/routes/probe.py +++ b/agentic_security/routes/probe.py @@ -3,7 +3,7 @@ import random from fastapi import APIRouter, File, Header, HTTPException, UploadFile from fastapi.responses import JSONResponse -from ..models.schemas import FileProbeResponse, Probe +from ..primitives import FileProbeResponse, Probe from ..probe_actor.refusal import REFUSAL_MARKS from ..probe_data import REGISTRY diff --git a/agentic_security/routes/proxy.py b/agentic_security/routes/proxy.py index d67bebf..8ef1d04 100644 --- a/agentic_security/routes/proxy.py +++ b/agentic_security/routes/proxy.py @@ -5,7 +5,7 @@ from fastapi import APIRouter from loguru import logger from ..core.app import get_current_run, get_tools_inbox -from ..models.schemas import CompletionRequest, Settings +from ..primitives import CompletionRequest, Settings from ..probe_actor.refusal import REFUSAL_MARKS router = APIRouter() diff --git a/agentic_security/routes/report.py b/agentic_security/routes/report.py index f45dbd3..6db0265 100644 --- a/agentic_security/routes/report.py +++ b/agentic_security/routes/report.py @@ -3,7 +3,7 @@ from pathlib import Path from fastapi import APIRouter, Response from fastapi.responses import FileResponse, StreamingResponse -from ..models.schemas import Table +from ..primitives import Table from ..report_chart import plot_security_report router = APIRouter() diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index e987863..12a345b 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -14,7 +14,7 @@ from fastapi.responses import StreamingResponse from ..core.app import get_stop_event, get_tools_inbox, set_current_run from ..dependencies import InMemorySecrets, get_in_memory_secrets from ..http_spec import LLMSpec -from ..models.schemas import LLMInfo, Scan +from ..primitives import LLMInfo, Scan from ..probe_actor import fuzzer router = APIRouter() diff --git a/agentic_security/routes/static.py b/agentic_security/routes/static.py index a919ae0..2abafd9 100644 --- a/agentic_security/routes/static.py +++ b/agentic_security/routes/static.py @@ -7,7 +7,7 @@ from fastapi.templating import Jinja2Templates from jinja2 import Environment, FileSystemLoader from starlette.responses import Response -from ..models.schemas import Settings +from ..primitives import Settings router = APIRouter() STATIC_DIR = Path(__file__).parent.parent / "static" diff --git a/agentic_security/routes/telemetry.py b/agentic_security/routes/telemetry.py index f7e7495..be48573 100644 --- a/agentic_security/routes/telemetry.py +++ b/agentic_security/routes/telemetry.py @@ -2,7 +2,7 @@ import sentry_sdk from loguru import logger from sentry_sdk.integrations.logging import ignore_logger -from ..models.schemas import Settings +from ..primitives import Settings def setup(app): diff --git a/agentic_security/routes/test_probe.py b/agentic_security/routes/test_probe.py index 8a3644b..7005976 100644 --- a/agentic_security/routes/test_probe.py +++ b/agentic_security/routes/test_probe.py @@ -6,7 +6,7 @@ import pytest from fastapi.testclient import TestClient from ..app import app -from ..models.schemas import Probe +from ..primitives import Probe from ..probe_actor.refusal import REFUSAL_MARKS from ..probe_data import REGISTRY diff --git a/agentic_security/routes/test_static.py b/agentic_security/routes/test_static.py index fc6ca90..d6811d0 100644 --- a/agentic_security/routes/test_static.py +++ b/agentic_security/routes/test_static.py @@ -4,7 +4,7 @@ import pytest from fastapi import HTTPException from fastapi.testclient import TestClient -from ..models.schemas import Settings +from ..primitives import Settings from .static import get_static_file, router client = TestClient(router)