yaab-sdk 0.1.0__py3-none-any.whl

yaab/__init__.py

@@ -0,0 +1,159 @@
+"""YAAB — Yet Another Agent Builder.
+
+A type-safe, governance-first agent SDK with a Rust performance core. Type-safe,
+optimizable, durable, and simple — the best ideas from across the agent
+ecosystem on one runtime, on a universal LiteLLM model layer.
+
+Quickstart::
+
+    from yaab import Agent
+
+    agent = Agent("assistant", model="openai/gpt-4o", instructions="Be concise.")
+    print(agent.run_sync("Say hello").output)
+
+Offline (no API key)::
+
+    from yaab import Agent
+    from yaab.testing import TestModel
+
+    agent = Agent("assistant", model=TestModel("hi!"))
+    assert agent.run_sync("hello").output == "hi!"
+"""
+
+from __future__ import annotations
+
+from . import _core
+from .agent import Agent
+from .artifacts.manager import ArtifactManager
+from .batch import batch_embed, batch_map, batch_run
+from .config import agent_from_dict, agent_from_yaml, runner_from_dict
+from .content import Content, Part, PartKind
+from .context import KeepAll, SummarizeHistory, TruncateMessages
+from .eval import available_metrics, get_metric, register_metric
+from .exceptions import (
+    ApprovalRequired,
+    GovernanceError,
+    MaxStepsExceeded,
+    ModelError,
+    OutputValidationError,
+    PolicyViolation,
+    RunCancelled,
+    ToolError,
+    UsageLimitExceeded,
+    YaabError,
+)
+from .extensions import available as available_components
+from .extensions import get as get_component
+from .extensions import register as register_component
+from .governance.eval import ToolTrajectoryMatch
+from .governance.evalset import EvalCase, EvalSet
+from .graph.state import RetryPolicy
+from .limits import CancellationToken, UsageLimits
+from .memory.extraction import MemoryExtractor
+from .memory.manager import MemoryManager
+from .models.router import ModelRouter
+from .multiagent import LoopAgent, MapAgent, ParallelAgent, SequentialAgent, Swarm
+from .prompts import PromptRegistry
+from .rag import Document, KnowledgeBase
+from .rag.memory_service import KnowledgeBaseMemory
+from .runner import Runner
+from .sessions.manager import SessionManager
+from .skills import Skill
+from .state import State
+from .tools import AgentTool, FunctionTool, tool
+from .tools.auth import ToolAuth, ToolAuthRequired, ToolCredential
+from .tools.openapi import OpenAPITool, openapi_toolset
+from .types import Event, EventType, Message, RunContext, RunResult, Usage
+
+__version__ = "0.1.0"
+
+#: Which performance backend is active: ``"rust"`` or ``"python"``.
+BACKEND = _core.backend()
+
+__all__ = [
+    "__version__",
+    "BACKEND",
+    "Agent",
+    "Runner",
+    "tool",
+    "FunctionTool",
+    "AgentTool",
+    # multi-agent workflow patterns
+    "SequentialAgent",
+    "ParallelAgent",
+    "MapAgent",
+    "LoopAgent",
+    "Swarm",
+    # managers
+    "SessionManager",
+    "MemoryManager",
+    "ArtifactManager",
+    "State",
+    # extensibility
+    "register_component",
+    "get_component",
+    "available_components",
+    # eval metrics (built-in + RAGAS/DeepEval adapters)
+    "register_metric",
+    "get_metric",
+    "available_metrics",
+    # reusable building blocks
+    "Skill",
+    "PromptRegistry",
+    # RAG
+    "KnowledgeBase",
+    "Document",
+    # memory intelligence
+    "MemoryExtractor",
+    "KnowledgeBaseMemory",
+    # model intelligence
+    "ModelRouter",
+    # graph retry policies
+    "RetryPolicy",
+    # eval depth (portable evalsets + trajectory metric)
+    "EvalSet",
+    "EvalCase",
+    "ToolTrajectoryMatch",
+    # OpenAPI toolset
+    "openapi_toolset",
+    "OpenAPITool",
+    # tool-level auth (credentials + OAuth2 consent)
+    "ToolAuth",
+    "ToolCredential",
+    "ToolAuthRequired",
+    # run controls
+    "UsageLimits",
+    "CancellationToken",
+    # context-window management
+    "TruncateMessages",
+    "SummarizeHistory",
+    "KeepAll",
+    # declarative config
+    "agent_from_dict",
+    "agent_from_yaml",
+    "runner_from_dict",
+    # batch / offline inference
+    "batch_run",
+    "batch_map",
+    "batch_embed",
+    "RunContext",
+    "RunResult",
+    "Message",
+    "Content",
+    "Part",
+    "PartKind",
+    "Event",
+    "EventType",
+    "Usage",
+    # exceptions
+    "YaabError",
+    "ModelError",
+    "ToolError",
+    "OutputValidationError",
+    "MaxStepsExceeded",
+    "UsageLimitExceeded",
+    "RunCancelled",
+    "GovernanceError",
+    "PolicyViolation",
+    "ApprovalRequired",
+]

