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

switch_trust/eval/core/models/model_http.py

@@ -0,0 +1,135 @@
+"""Shared HTTP-error handling for the aiohttp-backed model/agent types.
+
+`resp.raise_for_status()` raises `aiohttp.ClientResponseError` from just
+`status` + `reason` + `url` — the response body, where the useful detail
+lives (e.g. a Google quota error's `error.status`/`consumer`), is discarded.
+
+The body is never logged here, deliberately: `endpoint_url`/`db_model.host`
+(consumed by every caller of `raise_for_status()` below) is tenant-controlled
+(the same field `platform/ssrf.py` guards), so the body can carry arbitrary
+tenant/user content — echoed prompts, internal stack traces, anything the
+tenant's server chooses to put in an error response. There is no reliable way
+to tell a safe platform-template error apart from arbitrary tenant content
+from the outside, so instead of logging the body, this module extracts two
+narrow, content-free signals from it and discards the rest.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from urllib.parse import urlsplit, urlunsplit
+
+import aiohttp
+
+# Bodies are only read for their two content-free signals -- there is never a
+# reason to buffer more than a small error payload needs, and an unbounded
+# read lets a misbehaving tenant server exhaust worker memory with one huge
+# response.
+_MAX_BODY_BYTES = 64 * 1024
+
+# The fixed google.rpc.Code vocabulary. Only a status string that is an exact
+# match is ever surfaced -- this makes the field safe to log even though it
+# came from a tenant-reachable response: a server can put arbitrary text in a
+# field called "status", but it can't make arbitrary text collide with one of
+# these 17 fixed names, so anything else is treated as unrecognized and
+# dropped.
+_KNOWN_ERROR_STATUSES = frozenset(
+    {
+        "OK",
+        "CANCELLED",
+        "UNKNOWN",
+        "INVALID_ARGUMENT",
+        "DEADLINE_EXCEEDED",
+        "NOT_FOUND",
+        "ALREADY_EXISTS",
+        "PERMISSION_DENIED",
+        "UNAUTHENTICATED",
+        "RESOURCE_EXHAUSTED",
+        "FAILED_PRECONDITION",
+        "ABORTED",
+        "OUT_OF_RANGE",
+        "UNIMPLEMENTED",
+        "INTERNAL",
+        "UNAVAILABLE",
+        "DATA_LOSS",
+    }
+)
+_PROJECT_NUMBER_RE = re.compile(r"/projects/(\d+)/")
+
+
+class ModelHttpError(Exception):
+    """HTTP error from a model/agent endpoint.
+
+    Deliberately carries no response-body content (see module docstring).
+    `error_status` and `same_project_as_url` are the only signals derived
+    from the body, and both are content-free.
+
+    Duck-types `.status` like `aiohttp.ClientResponseError` so the existing
+    `_is_transient` (and any future `_is_terminal`) checks in
+    `model_retry.py`, which read `status_code`/`status`/`code` via
+    `getattr`, keep working unmodified.
+    """
+
+    def __init__(
+        self,
+        status: int,
+        url: str,
+        error_status: str | None,
+        same_project_as_url: bool | None,
+    ):
+        self.status = status
+        self.url = url
+        self.error_status = error_status
+        self.same_project_as_url = same_project_as_url
+        parts = [f"{status} from {url}"]
+        if error_status is not None:
+            parts.append(f"error_status={error_status}")
+        if same_project_as_url is not None:
+            parts.append(f"same_project_as_url={same_project_as_url}")
+        super().__init__(", ".join(parts))
+
+
+async def raise_for_status(resp: aiohttp.ClientResponse) -> None:
+    """Like `resp.raise_for_status()`, but derives two safe signals from the
+    body first. The body itself is never stored on the raised exception or
+    returned -- only `_extract_error_status()`/`_same_project_as_url()`'s
+    outputs are."""
+    if resp.status < 400:
+        return
+    raw = await resp.content.read(_MAX_BODY_BYTES)
+    body = raw.decode(resp.charset or "utf-8", errors="replace")
+    # Query/fragment can carry tenant-supplied secrets (e.g. an API key on a
+    # GenericHttpModel/LangServe endpoint URL) -- strip them before this URL
+    # flows into the exception string, logs, and the OTel span.
+    parts = urlsplit(str(resp.url))
+    url = urlunsplit((parts.scheme, parts.netloc, parts.path, "", ""))
+    raise ModelHttpError(
+        resp.status,
+        url,
+        _extract_error_status(body),
+        _same_project_as_url(body, url),
+    )
+
+
+def _extract_error_status(body: str) -> str | None:
+    """The fixed-vocabulary ``error.status`` enum (e.g. RESOURCE_EXHAUSTED),
+    if the body has that shape *and* the value is one of the known names.
+    Never returns ``error.message`` or any other free text -- that's exactly
+    the field a tenant server could use to smuggle arbitrary content."""
+    try:
+        status = json.loads(body)["error"]["status"]
+    except (json.JSONDecodeError, KeyError, TypeError):
+        return None
+    return status if status in _KNOWN_ERROR_STATUSES else None
+
+
+def _same_project_as_url(body: str, url: str) -> bool | None:
+    """Whether the GCP project number already visible in `url` also appears
+    as `project_number:<n>` in the body -- as a boolean only, never quoting
+    the body itself. `None` when `url` has no project number (the check
+    doesn't apply), distinct from `False` (checked, and it didn't match)."""
+    match = _PROJECT_NUMBER_RE.search(url)
+    if not match:
+        return None
+    return f"project_number:{match.group(1)}" in body

switch_trust/eval/core/models/model_huggingface.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from switch_trust.eval.common.schema import Content, Message, PartType, Role
+from switch_trust.eval.core.models.model import Model, ModelResponse
+from switch_trust.eval.core.optional_deps import require_transformers_pipeline
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+    from transformers import Pipeline
+
+
+class HuggingFaceModel(Model):
+    """Model adapter for Hugging Face ``text-generation`` pipelines.
+
+    Accepts either an already-constructed ``Pipeline`` or a model name /
+    path to load one on the fly.
+    """
+
+    _pipeline: Pipeline
+    _temperature: float
+
+    def __init__(
+        self,
+        model: str | Pipeline,
+        temperature: float = 0.0,
+        **pipeline_kwargs: Any,
+    ):
+        if isinstance(model, str):
+            pipeline = require_transformers_pipeline()
+            self._pipeline = pipeline(
+                "text-generation",
+                model=model,
+                **pipeline_kwargs,
+            )
+        else:
+            self._pipeline = model
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        # A local text-generation pipeline cannot enforce a JSON schema;
+        # output_schema is accepted for interface parity and ignored.
+        return await asyncio.to_thread(
+            self._generate_sync,
+            messages,
+            **kwargs,
+        )
+
+    def _generate_sync(self, messages: list[Message], **kwargs: Any) -> ModelResponse:
+        if len(messages) > 1:
+            raise ValueError("HuggingFaceModel does not support multiple messages")
+        prompt = _collect_text(messages[0].content)
+        kwargs.setdefault("max_new_tokens", 256)
+        kwargs.setdefault("return_full_text", False)
+        kwargs.setdefault("temperature", self._temperature or None)
+        outputs = self._pipeline(prompt, **kwargs)
+        generated_text = outputs[0]["generated_text"]
+        message = Message(
+            content=Content.text(Role.ASSISTANT, generated_text),
+        )
+        return ModelResponse(message)
+
+
+def _collect_text(content: Content) -> str:
+    return " ".join(
+        p.text for p in content.parts if p.part_type == PartType.TEXT and p.text
+    )

switch_trust/eval/core/models/model_langserve.py

@@ -0,0 +1,96 @@
+"""Model for LangServe endpoints.
+
+LangServe exposes LangChain runnables as REST APIs with
+a standard ``/{chain}/invoke`` endpoint.
+
+Request format::
+
+    POST /my-chain/invoke
+    {"input": "Hello"}
+
+Response format::
+
+    {"output": "Hi there!", "metadata": {...}}
+"""
+
+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 LangServeModel(Model):
+    _url: str
+    _headers: dict[str, str]
+    # Optional aiohttp connector factory; the worker passes the SSRF-guarded one
+    # (platform/ssrf.py). None keeps the default connector.
+    _connector_factory: Callable[[], aiohttp.BaseConnector] | None
+
+    def __init__(
+        self,
+        base_url: str,
+        chain_path: str = "",
+        headers: dict[str, str] | None = None,
+        connector_factory: Callable[[], aiohttp.BaseConnector] | None = None,
+    ):
+        base = base_url.rstrip("/")
+        path = chain_path.strip("/")
+        if path:
+            self._url = f"{base}/{path}/invoke"
+        else:
+            self._url = f"{base}/invoke"
+        self._headers = headers or {}
+        self._connector_factory = connector_factory
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        # A LangServe 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={"input": prompt_text},
+                headers=self._headers,
+            ) as resp:
+                await raise_for_status(resp)
+                data = await resp.json()
+
+        output = data.get("output")
+        if not output:
+            return ModelResponse(
+                message=None,
+                status=ResponseStatus.EMPTY_RESPONSE,
+            )
+
+        if isinstance(output, dict):
+            output = output.get("content", str(output))
+
+        message = Message(
+            content=Content(
+                role=Role.ASSISTANT,
+                parts=[Part.text_part(str(output))],
+            ),
+        )
+        return ModelResponse(message=message)

switch_trust/eval/core/models/model_litellm.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+import litellm
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_openai
+from switch_trust.eval.common.schema import Message
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+)
+
+
+class LiteLLMModel(Model):
+    """Model adapter that delegates to LiteLLM.
+
+    LiteLLM uses the OpenAI message format for all providers, so this
+    adapter reuses ``converter_openai`` for serialisation.
+    """
+
+    _model: str
+    _temperature: float
+
+    def __init__(
+        self,
+        model: str,
+        temperature: float = 0.0,
+    ):
+        self._model = model
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        openai_messages = [converter_openai.from_message(m) for m in messages]
+        if output_schema is not None and "response_format" not in kwargs:
+            # LiteLLM accepts a Pydantic class and translates it to each
+            # provider's structured-output mechanism.
+            kwargs["response_format"] = output_schema
+        response = await litellm.acompletion(
+            model=self._model,
+            messages=openai_messages,
+            temperature=kwargs.pop("temperature", self._temperature),
+            **kwargs,
+        )
+        choice = response.choices[0]
+        message = converter_openai.to_message(choice.message)
+        return ModelResponse(message)

switch_trust/eval/core/models/model_ollama.py

@@ -0,0 +1,65 @@
+from typing import Any
+
+import requests
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_openai
+from switch_trust.eval.common.schema import Message
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+)
+from switch_trust.eval.core.models.response_schema import to_json_schema
+
+
+class OllamaModel(Model):
+    _client: AsyncOpenAI
+    _model: str
+    _temperature: float
+
+    def __init__(
+        self,
+        model: str,
+        host: str = "http://localhost:11434",
+        temperature: float = 0.0,
+    ):
+        self._client = AsyncOpenAI(
+            base_url=f"{host}/v1",
+            api_key="ollama",
+        )
+        self._model = model
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        openai_messages = [converter_openai.from_message(m) for m in messages]
+        if output_schema is not None:
+            # Ollama constrains output via its native top-level ``format`` field
+            # (a JSON schema), passed through the OpenAI-compatible body.
+            extra_body = kwargs.setdefault("extra_body", {})
+            extra_body.setdefault("format", to_json_schema(output_schema))
+        response = await self._client.chat.completions.create(
+            model=self._model,
+            messages=openai_messages,
+            temperature=kwargs.pop("temperature", self._temperature),
+            **kwargs,
+        )
+        choice = response.choices[0]
+        message = converter_openai.to_message(choice.message)
+        return ModelResponse(message)
+
+
+def discover_ollama_models(
+    host: str = "http://localhost:11434",
+) -> list[str]:
+    """Return a list of model names available on an Ollama instance."""
+    response = requests.get(f"{host}/api/tags")
+    response.raise_for_status()
+    data = response.json()
+    return [m["name"] for m in data.get("models", [])]

switch_trust/eval/core/models/model_openai.py

@@ -0,0 +1,71 @@
+from typing import Any
+
+from openai import APIError, AsyncOpenAI, BadRequestError
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_openai
+from switch_trust.eval.common.schema import Message
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+    ResponseStatus,
+)
+from switch_trust.eval.core.models.response_schema import to_openai_response_format
+
+
+class OpenAIModel(Model):
+    _client: AsyncOpenAI
+    _model: str
+    _temperature: float
+
+    def __init__(
+        self,
+        client: AsyncOpenAI,
+        model: str,
+        temperature: float = 0.0,
+    ):
+        self._client = client
+        self._model = model
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        openai_messages = [converter_openai.from_message(m) for m in messages]
+        if output_schema is not None and "response_format" not in kwargs:
+            kwargs["response_format"] = to_openai_response_format(output_schema)
+        try:
+            response = await self._client.chat.completions.create(
+                model=self._model,
+                messages=openai_messages,
+                temperature=kwargs.pop("temperature", self._temperature),
+                **kwargs,
+            )
+        except APIError as e:
+            return ModelResponse(
+                None,
+                _classify_block_reason(e),
+            )
+        choice = response.choices[0]
+        message = converter_openai.to_message(choice.message)
+        return ModelResponse(message)
+
+
+def _classify_block_reason(api_error: APIError) -> ResponseStatus:
+    """Determine why a OpenAI response was blocked."""
+
+    # Check prompt-level blocking
+    if isinstance(api_error, BadRequestError):
+        code = (
+            api_error.body.get("code", "") if isinstance(api_error.body, dict) else ""
+        )
+        if code == "cyber_policy":
+            return ResponseStatus.BLOCKED_PROHIBITED
+        if code in ["invalid_prompt", "bio_policy"]:
+            return ResponseStatus.BLOCKED_SAFETY
+
+    return ResponseStatus.ERROR

switch_trust/eval/core/models/model_openai_agent.py

@@ -0,0 +1,92 @@
+"""Model implementation for OpenAI Agents SDK agents served over HTTP.
+
+OpenAI Agents SDK agents don't have a built-in HTTP server.
+The expected pattern is a FastAPI wrapper with a POST endpoint
+that accepts ``{"input": "<prompt>"}`` and returns
+``{"output": "<response>"}``.
+
+Example server::
+
+    from fastapi import FastAPI
+    from pydantic import BaseModel
+    from agents import Agent, Runner
+
+    app = FastAPI()
+    agent = Agent(name="my_agent", instructions="...")
+
+    class Request(BaseModel):
+        input: str
+
+    @app.post("/run")
+    async def run(req: Request):
+        result = await Runner.run(agent, req.input)
+        return {"output": result.final_output}
+"""
+
+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 OpenAIAgentModel(Model):
+    """Sends prompts to an OpenAI Agents 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("OpenAIAgentModel 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={"input": prompt_text},
+            ) as resp:
+                await raise_for_status(resp)
+                data = await resp.json()
+
+        output = data.get("output")
+        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_openai_compatible.py

