AgentAssure-ai 0.1.0__py3-none-any.whl

agentassure/__init__.py

@@ -0,0 +1,68 @@
+"""AgentAssure: pytest-based QA for enterprise AI agents.
+
+Write tests that talk to an agent, then check what it said, which tools it called, what it changed in the
+system of record and what it was allowed to do. Run them with `pytest` or `agentassure test`.
+
+Quick start (`agentassure init` scaffolds the config)::
+
+    import agentassure as aa
+
+    @aa.test(agent="service_agent", persona="service_rep")
+    def test_escalation(agent: aa.Conversation, sf: aa.StateProvider, expect: aa.Expect):
+        case = sf.create("Case", Subject="Delayed order", Priority="Medium")
+        with sf.watch() as diff:
+            reply = agent.ask("Escalate case {n}", n=case.CaseNumber)
+        expect(reply).tool_called("Escalate_Case").response_contains("High priority")
+        expect(diff).record(case).changed_only(Priority="High")
+
+The type hints are optional; add them to get autocomplete and hover docs on the fixtures.
+
+What you import from here:
+
+- Markers: `test`, `cases` (CSV-driven tests), `agent`, `tags`, `meta`.
+- Structure: `step` (named test steps), `Responder` / `when` (scripted replies to the agent's questions).
+- Fixture types: `Conversation` (`agent`), `StateProvider` (`state` / `sf`), `Expect` (`expect`),
+  `Row` (`row`), plus `Reply` and `StateDiff` for the values they return.
+- Constants: `metrics` (names for `expect(...).metric(...)`), `attacks` (prompt-injection payloads),
+  `Severity`.
+"""
+
+__version__ = "0.1.0"
+
+from agentassure import attacks, metrics
+from agentassure.authoring.cases import Row
+from agentassure.conversation.session import Conversation, Responder, when
+from agentassure.core.checks import Severity
+from agentassure.core.evidence import Reply
+from agentassure.dsl.decorators import agent, cases, meta, tags, test
+from agentassure.dsl.expect import Expect
+from agentassure.perf.engine import PerfExpect, PerfReply, PerfSession
+from agentassure.perf.scenarios import perf_scenario
+from agentassure.runtime.steps import step
+from agentassure.state.diff import StateDiff
+from agentassure.state.provider import StateProvider
+
+__all__ = [
+    "Conversation",
+    "Expect",
+    "PerfExpect",
+    "PerfReply",
+    "PerfSession",
+    "Reply",
+    "Responder",
+    "Row",
+    "Severity",
+    "StateDiff",
+    "StateProvider",
+    "__version__",
+    "agent",
+    "attacks",
+    "cases",
+    "meta",
+    "metrics",
+    "perf_scenario",
+    "step",
+    "tags",
+    "test",
+    "when",
+]

agentassure/adapters/__init__.py

@@ -0,0 +1,5 @@
+"""Adapters connect tests to agents: generic HTTP, in-process Python callables, and a lookup registry."""
+
+from agentassure.adapters.base import AdapterContext, AgentAdapter, AgentSession, PersonaContext, Resources
+
+__all__ = ["AdapterContext", "AgentAdapter", "AgentSession", "PersonaContext", "Resources"]

agentassure/adapters/base.py

