switch-trust-cli 1.7.0__py3-none-any.whl

switch_trust/eval/core/models/model_adk.py

@@ -0,0 +1,186 @@
+"""Model implementation for Google ADK agents.
+
+ADK agents are served via ``adk api_server`` or ``adk web`` and
+expose a REST API at a configurable host/port.  This model sends
+prompts to the ``/run`` endpoint and extracts the agent's text
+response from the returned events.
+
+Typical ADK server URL: ``http://localhost:8000``
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import aiohttp
+import requests
+
+from switch_trust.eval.common.schema import Content, Message, Part, Role
+from switch_trust.eval.core.models.model import Model, ModelResponse, ResponseStatus
+from switch_trust.eval.core.models.model_http import raise_for_status
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+
+class ADKModel(Model):
+    """Sends prompts to a Google ADK agent served over HTTP.
+
+    Args:
+        app_name: The ADK app/agent name.
+        host: The HTTP host where the ADK server is running.
+        user_id: User identifier for sessions.
+        immediate_result: If True, return the first event's
+            content (e.g. a tool call). If False (default),
+            return the final text response after all tool
+            calls complete.
+        headers: Extra HTTP headers sent on every request, e.g. an
+            ``Authorization`` header for a deployment gated by an
+            app-level bearer token rather than Cloud Run IAM.
+    """
+
+    def __init__(
+        self,
+        app_name: str,
+        host: str = "http://localhost:8000",
+        user_id: str = "aired",
+        immediate_result: bool = False,
+        headers: dict[str, str] | None = None,
+        connector_factory: Callable[[], aiohttp.BaseConnector] | None = None,
+    ):
+        self._app_name = app_name
+        self._host = host.rstrip("/")
+        self._user_id = user_id
+        self._immediate_result = immediate_result
+        self._headers = headers or {}
+        # Optional aiohttp connector factory; the worker passes the SSRF-guarded
+        # one (platform/ssrf.py). None keeps the default connector.
+        self._connector_factory = connector_factory
+
+    async def _create_session(
+        self,
+        session: aiohttp.ClientSession,
+    ) -> str:
+        """Create a new session for this request."""
+        url = f"{self._host}/apps/{self._app_name}/users/{self._user_id}/sessions"
+        async with session.post(url, json={}, headers=self._headers) as resp:
+            await raise_for_status(resp)
+            data = await resp.json()
+            return data["id"]
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **_kwargs: Any,
+    ) -> ModelResponse:
+        # An agent under test cannot enforce a JSON schema; output_schema is
+        # accepted for interface parity and ignored.
+        if len(messages) > 1:
+            raise ValueError(
+                "ADKModel does not support multiple messages; use session-based history"
+            )
+
+        connector = self._connector_factory() if self._connector_factory else None
+        async with aiohttp.ClientSession(connector=connector) as session:
+            session_id = await self._create_session(session)
+
+            text_parts = [p.text for p in messages[0].content.parts if p.text]
+            prompt_text = " ".join(text_parts)
+
+            payload = {
+                "appName": self._app_name,
+                "userId": self._user_id,
+                "sessionId": session_id,
+                "newMessage": {
+                    "role": "user",
+                    "parts": [{"text": prompt_text}],
+                },
+            }
+
+            async with session.post(
+                f"{self._host}/run",
+                json=payload,
+                headers=self._headers,
+            ) as resp:
+                await raise_for_status(resp)
+                events = await resp.json()
+
+        if self._immediate_result:
+            response_text = self._extract_first(events)
+        else:
+            response_text = self._extract_final(events)
+
+        if response_text is None:
+            return ModelResponse(
+                message=None,
+                status=ResponseStatus.EMPTY_RESPONSE,
+            )
+
+        message = Message(
+            content=Content(
+                role=Role.ASSISTANT,
+                parts=[Part.text_part(response_text)],
+            ),
+        )
+        return ModelResponse(message=message)
+
+    def _extract_final(
+        self,
+        events: list[dict],
+    ) -> str | None:
+        """Extract the last text response, skipping tool
+        call/response events.
+
+        FIXME: the skip below drops a *mixed* event — one turn carrying both
+        narration and a tool call, e.g. ``[{"text": "Let me look that up"},
+        {"functionCall": ...}]``, which Gemini emits routinely — losing the text
+        with it. The test is also redundant for pure-tool events, since those
+        carry no ``text`` key and the ``if texts`` check below already skips
+        them. Deleting the whole ``any(...)`` block both fixes and simplifies
+        this, as done in ``model_vertex_agent_runtime._extract_final``. Left
+        unchanged here only to keep that PR scoped; needs its own change plus a
+        mixed-event test.
+        """
+        for event in reversed(events):
+            content = event.get("content")
+            if content is None:
+                continue
+            parts = content.get("parts", [])
+            # Skip events that are tool calls or responses
+            if any(
+                "functionCall" in p or "functionResponse" in p
+                for p in parts
+                if isinstance(p, dict)
+            ):
+                continue
+            texts = [p["text"] for p in parts if isinstance(p, dict) and "text" in p]
+            if texts:
+                return " ".join(texts)
+        return None
+
+    def _extract_first(
+        self,
+        events: list[dict],
+    ) -> str | None:
+        """Extract the first event's text content."""
+        for event in events:
+            content = event.get("content")
+            if content is None:
+                continue
+            parts = content.get("parts", [])
+            texts = [p["text"] for p in parts if isinstance(p, dict) and "text" in p]
+            if texts:
+                return " ".join(texts)
+        return None
+
+
+def discover_adk_agents(
+    host: str = "http://localhost:8000",
+) -> list[str]:
+    """Return a list of agent app names on an ADK server."""
+    resp = requests.get(f"{host.rstrip('/')}/list-apps")
+    resp.raise_for_status()
+    return resp.json()