yaab/_core.py

@@ -0,0 +1,214 @@
+"""Performance-core shim.
+
+Prefers the compiled Rust extension (:mod:`yaab_core`) and transparently falls
+back to pure-Python implementations when the wheel is unavailable. The public
+functions here are the *only* way the rest of the SDK touches the hot paths, so
+the Rust/Python split is invisible to callers.
+
+Inspect :data:`RUST` to know which backend is active.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from typing import Any
+
+# Set YAAB_NO_RUST=1 to force the pure-Python fallback (used to keep the
+# fallback path green in CI even when the wheel is installed).
+if os.environ.get("YAAB_NO_RUST") == "1":
+    _rust = None
+    RUST = False
+else:
+    try:  # pragma: no cover - exercised by whichever backend is installed
+        import yaab_core as _rust
+
+        RUST = True
+    except ImportError:  # pragma: no cover
+        _rust = None
+        RUST = False
+
+__all__ = [
+    "RUST",
+    "backend",
+    "cosine_similarity",
+    "top_k",
+    "encode_checkpoint",
+    "decode_checkpoint",
+    "reduce_channel",
+    "advance_superstep",
+    "plan_supersteps",
+    "hash_event",
+    "verify_chain",
+    "aggregate_cost",
+]
+
+_MAGIC = b"YAAB"
+_VERSION = 1
+
+
+def backend() -> str:
+    """Return ``"rust"`` or ``"python"`` for diagnostics and tests."""
+    return "rust" if RUST else "python"
+
+
+def cosine_similarity(a: list[float], b: list[float]) -> float:
+    if RUST:
+        return _rust.cosine_similarity(list(a), list(b))
+    if len(a) != len(b) or not a:
+        return 0.0
+    dot = sum(x * y for x, y in zip(a, b, strict=False))
+    na = math.sqrt(sum(x * x for x in a))
+    nb = math.sqrt(sum(y * y for y in b))
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return dot / (na * nb)
+
+
+def top_k(query: list[float], matrix: list[list[float]], k: int) -> list[tuple[int, float]]:
+    if RUST:
+        return _rust.top_k(list(query), [list(r) for r in matrix], k)
+    scored = [(i, cosine_similarity(query, row)) for i, row in enumerate(matrix)]
+    scored.sort(key=lambda t: t[1], reverse=True)
+    return scored[:k]
+
+
+def encode_checkpoint(state: Any) -> bytes:
+    """Serialize a JSON-compatible state object into a framed checkpoint blob."""
+    json_str = json.dumps(state, separators=(",", ":"), sort_keys=True)
+    if RUST:
+        return bytes(_rust.encode_checkpoint(json_str))
+    return _MAGIC + bytes([_VERSION]) + json_str.encode("utf-8")
+
+
+def decode_checkpoint(blob: bytes) -> Any:
+    """Decode a framed checkpoint blob back into a Python object."""
+    if RUST:
+        return json.loads(_rust.decode_checkpoint(bytes(blob)))
+    if len(blob) < 5 or blob[0:4] != _MAGIC:
+        raise ValueError("invalid checkpoint: bad magic header")
+    if blob[4] != _VERSION:
+        raise ValueError(f"unsupported checkpoint version: {blob[4]}")
+    return json.loads(blob[5:].decode("utf-8"))
+
+
+def reduce_channel(reducer: str, current: Any, update: Any) -> Any:
+    """Merge ``update`` into ``current`` under the named channel reducer."""
+    if RUST:
+        out = _rust.reduce_channel(
+            reducer,
+            json.dumps(current, separators=(",", ":")),
+            json.dumps(update, separators=(",", ":")),
+        )
+        return json.loads(out)
+    if reducer == "last_value":
+        return update
+    if reducer == "append":
+        base = (
+            list(current) if isinstance(current, list) else ([] if current is None else [current])
+        )
+        if isinstance(update, list):
+            base.extend(update)
+        else:
+            base.append(update)
+        return base
+    if reducer == "add":
+        a = current if isinstance(current, (int, float)) else 0
+        b = update if isinstance(update, (int, float)) else 0
+        return a + b
+    return update
+
+
+def advance_superstep(
+    state: dict[str, Any],
+    reducers: dict[str, str],
+    updates: list[Any],
+) -> dict[str, Any]:
+    """Apply a whole superstep's node updates to ``state`` in one operation.
+
+    ``updates`` is a list of per-node update dicts (or ``None``), applied in
+    order; each key is reduced with the channel reducer named in ``reducers``
+    (defaulting to last-value). The Rust path does the entire fold in one call;
+    the pure-Python fallback reuses :func:`reduce_channel`.
+    """
+    if RUST:
+        out = _rust.advance_superstep(
+            json.dumps(state, separators=(",", ":")),
+            json.dumps(reducers, separators=(",", ":")),
+            json.dumps(updates, separators=(",", ":")),
+        )
+        return json.loads(out)
+    new_state = dict(state)
+    for update in updates:
+        if not update:
+            continue
+        for key, value in update.items():
+            reducer = reducers.get(key, "last_value")
+            new_state[key] = reduce_channel(reducer, new_state.get(key), value)
+    return new_state
+
+
+def plan_supersteps(nodes: list[str], edges: list[tuple[str, str]]) -> list[list[str]]:
+    """Group nodes into BSP supersteps; trailing cycle members form a final step."""
+    if RUST:
+        return _rust.plan_supersteps(list(nodes), [tuple(e) for e in edges])
+    node_set = set(nodes)
+    indegree = {n: 0 for n in nodes}
+    adj: dict[str, list[str]] = {n: [] for n in nodes}
+    for src, dst in edges:
+        if src not in node_set or dst not in node_set or src == dst:
+            continue
+        adj[src].append(dst)
+        indegree[dst] += 1
+    layers: list[list[str]] = []
+    ready = [n for n, d in indegree.items() if d == 0]
+    visited: set[str] = set()
+    while ready:
+        layer = sorted(n for n in ready if n not in visited)
+        visited.update(layer)
+        nxt: list[str] = []
+        for node in ready:
+            for child in adj.get(node, []):
+                indegree[child] -= 1
+                if indegree[child] == 0 and child not in visited:
+                    nxt.append(child)
+        if layer:
+            layers.append(layer)
+        ready = nxt
+    remaining = sorted(n for n in nodes if n not in visited)
+    if remaining:
+        layers.append(remaining)
+    return layers
+
+
+def hash_event(prev_hash: str, payload: str) -> str:
+    """Chained audit hash: ``sha256(prev_hash || payload)``."""
+    if RUST:
+        return _rust.hash_event(prev_hash, payload)
+    import hashlib
+
+    h = hashlib.sha256()
+    h.update(prev_hash.encode("utf-8"))
+    h.update(b"\x1f")
+    h.update(payload.encode("utf-8"))
+    return h.hexdigest()
+
+
+def verify_chain(genesis: str, entries: list[tuple[str, str]]) -> int | None:
+    """Return the index of the first broken hash-chain link, or ``None``."""
+    if RUST:
+        return _rust.verify_chain(genesis, [tuple(e) for e in entries])
+    prev = genesis
+    for i, (payload, recorded) in enumerate(entries):
+        expected = hash_event(prev, payload)
+        if expected != recorded:
+            return i
+        prev = expected
+    return None
+
+
+def aggregate_cost(values: list[float]) -> float:
+    if RUST:
+        return _rust.aggregate_cost([float(v) for v in values])
+    return float(sum(values))

yaab/a2a/__init__.py

@@ -0,0 +1,14 @@
+"""Agent-to-Agent (A2A) interop.
+
+The server side lives in :mod:`yaab.serve` (agent card + ``/a2a/tasks``). This
+package is the *client* side: :class:`RemoteAgent` discovers a remote agent via
+its Agent Card and delegates tasks to it. A ``RemoteAgent`` satisfies both the
+:class:`~yaab.tools.base.Tool` protocol (so it can be handed to a local agent as
+a tool) and the workflow ``run`` surface (so it composes like any agent).
+"""
+
+from __future__ import annotations
+
+from .client import RemoteAgent
+
+__all__ = ["RemoteAgent"]

yaab/a2a/client.py

@@ -0,0 +1,170 @@
+"""A2A client — discover and delegate to a remote agent over HTTP.
+
+``httpx`` is an optional dependency, imported lazily. An injectable
+``transport`` callable (``method, path, json -> dict``) keeps the client fully
+testable without a network — the test suite drives a real :func:`fastapi_server_app`
+server through an in-process transport.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from ..types import RunContext, RunResult, Usage
+
+Transport = Callable[[str, str, dict | None], Awaitable[dict]]
+
+
+class RemoteAgent:
+    """A handle to a remote A2A agent — usable as a tool or as an agent."""
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        name: str | None = None,
+        auth_token: str | None = None,
+        token_provider: Callable[[], str] | None = None,
+        transport: Transport | None = None,
+    ) -> None:
+        self.url = url.rstrip("/")
+        self.name = name or "remote_agent"
+        self.auth_token = auth_token
+        # token_provider lets an OAuth2 client supply a fresh bearer token per
+        # request (token exchange / refresh handled by the caller's IdP client).
+        self.token_provider = token_provider
+        self._transport = transport
+        self._card: dict[str, Any] | None = None
+
+    # --- transport -----------------------------------------------------
+    def _headers(self) -> dict[str, str]:
+        token = self.token_provider() if self.token_provider is not None else self.auth_token
+        return {"Authorization": f"Bearer {token}"} if token else {}
+
+    async def _request(self, method: str, path: str, json: dict | None = None) -> dict:
+        if self._transport is not None:
+            return await self._transport(method, path, json)
+        try:
+            import httpx
+        except ImportError as exc:  # pragma: no cover - optional extra
+            raise RuntimeError(
+                "httpx is required for the A2A client. `pip install httpx`."
+            ) from exc
+        async with httpx.AsyncClient() as client:
+            resp = await client.request(
+                method, f"{self.url}{path}", json=json, headers=self._headers()
+            )
+            resp.raise_for_status()
+            return resp.json()
+
+    # --- discovery -----------------------------------------------------
+    async def fetch_card(self, *, refresh: bool = False) -> dict[str, Any]:
+        """Fetch (and cache) the remote Agent Card."""
+        if self._card is None or refresh:
+            self._card = await self._request("GET", "/.well-known/agent.json", None)
+            if not self._card.get("name"):
+                self._card["name"] = self.name
+            else:
+                self.name = self.name or self._card["name"]
+        return self._card
+
+    # --- delegation ----------------------------------------------------
+    async def run(
+        self,
+        prompt: str,
+        *,
+        deps: Any = None,
+        session_id: str | None = None,
+        identity: str | None = None,
+    ) -> RunResult[str]:
+        """Send the prompt as an A2A task and return the remote agent's output."""
+        body = {"message": {"role": "user", "parts": [{"text": prompt}]}}
+        task = await self._request("POST", "/a2a/tasks", body)
+        text = _extract_text(task)
+        return RunResult(output=text, usage=Usage(requests=1), run_id=task.get("id", ""))
+
+    def run_sync(self, prompt: str, **kwargs: Any) -> RunResult[str]:
+        import asyncio
+
+        return asyncio.run(self.run(prompt, **kwargs))
+
+    async def get_task(self, task_id: str) -> dict[str, Any]:
+        """Poll a long-running task by id (returns the A2A task object)."""
+        return await self._request("GET", f"/a2a/tasks/{task_id}", None)
+
+    #: A2A task states considered terminal (stop polling).
+    TERMINAL_STATES = frozenset({"completed", "failed", "canceled", "cancelled", "rejected"})
+
+    async def poll_task(
+        self,
+        task_id: str,
+        *,
+        interval: float = 1.0,
+        timeout: float | None = 60.0,
+        terminal_states: set[str] | frozenset[str] | None = None,
+    ) -> dict[str, Any]:
+        """Poll a long-running remote task until it reaches a terminal state.
+
+        Returns the final task object. Raises :class:`TimeoutError` if the task
+        has not finished within ``timeout`` seconds (``None`` = no timeout).
+        ``interval`` is the delay between polls.
+        """
+        import asyncio
+        import time
+
+        terminal = terminal_states or self.TERMINAL_STATES
+        started = time.monotonic()
+        while True:
+            task = await self.get_task(task_id)
+            state = (task.get("status") or {}).get("state")
+            if state in terminal:
+                return task
+            if timeout is not None and time.monotonic() - started > timeout:
+                raise TimeoutError(
+                    f"A2A task {task_id!r} did not reach a terminal state within {timeout}s "
+                    f"(last state: {state!r})"
+                )
+            if interval:
+                await asyncio.sleep(interval)
+
+    # --- Tool protocol -------------------------------------------------
+    @property
+    def description(self) -> str:
+        if self._card:
+            return self._card.get("description", f"Remote agent at {self.url}")
+        return f"Remote A2A agent at {self.url}"
+
+    def schema(self) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "prompt": {"type": "string", "description": "Task for the remote agent."}
+                    },
+                    "required": ["prompt"],
+                },
+            },
+        }
+
+    async def execute(self, ctx: RunContext, *, prompt: str) -> Any:
+        result = await self.run(prompt, identity=ctx.identity)
+        ctx.usage.add(result.usage)
+        return result.output
+
+
+def _extract_text(task: dict[str, Any]) -> str:
+    """Pull the text out of an A2A task's artifacts."""
+    parts_text: list[str] = []
+    for artifact in task.get("artifacts", []):
+        for part in artifact.get("parts", []):
+            if "text" in part:
+                parts_text.append(part["text"])
+    return "\n".join(parts_text)
+
+
+__all__ = ["RemoteAgent"]