diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py index 09b6a0b..483369b 100644 --- a/agentic_security/http_spec.py +++ b/agentic_security/http_spec.py @@ -1,4 +1,5 @@ import base64 +import json from enum import Enum from urllib.parse import urlparse @@ -145,6 +146,18 @@ class LLMSpec(BaseModel): fn = probe + @property + def model_name(self) -> str: + """Extract the model name from the request body (JSON). + + Returns the value of the 'model' field if present, otherwise 'unknown'. + """ + try: + body_json = json.loads(self.body) + return body_json.get("model", "unknown") + except (json.JSONDecodeError, TypeError): + return "unknown" + @property def modality(self) -> Modality: if self.has_image: diff --git a/agentic_security/primitives/models.py b/agentic_security/primitives/models.py index b9bdbc9..05019c7 100644 --- a/agentic_security/primitives/models.py +++ b/agentic_security/primitives/models.py @@ -40,7 +40,7 @@ class Scan(BaseModel): class ScanResult(BaseModel): module: str tokens: float | int - cost: float + cost: float | None progress: float status: bool = False failureRate: float = 0.0 diff --git a/agentic_security/probe_actor/cost_module.py b/agentic_security/probe_actor/cost_module.py index 74db2dd..e103264 100644 --- a/agentic_security/probe_actor/cost_module.py +++ b/agentic_security/probe_actor/cost_module.py @@ -1,4 +1,7 @@ -def calculate_cost(tokens: int, model: str = "deepseek-chat") -> float: +from agentic_security.logutils import logger + + +def calculate_cost(tokens: int, model: str = "deepseek-chat") -> float | None: """Calculate API cost based on token count and model. Args: @@ -6,7 +9,7 @@ def calculate_cost(tokens: int, model: str = "deepseek-chat") -> float: model (str): Model name to calculate cost for Returns: - float: Cost in USD + float | None: Cost in USD, or None if the model pricing is unknown. """ # API pricing as of 2024-03-01 pricing = { @@ -49,7 +52,10 @@ def calculate_cost(tokens: int, model: str = "deepseek-chat") -> float: } if model not in pricing: - raise ValueError(f"Unknown model: {model}") + logger.warning( + f"Unknown model '{model}': pricing not available, cost will not be estimated." + ) + return None # For now, assume 1:1 input/output ratio input_cost = tokens * pricing[model]["input"] diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index ff7b3e0..a23a3aa 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -273,7 +273,7 @@ async def scan_module( failure_rate = module_failures / max(module_prompts, 1) failure_rates.append(failure_rate) - cost = calculate_cost(tokens) + cost = calculate_cost(tokens, model=getattr(request_factory, 'model_name', 'unknown')) response_text = fuzzer_state.get_last_output(prompt) or "" @@ -543,7 +543,7 @@ async def perform_many_shot_scan( failure_rate = module_failures / max(processed_prompts, 1) failure_rates.append(failure_rate) - cost = calculate_cost(tokens) + cost = calculate_cost(tokens, model=getattr(request_factory, 'model_name', 'unknown')) yield ScanResult( module=module.dataset_name, diff --git a/agentic_security/probe_data/audio_generator.py b/agentic_security/probe_data/audio_generator.py index 0b7a488..50babe5 100644 --- a/agentic_security/probe_data/audio_generator.py +++ b/agentic_security/probe_data/audio_generator.py @@ -131,6 +131,10 @@ class RequestAdapter: if not llm_spec.has_audio: raise ValueError("LLMSpec must have an image") + @property + def model_name(self) -> str: + return self.llm_spec.model_name + async def probe( self, prompt: str, encoded_image: str = "", encoded_audio: str = "", files={} ) -> httpx.Response: diff --git a/agentic_security/probe_data/image_generator.py b/agentic_security/probe_data/image_generator.py index 171f752..d986651 100644 --- a/agentic_security/probe_data/image_generator.py +++ b/agentic_security/probe_data/image_generator.py @@ -131,6 +131,10 @@ class RequestAdapter: if not llm_spec.has_image: raise ValueError("LLMSpec must have an image") + @property + def model_name(self) -> str: + return self.llm_spec.model_name + async def probe( self, prompt: str, encoded_image: str = "", encoded_audio: str = "", files={} ) -> httpx.Response: