From d646ecd61bbef6ea546ca6c32382f22a56a2a20f Mon Sep 17 00:00:00 2001 From: Alexander Myasoedov Date: Sat, 8 Mar 2025 12:35:16 +0200 Subject: [PATCH] feat(add logutils): --- agentic_security/agents/operator_crew.py | 4 +- agentic_security/agents/operator_pydantic.py | 6 +-- agentic_security/config.py | 3 +- agentic_security/logutils.py | 45 +++++++++++++++++++ agentic_security/middleware/logging.py | 3 +- agentic_security/probe_actor/fuzzer.py | 2 +- agentic_security/probe_actor/operator.py | 6 +-- agentic_security/probe_data/data.py | 2 +- .../probe_data/modules/adaptive_attacks.py | 3 +- .../probe_data/modules/fine_tuned.py | 3 +- .../probe_data/modules/garak_tool.py | 2 +- .../probe_data/modules/inspect_ai_tool.py | 2 +- .../probe_data/modules/rl_model.py | 3 +- agentic_security/routes/proxy.py | 3 +- agentic_security/routes/scan.py | 3 +- agentic_security/routes/telemetry.py | 3 +- 16 files changed, 68 insertions(+), 25 deletions(-) create mode 100644 agentic_security/logutils.py diff --git a/agentic_security/agents/operator_crew.py b/agentic_security/agents/operator_crew.py index 60b0e85..e836569 100644 --- a/agentic_security/agents/operator_crew.py +++ b/agentic_security/agents/operator_crew.py @@ -1,5 +1,4 @@ import asyncio -import logging import os from typing import Any @@ -10,12 +9,11 @@ from pydantic import BaseModel, ConfigDict, Field # Assuming LLMSpec is defined elsewhere; placeholder import from agentic_security.http_spec import LLMSpec +from agentic_security.logutils import logger LLM_SPECS = [] # Populate with LLM spec strings if needed # Configure logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) # Define AgentSpecification model diff --git a/agentic_security/agents/operator_pydantic.py b/agentic_security/agents/operator_pydantic.py index 7e2b670..6921768 100644 --- a/agentic_security/agents/operator_pydantic.py +++ b/agentic_security/agents/operator_pydantic.py @@ -1,5 +1,4 @@ import asyncio -import logging from typing import Any import httpx @@ -8,13 +7,10 @@ from pydantic_ai import Agent, RunContext, Tool # Assuming LLMSpec is defined elsewhere; placeholder import from agentic_security.http_spec import LLMSpec +from agentic_security.logutils import logger LLM_SPECS = [] # Populate this list with LLM spec strings if needed -# Configure logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - # Define AgentSpecification model class AgentSpecification(BaseModel): diff --git a/agentic_security/config.py b/agentic_security/config.py index f166a0b..29a6d23 100644 --- a/agentic_security/config.py +++ b/agentic_security/config.py @@ -1,7 +1,8 @@ from functools import lru_cache import tomli -from loguru import logger + +from agentic_security.logutils import logger SETTINGS_VERSION = 1 diff --git a/agentic_security/logutils.py b/agentic_security/logutils.py new file mode 100644 index 0000000..3e897fe --- /dev/null +++ b/agentic_security/logutils.py @@ -0,0 +1,45 @@ +import logging +from os import getenv + +from rich.logging import RichHandler + +LOGGER_NAME = "agentic_security" + + +def get_logger(logger_name: str) -> logging.Logger: + # https://rich.readthedocs.io/en/latest/reference/logging.html#rich.logging.RichHandler + # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions + rich_handler = RichHandler( + show_time=False, + rich_tracebacks=False, + show_path=True if getenv("API_RUNTIME") == "dev" else False, + tracebacks_show_locals=False, + ) + rich_handler.setFormatter( + logging.Formatter( + fmt="%(message)s", + datefmt="[%X]", + ) + ) + + _logger = logging.getLogger(logger_name) + _logger.addHandler(rich_handler) + _logger.setLevel(logging.INFO) + _logger.propagate = True + return _logger + + +logger: logging.Logger = get_logger(LOGGER_NAME) + + +def set_log_level_to_debug(): + _logger = logging.getLogger(LOGGER_NAME) + _logger.setLevel(logging.DEBUG) + + +def set_log_level_to_info(): + _logger = logging.getLogger(LOGGER_NAME) + _logger.setLevel(logging.INFO) + + +set_log_level_to_debug() diff --git a/agentic_security/middleware/logging.py b/agentic_security/middleware/logging.py index 0cf5be8..53b5042 100644 --- a/agentic_security/middleware/logging.py +++ b/agentic_security/middleware/logging.py @@ -1,7 +1,8 @@ from fastapi import Request -from loguru import logger from starlette.middleware.base import BaseHTTPMiddleware +from agentic_security.logutils import logger + class LogNon200ResponsesMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index bafaaf3..882fc7f 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -6,11 +6,11 @@ from json import JSONDecodeError import httpx import pandas as pd -from loguru import logger from skopt import Optimizer from skopt.space import Real from agentic_security.http_spec import Modality +from agentic_security.logutils import logger 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 diff --git a/agentic_security/probe_actor/operator.py b/agentic_security/probe_actor/operator.py index 99d65f0..77054c4 100644 --- a/agentic_security/probe_actor/operator.py +++ b/agentic_security/probe_actor/operator.py @@ -1,5 +1,4 @@ import asyncio -import logging from typing import Any import httpx @@ -7,13 +6,10 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext from agentic_security.http_spec import LLMSpec +from agentic_security.logutils import logger LLM_SPECS = [] -# Configure logging -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - class AgentSpecification(BaseModel): name: str | None = Field(None, description="Name of the LLM/agent") diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py index c0ef2d5..4c414d7 100644 --- a/agentic_security/probe_data/data.py +++ b/agentic_security/probe_data/data.py @@ -7,8 +7,8 @@ import httpx import pandas as pd from cache_to_disk import cache_to_disk from datasets import load_dataset -from loguru import logger +from agentic_security.logutils import logger from agentic_security.probe_data import stenography_fn from agentic_security.probe_data.models import ProbeDataset from agentic_security.probe_data.modules import ( diff --git a/agentic_security/probe_data/modules/adaptive_attacks.py b/agentic_security/probe_data/modules/adaptive_attacks.py index 82f7958..33a5dd9 100644 --- a/agentic_security/probe_data/modules/adaptive_attacks.py +++ b/agentic_security/probe_data/modules/adaptive_attacks.py @@ -3,7 +3,8 @@ import io import httpx import pandas as pd -from loguru import logger + +from agentic_security.logutils import logger url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harmful_behaviors/harmful_behaviors_pair.csv" diff --git a/agentic_security/probe_data/modules/fine_tuned.py b/agentic_security/probe_data/modules/fine_tuned.py index fa01951..0c2baa2 100644 --- a/agentic_security/probe_data/modules/fine_tuned.py +++ b/agentic_security/probe_data/modules/fine_tuned.py @@ -3,7 +3,8 @@ import os import uuid as U import httpx -from loguru import logger + +from agentic_security.logutils import logger AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115b") diff --git a/agentic_security/probe_data/modules/garak_tool.py b/agentic_security/probe_data/modules/garak_tool.py index 759c279..a68b6e7 100644 --- a/agentic_security/probe_data/modules/garak_tool.py +++ b/agentic_security/probe_data/modules/garak_tool.py @@ -4,7 +4,7 @@ import json import os import subprocess -from loguru import logger +from agentic_security.logutils import logger # TODO: add probes modules diff --git a/agentic_security/probe_data/modules/inspect_ai_tool.py b/agentic_security/probe_data/modules/inspect_ai_tool.py index a512351..d2058e7 100644 --- a/agentic_security/probe_data/modules/inspect_ai_tool.py +++ b/agentic_security/probe_data/modules/inspect_ai_tool.py @@ -2,7 +2,7 @@ import asyncio import importlib.util import os -from loguru import logger +from agentic_security.logutils import logger inspect_ai_task = ( __file__.replace("inspect_ai_tool.py", "inspect_ai_task.py") diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py index e1b71ed..e024976 100644 --- a/agentic_security/probe_data/modules/rl_model.py +++ b/agentic_security/probe_data/modules/rl_model.py @@ -8,7 +8,8 @@ from typing import Deque import numpy as np import requests -from loguru import logger + +from agentic_security.logutils import logger AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115b") diff --git a/agentic_security/routes/proxy.py b/agentic_security/routes/proxy.py index 8ef1d04..ff2a976 100644 --- a/agentic_security/routes/proxy.py +++ b/agentic_security/routes/proxy.py @@ -2,7 +2,8 @@ import random from asyncio import Event from fastapi import APIRouter -from loguru import logger + +from agentic_security.logutils import logger from ..core.app import get_current_run, get_tools_inbox from ..primitives import CompletionRequest, Settings diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index f5e2977..558bbb3 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -10,7 +10,8 @@ from fastapi import ( UploadFile, ) from fastapi.responses import StreamingResponse -from loguru import logger + +from agentic_security.logutils import logger from ..core.app import get_stop_event, get_tools_inbox, set_current_run from ..dependencies import InMemorySecrets, get_in_memory_secrets diff --git a/agentic_security/routes/telemetry.py b/agentic_security/routes/telemetry.py index be48573..4b3c885 100644 --- a/agentic_security/routes/telemetry.py +++ b/agentic_security/routes/telemetry.py @@ -1,7 +1,8 @@ import sentry_sdk -from loguru import logger from sentry_sdk.integrations.logging import ignore_logger +from agentic_security.logutils import logger + from ..primitives import Settings