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

switch_trust/eval/core/models/model_per_endpoint_limiter.py

@@ -0,0 +1,80 @@
+"""Per-endpoint concurrency cap and initial jitter for HTTP-backed models.
+
+Multiple :class:`Model` instances are built for the same tenant agent endpoint
+(the worker rebuilds a fresh ``Model`` per job) so the semaphore is held on the
+class rather than the instance: every instance sharing a URL shares one budget.
+
+asyncio is single-threaded, so no lock is needed around the dict lookup — the
+check-then-insert in :meth:`_get_semaphore` never awaits between the two ops.
+
+Bucketing is by ``endpoint_url`` as passed by the factory, which uses
+``db_model.host``. For HTTP model types whose real request URL is ``host +
+endpoint`` (GenericHttp, LangServe), calls to distinct paths on the same host
+share one budget — that matches "one tenant server" as a rate-limit unit.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from switch_trust.eval.core.models.model import Model, ModelResponse
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from switch_trust.eval.common.schema import Message
+
+
+class PerEndpointLimiter(Model):
+    """Bounds concurrent calls per endpoint URL and jitters the initial burst.
+
+    Wraps an inner :class:`Model`. Delegates to the inner model's public
+    :meth:`Model.generate` (not ``_generate``) so the outer ``generate``
+    wrapper's ``TokenUsage`` accumulator dedupe (see ``core/models/model.py``
+    ``TokenUsageAccumulator.record_once``) still runs at every layer.
+    """
+
+    _semaphores: ClassVar[dict[str, asyncio.Semaphore]] = {}
+
+    def __init__(
+        self,
+        model: Model,
+        endpoint_url: str,
+        limit: int = 8,
+        jitter_max: float = 1.0,
+    ):
+        self._model = model
+        self._endpoint_url = endpoint_url
+        self._limit = limit
+        self._jitter_max = jitter_max
+
+    def _get_semaphore(self) -> asyncio.Semaphore:
+        sem = PerEndpointLimiter._semaphores.get(self._endpoint_url)
+        if sem is None:
+            sem = asyncio.Semaphore(self._limit)
+            PerEndpointLimiter._semaphores[self._endpoint_url] = sem
+        return sem
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        if self._jitter_max > 0:
+            await asyncio.sleep(random.uniform(0, self._jitter_max))
+        async with self._get_semaphore():
+            return await self._model.generate(
+                messages, output_schema=output_schema, **kwargs
+            )
+
+    @classmethod
+    def _reset_state_for_testing(cls) -> None:
+        # asyncio.Semaphore binds to the loop it was created on, and
+        # IsolatedAsyncioTestCase spins up a fresh loop per test. Without
+        # clearing the class-level dict, later tests would await on a
+        # semaphore attached to a dead loop and hang.
+        cls._semaphores.clear()

switch_trust/eval/core/models/model_retry.py

@@ -0,0 +1,204 @@
+import asyncio
+import logging
+import random
+from typing import Any
+
+from pydantic import BaseModel
+
+from switch_trust.eval.common.schema import Message
+from switch_trust.eval.core.models.model import Model, ModelResponse
+
+logger = logging.getLogger(__name__)
+
+_TRANSIENT_EXCEPTIONS = (
+    ConnectionError,
+    TimeoutError,
+    OSError,
+)
+
+_TRANSIENT_STATUS_CODES = {429, 500, 502, 503, 504}
+# Terminal statuses stop retries immediately. 401/403 mean bad credentials
+# and 404 means the endpoint is wrong - none recover by waiting.
+_TERMINAL_STATUS_CODES = {401, 403, 404}
+# Upper bound on the delay we accept from a tenant-controlled Retry-After
+# header. Anything larger is clamped so a misconfigured or hostile endpoint
+# can't park a worker coroutine indefinitely (defeats the pacing goal).
+# 60s covers typical 429 recovery windows for the providers we've seen.
+_MAX_RETRY_AFTER_SECONDS = 60.0
+
+
+def _status_of(exc: Exception) -> int | None:
+    for attr in ("status_code", "status", "code"):
+        status = getattr(exc, attr, None)
+        if isinstance(status, int):
+            return status
+    return None
+
+
+def _is_terminal(exc: Exception) -> bool:
+    status = _status_of(exc)
+    return status is not None and status in _TERMINAL_STATUS_CODES
+
+
+def _is_transient(exc: Exception) -> bool:
+    if isinstance(exc, _TRANSIENT_EXCEPTIONS):
+        return True
+    status = _status_of(exc)
+    return status is not None and status in _TRANSIENT_STATUS_CODES
+
+
+def _extract_retry_after(exc: Exception) -> float | None:
+    """Parse a Retry-After (integer seconds) header off an exception.
+
+    Duck-types across ``aiohttp.ClientResponseError`` (``exc.headers``) and
+    SDK errors like ``openai.APIStatusError`` (``exc.response.headers``).
+    HTTP-date form is not supported. The parsed value is clamped to
+    ``[0, _MAX_RETRY_AFTER_SECONDS]`` because the endpoint is
+    tenant-controlled and cannot be trusted to name a sane delay.
+    """
+    for source in (exc, getattr(exc, "response", None)):
+        headers = getattr(source, "headers", None) if source is not None else None
+        if headers is None:
+            continue
+        try:
+            raw = headers.get("Retry-After") or headers.get("retry-after")
+        except AttributeError:
+            continue
+        if raw is None:
+            continue
+        try:
+            delay = float(int(str(raw).strip()))
+        except (ValueError, TypeError):
+            return None
+        return max(0.0, min(delay, _MAX_RETRY_AFTER_SECONDS))
+    return None
+
+
+class ExponentialRetryModel(Model):
+    """Wraps a Model and retries on transient errors with
+    exponential backoff."""
+
+    def __init__(
+        self,
+        model: Model,
+        max_retries: int = 5,
+        base_delay: float = 2.0,
+    ):
+        self._model = model
+        self._max_retries = max_retries
+        self._base_delay = base_delay
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        last_exc: Exception | None = None
+        for attempt in range(self._max_retries + 1):
+            try:
+                return await self._model.generate(
+                    messages, output_schema=output_schema, **kwargs
+                )
+            except Exception as exc:
+                if _is_terminal(exc):
+                    raise
+                if not _is_transient(exc) or attempt == self._max_retries:
+                    raise
+                last_exc = exc
+                retry_after = _extract_retry_after(exc)
+                if retry_after is not None:
+                    # Add jitter on top of the server-suggested delay: 100
+                    # concurrent 429s all carry the same Retry-After, so
+                    # sleeping exactly that value has every task fire again
+                    # at the same instant and hit the same limit. Same
+                    # shape as the exponential-backoff branch below.
+                    delay = retry_after + random.uniform(0, retry_after)
+                    delay_source = "retry-after"
+                else:
+                    base = self._base_delay * (2**attempt)
+                    delay = base + random.uniform(0, base)
+                    delay_source = "backoff"
+                logger.warning(
+                    "Transient error (attempt %d/%d), retrying in %.1fs [%s] (%s: %s)",
+                    attempt + 1,
+                    self._max_retries,
+                    delay,
+                    delay_source,
+                    type(exc).__name__,
+                    exc,
+                )
+                await asyncio.sleep(delay)
+        raise last_exc  # unreachable, but keeps type checker happy
+
+
+def _fibonacci_delays(max_value: float):
+    """Yield fibonacci-sequence delays capped at max_value."""
+    a, b = 1.0, 1.0
+    while True:
+        yield min(a, max_value)
+        a, b = b, a + b
+
+
+class FibonacciRetryModel(Model):
+    """Wraps a Model and retries on transient errors with
+    fibonacci backoff and jitter, capped at a maximum delay."""
+
+    def __init__(
+        self,
+        model: Model,
+        max_retries: int = 5,
+        max_delay: float = 70.0,
+    ):
+        self._model = model
+        self._max_retries = max_retries
+        self._max_delay = max_delay
+
+    async def _generate(
+        self,
+        messages: list[Message],
+        *,
+        output_schema: type[BaseModel] | None = None,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        last_exc: Exception | None = None
+        delays = _fibonacci_delays(self._max_delay)
+        for attempt in range(self._max_retries + 1):
+            try:
+                return await self._model.generate(
+                    messages, output_schema=output_schema, **kwargs
+                )
+            except Exception as exc:
+                if _is_terminal(exc):
+                    raise
+                if not _is_transient(exc) or attempt == self._max_retries:
+                    raise
+                last_exc = exc
+                base = next(delays)
+                retry_after = _extract_retry_after(exc)
+                if retry_after is not None:
+                    # Add jitter on top of the server-suggested delay: 100
+                    # concurrent 429s all carry the same Retry-After, so
+                    # sleeping exactly that value has every task fire again
+                    # at the same instant and hit the same limit. Same
+                    # shape as the exponential-backoff branch below.
+                    delay = retry_after + random.uniform(0, retry_after)
+                    delay_source = "retry-after"
+                else:
+                    delay = random.uniform(0, base)
+                    delay_source = "backoff"
+                logger.warning(
+                    "Transient error (attempt %d/%d), retrying in %.1fs [%s] (%s: %s)",
+                    attempt + 1,
+                    self._max_retries,
+                    delay,
+                    delay_source,
+                    type(exc).__name__,
+                    exc,
+                )
+                await asyncio.sleep(delay)
+        raise last_exc  # unreachable, but keeps type checker happy
+
+
+RetryModel = ExponentialRetryModel

switch_trust/eval/core/models/model_sync_wrapper.py

@@ -0,0 +1,38 @@
+"""Sync wrapper around an async Model for use in sync contexts.
+
+Used by garak integration which requires synchronous model calls.
+Creates its own event loop to bridge async-to-sync.
+"""
+
+import asyncio
+import logging
+from typing import Any
+
+from switch_trust.eval.core.models.model import Model, ModelContent, ModelResponse
+
+logger = logging.getLogger(__name__)
+
+
+class SyncModelWrapper:
+    """Wraps an async Model for use in sync code (e.g. garak)."""
+
+    def __init__(self, model: Model):
+        self._model = model
+        self._loop = asyncio.new_event_loop()
+
+    def generate(
+        self,
+        contents: ModelContent,
+        **kwargs: Any,
+    ) -> ModelResponse:
+        logger.debug("SyncModelWrapper: bridging async model call")
+        return self._loop.run_until_complete(
+            self._model.generate(contents, **kwargs),
+        )
+
+    def close(self):
+        self._loop.close()
+
+    def __del__(self):
+        if not self._loop.is_closed():
+            self._loop.close()

switch_trust/eval/core/models/model_vertex_agent_runtime.py

@@ -0,0 +1,315 @@
+"""Model implementation for Google Agent Runtime (Vertex AI reasoning engines).
+
+Agents deployed to Agent Runtime are invoked by POSTing to the
+engine resource with a ``:streamQuery`` suffix::
+
+    POST .../reasoningEngines/{engine_id}:streamQuery
+    {"class_method": "stream_query", "input": {"message": "…", "user_id": "…"}}
+
+The product has been renamed twice — Reasoning Engine, then Vertex AI Agent
+Engine, now Agent Runtime under Gemini Enterprise Agent Platform — but Google
+kept the old API identifiers for backwards compatibility. So ``reasoningEngines``
+in the URL is correct, and the service still returns ``Agent Engine Error: …``
+on a bad ``class_method``. See ``platform/MODEL_TYPES.md``.
+
+Two things make this incompatible with :class:`GenericHttpModel`, which is
+why it needs its own type:
+
+* the body is **nested** and carries a sibling ``class_method`` discriminator.
+  ``input`` must be a ``google.protobuf.Struct`` — passing the flat
+  ``{"input": "<prompt>"}`` that ``GenericHttpModel`` sends is rejected with
+  ``400 Invalid value at 'input' (…Struct)``;
+* the response is a **stream of newline-delimited JSON** events, not one JSON
+  object, so a single ``resp.json()`` cannot read it.
+
+The non-streaming ``:query`` sibling is not an alternative for ADK-deployed
+agents: their ``AdkApp`` wrapper marks every agent-invoking method
+(``stream_query``, ``async_stream_query``, ``streaming_agent_run_with_events``)
+as streaming, so ``:query`` exposes only session and memory management. Calling
+it for ``stream_query`` returns ``404 … method not found``.
+
+``session_id`` is deliberately omitted from the request: ``AdkApp`` creates a
+fresh session per call when it is absent, which is what an eval turn wants —
+each prompt is scored independently, with no history bleeding between them.
+
+Authentication accepts either form of ``endpoint_credential``:
+
+* a **service-account key** (raw or base64 JSON), from which a scoped access
+  token is minted and transparently refreshed — the durable option, since
+  Agent Runtime sits behind Google OAuth and a pasted token expires in an hour;
+* a **raw OAuth access token** (e.g. ``gcloud auth print-access-token``), sent
+  as-is. Convenient for a one-off local run, but it will start returning
+  ``401 Unauthorized`` once it expires.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+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
+
+# Agent Runtime is a Cloud Platform API; the broad scope is what service-account
+# tokens for aiplatform.googleapis.com are issued against.
+CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
+
+DEFAULT_CLASS_METHOD = "stream_query"
+DEFAULT_USER_ID = "eval"
+
+
+class VertexAgentRuntimeModel(Model):
+    """Sends prompts to a Google Agent Runtime reasoning engine.
+
+    Args:
+        url: The full engine URL including the ``:streamQuery`` suffix.
+        credential: A service-account key (raw or base64 JSON) or a raw OAuth
+            access token. When None, no ``Authorization`` header is added and
+            ``headers`` is expected to carry one.
+        headers: Extra headers sent with every request.
+        user_id: Identifies the caller to the agent; also scopes the
+            auto-created session.
+        class_method: The engine method to invoke. Engines that expose a
+            differently named streaming entry point (e.g.
+            ``streaming_agent_run_with_events``) can override it.
+        immediate_result: If True, return the first event's text (e.g. a tool
+            call). If False (default), return the final text response after all
+            tool calls complete.
+        connector_factory: Optional aiohttp connector factory; the worker passes
+            the SSRF-guarded one (platform/ssrf.py).
+    """
+
+    _url: str
+    _headers: dict[str, str]
+    _user_id: str
+    _class_method: str
+    _immediate_result: bool
+    _connector_factory: Callable[[], aiohttp.BaseConnector] | None
+    # Set when the credential was a service-account key; None when a raw access
+    # token (or nothing) was supplied.
+    _google_credentials: Any | None
+    _static_token: str | None
+    _refresh_lock: asyncio.Lock
+
+    def __init__(
+        self,
+        url: str,
+        credential: str | None = None,
+        headers: dict[str, str] | None = None,
+        user_id: str = DEFAULT_USER_ID,
+        class_method: str = DEFAULT_CLASS_METHOD,
+        immediate_result: bool = False,
+        connector_factory: Callable[[], aiohttp.BaseConnector] | None = None,
+    ):
+        self._url = url
+        self._headers = dict(headers or {})
+        self._user_id = user_id
+        self._class_method = class_method
+        self._immediate_result = immediate_result
+        self._connector_factory = connector_factory
+        self._google_credentials = _service_account_credentials(credential)
+        self._static_token = (
+            credential if credential and self._google_credentials is None else None
+        )
+        self._refresh_lock = asyncio.Lock()
+
+    def _refresh_credentials(self) -> None:
+        """Refresh the access token in place.
+
+        Blocking: google-auth's transport is synchronous, so callers run this
+        off the event loop.
+        """
+        from google.auth.transport.requests import (  # noqa: PLC0415 - blocking transport, imported at use
+            Request,
+        )
+
+        self._google_credentials.refresh(Request())
+
+    async def _auth_headers(self) -> dict[str, str]:
+        credentials = self._google_credentials
+        if credentials is not None:
+            # One model instance serves every prompt in a run concurrently, so
+            # the moment the token expires all of them see `valid == False` at
+            # once. Without the lock that is N simultaneous refreshes — N token
+            # requests, and a race on the token/expiry pair, which google-auth
+            # does not guard. Re-check inside the lock so only the first waiter
+            # refreshes; the rest fall through to the fresh token. The valid
+            # path stays a plain attribute read, off the lock and off a thread.
+            if not credentials.valid:
+                async with self._refresh_lock:
+                    if not credentials.valid:
+                        await asyncio.to_thread(self._refresh_credentials)
+            return {"Authorization": f"Bearer {credentials.token}"}
+        if self._static_token:
+            return {"Authorization": f"Bearer {self._static_token}"}
+        return {}
+
+    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.
+        prompt_text = flatten_messages(messages)
+
+        payload = {
+            "class_method": self._class_method,
+            "input": {
+                "message": prompt_text,
+                "user_id": self._user_id,
+            },
+        }
+        headers = {**self._headers, **await self._auth_headers()}
+
+        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=payload,
+                headers=headers,
+            ) as resp:
+                await raise_for_status(resp)
+                # Read the whole stream rather than parsing incrementally: an
+                # eval turn needs the complete event list to pick the final
+                # text, and responses are a handful of KB.
+                body = await resp.text()
+
+        events = _parse_events(body)
+        if self._immediate_result:
+            response_text = _extract_first(events)
+        else:
+            response_text = _extract_final(events)
+
+        if not response_text:
+            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 _service_account_credentials(credential: str | None) -> Any | None:
+    """Build google-auth credentials when ``credential`` is a service-account key.
+
+    Returns None for a raw access token (or no credential), which is the signal
+    to send it verbatim as a bearer token. Reuses the generator's parser so the
+    base64 handling and the structural/PEM validation stay in one place — a
+    newline-mangled key then fails with the same actionable message here.
+    """
+    if not credential:
+        return None
+
+    from google.oauth2 import (  # noqa: PLC0415 - only needed on the service-account path
+        service_account,
+    )
+
+    from switch_trust.eval.core.models.generator_model import (  # noqa: PLC0415 - avoids a module-level cycle
+        _parse_service_account,
+    )
+
+    try:
+        info = _parse_service_account(credential, source="endpoint_credential")
+    except ValueError:
+        # Not a service-account key — treat it as a raw access token.
+        return None
+
+    return service_account.Credentials.from_service_account_info(
+        info,
+        scopes=[CLOUD_PLATFORM_SCOPE],
+    )
+
+
+def _parse_events(body: str) -> list[dict]:
+    """Parse an Agent Runtime response into a list of event objects.
+
+    The success path is newline-delimited JSON, one event per line, but a
+    single-event response is indistinguishable from a bare JSON object and
+    errors come back as a JSON array — so try whole-body JSON first and fall
+    back to per-line parsing. Unparseable lines are skipped rather than
+    raising: a trailing partial line must not discard the events already read.
+    """
+    text = body.strip()
+    if not text:
+        return []
+
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        pass
+    else:
+        return _as_event_list(parsed)
+
+    events: list[dict] = []
+    for line in text.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            events.extend(_as_event_list(json.loads(line)))
+        except json.JSONDecodeError:
+            continue
+    return events
+
+
+def _as_event_list(parsed: Any) -> list[dict]:
+    if isinstance(parsed, dict):
+        return [parsed]
+    if isinstance(parsed, list):
+        return [item for item in parsed if isinstance(item, dict)]
+    return []
+
+
+def _texts(event: dict) -> list[str] | None:
+    """Return an event's text parts, or None if it carries no content."""
+    content = event.get("content")
+    if not isinstance(content, dict):
+        return None
+    parts = content.get("parts")
+    if not isinstance(parts, list):
+        return None
+    return [
+        part["text"]
+        for part in parts
+        if isinstance(part, dict) and isinstance(part.get("text"), str)
+    ]
+
+
+def _extract_final(events: list[dict]) -> str | None:
+    """Extract the last text response, ignoring tool call/response events."""
+    for event in reversed(events):
+        texts = _texts(event)
+        if texts:
+            return " ".join(texts)
+    return None
+
+
+def _extract_first(events: list[dict]) -> str | None:
+    """Extract the first event's text content."""
+    for event in events:
+        texts = _texts(event)
+        if texts:
+            return " ".join(texts)
+    return None