flock-core 0.4.543__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/utilities.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+
+"""Built-in utility components for metrics and logging."""
+
+import asyncio
+import contextlib
+import json
+import sys
+import time
+from collections.abc import Iterable, Mapping, MutableMapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from rich.console import Console
+from rich.json import JSON
+from rich.live import Live
+from rich.panel import Panel
+from rich.pretty import Pretty
+from rich.table import Table
+from rich.text import Text
+
+from flock.components import AgentComponent
+
+
+if TYPE_CHECKING:
+    from flock.runtime import Context, EvalInputs, EvalResult
+
+
+class MetricsUtility(AgentComponent):
+    """Records simple runtime metrics per agent execution."""
+
+    name: str | None = "metrics"
+
+    async def on_pre_evaluate(self, agent, ctx: Context, inputs: EvalInputs) -> EvalInputs:
+        ctx.state.setdefault("metrics", {})[f"{agent.name}:start"] = time.perf_counter()
+        return inputs
+
+    async def on_post_evaluate(
+        self, agent, ctx: Context, inputs: EvalInputs, result: EvalResult
+    ) -> EvalResult:
+        metrics = ctx.state.setdefault("metrics", {})
+        start = metrics.get(f"{agent.name}:start")
+        if start:
+            metrics[f"{agent.name}:duration_ms"] = (time.perf_counter() - start) * 1000
+        result.metrics.update({k: v for k, v in metrics.items() if k.endswith("duration_ms")})
+        return result
+
+
+class LoggingUtility(AgentComponent):
+    """Rich-powered logging with optional streaming previews."""
+
+    name: str | None = "logs"
+
+    def __init__(
+        self,
+        console: Console | None = None,
+        *,
+        highlight_json: bool = True,
+        stream_tokens: bool = True,
+    ) -> None:
+        super().__init__()
+        if console is None:
+            console = Console(
+                file=sys.stdout,
+                force_terminal=True,
+                highlight=False,
+                log_time=True,
+                log_path=False,
+            )
+        self._console = console
+        self._highlight_json = highlight_json
+        self._stream_tokens = stream_tokens
+        self._stream_context: dict[str, tuple[asyncio.Queue, asyncio.Task]] = {}
+
+    async def on_initialize(self, agent, ctx: Context) -> None:
+        self._console.log(f"[{agent.name}] start task={ctx.task_id}")
+        await super().on_initialize(agent, ctx)
+
+    async def on_pre_consume(self, agent, ctx: Context, inputs: list[Any]):
+        summary = ", ".join(self._summarize_artifact(art) for art in inputs) or "<none>"
+        self._console.log(f"[{agent.name}] consume n={len(inputs)} artifacts -> {summary}")
+        self._render_artifacts(agent.name, inputs, role="input")
+        return await super().on_pre_consume(agent, ctx, inputs)
+
+    async def on_pre_evaluate(self, agent, ctx: Context, inputs: EvalInputs) -> EvalInputs:
+        if self._stream_tokens:
+            self._maybe_start_stream(agent, ctx)
+        return await super().on_pre_evaluate(agent, ctx, inputs)
+
+    async def on_post_evaluate(
+        self, agent, ctx: Context, inputs: EvalInputs, result: EvalResult
+    ) -> EvalResult:
+        self._render_metrics(agent.name, result.metrics)
+        self._render_artifacts(agent.name, result.artifacts or inputs.artifacts, role="output")
+        if result.logs:
+            self._render_logs(agent.name, result.logs)
+        awaited = await super().on_post_evaluate(agent, ctx, inputs, result)
+        if self._stream_tokens:
+            await self._finalize_stream(agent, ctx)
+        return awaited
+
+    async def on_post_publish(self, agent, ctx: Context, artifact):
+        visibility = getattr(artifact.visibility, "kind", "Public")
+        subtitle = f"visibility={visibility}"
+        panel = self._build_artifact_panel(artifact, role="published", subtitle=subtitle)
+        self._console.print(panel)
+        await super().on_post_publish(agent, ctx, artifact)
+
+    async def on_error(self, agent, ctx: Context, error: Exception) -> None:
+        self._console.log(f"[{agent.name}] error {error!r}", style="bold red")
+        if self._stream_tokens:
+            await self._abort_stream(agent, ctx)
+        await super().on_error(agent, ctx, error)
+
+    async def on_terminate(self, agent, ctx: Context) -> None:
+        if self._stream_tokens:
+            await self._abort_stream(agent, ctx)
+        self._console.log(f"[{agent.name}] end task={ctx.task_id}")
+        await super().on_terminate(agent, ctx)
+
+    # ------------------------------------------------------------------
+    # Rendering helpers
+
+    def _render_artifacts(self, agent_name: str, artifacts: Sequence[Any], *, role: str) -> None:
+        for artifact in artifacts:
+            panel = self._build_artifact_panel(artifact, role=role)
+            self._console.print(panel)
+
+    def _build_artifact_panel(
+        self, artifact: Any, *, role: str, subtitle: str | None = None
+    ) -> Panel:
+        title = f"{role} • {self._summarize_artifact(artifact)}"
+        if subtitle is None:
+            produced_by = getattr(artifact, "produced_by", None)
+            visibility = getattr(artifact, "visibility", None)
+            visibility_name = getattr(visibility, "kind", None)
+            pieces = []
+            if produced_by:
+                pieces.append(f"from={produced_by}")
+            if visibility_name:
+                pieces.append(f"visibility={visibility_name}")
+            subtitle = " | ".join(pieces)
+
+        payload = getattr(artifact, "payload", None)
+        renderable = self._render_payload(payload)
+        return Panel(renderable, title=title, subtitle=subtitle, border_style="cyan")
+
+    def _render_payload(self, payload: Any):
+        if payload is None:
+            return Pretty(payload)
+        if isinstance(payload, Mapping):
+            if self._highlight_json:
+                try:
+                    return JSON.from_data(payload, indent=2, sort_keys=True)
+                except Exception:  # nosec B110 - pragma: no cover - serialization guard
+                    pass
+            return Pretty(dict(payload))
+        if isinstance(payload, (list, tuple, set)):
+            return Pretty(payload)
+        if hasattr(payload, "model_dump"):
+            model_dict = payload.model_dump()
+            return JSON.from_data(model_dict, indent=2, sort_keys=True)
+        return Pretty(payload)
+
+    def _render_metrics(self, agent_name: str, metrics: Mapping[str, Any]) -> None:
+        if not metrics:
+            return
+        table = Table(title=f"{agent_name} metrics", box=None, show_header=False)
+        for key, value in metrics.items():
+            display = f"{value:.2f}" if isinstance(value, (int, float)) else str(value)
+            table.add_row(key, display)
+        self._console.print(table)
+
+    def _render_logs(self, agent_name: str, logs: Iterable[str]) -> None:
+        if not logs:
+            return
+        textual: list[str] = []
+        json_sections: list[JSON] = []
+        for line in logs:
+            if line.startswith("dspy.output="):
+                _, _, payload = line.partition("=")
+                try:
+                    json_sections.append(
+                        JSON.from_data(json.loads(payload), indent=2, sort_keys=True)
+                    )
+                    continue
+                except json.JSONDecodeError:
+                    textual.append(line)
+            else:
+                textual.append(line)
+        for payload in json_sections:
+            panel = Panel(payload, title=f"{agent_name} ▸ dspy.output", border_style="green")
+            self._console.print(panel)
+        if textual:
+            body = Text("\n".join(textual) + "\n")
+            panel = Panel(body, title=f"{agent_name} logs", border_style="magenta")
+            self._console.print(panel)
+
+    def _summarize_artifact(self, artifact: Any) -> str:
+        try:
+            art_id = getattr(artifact, "id", None)
+            prefix = str(art_id)[:8] if art_id else "?"
+            art_type = getattr(artifact, "type", type(artifact).__name__)
+            return f"{art_type}@{prefix}"
+        except Exception:  # pragma: no cover - defensive
+            return repr(artifact)
+
+    # ------------------------------------------------------------------
+    # Streaming support
+
+    def _maybe_start_stream(self, agent, ctx: Context) -> None:
+        stream_key = self._stream_key(agent, ctx)
+        if stream_key in self._stream_context:
+            return
+        queue: asyncio.Queue = asyncio.Queue()
+        self._attach_stream_queue(ctx.state, queue)
+        task = asyncio.create_task(self._consume_stream(agent.name, stream_key, queue))
+        self._stream_context[stream_key] = (queue, task)
+
+    async def _finalize_stream(self, agent, ctx: Context) -> None:
+        stream_key = self._stream_key(agent, ctx)
+        record = self._stream_context.pop(stream_key, None)
+        self._detach_stream_queue(ctx.state)
+        if not record:
+            return
+        queue, task = record
+        if not task.done():
+            await queue.put({"kind": "end"})
+        try:
+            await asyncio.wait_for(task, timeout=2.0)
+        except asyncio.TimeoutError:  # pragma: no cover - defensive cancel
+            task.cancel()
+
+    async def _abort_stream(self, agent, ctx: Context) -> None:
+        stream_key = self._stream_key(agent, ctx)
+        record = self._stream_context.pop(stream_key, None)
+        self._detach_stream_queue(ctx.state)
+        if not record:
+            return
+        queue, task = record
+        if not task.done():
+            await queue.put({"kind": "end", "error": "aborted"})
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+    async def _consume_stream(self, agent_name: str, stream_key: str, queue: asyncio.Queue) -> None:
+        body = Text()
+        live: Live | None = None
+        try:
+            while True:
+                event = await queue.get()
+                if event is None or event.get("kind") == "end":
+                    break
+                kind = event.get("kind")
+                if live is None:
+                    live_panel = Panel(body, title=f"{agent_name} ▸ streaming", border_style="cyan")
+                    live = Live(
+                        live_panel,
+                        console=self._console,
+                        refresh_per_second=12,
+                        transient=True,
+                    )
+                    live.__enter__()
+                if kind == "chunk":
+                    chunk = event.get("chunk") or ""
+                    body.append(chunk)
+                elif kind == "status":
+                    message = event.get("message") or ""
+                    stage = event.get("stage")
+                    line = f"[{stage}] {message}" if stage else message
+                    body.append(f"\n{line}\n", style="dim")
+                elif kind == "error":
+                    message = event.get("message") or ""
+                    body.append(f"\n⚠ {message}\n", style="bold red")
+                if live is not None:
+                    live.update(Panel(body, title=f"{agent_name} ▸ streaming", border_style="cyan"))
+        finally:
+            if live is not None:
+                live.__exit__(None, None, None)
+        if body.plain:
+            self._console.print(
+                Panel(body, title=f"{agent_name} ▸ stream transcript", border_style="cyan")
+            )
+
+    def _stream_key(self, agent, ctx: Context) -> str:
+        return f"{ctx.task_id}:{agent.name}"
+
+    def _attach_stream_queue(self, state: MutableMapping[str, Any], queue: asyncio.Queue) -> None:
+        state.setdefault("_logging", {})["stream_queue"] = queue
+
+    def _detach_stream_queue(self, state: MutableMapping[str, Any]) -> None:
+        try:
+            logging_state = state.get("_logging")
+            if isinstance(logging_state, MutableMapping):
+                logging_state.pop("stream_queue", None)
+        except Exception:  # nosec B110 - pragma: no cover - defensive
+            pass
+
+
+__all__ = ["LoggingUtility", "MetricsUtility"]