switch_trust/eval/core/models/model_anthropic.py

@@ -0,0 +1,79 @@
+import json
+from typing import Any
+
+from anthropic import AsyncAnthropic
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_anthropic
+from switch_trust.eval.common.schema import Content, Message, Role
+from switch_trust.eval.core.models.model import Model, ModelResponse
+from switch_trust.eval.core.models.response_schema import to_anthropic_tool
+
+
+class AnthropicModel(Model):
+    _client: AsyncAnthropic
+    _model: str
+    _max_tokens: int
+    _temperature: float
+
+    def __init__(
+        self,
+        client: AsyncAnthropic,
+        model: str,
+        max_tokens: int = 1024,
+        temperature: float = 0.0,
+    ):
+        self._client = client
+        self._model = model
+        self._max_tokens = max_tokens
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        system_blocks = None
+        api_messages = []
+        for msg in messages:
+            anthropic_msg = converter_anthropic.from_message(msg)
+            if msg.content.role == Role.SYSTEM:
+                system_blocks = anthropic_msg["content"]
+            else:
+                api_messages.append(anthropic_msg)
+
+        if not api_messages:
+            api_messages = [{"role": "user", "content": ""}]
+
+        create_kwargs: dict[str, Any] = {
+            "model": self._model,
+            "max_tokens": kwargs.pop("max_tokens", self._max_tokens),
+            "temperature": kwargs.pop("temperature", self._temperature),
+            "messages": api_messages,
+            **kwargs,
+        }
+        if system_blocks is not None:
+            create_kwargs["system"] = system_blocks
+
+        if output_schema is not None and "tools" not in create_kwargs:
+            # Anthropic has no cross-version response_format: force a single
+            # tool whose input_schema is the target schema and read the
+            # tool_use input back as the JSON payload.
+            tool, tool_choice = to_anthropic_tool(output_schema)
+            create_kwargs["tools"] = [tool]
+            create_kwargs["tool_choice"] = tool_choice
+
+        response = await self._client.messages.create(**create_kwargs)
+
+        if output_schema is not None:
+            for block in response.content:
+                if getattr(block, "type", None) == "tool_use":
+                    payload = json.dumps(block.input)
+                    return ModelResponse(
+                        Message(content=Content.text(Role.ASSISTANT, payload)),
+                    )
+
+        message = converter_anthropic.to_message(response)
+        return ModelResponse(message)

switch_trust/eval/core/models/model_anthropic_agent.py

