flock-core 0.5.11__py3-none-any.whl → 0.5.21__py3-none-any.whl

flock/engines/examples/simple_batch_engine.py

@@ -4,9 +4,9 @@ from __future__ import annotations
 
 from pydantic import BaseModel, Field
 
-from flock.components import EngineComponent
+from flock.components.agent import EngineComponent
 from flock.registry import flock_type
-from flock.runtime import EvalInputs, EvalResult
+from flock.utils.runtime import EvalInputs, EvalResult
 
 
 @flock_type(name="BatchItem")

flock/engines/streaming/__init__.py

@@ -0,0 +1,3 @@
+"""Streaming engine shared utilities."""
+
+__all__ = ["sinks"]

flock/engines/streaming/sinks.py

@@ -0,0 +1,489 @@
+"""Shared streaming sink implementations for DSPy execution.
+
+This module provides a composable sink pattern for consuming streaming output
+from DSPy programs and routing it to different presentation layers (Rich terminal,
+WebSocket dashboard, etc.).
+
+Architecture
+------------
+The StreamSink protocol defines a minimal interface that all sinks must implement.
+Sinks receive normalized streaming events (status messages, tokens, final predictions)
+and handle presentation-specific logic (Rich display updates, WebSocket broadcasts, etc.).
+
+Sinks are designed to be:
+- Composable: Multiple sinks can consume the same stream in parallel
+- Isolated: Each sink maintains its own state and error handling
+- Testable: Sinks can be tested independently with mock dependencies
+
+Error Handling Contract
+-----------------------
+Sinks SHOULD NOT raise exceptions during normal streaming operations. Instead:
+- Log errors and continue processing remaining events
+- Use defensive programming (null checks, try/except where appropriate)
+- Only raise exceptions for unrecoverable errors (e.g., invalid configuration)
+
+The streaming loop treats sink exceptions as fatal and will abort the stream.
+For fault tolerance, sinks should catch and log their own errors.
+
+Example Usage
+-------------
+Basic WebSocket-only streaming:
+
+    async def ws_broadcast(event: StreamingOutputEvent) -> None:
+        await websocket_manager.broadcast(event)
+
+    def event_factory(output_type, content, seq, is_final):
+        return StreamingOutputEvent(
+            correlation_id="123",
+            agent_name="agent",
+            run_id="run-1",
+            output_type=output_type,
+            content=content,
+            sequence=seq,
+            is_final=is_final,
+        )
+
+    sink = WebSocketSink(ws_broadcast=ws_broadcast, event_factory=event_factory)
+
+    async for value in stream:
+        kind, text, field, final = normalize_value(value)
+        if kind == "status":
+            await sink.on_status(text)
+        elif kind == "token":
+            await sink.on_token(text, field)
+        elif kind == "prediction":
+            await sink.on_final(final, token_count)
+            break
+
+    await sink.flush()
+
+Dual-sink composition (CLI with WebSocket):
+
+    sinks = []
+    if rich_enabled:
+        sinks.append(RichSink(...))
+    if ws_enabled:
+        sinks.append(WebSocketSink(...))
+
+    # Dispatch to all sinks
+    for sink in sinks:
+        await sink.on_token(text, field)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, MutableMapping, Sequence
+from typing import (
+    Any,
+    Protocol,
+    runtime_checkable,
+)
+
+from pydantic import BaseModel
+
+from flock.dashboard.events import StreamingOutputEvent
+from flock.logging.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+@runtime_checkable
+class StreamSink(Protocol):
+    """Minimal sink protocol for consuming normalized stream events.
+
+    Sinks receive streaming events from DSPy execution and handle
+    presentation-specific logic (Rich display, WebSocket broadcast, etc.).
+
+    Implementations must be idempotent for on_final() to handle edge cases
+    where the stream loop might call it multiple times.
+
+    Error Handling
+    --------------
+    Implementations SHOULD catch and log their own errors rather than raising,
+    to prevent one sink failure from aborting the entire stream. Only raise
+    exceptions for unrecoverable errors during initialization/configuration.
+    """
+
+    async def on_status(self, text: str) -> None:
+        """Process a status message from the LLM.
+
+        Status messages are typically intermediate reasoning steps or
+        progress indicators (e.g., "Analyzing input...", "Generating response...").
+
+        Args:
+            text: Status message text (may include newlines)
+
+        Note:
+            Empty text should be ignored. Implementations should handle
+            this gracefully without raising.
+        """
+        ...
+
+    async def on_token(self, text: str, signature_field: str | None) -> None:
+        """Process a single token from the LLM output stream.
+
+        Tokens are emitted as the LLM generates text. signature_field indicates
+        which output field this token belongs to (for multi-field signatures).
+
+        Args:
+            text: Token text (typically a single word or word fragment)
+            signature_field: Name of the signature field being streamed,
+                or None if the token doesn't belong to a specific field
+
+        Note:
+            Empty text should be ignored. The field "description" is typically
+            skipped as it's the input prompt, not output.
+        """
+        ...
+
+    async def on_final(self, result: Any, tokens_emitted: int) -> None:
+        """Process the final prediction result.
+
+        Called once when streaming completes successfully. Contains the
+        complete DSPy Prediction object with all output fields populated.
+
+        Args:
+            result: DSPy Prediction object with output fields
+            tokens_emitted: Total number of tokens emitted during streaming
+
+        Note:
+            Implementations MUST be idempotent - this may be called multiple
+            times in edge cases. Use a finalization guard flag if necessary.
+        """
+        ...
+
+    async def flush(self) -> None:
+        """Flush any pending async operations.
+
+        Called after streaming completes to ensure all async tasks
+        (e.g., WebSocket broadcasts) complete before returning.
+
+        Implementations should await any background tasks and handle
+        errors gracefully (log but don't raise).
+
+        Note:
+            For synchronous sinks (e.g., Rich terminal), this is a no-op.
+        """
+        ...
+
+
+class RichSink(StreamSink):
+    """Rich terminal sink responsible for mutating live display data.
+
+    This sink updates a mutable display_data dictionary that represents
+    the artifact being streamed. It accumulates status messages and tokens
+    in buffers, then replaces them with final structured data when streaming
+    completes.
+
+    The sink integrates with Rich's Live display context, calling a refresh
+    callback after each update to trigger terminal re-rendering.
+
+    Display Data Flow
+    -----------------
+    1. Initialization: display_data contains empty payload fields and "streaming..." timestamp
+    2. on_status(): Accumulates status messages in a buffer, updates display_data["status"]
+    3. on_token(): Accumulates tokens in field-specific buffers, updates display_data["payload"]["_streaming"]
+    4. on_final(): Replaces streaming buffers with final Prediction fields, removes "status", adds real timestamp
+    5. flush(): No-op (Rich rendering is synchronous)
+
+    Error Handling
+    --------------
+    - refresh_panel() errors are caught and logged, never raised
+    - Idempotent: on_final() checks _finalized flag to prevent double-finalization
+    - Defensive: Uses setdefault() and get() to handle missing dictionary keys
+
+    Thread Safety
+    -------------
+    NOT thread-safe. Assumes single-threaded async execution within a single
+    Rich Live context. Multiple concurrent streams should use separate RichSink instances.
+
+    Example
+    -------
+        display_data = OrderedDict([("id", "artifact-123"), ("payload", {}), ...])
+        stream_buffers = defaultdict(list)
+
+        def refresh():
+            live.update(formatter.format_result(display_data, ...))
+
+        sink = RichSink(
+            display_data=display_data,
+            stream_buffers=stream_buffers,
+            status_field="_status",
+            signature_order=["output", "summary"],
+            formatter=formatter,
+            theme_dict=theme,
+            styles=styles,
+            agent_label="Agent - gpt-4",
+            refresh_panel=refresh,
+            timestamp_factory=lambda: datetime.now(UTC).isoformat(),
+        )
+
+        await sink.on_status("Processing...")
+        await sink.on_token("Hello", "output")
+        await sink.on_final(prediction, tokens_emitted=5)
+        await sink.flush()
+    """
+
+    def __init__(
+        self,
+        *,
+        display_data: MutableMapping[str, Any],
+        stream_buffers: MutableMapping[str, list[str]],
+        status_field: str,
+        signature_order: Sequence[str],
+        formatter: Any | None,
+        theme_dict: dict[str, Any] | None,
+        styles: dict[str, Any] | None,
+        agent_label: str | None,
+        refresh_panel: Callable[[], None],
+        timestamp_factory: Callable[[], str],
+    ) -> None:
+        self._display_data = display_data
+        self._stream_buffers = stream_buffers
+        self._status_field = status_field
+        self._signature_order = list(signature_order)
+        self._formatter = formatter
+        self._theme_dict = theme_dict
+        self._styles = styles
+        self._agent_label = agent_label
+        self._refresh_panel = refresh_panel
+        self._timestamp_factory = timestamp_factory
+        self._final_display = (
+            formatter,
+            display_data,
+            theme_dict,
+            styles,
+            agent_label,
+        )
+        # Ensure buffers exist for status updates
+        self._stream_buffers.setdefault(status_field, [])
+        self._finalized = False
+
+    def _refresh(self) -> None:
+        try:
+            self._refresh_panel()
+        except Exception:
+            logger.debug("Rich sink refresh panel callable failed", exc_info=True)
+
+    async def on_status(self, text: str) -> None:
+        if not text:
+            return
+
+        buffer = self._stream_buffers.setdefault(self._status_field, [])
+        buffer.append(f"{text}\n")
+        self._display_data["status"] = "".join(buffer)
+        self._refresh()
+
+    async def on_token(self, text: str, signature_field: str | None) -> None:
+        if not text:
+            return
+
+        if signature_field and signature_field != "description":
+            buffer_key = f"_stream_{signature_field}"
+            buffer = self._stream_buffers.setdefault(buffer_key, [])
+            buffer.append(str(text))
+            payload = self._display_data.setdefault("payload", {})
+            payload["_streaming"] = "".join(buffer)
+        else:
+            buffer = self._stream_buffers.setdefault(self._status_field, [])
+            buffer.append(str(text))
+            self._display_data["status"] = "".join(buffer)
+
+        self._refresh()
+
+    async def on_final(self, result: Any, tokens_emitted: int) -> None:  # noqa: ARG002
+        if self._finalized:
+            return
+
+        payload_section: MutableMapping[str, Any] = self._display_data.setdefault(
+            "payload", {}
+        )
+        payload_section.clear()
+
+        for field_name in self._signature_order:
+            if field_name == "description":
+                continue
+            if not hasattr(result, field_name):
+                continue
+
+            value = getattr(result, field_name)
+            if isinstance(value, list):
+                payload_section[field_name] = [
+                    item.model_dump() if isinstance(item, BaseModel) else item
+                    for item in value
+                ]
+            elif isinstance(value, BaseModel):
+                payload_section[field_name] = value.model_dump()
+            else:
+                payload_section[field_name] = value
+
+        self._display_data["created_at"] = self._timestamp_factory()
+        self._display_data.pop("status", None)
+        payload_section.pop("_streaming", None)
+        self._refresh()
+        self._finalized = True
+
+    async def flush(self) -> None:
+        # Rich sink has no async resources to drain.
+        return None
+
+    @property
+    def final_display_data(
+        self,
+    ) -> tuple[
+        Any,
+        MutableMapping[str, Any],
+        dict[str, Any] | None,
+        dict[str, Any] | None,
+        str | None,
+    ]:
+        return self._final_display
+
+
+class WebSocketSink(StreamSink):
+    """WebSocket-only sink that mirrors dashboard streaming behaviour.
+
+    This sink broadcasts StreamingOutputEvent messages via WebSocket for
+    real-time dashboard updates. It uses fire-and-forget task scheduling
+    to avoid blocking the streaming loop while ensuring all events are
+    delivered via flush().
+
+    Event Sequence
+    --------------
+    Each event gets a monotonically increasing sequence number for ordering:
+    - on_status("Loading"): seq=0, output_type="log", content="Loading\\n"
+    - on_token("Hello", field): seq=1, output_type="llm_token", content="Hello"
+    - on_token(" world", field): seq=2, output_type="llm_token", content=" world"
+    - on_final(pred, 2): seq=3, output_type="log", content="\\nAmount of output tokens: 2", is_final=True
+    -                    seq=4, output_type="log", content="--- End of output ---", is_final=True
+
+    The two terminal events are required for dashboard compatibility and must
+    appear in this exact order with is_final=True.
+
+    Task Management
+    ---------------
+    Events are broadcast using asyncio.create_task() to avoid blocking the
+    streaming loop. Tasks are tracked in a set and awaited during flush()
+    to ensure delivery before the stream completes.
+
+    Task lifecycle:
+    1. _schedule() creates task and adds to _tasks set
+    2. Task completion callback removes it from _tasks
+    3. flush() awaits remaining tasks with error handling
+
+    Error Handling
+    --------------
+    - Scheduling errors: Logged and ignored (event dropped)
+    - Broadcast errors: Caught during flush(), logged but don't raise
+    - Idempotent: on_final() checks _finalized flag to prevent duplicate terminal events
+
+    Thread Safety
+    -------------
+    NOT thread-safe. Assumes single-threaded async execution. Multiple
+    concurrent streams should use separate WebSocketSink instances.
+
+    Example
+    -------
+        async def broadcast(event: StreamingOutputEvent):
+            await websocket_manager.send_json(event.model_dump())
+
+        def event_factory(output_type, content, seq, is_final):
+            return StreamingOutputEvent(
+                correlation_id="corr-123",
+                agent_name="analyzer",
+                run_id="run-456",
+                output_type=output_type,
+                content=content,
+                sequence=seq,
+                is_final=is_final,
+                artifact_id="artifact-789",
+                artifact_type="Report",
+            )
+
+        sink = WebSocketSink(ws_broadcast=broadcast, event_factory=event_factory)
+
+        await sink.on_status("Processing input")
+        await sink.on_token("Analysis", "output")
+        await sink.on_final(prediction, tokens_emitted=1)
+        await sink.flush()  # Ensures all broadcasts complete
+    """
+
+    def __init__(
+        self,
+        *,
+        ws_broadcast: Callable[[StreamingOutputEvent], Awaitable[None]] | None,
+        event_factory: Callable[[str, str, int, bool], StreamingOutputEvent],
+    ) -> None:
+        self._ws_broadcast = ws_broadcast
+        self._event_factory = event_factory
+        self._sequence = 0
+        self._tasks: set[asyncio.Task[Any]] = set()
+        self._finalized = False
+
+    def _schedule(
+        self,
+        output_type: str,
+        content: str,
+        *,
+        is_final: bool,
+        advance_sequence: bool = True,
+    ) -> None:
+        if not self._ws_broadcast:
+            return
+
+        event = self._event_factory(output_type, content, self._sequence, is_final)
+        try:
+            task = asyncio.create_task(self._ws_broadcast(event))
+        except Exception as exc:  # pragma: no cover - scheduling should rarely fail
+            logger.warning(f"Failed to schedule streaming event: {exc}")
+            return
+
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+        if advance_sequence:
+            self._sequence += 1
+
+    async def on_status(self, text: str) -> None:
+        if not text:
+            return
+        self._schedule("log", f"{text}\n", is_final=False)
+
+    async def on_token(self, text: str, signature_field: str | None) -> None:  # noqa: ARG002
+        if not text:
+            return
+        self._schedule("llm_token", text, is_final=False)
+
+    async def on_final(self, result: Any, tokens_emitted: int) -> None:  # noqa: ARG002
+        if self._finalized:
+            return
+
+        self._schedule(
+            "log",
+            f"\nAmount of output tokens: {tokens_emitted}",
+            is_final=True,
+        )
+        self._schedule(
+            "log",
+            "--- End of output ---",
+            is_final=True,
+        )
+
+        self._finalized = True
+
+    async def flush(self) -> None:
+        if not self._tasks:
+            return
+
+        pending = list(self._tasks)
+        self._tasks.clear()
+
+        results = await asyncio.gather(*pending, return_exceptions=True)
+        for result in results:
+            if isinstance(result, Exception):
+                logger.warning(f"Streaming broadcast task failed: {result}")
+
+
+__all__ = ["RichSink", "StreamSink", "WebSocketSink"]

flock/examples.py

@@ -13,17 +13,17 @@ from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, Field
 
-from flock.artifacts import Artifact
-from flock.components import EngineComponent
-from flock.orchestrator import Flock
+from flock.components.agent import EngineComponent
+from flock.core import Flock
+from flock.core.artifacts import Artifact
+from flock.core.store import BlackboardStore
 from flock.registry import flock_tool, flock_type, type_registry
-from flock.runtime import EvalInputs, EvalResult
-from flock.store import BlackboardStore
-from flock.utilities import LoggingUtility, MetricsUtility
+from flock.utils.runtime import EvalInputs, EvalResult
+from flock.utils.utilities import LoggingUtility, MetricsUtility
 
 
 if TYPE_CHECKING:
-    from flock.agent import AgentBuilder
+    from flock.core import AgentBuilder
 
 
 @flock_type

flock/logging/logging.py

@@ -42,13 +42,6 @@ def in_workflow_context() -> bool:
     if successful. Otherwise, it returns False.
     """
     return False
-    # try:
-    #     workflow.logger.debug("Checking if in workflow context...")
-    #     # loguru_logger.debug("Checking if in workflow context...")
-    #     # This call will succeed only if we're in a workflow context.
-    #     return bool(hasattr(workflow.info(), "is_replaying"))
-    # except Exception:
-    #     return False
 
 
 def get_current_trace_id() -> str:
@@ -374,9 +367,6 @@ class FlockLogger:
         self.min_level_severity = initial_min_level_severity
 
     def _get_logger(self):
-        # if in_workflow_context():
-        #     # Use Temporal's workflow.logger inside a workflow context.
-        #     return workflow.logger
         # Bind our logger with category and trace_id
         return loguru_logger.bind(
             name=self.name,
@@ -562,12 +552,7 @@ def get_logger(name: str = "flock") -> FlockLogger:
 
 def get_module_loggers() -> list[FlockLogger]:
     """Return a cached FlockLogger instance for the given module name."""
-    result = []
-    for kvp in _LOGGER_CACHE:
-        if kvp.startswith("module."):
-            result.append(_LOGGER_CACHE[kvp])
-
-    return result
+    return [_LOGGER_CACHE[kvp] for kvp in _LOGGER_CACHE if kvp.startswith("module.")]
 
 
 def truncate_for_logging(obj, max_item_length=100, max_items=10):

flock/models/__init__.py

@@ -0,0 +1,10 @@
+"""System data models for Flock.
+
+This module contains system-level artifact types used by the orchestrator
+for error handling, workflow tracking, and internal communication.
+"""
+
+from flock.models.system_artifacts import WorkflowError
+
+
+__all__ = ["WorkflowError"]

flock/orchestrator/__init__.py

@@ -0,0 +1,45 @@
+"""Orchestrator implementation modules.
+
+This package contains internal implementation details for the Flock orchestrator.
+
+Phase 5A Additions:
+- ContextBuilder: Security-critical context building with provider resolution
+- EventEmitter: Dashboard WebSocket event emission
+- LifecycleManager: Background task coordination for batches and correlations
+
+Phase 3 Additions (Orchestrator Simplification):
+- TracingManager: Unified tracing and trace database management
+- ServerManager: HTTP service startup and dashboard integration
+- OrchestratorInitializer: Component setup and state initialization
+"""
+
+from flock.orchestrator.artifact_collector import ArtifactCollector
+from flock.orchestrator.artifact_manager import ArtifactManager
+from flock.orchestrator.batch_accumulator import BatchAccumulator
+from flock.orchestrator.component_runner import ComponentRunner
+from flock.orchestrator.context_builder import ContextBuilder
+from flock.orchestrator.correlation_engine import CorrelationEngine
+from flock.orchestrator.event_emitter import EventEmitter
+from flock.orchestrator.initialization import OrchestratorInitializer
+from flock.orchestrator.lifecycle_manager import LifecycleManager
+from flock.orchestrator.mcp_manager import MCPManager
+from flock.orchestrator.scheduler import AgentScheduler
+from flock.orchestrator.server_manager import ServerManager
+from flock.orchestrator.tracing import TracingManager
+
+
+__all__ = [
+    "AgentScheduler",
+    "ArtifactCollector",
+    "ArtifactManager",
+    "BatchAccumulator",
+    "ComponentRunner",
+    "ContextBuilder",
+    "CorrelationEngine",
+    "EventEmitter",
+    "LifecycleManager",
+    "MCPManager",
+    "OrchestratorInitializer",
+    "ServerManager",
+    "TracingManager",
+]

flock/{artifact_collector.py → orchestrator/artifact_collector.py}

@@ -17,9 +17,9 @@ from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
-    from flock.agent import Agent
-    from flock.artifacts import Artifact
-    from flock.subscription import Subscription
+    from flock.core import Agent
+    from flock.core.artifacts import Artifact
+    from flock.core.subscription import Subscription
 
 
 class ArtifactCollector: