flock-core 0.4.528__py3-none-any.whl → 0.5.0b0__py3-none-any.whl

flock/core/orchestration/flock_batch_processor.py

@@ -0,0 +1,94 @@
+# src/flock/core/orchestration/flock_batch_processor.py
+"""Batch processing functionality for Flock orchestrator."""
+
+from typing import TYPE_CHECKING, Any
+
+from box import Box
+from pandas import DataFrame
+
+from flock.core.logging.logging import get_logger
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock
+    from flock.core.flock_agent import FlockAgent
+
+logger = get_logger("flock.batch_processor")
+
+
+class FlockBatchProcessor:
+    """Handles batch processing functionality for Flock orchestrator."""
+
+    def __init__(self, flock: "Flock"):
+        self.flock = flock
+
+    async def run_batch_async(
+        self,
+        start_agent: "FlockAgent | str",
+        batch_inputs: list[dict[str, Any]] | DataFrame | str,
+        input_mapping: dict[str, str] | None = None,
+        static_inputs: dict[str, Any] | None = None,
+        parallel: bool = True,
+        max_workers: int = 5,
+        use_temporal: bool | None = None,
+        box_results: bool = True,
+        return_errors: bool = False,
+        silent_mode: bool = False,
+        write_to_csv: str | None = None,
+        hide_columns: list[str] | None = None,
+        delimiter: str = ",",
+    ) -> list[Box | dict | None | Exception]:
+        """Runs the specified agent/workflow for each item in a batch asynchronously (delegated)."""
+        # Import processor locally
+        from flock.core.execution.batch_executor import BatchProcessor
+
+        processor = BatchProcessor(self.flock)  # Pass flock instance
+        return await processor.run_batch_async(
+            start_agent=start_agent,
+            batch_inputs=batch_inputs,
+            input_mapping=input_mapping,
+            static_inputs=static_inputs,
+            parallel=parallel,
+            max_workers=max_workers,
+            use_temporal=use_temporal,
+            box_results=box_results,
+            return_errors=return_errors,
+            silent_mode=silent_mode,
+            write_to_csv=write_to_csv,
+            hide_columns=hide_columns,
+            delimiter=delimiter,
+        )
+
+    def run_batch(
+        self,
+        start_agent: "FlockAgent | str",
+        batch_inputs: list[dict[str, Any]] | DataFrame | str,
+        input_mapping: dict[str, str] | None = None,
+        static_inputs: dict[str, Any] | None = None,
+        parallel: bool = True,
+        max_workers: int = 5,
+        use_temporal: bool | None = None,
+        box_results: bool = True,
+        return_errors: bool = False,
+        silent_mode: bool = False,
+        write_to_csv: str | None = None,
+        hide_columns: list[str] | None = None,
+        delimiter: str = ",",
+    ) -> list[Box | dict | None | Exception]:
+        """Synchronous wrapper for batch processing."""
+        return self.flock._execution._run_sync(
+            self.run_batch_async(
+                start_agent=start_agent,
+                batch_inputs=batch_inputs,
+                input_mapping=input_mapping,
+                static_inputs=static_inputs,
+                parallel=parallel,
+                max_workers=max_workers,
+                use_temporal=use_temporal,
+                box_results=box_results,
+                return_errors=return_errors,
+                silent_mode=silent_mode,
+                write_to_csv=write_to_csv,
+                hide_columns=hide_columns,
+                delimiter=delimiter,
+            )
+        )

flock/core/orchestration/flock_evaluator.py

