kaboo-workflows 0.9.0__py3-none-any.whl

kaboo_workflows/__init__.py

@@ -0,0 +1,66 @@
+"""kaboo-workflows — YAML-driven multi-agent orchestration with AG-UI/CopilotKit support."""
+
+from __future__ import annotations
+
+from .config import (
+    AppConfig,
+    ConfigInput,
+    ResolvedConfig,
+    ResolvedInfra,
+    load,
+    load_config,
+    load_session,
+    resolve_infra,
+)
+from .config.resolvers.orchestrations import OrchestrationBuilder
+from .exceptions import (
+    CircularDependencyError,
+    ConfigurationError,
+    ImportResolutionError,
+    SchemaValidationError,
+    UnresolvedReferenceError,
+)
+from .hooks import EventPublisher, MaxToolCallsGuard, StopGuard, ToolNameSanitizer
+from .mcp import MCPLifecycle, create_mcp_client, create_mcp_server
+from .renderers import AnsiRenderer
+from .tools import (
+    node_as_async_tool,
+    node_as_tool,
+    serialize_multiagent_result,
+)
+from .types import EventType, StreamEvent
+from .utils import cli_errors
+from .wire import EventQueue, make_event_queue
+
+__all__ = [
+    "AnsiRenderer",
+    "AppConfig",
+    "CircularDependencyError",
+    "ConfigInput",
+    "ConfigurationError",
+    "EventPublisher",
+    "EventQueue",
+    "EventType",
+    "ImportResolutionError",
+    "MCPLifecycle",
+    "MaxToolCallsGuard",
+    "OrchestrationBuilder",
+    "ResolvedConfig",
+    "ResolvedInfra",
+    "SchemaValidationError",
+    "StopGuard",
+    "StreamEvent",
+    "ToolNameSanitizer",
+    "UnresolvedReferenceError",
+    "cli_errors",
+    "create_mcp_client",
+    "create_mcp_server",
+    "load",
+    "load_config",
+    "load_session",
+    "make_event_queue",
+    "node_as_async_tool",
+    "node_as_tool",
+    "resolve_infra",
+    "serialize_multiagent_result",
+]

kaboo_workflows/_context.py

@@ -0,0 +1,135 @@
+"""Per-conversation activity context.
+
+Carries the current ``thread_id`` / ``run_id`` through the async call stack so
+that :class:`~kaboo_workflows.hooks.EventPublisher` can stamp every
+:class:`~kaboo_workflows.types.StreamEvent` with the conversation it belongs to.
+
+This is what allows a single server process to serve multiple concurrent
+conversations without their activity streams clobbering one another. On
+AgentCore (one microVM per session) it is redundant but harmless; on a plain
+self-hosted process it is what makes per-conversation routing correct.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import contextvars
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from typing import Any
+
+_current_thread_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "kaboo_thread_id", default=None
+)
+_current_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "kaboo_run_id", default=None
+)
+_current_turn_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "kaboo_turn_id", default=None
+)
+_current_delegation_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "kaboo_delegation_id", default=None
+)
+
+
+@dataclass
+class HistoryExchange:
+    """Request-scoped carrier for client-driven, per-agent history.
+
+    ``inbound`` holds the transcripts the client sent this run (parsed from
+    ``RunAgentInput.state['kaboo_history']``), keyed per agent bucket.
+    ``outbound`` accumulates the transcripts agents produce this run; it is
+    merged back into the outgoing ``STATE_SNAPSHOT`` so the client persists
+    and replays them next turn.
+
+    One instance is created per request and shared by reference across the
+    request's tasks via :data:`_current_history`, so a sub-agent hook running
+    in a child task and the snapshot enricher running in the parent see the
+    same buffers.
+    """
+
+    inbound: dict[str, list[Any]] = field(default_factory=dict)
+    outbound: dict[str, list[Any]] = field(default_factory=dict)
+
+
+_current_history: contextvars.ContextVar[HistoryExchange | None] = contextvars.ContextVar(
+    "kaboo_history_exchange", default=None
+)
+
+
+def set_activity_context(
+    thread_id: str | None,
+    run_id: str | None = None,
+    turn_id: str | None = None,
+) -> None:
+    """Bind the current conversation's ``thread_id`` / ``run_id`` / ``turn_id``.
+
+    Call at the start of each request handler. Tasks created afterwards
+    (via ``asyncio.create_task``) inherit a copy of this context, so nested
+    agent/delegate activity is attributed to the right conversation.
+
+    ``turn_id`` identifies one logical user turn and is **stable across an
+    interrupt/resume** (unlike ``run_id``, which changes on the resume POST).
+    It lets the UI bind every group a turn produced — even ones that first
+    appear only after a resume — to that turn's single chat reply.
+    """
+    _current_thread_id.set(thread_id)
+    _current_run_id.set(run_id)
+    _current_turn_id.set(turn_id)
+
+
+def get_thread_id() -> str | None:
+    """Return the current conversation's ``thread_id`` (or ``None``)."""
+    return _current_thread_id.get()
+
+
+def get_run_id() -> str | None:
+    """Return the current conversation's ``run_id`` (or ``None``)."""
+    return _current_run_id.get()
+
+
+def get_turn_id() -> str | None:
+    """Return the current logical turn's id (stable across resume, or ``None``)."""
+    return _current_turn_id.get()
+
+
+@contextlib.contextmanager
+def delegation_context(tool_use_id: str | None) -> Iterator[None]:
+    """Bind the delegating tool-call id for the wrapped node invocation.
+
+    A coordinator delegates to a sub-agent by calling a tool; that tool call has
+    a ``toolUseId`` which is the same id the frontend (AG-UI / CopilotKit) sees
+    for the inline card. Binding it here lets
+    :class:`~kaboo_workflows.hooks.EventPublisher` stamp the sub-agent's stream
+    group with that id, so the card can be correlated to its activity group by a
+    stable key instead of by arrival order.
+
+    Nesting is token-based: a sub-agent that itself delegates temporarily
+    overrides the id and restores the parent's on exit.
+    """
+    token = _current_delegation_id.set(tool_use_id)
+    try:
+        yield
+    finally:
+        _current_delegation_id.reset(token)
+
+
+def get_delegation_id() -> str | None:
+    """Return the current delegating tool-call id (or ``None``)."""
+    return _current_delegation_id.get()
+
+
+def set_history_exchange(exchange: HistoryExchange | None) -> None:
+    """Bind the current request's :class:`HistoryExchange`.
+
+    Call at the start of each request handler, before creating the run task.
+    Tasks created afterwards inherit a copy of this context that references
+    the same exchange object, so hook writes made in child tasks are visible
+    to the parent snapshot enricher.
+    """
+    _current_history.set(exchange)
+
+
+def get_history_exchange() -> HistoryExchange | None:
+    """Return the current request's :class:`HistoryExchange` (or ``None``)."""
+    return _current_history.get()

