antz-audit 0.3.2__py3-none-any.whl

adversarial/adapters/http_agent.py

@@ -0,0 +1,139 @@
+"""
+HTTP Agent Adapter — wraps any external agent REST API as a LangChain
+BaseChatModel so the attack engine can audit it without code access.
+
+Usage:
+    from adversarial.adapters.http_agent import HTTPAgentAdapter
+
+    adapter = HTTPAgentAdapter(
+        url="https://my-agent.company.com/api/chat",
+        headers={"Authorization": "Bearer sk-..."},
+        payload_template={"session_id": "audit-001"},
+        message_field="message",
+        response_field="response",
+    )
+    attack = ConstraintBypassAttack(adapter, judge, constitution)
+
+Probe before attacking:
+    ok, info = await adapter.probe()
+    if not ok:
+        print(info)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import httpx
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.outputs import ChatGeneration, ChatResult
+
+
+class HTTPAgentAdapter(BaseChatModel):
+    """
+    Adapts any HTTP/REST agent endpoint to the LangChain BaseChatModel interface.
+
+    The adapter sends the last HumanMessage content to the configured URL and
+    extracts the response using a dot-path field selector.
+    """
+
+    url:              str
+    headers:          dict[str, str] = {}
+    payload_template: dict[str, Any] = {}
+    message_field:    str = "message"
+    response_field:   str = "response"
+    timeout:          float = 30.0
+    max_retries:      int = 2
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def _llm_type(self) -> str:
+        return "http_agent_adapter"
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return asyncio.get_event_loop().run_until_complete(
+            self._agenerate(messages, stop, **kwargs)
+        )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        user_text = next(
+            (m.content for m in reversed(messages) if m.type == "human"),
+            str(messages[-1].content) if messages else "",
+        )
+
+        payload = {**self.payload_template, self.message_field: user_text}
+
+        last_exc: Exception | None = None
+        for attempt in range(self.max_retries + 1):
+            try:
+                async with httpx.AsyncClient(timeout=self.timeout) as client:
+                    resp = await client.post(
+                        self.url,
+                        json=payload,
+                        headers=self.headers,
+                    )
+                    resp.raise_for_status()
+                    data = resp.json()
+                    text = self._extract(data, self.response_field)
+                    msg  = AIMessage(content=text)
+                    return ChatResult(generations=[ChatGeneration(message=msg)])
+            except Exception as exc:
+                last_exc = exc
+                if attempt < self.max_retries:
+                    await asyncio.sleep(2 ** attempt)
+
+        error_text = f"[HTTP_AGENT_ERROR after {self.max_retries + 1} attempts: {last_exc}]"
+        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=error_text))])
+
+    @staticmethod
+    def _extract(data: Any, field_path: str) -> str:
+        """Dot-path extractor: 'choices.0.message.content' style."""
+        parts = field_path.split(".")
+        current = data
+        for part in parts:
+            if isinstance(current, dict):
+                current = current.get(part, "")
+            elif isinstance(current, list) and part.isdigit():
+                current = current[int(part)]
+            else:
+                return str(current)
+        return str(current) if current is not None else ""
+
+    async def probe(self) -> tuple[bool, str]:
+        """
+        Connectivity check before running the full attack campaign.
+        Returns (reachable, info_message).
+        """
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                payload = {**self.payload_template, self.message_field: "ping"}
+                resp = await client.post(self.url, json=payload, headers=self.headers)
+                if resp.status_code < 500:
+                    return True, f"Agent reachable — HTTP {resp.status_code}"
+                return False, f"Agent returned HTTP {resp.status_code}"
+        except httpx.ConnectError:
+            return False, f"Cannot connect to {self.url}"
+        except httpx.TimeoutException:
+            return False, f"Timeout connecting to {self.url}"
+        except Exception as exc:
+            return False, f"Probe failed: {exc}"

adversarial/adapters/langgraph.py

@@ -0,0 +1,265 @@
+"""
+LangGraph and CrewAI Adapters — wrap existing agent graphs and crews
+as LangChain BaseChatModel so the attack engine can audit them directly.
+
+Usage (LangGraph):
+    from adversarial.adapters.langgraph import LangGraphAdapter
+
+    graph   = builder.compile()
+    adapter = LangGraphAdapter(graph, input_key="messages", output_key="messages")
+    attack  = ConstraintBypassAttack(adapter, judge, constitution)
+
+Usage (CrewAI):
+    from adversarial.adapters.crewai import CrewAIAdapter
+
+    crew    = Crew(agents=[...], tasks=[...])
+    adapter = CrewAIAdapter(crew, input_variable="customer_request")
+    attack  = ConstraintBypassAttack(adapter, judge, constitution)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langchain_core.outputs import ChatGeneration, ChatResult
+
+# ---------------------------------------------------------------------------
+# LangGraph Adapter
+# ---------------------------------------------------------------------------
+
+class LangGraphAdapter(BaseChatModel):
+    """
+    Wraps a compiled LangGraph (StateGraph) as a BaseChatModel.
+
+    The adapter invokes the graph with the last HumanMessage and extracts
+    the final AI response from the graph state.
+
+    Args:
+        graph:      A compiled LangGraph (result of builder.compile()).
+        input_key:  State key used to pass messages into the graph.
+        output_key: State key from which to extract the response.
+                    If the value is a list of messages, the last AIMessage
+                    content is returned.
+        config:     Optional LangGraph config dict (thread_id, etc.).
+    """
+
+    graph:      Any
+    input_key:  str = "messages"
+    output_key: str = "messages"
+    config:     dict[str, Any] = {}
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def _llm_type(self) -> str:
+        return "langgraph_adapter"
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return asyncio.get_event_loop().run_until_complete(
+            self._agenerate(messages, stop, **kwargs)
+        )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        user_text = next(
+            (m.content for m in reversed(messages) if m.type == "human"),
+            str(messages[-1].content) if messages else "",
+        )
+
+        graph_input = {self.input_key: [HumanMessage(content=user_text)]}
+
+        try:
+            if hasattr(self.graph, "ainvoke"):
+                state = await self.graph.ainvoke(graph_input, config=self.config)
+            else:
+                state = self.graph.invoke(graph_input, config=self.config)
+
+            response_text = self._extract_response(state)
+        except Exception as exc:
+            response_text = f"[LANGGRAPH_ERROR: {exc}]"
+
+        return ChatResult(
+            generations=[ChatGeneration(message=AIMessage(content=response_text))]
+        )
+
+    def _extract_response(self, state: Any) -> str:
+        value = state.get(self.output_key, "") if isinstance(state, dict) else state
+
+        if isinstance(value, list):
+            for msg in reversed(value):
+                if hasattr(msg, "content") and getattr(msg, "type", "") in ("ai", "assistant"):
+                    return str(msg.content)
+                if hasattr(msg, "content"):
+                    return str(msg.content)
+            return str(value[-1]) if value else ""
+
+        if isinstance(value, str):
+            return value
+
+        if hasattr(value, "content"):
+            return str(value.content)
+
+        return str(value)
+
+
+# ---------------------------------------------------------------------------
+# CrewAI Adapter
+# ---------------------------------------------------------------------------
+
+class CrewAIAdapter(BaseChatModel):
+    """
+    Wraps a CrewAI Crew as a BaseChatModel.
+
+    The adapter calls crew.kickoff() with the last HumanMessage injected
+    as the value of `input_variable`.
+
+    Args:
+        crew:           A CrewAI Crew instance.
+        input_variable: The input variable name the crew expects (default: "input").
+        extra_inputs:   Additional inputs to pass alongside the user message.
+    """
+
+    crew:           Any
+    input_variable: str = "input"
+    extra_inputs:   dict[str, Any] = {}
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def _llm_type(self) -> str:
+        return "crewai_adapter"
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return asyncio.get_event_loop().run_until_complete(
+            self._agenerate(messages, stop, **kwargs)
+        )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        user_text = next(
+            (m.content for m in reversed(messages) if m.type == "human"),
+            str(messages[-1].content) if messages else "",
+        )
+
+        inputs = {**self.extra_inputs, self.input_variable: user_text}
+
+        try:
+            if hasattr(self.crew, "kickoff_async"):
+                result = await self.crew.kickoff_async(inputs=inputs)
+            else:
+                loop   = asyncio.get_running_loop()
+                result = await loop.run_in_executor(
+                    None, lambda: self.crew.kickoff(inputs=inputs)
+                )
+
+            response_text = (
+                result.raw
+                if hasattr(result, "raw")
+                else str(result)
+            )
+        except Exception as exc:
+            response_text = f"[CREWAI_ERROR: {exc}]"
+
+        return ChatResult(
+            generations=[ChatGeneration(message=AIMessage(content=response_text))]
+        )
+
+
+# ---------------------------------------------------------------------------
+# AutoGen Adapter (basic)
+# ---------------------------------------------------------------------------
+
+class AutoGenAdapter(BaseChatModel):
+    """
+    Wraps an AutoGen ConversableAgent as a BaseChatModel.
+
+    Sends a single message and returns the last reply.
+
+    Args:
+        agent:          An AutoGen ConversableAgent or AssistantAgent.
+        initiator_name: Name of the human proxy sending the message.
+    """
+
+    agent:          Any
+    initiator_name: str = "adversarial_tester"
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def _llm_type(self) -> str:
+        return "autogen_adapter"
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        return asyncio.get_event_loop().run_until_complete(
+            self._agenerate(messages, stop, **kwargs)
+        )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        user_text = next(
+            (m.content for m in reversed(messages) if m.type == "human"),
+            str(messages[-1].content) if messages else "",
+        )
+
+        try:
+            loop = asyncio.get_running_loop()
+            chat_result = await loop.run_in_executor(
+                None,
+                lambda: self.agent.initiate_chat(
+                    self.agent,
+                    message=user_text,
+                    max_turns=1,
+                    silent=True,
+                ),
+            )
+            history  = chat_result.chat_history if hasattr(chat_result, "chat_history") else []
+            response = history[-1].get("content", "") if history else str(chat_result)
+        except Exception as exc:
+            response = f"[AUTOGEN_ERROR: {exc}]"
+
+        return ChatResult(
+            generations=[ChatGeneration(message=AIMessage(content=response))]
+        )