@@ -0,0 +1,104 @@
+"""Model implementation for Anthropic Agent SDK agents served over HTTP.
+
+Anthropic Agent SDK (claude_agent_sdk) agents don't have a
+built-in HTTP server.  The expected pattern is a FastAPI wrapper
+with a POST endpoint that accepts ``{"prompt": "<prompt>"}`` and
+returns ``{"response": "<response>"}``.
+
+Example server::
+
+    from fastapi import FastAPI
+    from pydantic import BaseModel
+    from claude_agent_sdk import (
+        query, ClaudeAgentOptions,
+        AssistantMessage, TextBlock, ResultMessage,
+    )
+
+    app = FastAPI()
+
+    class Request(BaseModel):
+        prompt: str
+
+    @app.post("/run")
+    async def run(req: Request):
+        options = ClaudeAgentOptions(
+            permission_mode="acceptEdits",
+        )
+        text = ""
+        async for msg in query(
+            prompt=req.prompt, options=options,
+        ):
+            if isinstance(msg, AssistantMessage):
+                for block in msg.content:
+                    if isinstance(block, TextBlock):
+                        text += block.text
+        return {"response": text}
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import aiohttp
+
+from switch_trust.eval.common.schema import Content, Message, Part, Role
+from switch_trust.eval.core.models.model import Model, ModelResponse, ResponseStatus
+from switch_trust.eval.core.models.model_http import raise_for_status
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+
+class AnthropicAgentModel(Model):
+    """Sends prompts to an Anthropic Agent SDK agent over HTTP."""
+
+    def __init__(
+        self,
+        host: str = "http://localhost:8000",
+        endpoint: str = "/run",
+        connector_factory: Callable[[], aiohttp.BaseConnector] | None = None,
+    ):
+        self._host = host.rstrip("/")
+        self._endpoint = endpoint
+        # Optional aiohttp connector factory; the worker passes the SSRF-guarded
+        # one (platform/ssrf.py). None keeps the default connector.
+        self._connector_factory = connector_factory
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        # An agent under test cannot enforce a JSON schema; output_schema is
+        # accepted for interface parity and ignored.
+        if len(messages) > 1:
+            raise ValueError("AnthropicAgentModel does not support multiple messages")
+        text_parts = [p.text for p in messages[0].content.parts if p.text]
+        prompt_text = " ".join(text_parts)
+
+        connector = self._connector_factory() if self._connector_factory else None
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.post(
+                f"{self._host}{self._endpoint}",
+                json={"prompt": prompt_text},
+            ) as resp:
+                await raise_for_status(resp)
+                data = await resp.json()
+
+        output = data.get("response")
+        if not output:
+            return ModelResponse(
+                message=None,
+                status=ResponseStatus.EMPTY_RESPONSE,
+            )
+
+        message = Message(
+            content=Content(
+                role=Role.ASSISTANT,
+                parts=[Part.text_part(output)],
+            ),
+        )
+        return ModelResponse(message=message)

switch_trust/eval/core/models/model_gemini.py

@@ -0,0 +1,146 @@
+from typing import Any
+
+from google.genai import Client
+from google.genai import types as genai_types
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_genai
+from switch_trust.eval.common.schema import Message, Role
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+    ResponseStatus,
+)
+
+# Safety settings that disable content blocking across every active harm
+# category. Red-teaming needs the attacker/judge model to see and produce
+# adversarial content without Vertex's filters returning an empty candidate — a
+# blocked judge yields no JSON to parse, which surfaces as a spurious eval error
+# rather than a real result. HARM_CATEGORY_CIVIC_INTEGRITY is deprecated and
+# omitted. Note this is only appropriate for the generator; a model *under test*
+# should keep its own safety behavior so blocking can be measured.
+SAFETY_OFF: list[genai_types.SafetySetting] = [
+    genai_types.SafetySetting(category=category, threshold="OFF")
+    for category in (
+        genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT,
+        genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+        genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+        genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+    )
+]
+
+
+class GeminiModel(Model):
+    _client: Client
+    _model: str
+    _temperature: float
+    _safety_settings: list[genai_types.SafetySetting] | None
+
+    def __init__(
+        self,
+        client: Client,
+        model: str,
+        temperature: float = 0.0,
+        safety_settings: list[genai_types.SafetySetting] | None = None,
+    ):
+        self._client = client
+        self._model = model
+        self._temperature = temperature
+        self._safety_settings = safety_settings
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        system_parts: list[dict[str, Any]] = []
+        contents: list[dict[str, Any]] = []
+        for msg in messages:
+            if msg.content.role == Role.SYSTEM:
+                system_parts.extend(
+                    converter_genai.from_content(
+                        msg.content,
+                    )["parts"]
+                )
+            else:
+                contents.append(
+                    converter_genai.from_content(msg.content),
+                )
+
+        config = kwargs.pop("config", None)
+        if config is None:
+            config = genai_types.GenerateContentConfig(
+                temperature=self._temperature,
+            )
+        elif config.temperature is None:
+            config.temperature = self._temperature
+        if system_parts:
+            system_instruction = genai_types.Content(
+                parts=[genai_types.Part(**p) for p in system_parts],
+            )
+            config.system_instruction = system_instruction
+        if self._safety_settings is not None and config.safety_settings is None:
+            config.safety_settings = self._safety_settings
+        if output_schema is not None and config.response_schema is None:
+            config.response_mime_type = "application/json"
+            config.response_schema = output_schema
+
+        generate_kwargs: dict[str, Any] = {
+            "model": self._model,
+            "contents": contents,
+            **kwargs,
+        }
+        if config is not None:
+            generate_kwargs["config"] = config
+
+        response = await self._client.aio.models.generate_content(
+            **generate_kwargs,
+        )
+        if (
+            response.candidates is not None
+            and len(response.candidates) == 1
+            and response.candidates[0].content is not None
+        ):
+            message = converter_genai.to_message(
+                response.candidates[0].content,
+            )
+            return ModelResponse(message)
+        return ModelResponse(
+            None,
+            _classify_block_reason(response),
+        )
+
+
+def _classify_block_reason(response) -> ResponseStatus:
+    """Determine why a Gemini response was blocked."""
+    # Check prompt-level blocking
+    reason = getattr(
+        response,
+        "prompt_feedback",
+        None,
+    )
+    if reason is not None:
+        block_reason = getattr(reason, "block_reason", None)
+        if block_reason is not None:
+            reason_str = str(block_reason).upper()
+            if "SAFETY" in reason_str:
+                return ResponseStatus.BLOCKED_SAFETY
+            if "PROHIBITED" in reason_str:
+                return ResponseStatus.BLOCKED_PROHIBITED
+
+    # Check candidate-level finish reason
+    candidates = getattr(response, "candidates", None)
+    if candidates:
+        finish = getattr(candidates[0], "finish_reason", None)
+        if finish is not None:
+            finish_str = str(finish).upper()
+            if "SAFETY" in finish_str:
+                return ResponseStatus.BLOCKED_SAFETY
+            if "RECITATION" in finish_str:
+                return ResponseStatus.BLOCKED_RECITATION
+            if "PROHIBITED" in finish_str:
+                return ResponseStatus.BLOCKED_PROHIBITED
+
+    return ResponseStatus.BLOCKED_SAFETY

switch_trust/eval/core/models/model_generic_http.py

@@ -0,0 +1,120 @@
+"""Model for any HTTP endpoint with configurable JSON fields.
+
+Connects to a generic REST API where the request sends
+a prompt in a configurable JSON field and the response
+returns the output in another configurable field.
+
+Default format::
+
+    POST /your-endpoint
+    {"input": "Hello"}
+
+    {"output": "Hi there!"}
+
+The ``output_path`` supports dot-separated paths for
+nested responses, e.g. ``"data.response.text"`` or
+``"choices.0.message.content"``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import aiohttp
+
+from switch_trust.eval.common.schema import Content, Message, Part, Role
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+    ResponseStatus,
+    flatten_messages,
+)
+from switch_trust.eval.core.models.model_http import raise_for_status
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+
+class GenericHttpModel(Model):
+    _url: str
+    _headers: dict[str, str]
+    _input_path: str
+    _output_path: str
+    # Optional aiohttp connector factory. The eval worker passes one that
+    # enforces the request-time SSRF guard (platform/ssrf.py); None keeps the
+    # default connector for the CLI and local use.
+    _connector_factory: Callable[[], aiohttp.BaseConnector] | None
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        input_path: str = "input",
+        output_path: str = "output",
+        connector_factory: Callable[[], aiohttp.BaseConnector] | None = None,
+    ):
+        self._url = url
+        self._headers = headers or {}
+        self._input_path = input_path
+        self._output_path = output_path
+        self._connector_factory = connector_factory
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        # A generic REST endpoint under test cannot enforce a JSON schema;
+        # output_schema is accepted for interface parity and ignored.
+        prompt_text = flatten_messages(messages)
+
+        connector = self._connector_factory() if self._connector_factory else None
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.post(
+                self._url,
+                json={self._input_path: prompt_text},
+                headers=self._headers,
+            ) as resp:
+                await raise_for_status(resp)
+                data = await resp.json()
+
+        output = _resolve_path(data, self._output_path)
+        if not output:
+            return ModelResponse(
+                message=None,
+                status=ResponseStatus.EMPTY_RESPONSE,
+            )
+
+        message = Message(
+            content=Content(
+                role=Role.ASSISTANT,
+                parts=[Part.text_part(str(output))],
+            ),
+        )
+        return ModelResponse(message=message)
+
+
+def _resolve_path(data: Any, path: str) -> Any:
+    """Traverse a nested dict/list by dot-separated path.
+
+    Supports integer indices for lists, e.g.
+    ``"choices.0.message.content"`` resolves
+    ``data["choices"][0]["message"]["content"]``.
+    """
+    current = data
+    for key in path.split("."):
+        if current is None:
+            return None
+        if isinstance(current, list):
+            try:
+                current = current[int(key)]
+            except (ValueError, IndexError):
+                return None
+        elif isinstance(current, dict):
+            current = current.get(key)
+        else:
+            return None
+    return current