@@ -0,0 +1,66 @@
+"""Model for any OpenAI-compatible chat completions endpoint.
+
+Works with vLLM, SGLang, Ollama, LiteLLM proxy, Azure
+OpenAI, Amazon Bedrock, TGI, and any other server that
+implements the ``/v1/chat/completions`` API.
+"""
+
+from typing import Any
+
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+from switch_trust.eval.common import converter_openai
+from switch_trust.eval.common.schema import Message
+from switch_trust.eval.core.models.model import (
+    Model,
+    ModelResponse,
+)
+from switch_trust.eval.core.models.response_schema import to_openai_response_format
+
+
+class OpenAICompatibleModel(Model):
+    _client: AsyncOpenAI
+    _model: str
+    _temperature: float
+
+    def __init__(
+        self,
+        base_url: str,
+        model: str,
+        api_key: str = "EMPTY",
+        temperature: float = 0.0,
+        headers: dict[str, str] | None = None,
+    ):
+        self._client = AsyncOpenAI(
+            base_url=base_url.rstrip("/"),
+            api_key=api_key,
+            default_headers=headers or {},
+        )
+        self._model = model
+        self._temperature = temperature
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        openai_messages = [converter_openai.from_message(m) for m in messages]
+        if output_schema is not None and "response_format" not in kwargs:
+            kwargs["response_format"] = to_openai_response_format(output_schema)
+        response = await self._client.chat.completions.create(
+            model=self._model,
+            messages=openai_messages,
+            temperature=kwargs.pop(
+                "temperature",
+                self._temperature,
+            ),
+            **kwargs,
+        )
+        choice = response.choices[0]
+        message = converter_openai.to_message(
+            choice.message,
+        )
+        return ModelResponse(message)