kaboo_workflows/adapters/__init__.py

@@ -0,0 +1,7 @@
+"""AG-UI adapter for kaboo-workflows."""
+
+from __future__ import annotations
+
+from .agui import create_agui_app
+
+__all__ = ["create_agui_app"]

kaboo_workflows/adapters/_activity.py

@@ -0,0 +1,172 @@
+"""Per-conversation activity state builder.
+
+Folds kaboo ``StreamEvent``s into a hierarchical activity snapshot
+(``{groups: {...}}``) keyed by ``thread_id``. The snapshot is emitted onto the
+AG-UI run stream as ``ACTIVITY_SNAPSHOT`` events (see ``agui.py``), so there is
+no separate SSE endpoint or subscriber machinery here — this is state only.
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+import threading
+from typing import Any
+
+from kaboo_workflows.types import EventType, StreamEvent
+
+logger = logging.getLogger(__name__)
+
+
+def _update_activity_state(state: dict[str, Any], event: StreamEvent) -> bool:
+    """Mutate *state* based on an incoming kaboo StreamEvent.
+
+    Maintains a ``timeline`` list that preserves the chronological order of
+    text segments and tool calls as they arrive. Adjacent TOKEN events are
+    coalesced into a single text entry.
+
+    Returns:
+        ``True`` if the state changed (and a fresh snapshot should be emitted),
+        ``False`` for events that carry no stream group (nothing to render).
+    """
+    group = event.data.get("stream_group", "")
+    if not group:
+        return False
+
+    groups = state["groups"]
+
+    if event.type == EventType.STREAM_GROUP_START:
+        existing = groups.get(group)
+        if existing is not None:
+            existing["status"] = "active"
+            existing.pop("interrupt", None)
+        else:
+            groups[group] = {
+                "title": event.data.get("stream_title", ""),
+                "agentName": event.data.get("agent_name", ""),
+                "parentGroup": event.data.get("parent_group"),
+                "toolCallId": event.data.get("tool_call_id"),
+                "runId": event.data.get("run_id"),
+                "turnId": event.data.get("turn_id"),
+                "task": event.data.get("task"),
+                "isChatReply": bool(event.data.get("is_chat_reply", False)),
+                "inlineChatOwner": bool(event.data.get("inline_chat_owner", False)),
+                "status": "active",
+                "tools": [],
+                "tokens": "",
+                "timeline": [],
+            }
+            logger.info(
+                "activity.STREAM_GROUP_START group=%s agent=%s parent=%s "
+                "tool_call_id=%s run_id=%s is_chat_reply=%s task=%r thread_id=%s",
+                group,
+                event.data.get("agent_name"),
+                event.data.get("parent_group"),
+                event.data.get("tool_call_id"),
+                event.data.get("run_id"),
+                event.data.get("is_chat_reply"),
+                event.data.get("task"),
+                event.data.get("thread_id"),
+            )
+    elif event.type == EventType.AGENT_START:
+        if group in groups:
+            groups[group]["status"] = "active"
+            groups[group].pop("interrupt", None)
+    elif event.type == EventType.STREAM_GROUP_END:
+        if group in groups:
+            groups[group]["status"] = event.data.get("status", "completed")
+    elif event.type == EventType.TOOL_START:
+        if group in groups:
+            tool_use_id = event.data.get("tool_use_id", "")
+            existing_tool = next(
+                (t for t in groups[group]["tools"] if t["toolUseId"] == tool_use_id),
+                None,
+            )
+            if existing_tool is not None:
+                existing_tool["status"] = "running"
+            else:
+                tool_entry = {
+                    "toolUseId": tool_use_id,
+                    "toolName": event.data.get("tool_name", ""),
+                    "toolLabel": event.data.get("tool_label", ""),
+                    "toolInput": event.data.get("tool_input"),
+                    "status": "running",
+                }
+                groups[group]["tools"].append(tool_entry)
+                groups[group]["timeline"].append({"type": "tool", "tool": tool_entry})
+    elif event.type == EventType.TOOL_END:
+        if group in groups:
+            tool_use_id = event.data.get("tool_use_id")
+            status = event.data.get("status", "done")
+            result = event.data.get("tool_result", "")
+            for tool in groups[group]["tools"]:
+                if tool["toolUseId"] == tool_use_id:
+                    tool["status"] = status
+                    tool["toolResult"] = result
+                    break
+    elif event.type == EventType.TOKEN:
+        if group in groups:
+            text = event.data.get("text", "")
+            groups[group]["tokens"] += text
+            timeline = groups[group]["timeline"]
+            if timeline and timeline[-1]["type"] == "text":
+                timeline[-1]["text"] += text
+            else:
+                timeline.append({"type": "text", "text": text})
+    elif event.type == EventType.AGENT_COMPLETE:
+        if group in groups and "structured_output" in event.data:
+            groups[group]["structuredOutput"] = event.data["structured_output"]
+            groups[group]["outputSchemaName"] = event.data.get("output_schema_name")
+    elif event.type == EventType.INTERRUPT:
+        if group in groups:
+            groups[group]["status"] = "interrupted"
+            groups[group]["interrupt"] = {
+                "id": event.data.get("interrupt_id"),
+                "name": event.data.get("name"),
+                "reason": event.data.get("reason"),
+            }
+    elif event.type == EventType.ERROR:
+        if group in groups:
+            groups[group]["status"] = "error"
+
+    return True
+
+
+class ActivityRegistry:
+    """Thread-scoped builder of hierarchical activity state.
+
+    One instance is created per :func:`create_agui_app` and captured in the
+    endpoint closures. All mutation is guarded by a lock so concurrent request
+    tasks stay consistent. Snapshots are read via :meth:`snapshot` and emitted
+    on the run stream as ``ACTIVITY_SNAPSHOT`` events.
+    """
+
+    def __init__(self) -> None:
+        self._states: dict[str, dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+    def _state_for(self, thread_id: str) -> dict[str, Any]:
+        return self._states.setdefault(thread_id, {"groups": {}})
+
+    def apply(self, thread_id: str, event: StreamEvent) -> bool:
+        """Fold *event* into *thread_id*'s state.
+
+        Returns ``True`` when the state changed (caller should emit a fresh
+        snapshot), ``False`` for events that carry no stream group.
+        """
+        with self._lock:
+            state = self._state_for(thread_id)
+            return _update_activity_state(state, event)
+
+    def snapshot(self, thread_id: str) -> dict[str, Any]:
+        """Return a deep copy of the activity state for *thread_id* (``{groups: {...}}``).
+
+        Deep-copied because the returned snapshot is encoded later on the SSE
+        loop, after the lock is released, while subsequent events may still be
+        mutating the live state.
+        """
+        with self._lock:
+            state = self._states.get(thread_id)
+            if state is None:
+                return {"groups": {}}
+            return copy.deepcopy(state)