flock/utility/output_utility_component.py

@@ -0,0 +1,226 @@
+# src/flock/components/utility/output_utility_component.py
+"""Output formatting and display functionality for agents using unified component architecture."""
+
+import re
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field
+
+from flock.components import AgentComponent, AgentComponentConfig
+from flock.logging.formatters.themed_formatter import (
+    ThemedAgentResultFormatter,
+)
+from flock.logging.formatters.themes import OutputTheme
+from flock.logging.logging import get_logger
+from flock.runtime import Context, EvalInputs, EvalResult
+
+
+if TYPE_CHECKING:  # pragma: no cover - type checking only
+    from flock.agent import Agent
+
+
+logger = get_logger("components.utility.output")
+
+
+class OutputUtilityConfig(AgentComponentConfig):
+    """Configuration for output formatting and display."""
+
+    theme: OutputTheme = Field(
+        default=OutputTheme.catppuccin_mocha, description="Theme for output formatting"
+    )
+    render_table: bool = Field(default=True, description="Whether to render output as a table")
+    max_length: int = Field(default=1000, description="Maximum length for displayed output")
+    truncate_long_values: bool = Field(
+        default=True, description="Whether to truncate long values in display"
+    )
+    show_metadata: bool = Field(
+        default=True, description="Whether to show metadata like timestamps"
+    )
+    format_code_blocks: bool = Field(
+        default=True,
+        description="Whether to apply syntax highlighting to code blocks",
+    )
+    custom_formatters: dict[str, str] = Field(
+        default_factory=dict,
+        description="Custom formatters for specific output types",
+    )
+    no_output: bool = Field(
+        default=False,
+        description="Whether to suppress output",
+    )
+    print_context: bool = Field(
+        default=False,
+        description="Whether to print the context",
+    )
+
+
+class OutputUtilityComponent(AgentComponent):
+    """Utility component that handles output formatting and display."""
+
+    config: OutputUtilityConfig = Field(
+        default_factory=OutputUtilityConfig, description="Output configuration"
+    )
+
+    def __init__(self, name: str = "output", config: OutputUtilityConfig | None = None, **data):
+        if config is None:
+            config = OutputUtilityConfig()
+        super().__init__(name=name, config=config, **data)
+        self._formatter = ThemedAgentResultFormatter(
+            theme=self.config.theme,
+            max_length=self.config.max_length,
+            render_table=self.config.render_table,
+        )
+
+    def _format_value(self, value: Any, key: str) -> str:
+        """Format a single value based on its type and configuration."""
+        # Check for custom formatter
+        if key in self.config.custom_formatters:
+            formatter_name = self.config.custom_formatters[key]
+            if hasattr(self, f"_format_{formatter_name}"):
+                return getattr(self, f"_format_{formatter_name}")(value)
+
+        # Default formatting based on type
+        if isinstance(value, dict):
+            return self._format_dict(value)
+        if isinstance(value, list):
+            return self._format_list(value)
+        if isinstance(value, str) and self.config.format_code_blocks:
+            return self._format_potential_code(value)
+        return str(value)
+
+    def _format_dict(self, d: dict[str, Any], indent: int = 0) -> str:
+        """Format a dictionary with proper indentation."""
+        if not d:
+            return "{}"
+
+        items = []
+        prefix = "  " * indent
+        for key, value in d.items():
+            if self.config.truncate_long_values and isinstance(value, str) and len(value) > 100:
+                value = value[:97] + "..."
+            formatted_value = self._format_value(value, key)
+            items.append(f"{prefix}  {key}: {formatted_value}")
+
+        return "{\n" + "\n".join(items) + f"\n{prefix}}}"
+
+    def _format_list(self, lst: list[Any]) -> str:
+        """Format a list with proper structure."""
+        if not lst:
+            return "[]"
+
+        if len(lst) <= 3:
+            return str(lst)
+
+        # For longer lists, show first few items and count
+        preview = [str(item) for item in lst[:3]]
+        return f"[{', '.join(preview)}, ... ({len(lst)} total)]"
+
+    def _format_potential_code(self, text: str) -> str:
+        """Apply syntax highlighting to potential code blocks."""
+
+        # Simple pattern matching for code blocks
+        def replace_code_block(match):
+            language = match.group(1) or "text"
+            code = match.group(2)
+            return f"[CODE:{language}]\n{code}\n[/CODE]"
+
+        # Replace markdown-style code blocks
+        return re.sub(r"```(\w+)?\n(.*?)\n```", replace_code_block, text, flags=re.DOTALL)
+
+    async def on_post_evaluate(
+        self, agent: "Agent", ctx: Context, inputs: EvalInputs, result: EvalResult
+    ) -> dict[str, Any]:
+        """Format and display the output."""
+        logger.debug("Formatting and displaying output")
+
+        streaming_live_handled = False
+        output_queued = False
+        streamed_artifact_id = None
+
+        if ctx:
+            streaming_live_handled = bool(ctx.get_variable("_flock_stream_live_active", False))
+            output_queued = bool(ctx.get_variable("_flock_output_queued", False))
+            streamed_artifact_id = ctx.get_variable("_flock_streamed_artifact_id")
+
+            if streaming_live_handled:
+                ctx.state.pop("_flock_stream_live_active", None)
+
+            if output_queued:
+                ctx.state.pop("_flock_output_queued", None)
+
+            if streamed_artifact_id:
+                ctx.state.pop("_flock_streamed_artifact_id", None)
+
+        # If streaming was handled, we need to update the final display with the real artifact ID
+        if streaming_live_handled and streamed_artifact_id:
+            logger.debug(
+                f"Updating streamed display with final artifact ID: {streamed_artifact_id}"
+            )
+            # The streaming display already showed everything, we just need to update the ID
+            # This is handled by a final refresh in the streaming code
+            return result
+
+        # Skip output if streaming already handled it (and no ID to update)
+        if streaming_live_handled:
+            logger.debug("Skipping static table because streaming rendered live output.")
+            return result
+
+        # If output was queued due to concurrent stream, wait and then display
+        if output_queued:
+            # Wait for active streams to complete
+            orchestrator = getattr(ctx, "orchestrator", None)
+            if orchestrator:
+                import asyncio
+
+                # Wait until no streams are active
+                max_wait = 30  # seconds
+                waited = 0
+                while getattr(orchestrator, "_active_streams", 0) > 0 and waited < max_wait:
+                    await asyncio.sleep(0.1)
+                    waited += 0.1
+                logger.debug(
+                    f"Queued output displayed after waiting {waited:.1f}s for streams to complete."
+                )
+
+        logger.debug("Formatting and displaying output to console.")
+
+        if self.config.print_context and ctx:
+            # Add context snapshot if requested (be careful with large contexts)
+            try:
+                # Create a copy or select relevant parts to avoid modifying original result dict directly
+                display_result = result.copy()
+                display_result["context_snapshot"] = ctx.to_dict()  # Potential performance hit
+            except Exception:
+                display_result = result.copy()
+                display_result["context_snapshot"] = "[Error serializing context]"
+            result_to_display = display_result
+        else:
+            result_to_display = result
+
+        if not hasattr(self, "_formatter") or self._formatter is None:
+            self._formatter = ThemedAgentResultFormatter(
+                theme=self.config.theme,
+                max_length=self.config.max_length,
+                render_table=self.config.render_table,
+            )
+        model = agent.model if agent.model else ctx.get_variable("model")
+        # Handle None model gracefully
+        model_display = model if model is not None else "default"
+        self._formatter.display_result(
+            result_to_display.artifacts, agent.name + " - " + model_display
+        )
+
+        return result  # Return the original, unmodified result
+
+    def update_theme(self, new_theme: OutputTheme) -> None:
+        """Update the output theme."""
+        self.config.theme = new_theme
+        self._formatter = ThemedAgentResultFormatter(
+            theme=self.config.theme,
+            max_length=self.config.max_length,
+            render_table=self.config.render_table,
+        )
+
+    def add_custom_formatter(self, key: str, formatter_name: str) -> None:
+        """Add a custom formatter for a specific output key."""
+        self.config.custom_formatters[key] = formatter_name