@@ -0,0 +1,210 @@
+"""Adapter extension point: `AgentAdapter` / `AgentSession` base classes and shared adapter context types."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, Protocol, TypeVar
+
+from agentassure.config.schema import AgentConfig, ResolvedConfig
+from agentassure.core.evidence import Capabilities, Exchange, Reply
+from agentassure.state.provider import StateProvider
+
+T = TypeVar("T")
+
+
+class Resources:
+    """Session-wide shared objects (connections, token providers) keyed by name."""
+
+    def __init__(self) -> None:
+        self._items: dict[str, Any] = {}
+
+    def get_or_create(self, key: str, factory: Callable[[], T]) -> T:
+        """Return the object stored under `key`, creating and storing it with `factory` on first use."""
+        if key not in self._items:
+            self._items[key] = factory()
+        return self._items[key]
+
+    def get(self, key: str) -> Any:
+        """Return the object stored under `key`, or None if absent."""
+        return self._items.get(key)
+
+    def set(self, key: str, value: Any) -> None:
+        """Store `value` under `key`, replacing any existing entry."""
+        self._items[key] = value
+
+    def close(self) -> None:
+        """Close stored objects that have a `close()` method, in reverse insertion order."""
+        for item in reversed(list(self._items.values())):
+            close = getattr(item, "close", None)
+            if callable(close):
+                close()
+        self._items.clear()
+
+
+@dataclass
+class AdapterContext:
+    """Wiring handed to every adapter: the resolved run config and shared session resources."""
+
+    config: ResolvedConfig
+    resources: Resources
+
+
+@dataclass
+class PersonaContext:
+    """Runtime persona for a session: test-user identity, context variables, and adapter options."""
+
+    name: str
+    username: str | None = None
+    description: str | None = None
+    context_variables: dict[str, Any] = field(default_factory=dict)
+    options: dict[str, Any] = field(default_factory=dict)
+
+
+class PermissionChecker(Protocol):
+    """Platform access checker: which records a persona's test user may read or edit.
+
+    Return an implementation from `AgentAdapter.permission_checker()`.
+    """
+
+    def extract_record_ids(self, text: str) -> list[str]:
+        """Return the record IDs mentioned in `text` that permission assertions should check."""
+        ...
+
+    def access(self, persona: PersonaContext, record_ids: list[str]) -> dict[str, dict[str, bool]]:
+        """Return {record_id: {"read": bool, "edit": bool}} evaluated as the persona user. Leave out IDs
+        that can't be evaluated (for example records the checking user can't see): they are reported
+        as unknown, not as violations."""
+        ...
+
+
+class AgentSession(ABC):
+    """One live conversation with the agent under test.
+
+    Subclass this in your adapter and implement `send`. Override `set_variable`, `close`, and
+    `finalize` only if your platform needs them. Set `capabilities` to declare what evidence the
+    agent reports (tool calls, citations, ...).
+    """
+
+    capabilities: Capabilities
+
+    @abstractmethod
+    def send(self, text: str) -> Exchange:
+        """Send one user message and return the recorded exchange."""
+        ...
+
+    def set_variable(self, name: str, value: Any) -> None:
+        """Set a session variable visible to the agent; unsupported adapters raise NotImplementedError."""
+        raise NotImplementedError(f"{type(self).__name__} does not support session variables")
+
+    def close(self) -> None:
+        """End the session and release any live resources."""
+        return None
+
+    def finalize(self, replies: list[Reply]) -> None:
+        """Called after `close()`: fill evidence that only exists once the session has ended (into the
+        replies' exchanges), then clear `capabilities.deferred`."""
+        return None
+
+
+class StreamObserver(Protocol):
+    """Receives timing signals from an `AsyncAgentSession.send` (used by performance testing).
+
+    Call `dispatched()` once the request has been handed to the network, and `content(text)` for every
+    chunk that carries reply text (empty strings are ignored). Both are cheap and must not block.
+    """
+
+    def dispatched(self) -> None: ...
+
+    def content(self, text: str) -> None: ...
+
+
+class AsyncAgentSession(ABC):
+    """Optional async, stream-capable conversation (see `AgentAdapter.open_session_async`).
+
+    `send` must call `observer.dispatched()` when the request leaves the process and `observer.content()`
+    for each text chunk as it arrives; that is how first-content latency and stream gaps are measured.
+    """
+
+    capabilities: Capabilities
+    streams: bool = False
+    """True when replies arrive in chunks. Only streaming sessions get first-content and stream-silence
+    timeouts; non-streaming replies are bounded by the turn timeout alone."""
+
+    @abstractmethod
+    async def send(self, text: str, observer: StreamObserver | None = None) -> Exchange:
+        """Send one user message, report chunks to `observer`, and return the complete exchange."""
+        ...
+
+    async def aclose(self) -> None:
+        """End the session and release live resources."""
+        return None
+
+
+class AgentAdapter(ABC):
+    """Connects AgentAssure to one agent. Instantiated once per agent per test session.
+
+    To support a new agent platform, subclass this, implement `open_session` (returning an
+    `AgentSession`), and register the class under a name so `adapter: <name>` works in agentassure.yaml
+    (see `agentassure.adapters.registry.register`, or the `agentassure.adapters` entry-point group).
+    Optionally override `preflight`, `state_provider`, `permission_checker`, `instructions`, `describe`,
+    `doctor`, and `close`. Set `name` and, if the agent reports tool calls etc., `capabilities`.
+    """
+
+    name: str = "adapter"
+    capabilities: Capabilities = Capabilities()
+    supports_async: bool = False
+    """True when `open_session_async` is implemented (perf runs then stream instead of using threads)."""
+    usage_model: str | None = None
+    """`provider/model` key used to price reported token usage (perf `pricing`), if known."""
+
+    def __init__(self, agent_name: str, agent_config: AgentConfig, context: AdapterContext):
+        self.agent_name = agent_name
+        self.agent_config = agent_config
+        self.context = context
+
+    @property
+    def options(self) -> dict[str, Any]:
+        """Adapter-specific options: the extra keys of this agent's config entry."""
+        return self.agent_config.options
+
+    def preflight(self) -> None:
+        """Validate connectivity and safety before any test runs. Raise to abort the session."""
+
+    @abstractmethod
+    def open_session(self, persona: PersonaContext | None) -> AgentSession:
+        """Start a new conversation, optionally running as `persona`'s test user."""
+        ...
+
+    async def open_session_async(self, persona: PersonaContext | None) -> AsyncAgentSession:
+        """Start an async, stream-capable conversation. Only called when `supports_async` is True."""
+        raise NotImplementedError(f"{type(self).__name__} has no async sessions")
+
+    async def aclose(self) -> None:
+        """Release async resources (e.g. an `httpx.AsyncClient`) created by `open_session_async`."""
+        return None
+
+    def state_provider(self) -> StateProvider | None:
+        """Default state provider for this agent, if the platform has one."""
+        return None
+
+    def permission_checker(self) -> PermissionChecker | None:
+        """Platform access checks used by persona permission assertions."""
+        return None
+
+    def instructions(self) -> list[str]:
+        """The agent's own instructions, for leak checks. Defaults to `instructions:` in the agent config."""
+        return list(self.agent_config.instructions)
+
+    def describe(self) -> dict[str, Any]:
+        """Return a short dict identifying this adapter for reports and diagnostics."""
+        return {"adapter": self.name}
+
+    def doctor(self) -> list[tuple[str, bool, str]]:
+        """Diagnostics for `agentassure doctor`: (check, ok, detail)."""
+        return []
+
+    def close(self) -> None:
+        """Release adapter-level resources such as open connections."""
+        return None

agentassure/adapters/http.py

@@ -0,0 +1,326 @@
+"""Generic HTTP adapter (`adapter: http`): test any agent that exposes a JSON chat endpoint.
+
+Sync sessions (functional tests) POST and map the JSON reply. Async sessions (performance tests) can also
+consume a streamed reply: `stream: sse` (Server-Sent Events) or `stream: ndjson`, reporting each text
+chunk so first-content latency and stream gaps are measured. See the GitBook "Agents and adapters" page.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+import uuid
+from typing import Any
+
+import httpx
+
+from agentassure.adapters.base import (
+    AgentAdapter,
+    AgentSession,
+    AsyncAgentSession,
+    PersonaContext,
+    StreamObserver,
+)
+from agentassure.adapters.mapping import dig, exchange_from_payload, usage_from_payload
+from agentassure.adapters.streaming import StreamProtocolError, iter_ndjson, iter_sse
+from agentassure.config.loader import require
+from agentassure.core.evidence import AgentMessage, Capabilities, Exchange
+from agentassure.errors import AdapterError, ConfigError
+
+STREAM_MODES = ("none", "sse", "ndjson")
+
+
+class HttpSession(AgentSession):
+    """One conversation with an HTTP agent: tracks the server-side session ID and mutable variables."""
+
+    def __init__(self, adapter: HttpAdapter, persona: PersonaContext | None):
+        self.adapter = adapter
+        self.persona = persona
+        self.capabilities = adapter.capabilities
+        self.session_id = str(uuid.uuid4())
+        self.variables: dict[str, Any] = dict(persona.context_variables) if persona else {}
+        start_path = adapter.options.get("session_start_path")
+        if start_path:
+            data = adapter.post(
+                start_path, {"persona": persona.name if persona else None, "variables": self.variables}
+            )
+            self.session_id = str(
+                dig(data, adapter.options.get("session_id_field", "session_id"), self.session_id)
+            )
+
+    def set_variable(self, name: str, value: Any) -> None:
+        """Set a session variable sent with subsequent requests."""
+        self.variables[name] = value
+
+    def send(self, text: str) -> Exchange:
+        """POST the user message to the chat endpoint and map the JSON response to an Exchange."""
+        req = self.adapter.options.get("request", {})
+        body = {
+            req.get("message_field", "message"): text,
+            req.get("session_field", "session_id"): self.session_id,
+            req.get("persona_field", "persona"): self.persona.name if self.persona else None,
+            req.get("variables_field", "variables"): self.variables,
+        }
+        started = time.perf_counter()
+        data = self.adapter.post(self.adapter.options.get("path", "/chat"), body)
+        exchange = self.adapter.to_exchange(text, data)
+        exchange.latency_ms = (time.perf_counter() - started) * 1000
+        return exchange
+
+    def close(self) -> None:
+        """Call the optional `session_end_path` endpoint, ignoring request failures."""
+        end_path = self.adapter.options.get("session_end_path")
+        if end_path:
+            with contextlib.suppress(AdapterError):
+                self.adapter.post(end_path, {"session_id": self.session_id})
+
+
+def _body(session: HttpSession | HttpAsyncSession, text: str) -> dict[str, Any]:
+    req = session.adapter.options.get("request", {})
+    return {
+        req.get("message_field", "message"): text,
+        req.get("session_field", "session_id"): session.session_id,
+        req.get("persona_field", "persona"): session.persona.name if session.persona else None,
+        req.get("variables_field", "variables"): session.variables,
+    }
+
+
+_SEND_STARTED = ("http11.send_request_headers.started", "http2.send_request_headers.started")
+
+
+def trace_extensions(observer: StreamObserver | None) -> dict[str, Any]:
+    """httpx `extensions` that call `observer.dispatched()` when the request headers start going out
+    (after connection setup), so connect/TLS time counts as client wait, not agent latency."""
+    if observer is None:
+        return {}
+
+    async def trace(event: str, info: dict[str, Any]) -> None:
+        if event in _SEND_STARTED:
+            observer.dispatched()
+
+    return {"trace": trace}
+
+
+class HttpAsyncSession(AsyncAgentSession):
+    """Async conversation with an HTTP agent; streams the reply when `stream` is `sse` or `ndjson`."""
+
+    def __init__(self, adapter: HttpAdapter, persona: PersonaContext | None):
+        self.adapter = adapter
+        self.persona = persona
+        self.capabilities = adapter.capabilities
+        self.session_id = str(uuid.uuid4())
+        self.variables: dict[str, Any] = dict(persona.context_variables) if persona else {}
+        self.streams = adapter.stream_mode != "none"
+
+    async def start(self) -> None:
+        start_path = self.adapter.options.get("session_start_path")
+        if start_path:
+            data = await self.adapter.apost(
+                start_path,
+                {"persona": self.persona.name if self.persona else None, "variables": self.variables},
+            )
+            self.session_id = str(
+                dig(data, self.adapter.options.get("session_id_field", "session_id"), self.session_id)
+            )
+
+    async def send(self, text: str, observer: StreamObserver | None = None) -> Exchange:
+        """Send one message; with streaming, report each text chunk to `observer` as it arrives."""
+        opts = self.adapter.options
+        mode = self.adapter.stream_mode
+        started = time.perf_counter()
+        if mode == "none":
+            data = await self.adapter.apost(opts.get("path", "/chat"), _body(self, text), observer=observer)
+            exchange = self.adapter.to_exchange(text, data)
+        else:
+            exchange = await self._stream(text, observer)
+        exchange.latency_ms = (time.perf_counter() - started) * 1000
+        return exchange
+
+    async def _stream(self, text: str, observer: StreamObserver | None) -> Exchange:
+        opts = self.adapter.options
+        path = opts.get("stream_path") or opts.get("path", "/chat")
+        text_path = opts.get("stream_text_path", "text")
+        done_marker = opts.get("stream_done", "[DONE]")
+        final_event = opts.get("stream_final_event", "final")
+        usage_path = opts.get("usage_path", "usage")
+        parts: list[str] = []
+        final_payload: dict[str, Any] | None = None
+        usage_raw: Any = None
+        done = False
+        try:
+            async with self.adapter.async_client().stream(
+                "POST", path, json=_body(self, text), extensions=trace_extensions(observer)
+            ) as response:
+                response.raise_for_status()
+                lines = response.aiter_lines()
+                if self.adapter.stream_mode == "sse":
+                    async for event in iter_sse(lines):
+                        if event.data.strip() == done_marker or event.event == done_marker:
+                            done = True
+                            break
+                        try:
+                            payload = event.json()
+                        except ValueError as exc:
+                            raise StreamProtocolError(
+                                f"{self.adapter.agent_name}: SSE data isn't JSON: {event.data[:200]!r}"
+                            ) from exc
+                        if event.event == final_event:
+                            final_payload = payload if isinstance(payload, dict) else None
+                            done = True
+                            break
+                        chunk = dig(payload, text_path)
+                        if isinstance(chunk, str) and chunk:
+                            parts.append(chunk)
+                            if observer:
+                                observer.content(chunk)
+                        usage_raw = dig(payload, usage_path, usage_raw)
+                else:
+                    async for payload in iter_ndjson(lines):
+                        if isinstance(payload, dict) and payload.get("done") is True:
+                            final_payload = payload
+                            done = True
+                            break
+                        chunk = dig(payload, text_path)
+                        if isinstance(chunk, str) and chunk:
+                            parts.append(chunk)
+                            if observer:
+                                observer.content(chunk)
+                        usage_raw = dig(payload, usage_path, usage_raw)
+        except ValueError as exc:
+            if isinstance(exc, StreamProtocolError):
+                raise
+            raise StreamProtocolError(f"{self.adapter.agent_name}: malformed stream: {exc}") from exc
+        except httpx.HTTPError as exc:
+            raise AdapterError(f"{self.adapter.agent_name}: HTTP stream to {path} failed: {exc}") from exc
+        if not done and opts.get("stream_require_done", True):
+            raise StreamProtocolError(f"{self.adapter.agent_name}: stream ended without a completion marker")
+        streamed = "".join(parts)
+        if final_payload is not None:
+            exchange = self.adapter.to_exchange(text, final_payload)
+            if not exchange.response and streamed:
+                exchange.messages = [AgentMessage(text=streamed, agent=self.adapter.agent_name)]
+            if exchange.usage is None:
+                exchange.usage = usage_from_payload(usage_raw, opts.get("usage_fields"))
+            return exchange
+        exchange = self.adapter.to_exchange(text, {"text": streamed})
+        exchange.usage = usage_from_payload(usage_raw, opts.get("usage_fields"))
+        return exchange
+
+    async def aclose(self) -> None:
+        end_path = self.adapter.options.get("session_end_path")
+        if end_path:
+            with contextlib.suppress(AdapterError):
+                await self.adapter.apost(end_path, {"session_id": self.session_id})
+
+
+class HttpAdapter(AgentAdapter):
+    """Generic JSON-over-HTTP agent.
+
+    Options: base_url, path, headers, timeout_s, request{...field names}, response{...dotted paths},
+    capabilities{...}, session_start_path, session_end_path, session_id_field, usage_path, usage_fields;
+    streaming (async/perf only): stream (none|sse|ndjson), stream_path, stream_text_path, stream_done,
+    stream_final_event, stream_require_done.
+    """
+
+    name = "http"
+    supports_async = True
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        base_url = require(self.options.get("base_url"), f"agents.{self.agent_name}.base_url")
+        self.client = httpx.Client(
+            base_url=base_url,
+            headers=self.options.get("headers") or {},
+            timeout=float(self.options.get("timeout_s", self.context.config.defaults.timeout_s)),
+        )
+        declared = self.options.get("capabilities") or {}
+        response_map = self.options.get("response") or {}
+        inferred = {
+            "tool_calls": "tool_calls" in response_map,
+            "tool_arguments": "tool_calls" in response_map,
+            "tool_outputs": "tool_calls" in response_map,
+            "routing": "routing" in response_map,
+            "handoffs": "handoffs" in response_map,
+            "escalation": "escalation" in response_map,
+            "retrieved_contexts": "contexts" in response_map,
+        }
+        self.capabilities = Capabilities(**(inferred | declared))
+        self.stream_mode = str(self.options.get("stream", "none"))
+        if self.stream_mode not in STREAM_MODES:
+            raise ConfigError(
+                f"agents.{self.agent_name}.stream must be one of {', '.join(STREAM_MODES)}, got {self.stream_mode!r}"
+            )
+        self.usage_model = self.options.get("usage_model")
+        self._aclient: httpx.AsyncClient | None = None
+
+    def to_exchange(self, text: str, data: Any) -> Exchange:
+        """Map a JSON reply to an Exchange, including token usage at `usage_path` (default `usage`)."""
+        payload = data if isinstance(data, dict) else {"text": data}
+        exchange = exchange_from_payload(text, payload, self.options.get("response", {}), self.agent_name)
+        exchange.usage = usage_from_payload(
+            dig(payload, self.options.get("usage_path", "usage")), self.options.get("usage_fields")
+        )
+        return exchange
+
+    def async_client(self) -> httpx.AsyncClient:
+        """The adapter's async client, created on first use inside the running event loop."""
+        if self._aclient is None:
+            self._aclient = httpx.AsyncClient(
+                base_url=str(self.client.base_url),
+                headers=self.options.get("headers") or {},
+                timeout=self.client.timeout,
+                limits=httpx.Limits(max_connections=None, max_keepalive_connections=None),
+            )
+        return self._aclient
+
+    async def apost(self, path: str, body: dict[str, Any], observer: StreamObserver | None = None) -> Any:
+        """Async POST returning the decoded JSON body.
+
+        Raises:
+            AdapterError: the request failed or returned a non-2xx status (the httpx error is the cause).
+        """
+        try:
+            response = await self.async_client().post(path, json=body, extensions=trace_extensions(observer))
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPError as exc:
+            raise AdapterError(f"{self.agent_name}: HTTP call to {path} failed: {exc}") from exc
+        except ValueError as exc:
+            raise StreamProtocolError(f"{self.agent_name}: response from {path} isn't JSON") from exc
+
+    async def open_session_async(self, persona: PersonaContext | None) -> AsyncAgentSession:
+        """Start an async session (calls `session_start_path` when configured)."""
+        session = HttpAsyncSession(self, persona)
+        await session.start()
+        return session
+
+    async def aclose(self) -> None:
+        """Close the async client."""
+        if self._aclient is not None:
+            await self._aclient.aclose()
+            self._aclient = None
+
+    def post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
+        """POST JSON to `path` and return the decoded response body.
+
+        Raises:
+            AdapterError: If the request fails or returns a non-2xx status.
+        """
+        try:
+            response = self.client.post(path, json=body)
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPError as exc:
+            raise AdapterError(f"{self.agent_name}: HTTP call to {path} failed: {exc}") from exc
+
+    def open_session(self, persona: PersonaContext | None) -> AgentSession:
+        """Start a new HTTP session, optionally as `persona`'s user."""
+        return HttpSession(self, persona)
+
+    def describe(self) -> dict[str, Any]:
+        """Return adapter name and base URL for reports and diagnostics."""
+        return {"adapter": self.name, "base_url": str(self.client.base_url)}
+
+    def close(self) -> None:
+        """Close the underlying httpx client."""
+        self.client.close()