flock-core 0.5.0b27__py3-none-any.whl → 0.5.0b50__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/{components/utility → utility}/output_utility_component.py

@@ -6,19 +6,18 @@ from typing import TYPE_CHECKING, Any
 
 from pydantic import Field
 
-from flock.core.component.agent_component_base import AgentComponentConfig
-from flock.core.component.utility_component import UtilityComponent
-from flock.core.context.context import FlockContext
-from flock.core.context.context_vars import FLOCK_BATCH_SILENT_MODE
-from flock.core.logging.formatters.themed_formatter import (
+from flock.components import AgentComponent, AgentComponentConfig
+from flock.logging.formatters.themed_formatter import (
     ThemedAgentResultFormatter,
 )
-from flock.core.logging.formatters.themes import OutputTheme
-from flock.core.logging.logging import get_logger
-from flock.core.registry import flock_component
+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
 
-if TYPE_CHECKING:
-    from flock.core.flock_agent import FlockAgent
 
 logger = get_logger("components.utility.output")
 
@@ -27,14 +26,10 @@ class OutputUtilityConfig(AgentComponentConfig):
     """Configuration for output formatting and display."""
 
     theme: OutputTheme = Field(
-        default=OutputTheme.afterglow, description="Theme for output formatting"
-    )
-    render_table: bool = Field(
-        default=False, description="Whether to render output as a table"
-    )
-    max_length: int = Field(
-        default=1000, description="Maximum length for displayed output"
+        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"
     )
@@ -59,8 +54,7 @@ class OutputUtilityConfig(AgentComponentConfig):
     )
 
 
-@flock_component(config_class=OutputUtilityConfig)
-class OutputUtilityComponent(UtilityComponent):
+class OutputUtilityComponent(AgentComponent):
     """Utility component that handles output formatting and display."""
 
     config: OutputUtilityConfig = Field(
@@ -88,12 +82,11 @@ class OutputUtilityComponent(UtilityComponent):
         # Default formatting based on type
         if isinstance(value, dict):
             return self._format_dict(value)
-        elif isinstance(value, list):
+        if isinstance(value, list):
             return self._format_list(value)
-        elif isinstance(value, str) and self.config.format_code_blocks:
+        if isinstance(value, str) and self.config.format_code_blocks:
             return self._format_potential_code(value)
-        else:
-            return str(value)
+        return str(value)
 
     def _format_dict(self, d: dict[str, Any], indent: int = 0) -> str:
         """Format a dictionary with proper indentation."""
@@ -124,6 +117,7 @@ class OutputUtilityComponent(UtilityComponent):
 
     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"
@@ -131,57 +125,74 @@ class OutputUtilityComponent(UtilityComponent):
             return f"[CODE:{language}]\n{code}\n[/CODE]"
 
         # Replace markdown-style code blocks
-        text = re.sub(
-            r"```(\w+)?\n(.*?)\n```", replace_code_block, text, flags=re.DOTALL
-        )
-        return text
+        return re.sub(r"```(\w+)?\n(.*?)\n```", replace_code_block, text, flags=re.DOTALL)
 
     async def on_post_evaluate(
-        self,
-        agent: "FlockAgent",
-        inputs: dict[str, Any],
-        context: FlockContext | None = None,
-        result: dict[str, Any] | None = None,
+        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
-        if context:
-            streaming_live_handled = bool(
-                context.get_variable("_flock_stream_live_active", 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:
-                context.state.pop("_flock_stream_live_active", None)
+                ctx.state.pop("_flock_stream_live_active", None)
 
-        # Determine if output should be suppressed
-        is_silent = self.config.no_output or (
-            context and context.get_variable(FLOCK_BATCH_SILENT_MODE, False)
-        )
+            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 is_silent:
-            logger.debug("Output suppressed (config or batch silent mode).")
-            return result  # Skip console output
+        # 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 context:
+        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"] = (
-                    context.to_dict()
-                )  # Potential performance hit
+                display_result["context_snapshot"] = ctx.to_dict()  # Potential performance hit
             except Exception:
                 display_result = result.copy()
-                display_result["context_snapshot"] = (
-                    "[Error serializing context]"
-                )
+                display_result["context_snapshot"] = "[Error serializing context]"
             result_to_display = display_result
         else:
             result_to_display = result
@@ -192,8 +203,12 @@ class OutputUtilityComponent(UtilityComponent):
                 max_length=self.config.max_length,
                 render_table=self.config.render_table,
             )
-        model = agent.model if agent.model else context.get_variable("model")
-        self._formatter.display_result(result_to_display, agent.name + " - " + model)
+        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
 

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",
+]