flock/visibility.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+
+"""Artifact visibility policies."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+class AgentIdentity(BaseModel):
+    """Minimal identity information about an agent for visibility checks."""
+
+    name: str
+    labels: set[str] = Field(default_factory=set)
+    tenant_id: str | None = None
+
+
+class Visibility(BaseModel):
+    """Base visibility contract."""
+
+    kind: Literal["Public", "Private", "Labelled", "Tenant", "After"]
+
+    def allows(self, agent: AgentIdentity, *, now: datetime | None = None) -> bool:
+        raise NotImplementedError
+
+
+class PublicVisibility(Visibility):
+    kind: Literal["Public"] = "Public"
+
+    def allows(
+        self, agent: AgentIdentity, *, now: datetime | None = None
+    ) -> bool:  # pragma: no cover - trivial
+        return True
+
+
+class PrivateVisibility(Visibility):
+    kind: Literal["Private"] = "Private"
+    agents: set[str] = Field(default_factory=set)
+
+    def allows(self, agent: AgentIdentity, *, now: datetime | None = None) -> bool:
+        return agent.name in self.agents
+
+
+class LabelledVisibility(Visibility):
+    kind: Literal["Labelled"] = "Labelled"
+    required_labels: set[str] = Field(default_factory=set)
+
+    def allows(self, agent: AgentIdentity, *, now: datetime | None = None) -> bool:
+        return self.required_labels.issubset(agent.labels)
+
+
+class TenantVisibility(Visibility):
+    kind: Literal["Tenant"] = "Tenant"
+    tenant_id: str | None = None
+
+    def allows(self, agent: AgentIdentity, *, now: datetime | None = None) -> bool:
+        if self.tenant_id is None:
+            return True
+        return agent.tenant_id == self.tenant_id
+
+
+class AfterVisibility(Visibility):
+    kind: Literal["After"] = "After"
+    ttl: timedelta = Field(default=timedelta())
+    then: Visibility | None = None
+    _created_at: datetime = PrivateAttr(default_factory=lambda: datetime.now(timezone.utc))
+
+    def allows(self, agent: AgentIdentity, *, now: datetime | None = None) -> bool:
+        now = now or datetime.now(timezone.utc)
+        if now - self._created_at >= self.ttl:
+            if self.then:
+                return self.then.allows(agent, now=now)
+            return True
+        return False
+
+
+def ensure_visibility(value: Visibility | None) -> Visibility:
+    if value is None:
+        return PublicVisibility()
+    return value
+
+
+def only_for(*agent_names: str) -> PrivateVisibility:
+    return PrivateVisibility(agents=set(agent_names))
+
+
+def agents_from_names(names: Iterable[str]) -> set[str]:  # pragma: no cover - helper
+    return set(names)
+
+
+__all__ = [
+    "AfterVisibility",
+    "AgentIdentity",
+    "LabelledVisibility",
+    "PrivateVisibility",
+    "PublicVisibility",
+    "TenantVisibility",
+    "Visibility",
+    "ensure_visibility",
+    "only_for",
+]