prompture 0.0.36.dev1__py3-none-any.whl → 0.0.37.dev1__py3-none-any.whl

prompture/groups.py

@@ -0,0 +1,530 @@
+"""Synchronous multi-agent group coordination.
+
+Provides :class:`SequentialGroup`, :class:`LoopGroup`,
+:class:`RouterAgent`, and :class:`GroupAsAgent` for composing
+multiple agents into deterministic workflows.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+from collections.abc import Callable
+from typing import Any
+
+from .agent_types import AgentResult, AgentState
+from .group_types import (
+    AgentError,
+    ErrorPolicy,
+    GroupCallbacks,
+    GroupResult,
+    GroupStep,
+    _aggregate_usage,
+)
+
+logger = logging.getLogger("prompture.groups")
+
+
+# ------------------------------------------------------------------
+# State injection helper
+# ------------------------------------------------------------------
+
+
+def _inject_state(template: str, state: dict[str, Any]) -> str:
+    """Replace ``{key}`` placeholders with state values.
+
+    Unknown keys pass through unchanged so downstream agents can
+    still see the literal placeholder.
+    """
+
+    def _replacer(m: re.Match[str]) -> str:
+        key = m.group(1)
+        if key in state:
+            return str(state[key])
+        return m.group(0)  # leave unchanged
+
+    return re.sub(r"\{(\w+)\}", _replacer, template)
+
+
+# ------------------------------------------------------------------
+# Agent entry normalisation
+# ------------------------------------------------------------------
+
+AgentEntry = Any  # Agent | tuple[Agent, str]
+
+
+def _normalise_agents(agents: list[Any]) -> list[tuple[Any, str | None]]:
+    """Convert a mixed list of ``Agent`` or ``(Agent, prompt_template)`` to uniform tuples."""
+    result: list[tuple[Any, str | None]] = []
+    for item in agents:
+        if isinstance(item, tuple):
+            result.append((item[0], item[1]))
+        else:
+            result.append((item, None))
+    return result
+
+
+def _agent_name(agent: Any, index: int) -> str:
+    """Determine a display name for an agent."""
+    name = getattr(agent, "name", "") or ""
+    return name if name else f"agent_{index}"
+
+
+# ------------------------------------------------------------------
+# SequentialGroup
+# ------------------------------------------------------------------
+
+
+class SequentialGroup:
+    """Execute agents in sequence, passing state between them.
+
+    Each agent's ``output_key`` (if set) writes its output text
+    into the shared state dict, making it available as ``{key}``
+    in subsequent agent prompts.
+
+    Args:
+        agents: List of agents or ``(agent, prompt_template)`` tuples.
+        state: Initial shared state dict.
+        error_policy: How to handle agent failures.
+        max_total_turns: Limit on total agent runs across the sequence.
+        callbacks: Observability hooks.
+        max_total_cost: Budget cap in USD.
+    """
+
+    def __init__(
+        self,
+        agents: list[Any],
+        *,
+        state: dict[str, Any] | None = None,
+        error_policy: ErrorPolicy = ErrorPolicy.fail_fast,
+        max_total_turns: int | None = None,
+        callbacks: GroupCallbacks | None = None,
+        max_total_cost: float | None = None,
+    ) -> None:
+        self._agents = _normalise_agents(agents)
+        self._state: dict[str, Any] = dict(state) if state else {}
+        self._error_policy = error_policy
+        self._max_total_turns = max_total_turns
+        self._callbacks = callbacks or GroupCallbacks()
+        self._max_total_cost = max_total_cost
+        self._stop_requested = False
+
+    def stop(self) -> None:
+        """Request graceful shutdown after the current agent finishes."""
+        self._stop_requested = True
+
+    def save(self, path: str) -> None:
+        """Run and save result to file. Convenience wrapper."""
+        result = self.run()
+        result.save(path)
+
+    def run(self, prompt: str = "") -> GroupResult:
+        """Execute all agents in order."""
+        self._stop_requested = False
+        t0 = time.perf_counter()
+        timeline: list[GroupStep] = []
+        agent_results: dict[str, Any] = {}
+        errors: list[AgentError] = []
+        usage_summaries: list[dict[str, Any]] = []
+        turns = 0
+
+        for idx, (agent, custom_prompt) in enumerate(self._agents):
+            if self._stop_requested:
+                break
+
+            name = _agent_name(agent, idx)
+
+            # Build effective prompt
+            if custom_prompt is not None:
+                effective = _inject_state(custom_prompt, self._state)
+            elif prompt:
+                effective = _inject_state(prompt, self._state)
+            else:
+                effective = ""
+
+            # Check budget
+            if self._max_total_cost is not None:
+                total_so_far = sum(s.get("total_cost", 0.0) for s in usage_summaries)
+                if total_so_far >= self._max_total_cost:
+                    logger.debug("Budget exceeded, stopping group")
+                    break
+
+            # Check max turns
+            if self._max_total_turns is not None and turns >= self._max_total_turns:
+                logger.debug("Max total turns reached")
+                break
+
+            # Fire callback
+            if self._callbacks.on_agent_start:
+                self._callbacks.on_agent_start(name, effective)
+
+            step_t0 = time.perf_counter()
+            try:
+                result = agent.run(effective)
+                duration_ms = (time.perf_counter() - step_t0) * 1000
+                turns += 1
+
+                agent_results[name] = result
+                usage = getattr(result, "run_usage", {})
+                usage_summaries.append(usage)
+
+                # Write to shared state
+                output_key = getattr(agent, "output_key", None)
+                if output_key:
+                    self._state[output_key] = result.output_text
+                    if self._callbacks.on_state_update:
+                        self._callbacks.on_state_update(output_key, result.output_text)
+
+                timeline.append(
+                    GroupStep(
+                        agent_name=name,
+                        step_type="agent_run",
+                        timestamp=step_t0,
+                        duration_ms=duration_ms,
+                        usage_delta=usage,
+                    )
+                )
+
+                if self._callbacks.on_agent_complete:
+                    self._callbacks.on_agent_complete(name, result)
+
+            except Exception as exc:
+                duration_ms = (time.perf_counter() - step_t0) * 1000
+                turns += 1
+                err = AgentError(
+                    agent_name=name,
+                    error=exc,
+                    output_key=getattr(agent, "output_key", None),
+                )
+                errors.append(err)
+                timeline.append(
+                    GroupStep(
+                        agent_name=name,
+                        step_type="agent_error",
+                        timestamp=step_t0,
+                        duration_ms=duration_ms,
+                        error=str(exc),
+                    )
+                )
+
+                if self._callbacks.on_agent_error:
+                    self._callbacks.on_agent_error(name, exc)
+
+                if self._error_policy == ErrorPolicy.fail_fast:
+                    break
+                # continue_on_error / retry_failed: continue to next agent
+
+        elapsed_ms = (time.perf_counter() - t0) * 1000
+        return GroupResult(
+            agent_results=agent_results,
+            aggregate_usage=_aggregate_usage(*usage_summaries),
+            shared_state=dict(self._state),
+            elapsed_ms=elapsed_ms,
+            timeline=timeline,
+            errors=errors,
+            success=len(errors) == 0,
+        )
+
+
+# ------------------------------------------------------------------
+# LoopGroup
+# ------------------------------------------------------------------
+
+
+class LoopGroup:
+    """Repeat a sequence of agents until an exit condition is met.
+
+    Args:
+        agents: List of agents or ``(agent, prompt_template)`` tuples.
+        exit_condition: Callable ``(state, iteration) -> bool``.
+            When it returns ``True`` the loop stops.
+        max_iterations: Hard cap on loop iterations.
+        state: Initial shared state dict.
+        error_policy: How to handle agent failures.
+        callbacks: Observability hooks.
+    """
+
+    def __init__(
+        self,
+        agents: list[Any],
+        *,
+        exit_condition: Callable[[dict[str, Any], int], bool],
+        max_iterations: int = 10,
+        state: dict[str, Any] | None = None,
+        error_policy: ErrorPolicy = ErrorPolicy.fail_fast,
+        callbacks: GroupCallbacks | None = None,
+    ) -> None:
+        self._agents = _normalise_agents(agents)
+        self._exit_condition = exit_condition
+        self._max_iterations = max_iterations
+        self._state: dict[str, Any] = dict(state) if state else {}
+        self._error_policy = error_policy
+        self._callbacks = callbacks or GroupCallbacks()
+        self._stop_requested = False
+
+    def stop(self) -> None:
+        """Request graceful shutdown."""
+        self._stop_requested = True
+
+    def run(self, prompt: str = "") -> GroupResult:
+        """Execute the loop."""
+        self._stop_requested = False
+        t0 = time.perf_counter()
+        timeline: list[GroupStep] = []
+        agent_results: dict[str, Any] = {}
+        errors: list[AgentError] = []
+        usage_summaries: list[dict[str, Any]] = []
+
+        for iteration in range(self._max_iterations):
+            if self._stop_requested:
+                break
+            if self._exit_condition(self._state, iteration):
+                break
+
+            for idx, (agent, custom_prompt) in enumerate(self._agents):
+                if self._stop_requested:
+                    break
+
+                name = _agent_name(agent, idx)
+                result_key = f"{name}_iter{iteration}"
+
+                if custom_prompt is not None:
+                    effective = _inject_state(custom_prompt, self._state)
+                elif prompt:
+                    effective = _inject_state(prompt, self._state)
+                else:
+                    effective = ""
+
+                if self._callbacks.on_agent_start:
+                    self._callbacks.on_agent_start(name, effective)
+
+                step_t0 = time.perf_counter()
+                try:
+                    result = agent.run(effective)
+                    duration_ms = (time.perf_counter() - step_t0) * 1000
+
+                    agent_results[result_key] = result
+                    usage = getattr(result, "run_usage", {})
+                    usage_summaries.append(usage)
+
+                    output_key = getattr(agent, "output_key", None)
+                    if output_key:
+                        self._state[output_key] = result.output_text
+                        if self._callbacks.on_state_update:
+                            self._callbacks.on_state_update(output_key, result.output_text)
+
+                    timeline.append(
+                        GroupStep(
+                            agent_name=name,
+                            step_type="agent_run",
+                            timestamp=step_t0,
+                            duration_ms=duration_ms,
+                            usage_delta=usage,
+                        )
+                    )
+
+                    if self._callbacks.on_agent_complete:
+                        self._callbacks.on_agent_complete(name, result)
+
+                except Exception as exc:
+                    duration_ms = (time.perf_counter() - step_t0) * 1000
+                    err = AgentError(
+                        agent_name=name,
+                        error=exc,
+                        output_key=getattr(agent, "output_key", None),
+                    )
+                    errors.append(err)
+                    timeline.append(
+                        GroupStep(
+                            agent_name=name,
+                            step_type="agent_error",
+                            timestamp=step_t0,
+                            duration_ms=duration_ms,
+                            error=str(exc),
+                        )
+                    )
+
+                    if self._callbacks.on_agent_error:
+                        self._callbacks.on_agent_error(name, exc)
+
+                    if self._error_policy == ErrorPolicy.fail_fast:
+                        break
+
+            # Check if error caused early exit
+            if errors and self._error_policy == ErrorPolicy.fail_fast:
+                break
+
+        elapsed_ms = (time.perf_counter() - t0) * 1000
+        return GroupResult(
+            agent_results=agent_results,
+            aggregate_usage=_aggregate_usage(*usage_summaries),
+            shared_state=dict(self._state),
+            elapsed_ms=elapsed_ms,
+            timeline=timeline,
+            errors=errors,
+            success=len(errors) == 0,
+        )
+
+
+# ------------------------------------------------------------------
+# RouterAgent
+# ------------------------------------------------------------------
+
+_DEFAULT_ROUTING_PROMPT = """Given these specialists:
+{agent_list}
+
+Which should handle this? Reply with ONLY the name.
+
+Request: {prompt}"""
+
+
+class RouterAgent:
+    """LLM-driven router that delegates to the best-matching agent.
+
+    Args:
+        model: Model string for the routing LLM call.
+        agents: List of agents to route between.
+        routing_prompt: Custom prompt template (must include ``{agent_list}``
+            and ``{prompt}`` placeholders).
+        fallback: Agent to use when routing fails.
+        driver: Pre-built driver instance for the routing call.
+    """
+
+    def __init__(
+        self,
+        model: str = "",
+        *,
+        agents: list[Any],
+        routing_prompt: str | None = None,
+        fallback: Any | None = None,
+        driver: Any | None = None,
+        name: str = "",
+        description: str = "",
+        output_key: str | None = None,
+    ) -> None:
+        self._model = model
+        self._driver = driver
+        self._agents = {_agent_name(a, i): a for i, a in enumerate(agents)}
+        self._routing_prompt = routing_prompt or _DEFAULT_ROUTING_PROMPT
+        self._fallback = fallback
+        self.name = name
+        self.description = description
+        self.output_key = output_key
+
+    def run(self, prompt: str, *, deps: Any = None) -> AgentResult:
+        """Route the prompt to the best agent and return its result."""
+        from .conversation import Conversation
+
+        # Build agent list for routing prompt
+        agent_lines = []
+        for name, agent in self._agents.items():
+            desc = getattr(agent, "description", "") or ""
+            agent_lines.append(f"- {name}: {desc}" if desc else f"- {name}")
+        agent_list = "\n".join(agent_lines)
+
+        routing_text = self._routing_prompt.replace("{agent_list}", agent_list).replace("{prompt}", prompt)
+
+        # Single LLM call for routing
+        kwargs: dict[str, Any] = {}
+        if self._driver is not None:
+            kwargs["driver"] = self._driver
+        else:
+            kwargs["model_name"] = self._model
+
+        conv = Conversation(**kwargs)
+        route_response = conv.ask(routing_text)
+
+        # Fuzzy match against known agent names
+        selected = self._fuzzy_match(route_response.strip())
+
+        if selected is not None:
+            return selected.run(prompt, deps=deps) if deps is not None else selected.run(prompt)
+        elif self._fallback is not None:
+            return self._fallback.run(prompt, deps=deps) if deps is not None else self._fallback.run(prompt)
+        else:
+            # Return routing response as fallback
+            return AgentResult(
+                output=route_response,
+                output_text=route_response,
+                messages=conv.messages,
+                usage=conv.usage,
+                state=AgentState.idle,
+            )
+
+    def _fuzzy_match(self, response: str) -> Any | None:
+        """Find the best matching agent name in the LLM response."""
+        response_lower = response.lower().strip()
+
+        # Exact match
+        for name, agent in self._agents.items():
+            if name.lower() == response_lower:
+                return agent
+
+        # Substring match
+        for name, agent in self._agents.items():
+            if name.lower() in response_lower:
+                return agent
+
+        # Word-level match
+        response_words = set(response_lower.split())
+        for name, agent in self._agents.items():
+            name_words = set(name.lower().replace("_", " ").split())
+            if name_words & response_words:
+                return agent
+
+        return None
+
+
+# ------------------------------------------------------------------
+# GroupAsAgent
+# ------------------------------------------------------------------
+
+
+class GroupAsAgent:
+    """Adapter that makes a group behave like an Agent for composability.
+
+    Allows nesting groups inside other groups by presenting the same
+    ``run(prompt) -> AgentResult`` interface.
+
+    Args:
+        group: The group to wrap (SequentialGroup, LoopGroup, etc.).
+        name: Agent identity name.
+        output_key: Shared state key for writing output.
+    """
+
+    def __init__(
+        self,
+        group: Any,
+        *,
+        name: str = "",
+        output_key: str | None = None,
+    ) -> None:
+        self._group = group
+        self.name = name
+        self.output_key = output_key
+        self.description = ""
+
+    def run(self, prompt: str, **kwargs: Any) -> AgentResult:
+        """Run the wrapped group and return an AgentResult."""
+        group_result = self._group.run(prompt)
+
+        # Use the last agent's output text, or the shared state
+        output_text = ""
+        if group_result.agent_results:
+            last_result = list(group_result.agent_results.values())[-1]
+            output_text = getattr(last_result, "output_text", str(last_result))
+
+        return AgentResult(
+            output=output_text,
+            output_text=output_text,
+            messages=[],
+            usage=group_result.aggregate_usage,
+            state=AgentState.idle,
+            run_usage=group_result.aggregate_usage,
+        )
+
+    def stop(self) -> None:
+        """Propagate stop to the wrapped group."""
+        if hasattr(self._group, "stop"):
+            self._group.stop()

prompture/image.py

@@ -0,0 +1,180 @@
+"""Image handling utilities for vision-capable LLM drivers."""
+
+from __future__ import annotations
+
+import base64
+import mimetypes
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+
+@dataclass(frozen=True)
+class ImageContent:
+    """Normalized image representation for vision-capable drivers.
+
+    Attributes:
+        data: Base64-encoded image data.
+        media_type: MIME type (e.g. ``"image/png"``, ``"image/jpeg"``).
+        source_type: How the image is delivered — ``"base64"`` or ``"url"``.
+        url: Original URL when ``source_type`` is ``"url"``.
+    """
+
+    data: str
+    media_type: str
+    source_type: str = "base64"
+    url: str | None = None
+
+
+# Public type alias accepted by all image-aware APIs.
+ImageInput = Union[bytes, str, Path, ImageContent]
+
+# Known data-URI prefix pattern
+_DATA_URI_RE = re.compile(r"^data:(image/[a-zA-Z0-9.+-]+);base64,(.+)$", re.DOTALL)
+
+# Base64 detection heuristic — must look like pure base64 of reasonable length
+_BASE64_RE = re.compile(r"^[A-Za-z0-9+/\n\r]+=*$")
+
+_MIME_FROM_EXT: dict[str, str] = {
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".png": "image/png",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+    ".bmp": "image/bmp",
+    ".svg": "image/svg+xml",
+    ".tiff": "image/tiff",
+    ".tif": "image/tiff",
+}
+
+_MAGIC_BYTES: list[tuple[bytes, str]] = [
+    (b"\x89PNG", "image/png"),
+    (b"\xff\xd8\xff", "image/jpeg"),
+    (b"GIF87a", "image/gif"),
+    (b"GIF89a", "image/gif"),
+    (b"RIFF", "image/webp"),  # WebP starts with RIFF...WEBP
+    (b"BM", "image/bmp"),
+]
+
+
+def _guess_media_type_from_bytes(data: bytes) -> str:
+    """Guess MIME type from the first few bytes of image data."""
+    for magic, mime in _MAGIC_BYTES:
+        if data[: len(magic)] == magic:
+            return mime
+    return "image/png"  # safe fallback
+
+
+def _guess_media_type(path: str) -> str:
+    """Guess MIME type from a file path or URL."""
+    # Strip query strings for URLs
+    clean = path.split("?")[0].split("#")[0]
+    ext = Path(clean).suffix.lower()
+    if ext in _MIME_FROM_EXT:
+        return _MIME_FROM_EXT[ext]
+    guessed = mimetypes.guess_type(clean)[0]
+    return guessed or "image/png"
+
+
+# ------------------------------------------------------------------
+# Constructor functions
+# ------------------------------------------------------------------
+
+
+def image_from_bytes(data: bytes, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` from raw bytes.
+
+    Args:
+        data: Raw image bytes.
+        media_type: MIME type.  Auto-detected from magic bytes when *None*.
+    """
+    if not data:
+        raise ValueError("Image data cannot be empty")
+    b64 = base64.b64encode(data).decode("ascii")
+    mt = media_type or _guess_media_type_from_bytes(data)
+    return ImageContent(data=b64, media_type=mt)
+
+
+def image_from_base64(b64: str, media_type: str = "image/png") -> ImageContent:
+    """Create an :class:`ImageContent` from a base64-encoded string.
+
+    Accepts both raw base64 and ``data:`` URIs.
+    """
+    m = _DATA_URI_RE.match(b64)
+    if m:
+        return ImageContent(data=m.group(2), media_type=m.group(1))
+    return ImageContent(data=b64, media_type=media_type)
+
+
+def image_from_file(path: str | Path, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` by reading a local file.
+
+    Args:
+        path: Path to an image file.
+        media_type: MIME type.  Guessed from extension when *None*.
+    """
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Image file not found: {p}")
+    raw = p.read_bytes()
+    mt = media_type or _guess_media_type(str(p))
+    return image_from_bytes(raw, mt)
+
+
+def image_from_url(url: str, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` referencing a remote URL.
+
+    The image is **not** downloaded — the URL is stored directly so
+    drivers that accept URL-based images can pass it through.  For
+    drivers that require base64, the URL is embedded as a data URI.
+
+    Args:
+        url: Publicly-accessible image URL.
+        media_type: MIME type.  Guessed from the URL when *None*.
+    """
+    mt = media_type or _guess_media_type(url)
+    return ImageContent(data="", media_type=mt, source_type="url", url=url)
+
+
+# ------------------------------------------------------------------
+# Smart constructor
+# ------------------------------------------------------------------
+
+
+def make_image(source: ImageInput) -> ImageContent:
+    """Auto-detect the source type and return an :class:`ImageContent`.
+
+    Accepts:
+    - ``ImageContent`` — returned as-is.
+    - ``bytes`` — base64-encoded with auto-detected MIME.
+    - ``str`` — tries (in order): data URI, URL, file path, raw base64.
+    - ``pathlib.Path`` — read from disk.
+    """
+    if isinstance(source, ImageContent):
+        return source
+
+    if isinstance(source, bytes):
+        return image_from_bytes(source)
+
+    if isinstance(source, Path):
+        return image_from_file(source)
+
+    if isinstance(source, str):
+        # 1. data URI
+        if source.startswith("data:"):
+            return image_from_base64(source)
+
+        # 2. URL
+        if source.startswith(("http://", "https://")):
+            return image_from_url(source)
+
+        # 3. File path (if exists on disk)
+        p = Path(source)
+        if p.exists():
+            return image_from_file(p)
+
+        # 4. Assume raw base64
+        return image_from_base64(source)
+
+    raise TypeError(f"Unsupported image source type: {type(source).__name__}")