feat(add logutils):

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

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):

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
 

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()

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):

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

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")

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 (

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"
 

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")
 

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
 

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")

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")
 

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

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

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