@@ -0,0 +1,113 @@
+# src/flock/core/orchestration/flock_evaluator.py
+"""Evaluation functionality for Flock orchestrator."""
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+from datasets import Dataset
+from pandas import DataFrame
+
+from flock.core.flock_agent import FlockAgent
+from flock.core.logging.logging import get_logger
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock
+
+
+logger = get_logger("flock.evaluator")
+
+
+class FlockEvaluator:
+    """Handles evaluation functionality for Flock orchestrator."""
+
+    def __init__(self, flock: "Flock"):
+        self.flock = flock
+
+    async def evaluate_async(
+        self,
+        dataset: str | Path | list[dict[str, Any]] | DataFrame | Dataset,
+        start_agent: FlockAgent | str,
+        input_mapping: dict[str, str],
+        answer_mapping: dict[str, str],
+        metrics: list[
+            str
+            | Callable[[Any, Any], bool | float | dict[str, Any]]
+            | FlockAgent
+        ],
+        metric_configs: dict[str, dict[str, Any]] | None = None,
+        static_inputs: dict[str, Any] | None = None,
+        parallel: bool = True,
+        max_workers: int = 5,
+        use_temporal: bool | None = None,
+        error_handling: Literal["raise", "skip", "log"] = "log",
+        output_file: str | Path | None = None,
+        return_dataframe: bool = True,
+        silent_mode: bool = False,
+        metadata_columns: list[str] | None = None,
+    ) -> "DataFrame | list[dict[str, Any]]":
+        """Evaluates the Flock's performance against a dataset (delegated)."""
+        # Import processor locally
+        from flock.core.execution.evaluation_executor import EvaluationExecutor
+
+        processor = EvaluationExecutor(self.flock)  # Pass flock instance
+        return await processor.evaluate_async(
+            dataset=dataset,
+            start_agent=start_agent,
+            input_mapping=input_mapping,
+            answer_mapping=answer_mapping,
+            metrics=metrics,
+            metric_configs=metric_configs,
+            static_inputs=static_inputs,
+            parallel=parallel,
+            max_workers=max_workers,
+            use_temporal=use_temporal,
+            error_handling=error_handling,
+            output_file=output_file,
+            return_dataframe=return_dataframe,
+            silent_mode=silent_mode,
+            metadata_columns=metadata_columns,
+        )
+
+    def evaluate(
+        self,
+        dataset: str | Path | list[dict[str, Any]] | DataFrame | Dataset,
+        start_agent: FlockAgent | str,
+        input_mapping: dict[str, str],
+        answer_mapping: dict[str, str],
+        metrics: list[
+            str
+            | Callable[[Any, Any], bool | float | dict[str, Any]]
+            | FlockAgent
+        ],
+        metric_configs: dict[str, dict[str, Any]] | None = None,
+        static_inputs: dict[str, Any] | None = None,
+        parallel: bool = True,
+        max_workers: int = 5,
+        use_temporal: bool | None = None,
+        error_handling: Literal["raise", "skip", "log"] = "log",
+        output_file: str | Path | None = None,
+        return_dataframe: bool = True,
+        silent_mode: bool = False,
+        metadata_columns: list[str] | None = None,
+    ) -> "DataFrame | list[dict[str, Any]]":
+        """Synchronous wrapper for evaluation."""
+        return self.flock._execution._run_sync(
+            self.evaluate_async(
+                dataset=dataset,
+                start_agent=start_agent,
+                input_mapping=input_mapping,
+                answer_mapping=answer_mapping,
+                metrics=metrics,
+                metric_configs=metric_configs,
+                static_inputs=static_inputs,
+                parallel=parallel,
+                max_workers=max_workers,
+                use_temporal=use_temporal,
+                error_handling=error_handling,
+                output_file=output_file,
+                return_dataframe=return_dataframe,
+                silent_mode=silent_mode,
+                metadata_columns=metadata_columns,
+            )
+        )

flock/core/orchestration/flock_execution.py

