From cde9969c6135de2a692f141b80118125492d6221 Mon Sep 17 00:00:00 2001 From: Dev Gajjar Date: Fri, 21 Mar 2025 17:36:58 +0530 Subject: [PATCH 1/3] errorhandling in httpx request --- agentic_security/http_spec.py | 64 ++++++++++++++++++++------ agentic_security/mcp/main.py | 84 +++++++++++++++++++++++++++-------- 2 files changed, 116 insertions(+), 32 deletions(-) diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py index dcb938d..dfa9ecc 100644 --- a/agentic_security/http_spec.py +++ b/agentic_security/http_spec.py @@ -18,16 +18,33 @@ class Modality(Enum): def encode_image_base64_by_url(url: str = "https://github.com/fluidicon.png") -> str: """Encode image data to base64 from a URL""" - response = httpx.get(url) - encoded_content = base64.b64encode(response.content).decode("utf-8") - return "data:image/jpeg;base64," + encoded_content - + try: + response = httpx.get(url) + encoded_content = base64.b64encode(response.content).decode("utf-8") + return "data:image/jpeg;base64," + encoded_content + except httpx.TimeoutException as e: + raise ValueError(f"Image fetch timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Image fetch failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error fetching image: {str(e)}") def encode_audio_base64_by_url(url: str) -> str: """Encode audio data to base64 from a URL""" - response = httpx.get(url) - encoded_content = base64.b64encode(response.content).decode("utf-8") - return "data:audio/mpeg;base64," + encoded_content + try: + response = httpx.get(url) + encoded_content = base64.b64encode(response.content).decode("utf-8") + return "data:audio/mpeg;base64," + encoded_content + except httpx.TimeoutException as e: + raise ValueError(f"Audio fetch timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Audio fetch failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error fetching audio: {str(e)}") class InvalidHTTPSpecError(Exception): @@ -58,7 +75,8 @@ def timeout(self): async def _probe_with_files(self, files): transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) - async with httpx.AsyncClient(transport=transport) as client: + try: + async with httpx.AsyncClient(transport=transport) as client: response = await client.request( method=self.method, url=self.url, @@ -66,8 +84,16 @@ async def _probe_with_files(self, files): files=files, timeout=self.timeout(), ) - - return response + response.raise_for_status() + return response + except httpx.TimeoutException as e: + raise ValueError(f"Request timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Request failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error: {str(e)}") def validate(self, prompt, encoded_image, encoded_audio, files) -> None: if self.has_files and not files: @@ -102,16 +128,26 @@ async def probe( content = content.replace("<>", encoded_audio) transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) - async with httpx.AsyncClient(transport=transport) as client: - response = await client.request( + try: + async with httpx.AsyncClient(transport=transport) as client: + response = await client.request( method=self.method, url=self.url, headers=self.headers, content=content, timeout=self.timeout(), ) - - return response + response.raise_for_status() + return response + + except httpx.TimeoutException as e: + raise ValueError(f"Request timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Request failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error: {str(e)}") async def verify(self) -> httpx.Response: match self: diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 0e754dd..dd5b05c 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -16,10 +16,19 @@ async def verify_llm(spec: str) -> dict: """Verify an LLM model specification using the FastAPI server.""" url = f"{AGENTIC_SECURITY}/verify" - async with httpx.AsyncClient() as client: - response = await client.post(url, json={"spec": spec}) - return response.json() - + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url, json={"spec": spec}) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def start_scan( @@ -39,36 +48,75 @@ async def start_scan( "probe_datasets": [], "secrets": {}, } - async with httpx.AsyncClient() as client: - response = await client.post(url, json=payload) - return response.json() - + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url, json=payload) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def stop_scan() -> dict: """Stop an ongoing scan via the FastAPI server.""" url = f"{AGENTIC_SECURITY}/stop" - async with httpx.AsyncClient() as client: - response = await client.post(url) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def get_data_config() -> list: """Retrieve data configuration from the FastAPI server.""" url = f"{AGENTIC_SECURITY}/v1/data-config" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() - + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get(url) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def get_spec_templates() -> list: """Retrieve data configuration from the FastAPI server.""" url = f"{AGENTIC_SECURITY}/v1/llm-specs" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get(url) + response.raise_for_status() + return response.json() + + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} # Run the MCP server From a7feb540958c18ad50f68db67e21b86879308988 Mon Sep 17 00:00:00 2001 From: Dev Gajjar Date: Fri, 21 Mar 2025 18:29:26 +0530 Subject: [PATCH 2/3] fixe black --- agentic_security/agents/operator_crew.py | 6 +-- agentic_security/http_spec.py | 39 ++++++++++--------- agentic_security/integrations/__init__.py | 6 +-- agentic_security/mcp/main.py | 7 +++- .../probe_data/modules/rl_model.py | 3 +- pip | 0 6 files changed, 31 insertions(+), 30 deletions(-) create mode 100644 pip diff --git a/agentic_security/agents/operator_crew.py b/agentic_security/agents/operator_crew.py index e836569..72b706a 100644 --- a/agentic_security/agents/operator_crew.py +++ b/agentic_security/agents/operator_crew.py @@ -246,9 +246,9 @@ async def run_crew(): os.environ["OPENAI_API_KEY"] = os.environ.get( "DEEPSEEK_API_KEY", "" ) # CrewAI uses OPENAI_API_KEY -os.environ[ - "OPENAI_MODEL_NAME" -] = "deepseek:chat" # Specify DeepSeek model (adjust if needed) +os.environ["OPENAI_MODEL_NAME"] = ( + "deepseek:chat" # Specify DeepSeek model (adjust if needed) +) if __name__ == "__main__": asyncio.run(run_crew()) diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py index dfa9ecc..2ef594e 100644 --- a/agentic_security/http_spec.py +++ b/agentic_security/http_spec.py @@ -31,6 +31,7 @@ def encode_image_base64_by_url(url: str = "https://github.com/fluidicon.png") -> except Exception as e: raise ValueError(f"Unexpected error fetching image: {str(e)}") + def encode_audio_base64_by_url(url: str) -> str: """Encode audio data to base64 from a URL""" try: @@ -76,16 +77,16 @@ def timeout(self): async def _probe_with_files(self, files): transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) try: - async with httpx.AsyncClient(transport=transport) as client: - response = await client.request( - method=self.method, - url=self.url, - headers=self.headers, - files=files, - timeout=self.timeout(), - ) - response.raise_for_status() - return response + async with httpx.AsyncClient(transport=transport) as client: + response = await client.request( + method=self.method, + url=self.url, + headers=self.headers, + files=files, + timeout=self.timeout(), + ) + response.raise_for_status() + return response except httpx.TimeoutException as e: raise ValueError(f"Request timed out: {str(e)}") except httpx.HTTPStatusError as e: @@ -128,18 +129,18 @@ async def probe( content = content.replace("<>", encoded_audio) transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) - try: - async with httpx.AsyncClient(transport=transport) as client: + try: + async with httpx.AsyncClient(transport=transport) as client: response = await client.request( - method=self.method, - url=self.url, - headers=self.headers, - content=content, - timeout=self.timeout(), - ) + method=self.method, + url=self.url, + headers=self.headers, + content=content, + timeout=self.timeout(), + ) response.raise_for_status() return response - + except httpx.TimeoutException as e: raise ValueError(f"Request timed out: {str(e)}") except httpx.HTTPStatusError as e: diff --git a/agentic_security/integrations/__init__.py b/agentic_security/integrations/__init__.py index 4b2a360..dc138c1 100644 --- a/agentic_security/integrations/__init__.py +++ b/agentic_security/integrations/__init__.py @@ -5,8 +5,6 @@ class IntegrationProto(Protocol): def __init__( self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {} - ): - ... + ): ... - async def apply(self) -> list: - ... + async def apply(self) -> list: ... diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index dd5b05c..5397fb8 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -28,7 +28,8 @@ async def verify_llm(spec: str) -> dict: except httpx.RequestError as e: return {"error": f"Request failed: {str(e)}"} except Exception as e: - return {"error": f"Unexpected error: {str(e)}"} + return {"error": f"Unexpected error: {str(e)}"} + @mcp.tool() async def start_scan( @@ -62,6 +63,7 @@ async def start_scan( except Exception as e: return {"error": f"Unexpected error: {str(e)}"} + @mcp.tool() async def stop_scan() -> dict: """Stop an ongoing scan via the FastAPI server.""" @@ -99,6 +101,7 @@ async def get_data_config() -> list: except Exception as e: return {"error": f"Unexpected error: {str(e)}"} + @mcp.tool() async def get_spec_templates() -> list: """Retrieve data configuration from the FastAPI server.""" @@ -108,7 +111,7 @@ async def get_spec_templates() -> list: response = await client.get(url) response.raise_for_status() return response.json() - + except httpx.TimeoutException as e: return {"error": f"Request timed out: {str(e)}"} except httpx.HTTPStatusError as e: diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py index 1b5befd..002158c 100644 --- a/agentic_security/probe_data/modules/rl_model.py +++ b/agentic_security/probe_data/modules/rl_model.py @@ -121,8 +121,7 @@ def update_rewards( current_prompt: str, reward: float, passed_guard: bool, - ) -> None: - ... + ) -> None: ... class QLearningPromptSelector(PromptSelectionInterface): diff --git a/pip b/pip new file mode 100644 index 0000000..e69de29 From 45a6054438db2006345a7f87e3858552c925c2a9 Mon Sep 17 00:00:00 2001 From: Dev Gajjar Date: Fri, 21 Mar 2025 18:40:22 +0530 Subject: [PATCH 3/3] style: reformat code with black --- agentic_security/agents/operator_crew.py | 6 +++--- agentic_security/integrations/__init__.py | 6 ++++-- agentic_security/probe_data/modules/rl_model.py | 3 ++- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/agentic_security/agents/operator_crew.py b/agentic_security/agents/operator_crew.py index 72b706a..e836569 100644 --- a/agentic_security/agents/operator_crew.py +++ b/agentic_security/agents/operator_crew.py @@ -246,9 +246,9 @@ async def run_crew(): os.environ["OPENAI_API_KEY"] = os.environ.get( "DEEPSEEK_API_KEY", "" ) # CrewAI uses OPENAI_API_KEY -os.environ["OPENAI_MODEL_NAME"] = ( - "deepseek:chat" # Specify DeepSeek model (adjust if needed) -) +os.environ[ + "OPENAI_MODEL_NAME" +] = "deepseek:chat" # Specify DeepSeek model (adjust if needed) if __name__ == "__main__": asyncio.run(run_crew()) diff --git a/agentic_security/integrations/__init__.py b/agentic_security/integrations/__init__.py index dc138c1..4b2a360 100644 --- a/agentic_security/integrations/__init__.py +++ b/agentic_security/integrations/__init__.py @@ -5,6 +5,8 @@ class IntegrationProto(Protocol): def __init__( self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {} - ): ... + ): + ... - async def apply(self) -> list: ... + async def apply(self) -> list: + ... diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py index 002158c..1b5befd 100644 --- a/agentic_security/probe_data/modules/rl_model.py +++ b/agentic_security/probe_data/modules/rl_model.py @@ -121,7 +121,8 @@ def update_rewards( current_prompt: str, reward: float, passed_guard: bool, - ) -> None: ... + ) -> None: + ... class QLearningPromptSelector(PromptSelectionInterface):