agentlatch 0.1.0__py3-none-any.whl

agentlatch/__init__.py

@@ -0,0 +1,34 @@
+"""AgentLatch — Terminal-native agent resilience middleware.
+
+Two decorators are all you need::
+
+    from agentlatch import profile_agent, safe_tool
+
+    @safe_tool
+    def query_db(sql: str) -> str:
+        ...
+
+    @profile_agent
+    def run_agent():
+        result = query_db("SELECT ...")
+        ...
+"""
+
+from agentlatch.config import is_dev_mode, set_dev_mode
+from agentlatch.decorators import profile_agent, safe_tool
+from agentlatch.renderer import render_flamegraph
+from agentlatch.sampler import sample_response
+from agentlatch.tracker import TraceEvent, get_trace
+
+__version__ = "0.1.0"
+__all__ = [
+    "profile_agent",
+    "safe_tool",
+    "render_flamegraph",
+    "sample_response",
+    "is_dev_mode",
+    "set_dev_mode",
+    "TraceEvent",
+    "get_trace",
+    "__version__",
+]

agentlatch/_types.py

@@ -0,0 +1,26 @@
+"""Shared types and enumerations for AgentLatch."""
+
+from __future__ import annotations
+
+import enum
+from typing import Any, TypeAlias
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class EventStatus(enum.Enum):
+    """Status of a traced execution event."""
+
+    SUCCESS = "success"
+    ERROR = "error"
+    TIMEOUT = "timeout"
+
+
+# ---------------------------------------------------------------------------
+# Type Aliases
+# ---------------------------------------------------------------------------
+
+ErrorPayload: TypeAlias = dict[str, Any]
+"""Structured error information returned to the LLM instead of raising."""

agentlatch/banner.py