@@ -0,0 +1,288 @@
+# src/flock/core/orchestration/flock_execution.py
+"""Execution management functionality for Flock orchestrator."""
+
+import asyncio
+import contextvars
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from box import Box
+from opentelemetry import trace
+from opentelemetry.baggage import set_baggage
+
+from flock.config import DEFAULT_MODEL
+from flock.core.context.context import FlockContext
+from flock.core.context.context_manager import initialize_context
+from flock.core.execution.local_executor import run_local_workflow
+from flock.core.execution.temporal_executor import run_temporal_workflow
+from flock.core.logging.logging import get_logger
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock
+    from flock.core.flock_agent import FlockAgent
+
+logger = get_logger("flock.execution")
+tracer = trace.get_tracer(__name__)
+_R = TypeVar("_R")
+
+
+class FlockExecution:
+    """Handles execution management for Flock including run, run_async, and execution coordination."""
+
+    def __init__(self, flock: "Flock"):
+        self.flock = flock
+
+    def _run_sync(self, coro) -> _R:
+        """Execute *coro* synchronously.
+
+        * If no loop is running → ``asyncio.run``.
+        * Otherwise run ``asyncio.run`` inside a fresh thread **with**
+          context-vars propagation.
+        """
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:  # no loop → simple
+            return asyncio.run(coro)
+
+        # A loop is already running – Jupyter / ASGI / etc.
+        ctx = contextvars.copy_context()  # propagate baggage
+        with ThreadPoolExecutor(max_workers=1) as pool:
+            future = pool.submit(ctx.run, asyncio.run, coro)
+            try:
+                return future.result()
+            finally:
+                if not future.done():
+                    future.cancel()
+
+    def run(
+        self,
+        agent: "FlockAgent | str | None" = None,
+        input: dict | None = None,
+        context: FlockContext | None = None,
+        run_id: str = "",
+        box_result: bool = True,
+        agents: list["FlockAgent"] | None = None,
+        servers: list[Any] | None = None,
+        memo: dict[str, Any] | None = None,
+    ) -> Box | dict:
+        """Synchronous execution wrapper."""
+        return self._run_sync(
+            self.run_async(
+                agent=agent,
+                input=input,
+                context=context,
+                run_id=run_id,
+                box_result=box_result,
+                agents=agents,
+                servers=servers,
+                memo=memo,
+            )
+        )
+
+    async def run_async(
+        self,
+        agent: "FlockAgent | str | None" = None,
+        input: dict | None = None,
+        context: FlockContext | None = None,
+        run_id: str = "",
+        box_result: bool = True,
+        agents: list["FlockAgent"] | None = None,
+        servers: list[Any] | None = None,
+        memo: dict[str, Any] | None = None,
+    ) -> Box | dict:
+        """Entry point for running an agent system asynchronously."""
+        # Import here to allow forward reference resolution
+        from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServer as ConcreteFlockServer,
+        )
+
+        with tracer.start_as_current_span("flock.run_async") as span:
+            # Add passed servers so that agents have access to them.
+            if servers:
+                for server_obj in servers:
+                    if isinstance(server_obj, ConcreteFlockServer):
+                        self.flock.add_server(server=server_obj)
+                    else:
+                        logger.warning(
+                            f"Item in 'servers' list is not a FlockMCPServer: {type(server_obj)}"
+                        )
+
+            # Add passed agents
+            if agents:
+                for agent_obj in agents:
+                    if isinstance(agent_obj, ConcreteFlockAgent):
+                        self.flock.add_agent(agent_obj)
+                    else:
+                        logger.warning(
+                            f"Item in 'agents' list is not a FlockAgent: {type(agent_obj)}"
+                        )
+
+            # Determine starting agent name
+            start_agent_name = self._resolve_start_agent(agent)
+
+            # Setup execution context and input
+            run_input = input if input is not None else self.flock._start_input
+            effective_run_id = run_id or f"flockrun_{uuid.uuid4().hex[:8]}"
+
+            # Set span attributes
+            span.set_attribute("start_agent", start_agent_name)
+            span.set_attribute("input", str(run_input))
+            span.set_attribute("run_id", effective_run_id)
+            span.set_attribute("enable_temporal", self.flock.enable_temporal)
+
+            logger.info(
+                f"Initiating Flock run '{self.flock.name}'. Start Agent: '{start_agent_name}'. Temporal: {self.flock.enable_temporal}."
+            )
+
+            try:
+                # Setup execution context
+                run_context = self._setup_execution_context(
+                    context, start_agent_name, run_input, effective_run_id
+                )
+
+                # Execute workflow with server management
+                async with self.flock._mgr:
+                    logger.info("Entering managed server context. Servers starting up.")
+                    logger.info(
+                        "Starting agent execution",
+                        agent=start_agent_name,
+                        enable_temporal=self.flock.enable_temporal,
+                    )
+
+                    # Execute workflow using appropriate engine
+                    result = await self._execute_workflow(run_context, memo)
+
+                    # Set result attributes on span
+                    span.set_attribute("result.type", str(type(result)))
+                    result_str = str(result)
+                    span.set_attribute(
+                        "result.preview",
+                        result_str[:1000] + ("..." if len(result_str) > 1000 else ""),
+                    )
+
+                    # Format and return result
+                    return self._format_result(result, box_result)
+
+            except Exception as e:
+                logger.error(f"Flock run '{self.flock.name}' failed: {e}", exc_info=True)
+                span.record_exception(e)
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
+
+                # Return a consistent error structure
+                error_output = {
+                    "error": str(e),
+                    "details": f"Flock run '{self.flock.name}' failed.",
+                    "run_id": effective_run_id,
+                    "start_agent": start_agent_name,
+                }
+                return Box(error_output) if box_result else error_output
+
+    def _resolve_start_agent(self, agent: "FlockAgent | str | None") -> str:
+        """Resolve the start agent name from various input types."""
+        from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+        from flock.core.registry import get_registry
+
+        registry = get_registry()
+
+        # Determine starting agent name
+        start_agent_name: str | None = None
+        if isinstance(agent, ConcreteFlockAgent):
+            start_agent_name = agent.name
+            if start_agent_name not in self.flock._agents:  # Add if not already present
+                self.flock.add_agent(agent)
+        elif isinstance(agent, str):
+            start_agent_name = agent
+        else:  # start_agent is None
+            start_agent_name = self.flock._start_agent_name
+
+        # Default to first agent if only one exists and none specified
+        if not start_agent_name and len(self.flock._agents) == 1:
+            start_agent_name = next(iter(self.flock._agents.keys()))
+        elif not start_agent_name:
+            raise ValueError(
+                "No start_agent specified and multiple/no agents exist in the Flock instance."
+            )
+
+        # Check if start_agent is in agents
+        if start_agent_name not in self.flock._agents:
+            # Try loading from registry if not found locally yet
+            reg_agent = registry.get_agent(start_agent_name)
+            if reg_agent:
+                self.flock.add_agent(reg_agent)
+                logger.info(f"Loaded start agent '{start_agent_name}' from registry.")
+            else:
+                raise ValueError(
+                    f"Start agent '{start_agent_name}' not found locally or in registry."
+                )
+
+        return start_agent_name
+
+    def _setup_execution_context(
+        self,
+        context: FlockContext | None,
+        start_agent_name: str,
+        run_input: dict,
+        run_id: str,
+    ) -> FlockContext:
+        """Setup the execution context for the workflow."""
+        resolved_start_agent = self.flock._agents.get(start_agent_name)
+        if not resolved_start_agent:  # Should have been handled by now
+            raise ValueError(f"Start agent '{start_agent_name}' not found after checks.")
+
+        run_context = context if context else FlockContext()
+        set_baggage("run_id", run_id)  # Set for OpenTelemetry
+
+        initialize_context(
+            run_context,
+            start_agent_name,
+            run_input,
+            run_id,
+            not self.flock.enable_temporal,  # local_debug is inverse of enable_temporal
+            self.flock.model or resolved_start_agent.model or DEFAULT_MODEL,
+        )
+
+        # Add agent definitions to context for routing/serialization within workflow
+        for agent_name_iter, agent_instance_iter in self.flock.agents.items():
+            agent_dict_repr = agent_instance_iter.to_dict()  # Agents handle their own serialization
+            run_context.add_agent_definition(
+                agent_type=type(agent_instance_iter),
+                agent_name=agent_name_iter,
+                agent_data=agent_dict_repr,
+            )
+
+        # Add temporal config to context if enabled
+        if self.flock.enable_temporal and self.flock.temporal_config:
+            run_context.set_variable(
+                "flock.temporal_workflow_config",
+                self.flock.temporal_config.model_dump(mode="json"),
+            )
+
+        return run_context
+
+    async def _execute_workflow(
+        self, run_context: FlockContext, memo: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        """Execute the workflow using the appropriate execution engine."""
+        if not self.flock.enable_temporal:
+            return await run_local_workflow(run_context, box_result=False)
+        else:
+            return await run_temporal_workflow(
+                self.flock,  # Pass the Flock instance
+                run_context,
+                box_result=False,
+                memo=memo,
+            )
+
+    def _format_result(self, result: dict[str, Any], box_result: bool) -> Box | dict:
+        """Format the execution result."""
+        if box_result:
+            try:
+                logger.debug("Boxing final result.")
+                return Box(result)
+            except ImportError:
+                logger.warning("Box library not installed, returning raw dict.")
+                return result
+        else:
+            return result

flock/core/orchestration/flock_initialization.py

@@ -0,0 +1,125 @@
+# src/flock/core/orchestration/flock_initialization.py
+"""Initialization functionality for Flock orchestrator."""
+
+import os
+import uuid
+from typing import TYPE_CHECKING
+
+from opentelemetry.baggage import get_baggage, set_baggage
+
+from flock.core.logging.logging import get_logger
+from flock.core.registry import get_registry
+from flock.core.util.cli_helper import init_console
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock
+    from flock.core.flock_agent import FlockAgent
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
+
+logger = get_logger("flock.initialization")
+
+
+class FlockInitialization:
+    """Handles initialization logic for Flock orchestrator."""
+
+    def __init__(self, flock: "Flock"):
+        self.flock = flock
+
+    def setup(
+        self,
+        agents: list["FlockAgent"] | None = None,
+        servers: list["FlockMCPServer"] | None = None,
+    ) -> None:
+        """Handle all initialization side effects and setup."""
+        # Register passed servers first (agents may depend on them)
+        if servers:
+            self._register_servers(servers)
+
+        # Register passed agents
+        if agents:
+            self._register_agents(agents)
+
+        # Initialize console if needed for banner
+        if self.flock.show_flock_banner:
+            init_console(clear_screen=True, show_banner=self.flock.show_flock_banner)
+
+        # Set Temporal debug environment variable
+        self._set_temporal_debug_flag()
+
+        # Ensure session ID exists in baggage
+        self._ensure_session_id()
+
+        # Auto-discover components
+        registry = get_registry()
+        registry.discover_and_register_components()
+
+        # Setup Opik if enabled
+        if self.flock.enable_opik:
+            self._setup_opik()
+
+        logger.info(
+            "Flock instance initialized",
+            name=self.flock.name,
+            model=self.flock.model,
+            enable_temporal=self.flock.enable_temporal,
+        )
+
+    def _register_servers(self, servers: list["FlockMCPServer"]) -> None:
+        """Register servers with the Flock instance."""
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServer as ConcreteFlockMCPServer,
+        )
+
+        for server in servers:
+            if isinstance(server, ConcreteFlockMCPServer):
+                self.flock.add_server(server)
+            else:
+                logger.warning(
+                    f"Item provided in 'servers' list is not a FlockMCPServer: {type(server)}"
+                )
+
+    def _register_agents(self, agents: list["FlockAgent"]) -> None:
+        """Register agents with the Flock instance."""
+        from flock.core.flock_agent import FlockAgent as ConcreteFlockAgent
+
+        for agent in agents:
+            if isinstance(agent, ConcreteFlockAgent):
+                self.flock.add_agent(agent)
+            else:
+                logger.warning(
+                    f"Item provided in 'agents' list is not a FlockAgent: {type(agent)}"
+                )
+
+    def _set_temporal_debug_flag(self) -> None:
+        """Set or remove LOCAL_DEBUG env var based on enable_temporal."""
+        if not self.flock.enable_temporal:
+            if "LOCAL_DEBUG" not in os.environ:
+                os.environ["LOCAL_DEBUG"] = "1"
+                logger.debug("Set LOCAL_DEBUG environment variable for local execution.")
+        elif "LOCAL_DEBUG" in os.environ:
+            del os.environ["LOCAL_DEBUG"]
+            logger.debug("Removed LOCAL_DEBUG environment variable for Temporal execution.")
+
+    def _ensure_session_id(self) -> None:
+        """Ensure a session_id exists in the OpenTelemetry baggage."""
+        session_id = get_baggage("session_id")
+        if not session_id:
+            session_id = str(uuid.uuid4())
+            set_baggage("session_id", session_id)
+            logger.debug(f"Generated new session_id: {session_id}")
+
+    def _setup_opik(self) -> None:
+        """Setup Opik integration."""
+        try:
+            import dspy
+            import opik
+            from opik.integrations.dspy.callback import OpikCallback
+
+            opik.configure(use_local=True, automatic_approvals=True)
+            opik_callback = OpikCallback(project_name=self.flock.name, log_graph=True)
+            dspy.settings.configure(callbacks=[opik_callback])
+            logger.info("Opik integration enabled")
+        except ImportError as e:
+            logger.error(f"Failed to setup Opik integration: {e}")
+            logger.warning("Disabling Opik integration")
+            self.flock.enable_opik = False

flock/core/orchestration/flock_server_manager.py

@@ -0,0 +1,67 @@
+# src/flock/core/orchestration/flock_server_manager.py
+"""Server management functionality for Flock orchestrator."""
+
+from typing import TYPE_CHECKING
+
+from flock.core.flock_server_manager import (
+    FlockServerManager as InternalServerManager,
+)
+from flock.core.logging.logging import get_logger
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock
+    from flock.core.mcp.flock_mcp_server import FlockMCPServer
+
+logger = get_logger("flock.server_manager")
+
+
+class FlockServerManager:
+    """Handles server lifecycle management for Flock orchestrator."""
+
+    def __init__(self, flock: "Flock"):
+        self.flock = flock
+        # Use the existing internal server manager
+        self._internal_mgr = InternalServerManager()
+
+    def add_server(self, server: "FlockMCPServer") -> "FlockMCPServer":
+        """Adds a server instance to this Flock configuration and registry as well as set it up to be managed by internal manager."""
+        from flock.core.mcp.flock_mcp_server import (
+            FlockMCPServer as ConcreteFlockMCPServer,
+        )
+        from flock.core.registry import get_registry
+
+        registry = get_registry()
+
+        if not isinstance(server, ConcreteFlockMCPServer):
+            raise TypeError("Provided object is not a FlockMCPServer instance.")
+        if not server.config.name:
+            raise ValueError("Server must have a name.")
+
+        if server.config.name in self.flock.servers:
+            raise ValueError(
+                f"Server with this name already exists. Name: '{server.config.name}'"
+            )
+
+        self.flock._servers[server.config.name] = server
+        registry.register_server(server)  # Register globally.
+
+        # Prepare server to be managed by the FlockServerManager
+        logger.info(f"Adding server '{server.config.name}' to managed list.")
+        self._internal_mgr.add_server_sync(server=server)
+        logger.info(f"Server '{server.config.name}' is now on managed list.")
+
+        logger.info(f"Server '{server.config.name}' added to Flock '{self.flock.name}'")
+        return server
+
+    async def __aenter__(self):
+        """Start all managed servers."""
+        return await self._internal_mgr.__aenter__()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Cleanup all managed servers."""
+        return await self._internal_mgr.__aexit__(exc_type, exc_val, exc_tb)
+
+    @property
+    def servers(self) -> dict[str, "FlockMCPServer"]:
+        """Returns the dictionary of servers managed by this Flock instance."""
+        return self.flock._servers