@@ -0,0 +1,336 @@
+"""Startup banner animation — Claude Code-style cosmic reveal.
+
+Displays a large ASCII art banner spelling AGENT / LATCH in bold block
+letters, surrounded by cosmic cloud formations and scattered stars.
+The entire scene progressively "decrypts" from noise into the final art
+via a diagonal sweep, followed by a typing effect for the welcome message.
+
+Gracefully degrades in non-TTY environments (CI, Docker, piped output).
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import sys
+import time
+
+from rich.console import Console
+from rich.live import Live
+from rich.text import Text
+
+# ---------------------------------------------------------------------------
+# Block Letter Definitions  (each char is exactly 6 columns wide)
+# ---------------------------------------------------------------------------
+
+_LETTER_DATA: dict[str, list[str]] = {
+    "A": [" ████ ", "██  ██", "██████", "██  ██", "██  ██"],
+    "G": [" █████", "██    ", "██ ███", "██  ██", " █████"],
+    "E": ["██████", "██    ", "████  ", "██    ", "██████"],
+    "N": ["██  ██", "███ ██", "██████", "██ ███", "██  ██"],
+    "T": ["██████", "  ██  ", "  ██  ", "  ██  ", "  ██  "],
+    "L": ["██    ", "██    ", "██    ", "██    ", "██████"],
+    "C": [" █████", "██    ", "██    ", "██    ", " █████"],
+    "H": ["██  ██", "██  ██", "██████", "██  ██", "██  ██"],
+}
+
+
+def _build_word_lines(word: str, indent: int = 4, gap: int = 2) -> list[str]:
+    """Render a word as 5 lines of block-letter ASCII art."""
+    pad = " " * indent
+    sep = " " * gap
+    return [pad + sep.join(_LETTER_DATA[ch][row] for ch in word) for row in range(5)]
+
+
+# ---------------------------------------------------------------------------
+# Build the Full Art Scene
+# ---------------------------------------------------------------------------
+
+_ATMOS_TOP: list[str] = [
+    "...............................................................",
+    "",
+    "  *        ░░░░░░                              █████▓▓░       ",
+    "         ░░░░░░░░░░              *           ███▓░   ░░       ",
+    "       ░░░░░░░░░░░░░░░░                      ███▓░            ",
+    "    *                        *                 ██▓░░    ▓     ",
+    "                                                ░▓▓██▓▓░     ",
+    "",
+]
+
+_AGENT_LINES: list[str] = _build_word_lines("AGENT")
+_DIVIDER: list[str] = ["                       ⚡"]
+_LATCH_LINES: list[str] = _build_word_lines("LATCH")
+
+_ATMOS_BOTTOM: list[str] = [
+    "",
+    " *                       ░░░░                          *      ",
+    "                       ░░░░░░░░                               ",
+    "                     ░░░░░░░░░░░░░░                           ",
+    "                                                *             ",
+    "...............................................................",
+]
+
+_ART_LINES: list[str] = (
+    _ATMOS_TOP + _AGENT_LINES + _DIVIDER + _LATCH_LINES + _ATMOS_BOTTOM
+)
+
+# Pre-compute which rows are ASCII-art text (not atmosphere).
+_TEXT_ROW_START = len(_ATMOS_TOP)
+_TEXT_ROW_END = _TEXT_ROW_START + len(_AGENT_LINES) + len(_DIVIDER) + len(_LATCH_LINES)
+_TEXT_ROWS: set[int] = set(range(_TEXT_ROW_START, _TEXT_ROW_END))
+
+# ---------------------------------------------------------------------------
+# Welcome / Tagline
+# ---------------------------------------------------------------------------
+
+_WELCOME_LINE = "  Terminal-native agent resilience middleware  v0.1.0"
+_READY_LINE = "  Let's get started."
+
+# ---------------------------------------------------------------------------
+# Noise / Animation Config
+# ---------------------------------------------------------------------------
+
+_NOISE_CHARS: list[str] = [
+    "█",
+    "▓",
+    "▒",
+    "░",
+    "#",
+    "@",
+    "%",
+    "&",
+    "╬",
+    "╠",
+    "╣",
+    "╋",
+    "┃",
+    "┫",
+]
+
+_TOTAL_FRAMES = 32  # total decryption frames
+_FRAME_DELAY = 0.015  # ~15ms per frame → ~480ms total
+_TYPING_DELAY = 0.025  # per character for welcome text
+
+# ---------------------------------------------------------------------------
+# State
+# ---------------------------------------------------------------------------
+
+_banner_shown: bool = False
+
+# ---------------------------------------------------------------------------
+# Style Mapping
+# ---------------------------------------------------------------------------
+
+
+def _char_style(char: str, row: int) -> str:
+    """Determine the Rich style for a resolved character."""
+    if char == "⚡":
+        return "bold bright_yellow"
+    if char == "*":
+        return "bright_yellow"
+    if char == ".":
+        return "dim white"
+
+    # ASCII-art text rows → bold block letters
+    if row in _TEXT_ROWS:
+        return "bold bright_cyan"
+
+    # Atmosphere gradient (clouds with depth)
+    if char == "█":
+        return "bright_cyan"
+    if char == "▓":
+        return "cyan"
+    if char == "▒":
+        return "dim cyan"
+    if char == "░":
+        return "dim bright_white"
+
+    return "white"
+
+
+# ---------------------------------------------------------------------------
+# Resolve Map — when each character "decrypts"
+# ---------------------------------------------------------------------------
+
+
+def _build_resolve_map(
+    lines: list[str],
+    total_frames: int,
+) -> list[list[int]]:
+    """Assign a resolve-frame to every non-space character.
+
+    Dot borders flicker in first (frames 0–3).
+    The rest decrypts in a diagonal sweep (top-left → bottom-right)
+    with random jitter for an organic feel.
+    """
+    max_row = len(lines)
+    max_col = max((len(ln) for ln in lines), default=1)
+
+    rmap: list[list[int]] = []
+    for row, line in enumerate(lines):
+        row_map: list[int] = []
+        for col, ch in enumerate(line):
+            if ch == " ":
+                row_map.append(-1)  # always transparent
+            elif ch == ".":
+                row_map.append(random.randint(0, 3))  # dots early
+            else:
+                row_pct = row / max(max_row - 1, 1)
+                col_pct = col / max(max_col - 1, 1)
+                progress = row_pct * 0.55 + col_pct * 0.45
+                base = int(progress * (total_frames - 6)) + 3
+                jitter = random.randint(-3, 3)
+                frame = max(2, min(total_frames - 1, base + jitter))
+                row_map.append(frame)
+        rmap.append(row_map)
+    return rmap
+
+
+# ---------------------------------------------------------------------------
+# Frame Renderer
+# ---------------------------------------------------------------------------
+
+
+def _render_frame(
+    frame: int,
+    lines: list[str],
+    resolve_map: list[list[int]],
+) -> Text:
+    """Produce one animation frame as a Rich Text object."""
+    output = Text()
+
+    for row, line in enumerate(lines):
+        for col, ch in enumerate(line):
+            if ch == " ":
+                output.append(" ")
+                continue
+
+            resolve_at = resolve_map[row][col]
+
+            if frame >= resolve_at:
+                style = _char_style(ch, row)
+                output.append(ch, style=style)
+            else:
+                noise = random.choice(_NOISE_CHARS)
+                output.append(noise, style="dim bright_white")
+
+        output.append("\n")
+
+    return output
+
+
+# ---------------------------------------------------------------------------
+# Typing Effect
+# ---------------------------------------------------------------------------
+
+
+def _type_text(console: Console, text: str, style: str, delay: float) -> None:
+    """Print text character-by-character with a typing effect."""
+    for ch in text:
+        console.print(ch, end="", style=style, highlight=False)
+        time.sleep(delay)
+    console.print()  # newline
+
+
+# ---------------------------------------------------------------------------
+# Interactive Animation
+# ---------------------------------------------------------------------------
+
+
+def _play_animation(console: Console) -> None:
+    """Run the full cosmic reveal sequence."""
+    resolve_map = _build_resolve_map(_ART_LINES, _TOTAL_FRAMES)
+
+    # Phase 1: Decryption sweep of the ASCII art.
+    with Live(
+        _render_frame(0, _ART_LINES, resolve_map),
+        console=console,
+        refresh_per_second=62,
+        transient=True,
+    ) as live:
+        for frame in range(_TOTAL_FRAMES):
+            live.update(_render_frame(frame, _ART_LINES, resolve_map))
+            time.sleep(_FRAME_DELAY)
+
+    # Print the final, fully-resolved art (stays on screen).
+    final_art = _render_frame(_TOTAL_FRAMES, _ART_LINES, resolve_map)
+    console.print(final_art, end="")
+
+    # Phase 2: Welcome text types in.
+    console.print()
+    _type_text(console, _WELCOME_LINE, "bright_white", _TYPING_DELAY)
+    time.sleep(0.12)
+    _type_text(console, _READY_LINE, "dim bright_green", _TYPING_DELAY * 0.7)
+    console.print()
+
+
+# ---------------------------------------------------------------------------
+# Non-TTY Fallback
+# ---------------------------------------------------------------------------
+
+
+def _print_fallback(console: Console) -> None:
+    """Clean output for CI / Docker / piped environments."""
+    # Still show the ASCII art, just without animation.
+    for row, line in enumerate(_ART_LINES):
+        styled = Text()
+        for ch in line:
+            if ch == " ":
+                styled.append(" ")
+            else:
+                styled.append(ch, style=_char_style(ch, row))
+        console.print(styled)
+
+    console.print()
+    console.print(_WELCOME_LINE, style="bright_white")
+    console.print(_READY_LINE, style="dim bright_green")
+    console.print()
+
+
+# ---------------------------------------------------------------------------
+# Environment Detection
+# ---------------------------------------------------------------------------
+
+
+def _is_interactive() -> bool:
+    """Return True only if stdout is a real interactive terminal."""
+    if not sys.stdout.isatty():
+        return False
+    if os.environ.get("CI", "").lower() in ("true", "1"):
+        return False
+    if os.environ.get("TERM", "") == "dumb":
+        return False
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def initialize_latch(console: Console | None = None) -> None:
+    """Display the AgentLatch startup banner (once per process).
+
+    - **Interactive terminal**: plays the full cosmic decryption animation
+      with block-letter ASCII art reveal and typing effect.
+    - **Non-TTY / CI / Docker**: prints the static colored art without animation.
+
+    Calling this more than once per process is a no-op.
+    """
+    global _banner_shown
+    if _banner_shown:
+        return
+    _banner_shown = True
+
+    con = console or Console()
+
+    if _is_interactive():
+        _play_animation(con)
+    else:
+        _print_fallback(con)
+
+
+def reset_banner() -> None:
+    """Reset the banner flag.  Useful for tests."""
+    global _banner_shown
+    _banner_shown = False

agentlatch/config.py

@@ -0,0 +1,45 @@
+"""Environment detection and runtime configuration.
+
+Controls whether AgentLatch renders ASCII visuals (banner, flamegraph) to
+the terminal.  In production / server environments, visuals are suppressed
+automatically — only the structured data (headers, JSON profile) is emitted.
+
+**Rules:**
+
+*  ``AGENTLATCH_ENV=production``  → visuals OFF
+*  ``AGENTLATCH_ENV=development`` → visuals ON  (default when unset)
+*  Programmatic override via ``set_dev_mode(True/False)``
+"""
+
+from __future__ import annotations
+
+import os
+
+_dev_mode_override: bool | None = None
+
+
+def is_dev_mode() -> bool:
+    """Return ``True`` if ASCII visuals should be rendered.
+
+    Checks (in priority order):
+    1. Programmatic override via :func:`set_dev_mode`.
+    2. ``AGENTLATCH_ENV`` environment variable (``production`` → False).
+    3. Default: ``True`` (development mode).
+    """
+    if _dev_mode_override is not None:
+        return _dev_mode_override
+
+    env = os.environ.get("AGENTLATCH_ENV", "development").lower().strip()
+    return env != "production"
+
+
+def set_dev_mode(enabled: bool) -> None:
+    """Programmatically force dev mode on or off."""
+    global _dev_mode_override
+    _dev_mode_override = enabled
+
+
+def reset_dev_mode() -> None:
+    """Clear the programmatic override.  Useful for tests."""
+    global _dev_mode_override
+    _dev_mode_override = None

agentlatch/decorators.py

@@ -0,0 +1,250 @@
+"""Decorators that form AgentLatch's public API.
+
+* ``@safe_tool``     — wraps tool functions with error interception & timing.
+* ``@profile_agent`` — wraps the outer agent loop with tracing & visualization.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import functools
+import inspect
+import json
+from collections.abc import Callable
+from typing import Any, TypeVar, overload
+
+from agentlatch._types import ErrorPayload, EventStatus
+from agentlatch.sampler import sample_response
+from agentlatch.tracker import (
+    end_child,
+    finalize_trace,
+    get_trace,
+    init_trace,
+    start_child,
+)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+# =====================================================================
+# @safe_tool
+# =====================================================================
+
+
+def _build_error_payload(exc: Exception) -> ErrorPayload:
+    """Translate a raw Python exception into an LLM-friendly JSON dict."""
+    return {
+        "status": "error",
+        "error_type": type(exc).__name__,
+        "message": str(exc),
+        "instruction": (
+            "The tool execution failed. Review your parameters and "
+            "retry with corrected inputs."
+        ),
+    }
+
+
+def _build_timeout_payload(name: str, timeout: float) -> ErrorPayload:
+    """Payload returned when a tool exceeds its allowed time budget."""
+    return {
+        "status": "error",
+        "error_type": "TimeoutError",
+        "message": f"Tool '{name}' exceeded the {timeout}s timeout.",
+        "instruction": (
+            "The tool timed out. Consider simplifying the request or "
+            "breaking it into smaller steps."
+        ),
+    }
+
+
+@overload
+def safe_tool(func: F) -> F: ...
+
+
+@overload
+def safe_tool(
+    *,
+    timeout: float | None = ...,
+    on_fail: str = ...,
+    max_response_tokens: int | None = ...,
+    sample_rows: int | None = ...,
+) -> Callable[[F], F]: ...
+
+
+def safe_tool(
+    func: F | None = None,
+    *,
+    timeout: float | None = None,
+    on_fail: str = "instruct_llm",
+    max_response_tokens: int | None = None,
+    sample_rows: int | None = None,
+) -> F | Callable[[F], F]:
+    """Decorator that makes a tool function resilient and observable.
+
+    Can be used bare (``@safe_tool``) or with arguments
+    (``@safe_tool(timeout=5.0, sample_rows=10)``).
+
+    On exception the decorated function returns a JSON error string to the
+    caller (typically the LLM) instead of raising, preventing silent crashes.
+
+    Args:
+        timeout:              Optional wall-clock budget in seconds.
+        on_fail:              Error strategy (currently ``"instruct_llm"``).
+        max_response_tokens:  Approximate token ceiling for responses.
+        sample_rows:          If the response contains a list, keep only
+                              the first N elements.
+    """
+
+    def decorator(fn: F) -> F:
+        if inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                # If no active trace, still protect — just skip timing.
+                trace_active = get_trace() is not None
+                event = start_child(fn.__name__) if trace_active else None
+
+                try:
+                    if timeout is not None:
+                        result = await asyncio.wait_for(
+                            fn(*args, **kwargs), timeout=timeout
+                        )
+                    else:
+                        result = await fn(*args, **kwargs)
+                except asyncio.TimeoutError:
+                    payload = _build_timeout_payload(fn.__name__, timeout or 0.0)
+                    if event:
+                        end_child(event, EventStatus.TIMEOUT, payload)
+                    return json.dumps(payload)
+                except Exception as exc:
+                    payload = _build_error_payload(exc)
+                    if event:
+                        end_child(event, EventStatus.ERROR, payload)
+                    return json.dumps(payload)
+                else:
+                    result = sample_response(
+                        result,
+                        max_tokens=max_response_tokens,
+                        sample_rows=sample_rows,
+                    )
+                    if event:
+                        end_child(event, EventStatus.SUCCESS)
+                    return result
+
+            return async_wrapper  # type: ignore[return-value]
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                trace_active = get_trace() is not None
+                event = start_child(fn.__name__) if trace_active else None
+
+                try:
+                    if timeout is not None:
+                        with concurrent.futures.ThreadPoolExecutor(
+                            max_workers=1
+                        ) as pool:
+                            future = pool.submit(fn, *args, **kwargs)
+                            result = future.result(timeout=timeout)
+                    else:
+                        result = fn(*args, **kwargs)
+                except concurrent.futures.TimeoutError:
+                    payload = _build_timeout_payload(fn.__name__, timeout or 0.0)
+                    if event:
+                        end_child(event, EventStatus.TIMEOUT, payload)
+                    return json.dumps(payload)
+                except Exception as exc:
+                    payload = _build_error_payload(exc)
+                    if event:
+                        end_child(event, EventStatus.ERROR, payload)
+                    return json.dumps(payload)
+                else:
+                    result = sample_response(
+                        result,
+                        max_tokens=max_response_tokens,
+                        sample_rows=sample_rows,
+                    )
+                    if event:
+                        end_child(event, EventStatus.SUCCESS)
+                    return result
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    # Handle bare @safe_tool (no parentheses) vs @safe_tool(timeout=5)
+    if func is not None:
+        return decorator(func)
+    return decorator  # type: ignore[return-value]
+
+
+# =====================================================================
+# @profile_agent
+# =====================================================================
+
+
+def profile_agent(
+    func: F | None = None,
+    *,
+    name: str | None = None,
+) -> F | Callable[[F], F]:
+    """Decorator for the main agent loop.
+
+    Initializes the trace context, runs the agent, then renders the
+    execution flamegraph to the terminal.
+
+    Can be used bare (``@profile_agent``) or with a custom name
+    (``@profile_agent(name="MyAgent")``).
+    """
+
+    def decorator(fn: F) -> F:
+        label = name or fn.__name__
+
+        if inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                # Import here to avoid circular imports at module level.
+                from agentlatch.banner import initialize_latch
+                from agentlatch.config import is_dev_mode
+                from agentlatch.renderer import render_flamegraph
+
+                if is_dev_mode():
+                    initialize_latch()
+                init_trace(label)
+
+                try:
+                    result = await fn(*args, **kwargs)
+                finally:
+                    trace = finalize_trace()
+                    if is_dev_mode():
+                        render_flamegraph(trace)
+
+                return result
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                from agentlatch.banner import initialize_latch
+                from agentlatch.config import is_dev_mode
+                from agentlatch.renderer import render_flamegraph
+
+                if is_dev_mode():
+                    initialize_latch()
+                init_trace(label)
+
+                try:
+                    result = fn(*args, **kwargs)
+                finally:
+                    trace = finalize_trace()
+                    if is_dev_mode():
+                        render_flamegraph(trace)
+
+                return result
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    if func is not None:
+        return decorator(func)
+    return decorator  # type: ignore[return-value]