flock-core 0.1.2__py3-none-any.whl → 0.2.2__py3-none-any.whl

flock/core/flock.py

@@ -1,208 +1,232 @@
-# src/your_package/core/manager.py
-import contextlib
-import os
-import uuid
-from typing import TypeVar
-
-from flock.core.agent_registry import Registry
-from flock.core.context import FlockContext
-from flock.core.context_vars import FLOCK_CURRENT_AGENT, FLOCK_INITIAL_INPUT, FLOCK_LOCAL_DEBUG, FLOCK_RUN_ID
-from flock.core.logging import flock_logger, performance_handler
-from flock.core.logging.formatters import AgentResultFormatter
-from flock.core.logging.logger import FlockLogger, LogLevel
-from flock.workflow.activities import run_agent
-from flock.workflow.temporal_setup import create_temporal_client, setup_worker
-from flock.workflow.workflow import FlockWorkflow
-
-from .agent import Agent
-
-T = TypeVar("T", bound=Agent)
-
-
-class Flock:
-    """Manages creation and execution of agents.
-    This is a high-level class that the user interacts with directly.
-    """
-
-    def __init__(
-        self,
-        model: str = "openai/gpt-4o",
-        log_level: int = LogLevel.MINIMAL,
-        enable_performance_tracking: bool = False,
-        local_debug: bool = False,
-    ):
-        """Initialize the Flock.
-
-        Args:
-            model: The default model to use for agents.
-            log_level: The logging verbosity level:
-                - LogLevel.NONE: No logging except errors
-                - LogLevel.MINIMAL: Only agent outputs (default)
-                - LogLevel.BASIC: Agent outputs + basic workflow info
-                - LogLevel.VERBOSE: Everything including debug info
-            log_format: The logging format to use:
-                - LogFormat.SIMPLE: Simple message format (default)
-                - LogFormat.STRUCTURED: Table-like structured format
-            enable_performance_tracking: Whether to track and display performance metrics.
-            local_debug: Whether to run in local debug mode (default: False).
-        """
-        self.agents: dict[str, Agent] = {}
-        self.registry = Registry()
-        self.context = FlockContext()
-        self.model = model
-        self.enable_performance_tracking = enable_performance_tracking
-        self.local_debug = local_debug
-
-        # Configure performance tracking
-        if enable_performance_tracking:
-            performance_handler.enable()
-        else:
-            performance_handler.disable()
-
-        # Set LOCAL_DEBUG environment variable for console detection
-        if local_debug:
-            os.environ["LOCAL_DEBUG"] = "1"
-        elif "LOCAL_DEBUG" in os.environ:
-            del os.environ["LOCAL_DEBUG"]
-
-        # Configure logging
-        global flock_logger
-        flock_logger = FlockLogger(level=log_level)
-        flock_logger.info("Initialized Flock", model=model, local_debug=local_debug)
-
-    def add_agent(self, agent: T) -> T:
-        if not agent.model:
-            agent.model = self.model
-
-        if agent.name in self.agents:
-            flock_logger.debug(f"Agent {agent.name} already exists, returning existing instance")
-            return self.agents[agent.name]
-
-        self.agents[agent.name] = agent
-        self.registry.register_agent(agent)
-        self.context.add_agent_definition(type(agent), agent.name, agent.to_dict())
-        if hasattr(agent, "tools") and agent.tools:
-            for tool in agent.tools:
-                self.registry.register_tool(tool.__name__, tool)
-                flock_logger.debug(f"Registered tool: {tool.__name__}")
-
-        flock_logger.info(f"Added agent: {agent.name}", model=agent.model)
-        return agent
-
-    def add_tool(self, tool_name: str, tool: callable):
-        self.registry.register_tool(tool_name, tool)
-        flock_logger.debug(f"Registered tool: {tool_name}")
-
-    async def _run_local_debug_workflow(self, box_result: bool = True):
-        if self.enable_performance_tracking:
-            perf_context = performance_handler.track_time("local_workflow")
-        else:
-            perf_context = contextlib.nullcontext()
-
-        with perf_context:
-            flock_logger.info("Running workflow in local debug mode")
-            result = await run_agent(self.context)
-
-            # Format and display the result (always shown regardless of log level)
-            agent_name = self.context.get_variable(FLOCK_CURRENT_AGENT)
-            AgentResultFormatter.print_result(result, agent_name)
-
-            if self.enable_performance_tracking:
-                performance_handler.display_timings()
-
-            if box_result:
-                from box import Box
-
-                flock_logger.debug("Boxing result")
-                return Box(result)
-            return result
-
-    async def _run_temporal_workflow(self, box_result: bool = True):
-        if self.enable_performance_tracking:
-            perf_context = performance_handler.track_time("temporal_workflow")
-        else:
-            perf_context = contextlib.nullcontext()
-
-        with perf_context:
-            flock_logger.info("Connecting to Temporal and starting worker...")
-
-            await setup_worker(workflow=FlockWorkflow, activity=run_agent)
-
-            flock_logger.info("Starting workflow execution")
-
-            flock_client = await create_temporal_client()
-            result = await flock_client.execute_workflow(
-                FlockWorkflow.run,
-                self.context.to_dict(),
-                id=self.context.get_variable(FLOCK_RUN_ID),
-                task_queue="flock-queue",
-            )
-
-            # Format and display the result (always shown regardless of log level)
-            agent_name = self.context.get_variable(FLOCK_CURRENT_AGENT)
-            flock_logger.result(result, agent_name)
-
-            if self.enable_performance_tracking:
-                performance_handler.display_timings()
-
-            if box_result:
-                from box import Box
-
-                flock_logger.debug("Boxing result")
-                return Box(result)
-            return result
-
-    async def run_async(
-        self,
-        start_agent: Agent | str,
-        input: str,
-        context: FlockContext = None,
-        run_id: str = "",
-        box_result: bool = True,
-    ) -> dict:
-        """Entry point for running an agent system."""
-        try:
-            if self.enable_performance_tracking:
-                perf_context = performance_handler.track_time("run_setup")
-            else:
-                perf_context = contextlib.nullcontext()
-
-            with perf_context:
-                if isinstance(start_agent, str):
-                    flock_logger.debug(f"Looking up agent by name: {start_agent}")
-                    start_agent = self.registry.get_agent(start_agent)
-                    if not start_agent:
-                        raise ValueError(f"Agent '{start_agent}' not found in registry")
-
-                if context:
-                    flock_logger.debug("Using provided context")
-                    self.context = context
-
-                flock_logger.info("Setting up run context", agent=start_agent.name, input=input)
-                self.context.set_variable(FLOCK_CURRENT_AGENT, start_agent.name)
-                self.context.set_variable(FLOCK_INITIAL_INPUT, input)
-                self.context.set_variable(FLOCK_LOCAL_DEBUG, self.local_debug)
-
-                if not run_id:
-                    run_id = start_agent.name + "_" + uuid.uuid4().hex[:4]
-                    flock_logger.debug(f"Generated run ID: {run_id}")
-
-                self.context.set_variable(FLOCK_RUN_ID, run_id)
-
-            if self.enable_performance_tracking:
-                performance_handler.display_timings()
-
-            # we can run the flow locally for debugging purposes
-            if self.local_debug:
-                return await self._run_local_debug_workflow(box_result)
-            else:
-                return await self._run_temporal_workflow(box_result)
-
-        except Exception as e:
-            flock_logger.error(
-                "Run failed",
-                error=str(e),
-                agent=getattr(start_agent, "name", str(start_agent)),
-                input=input,
-            )
-            raise
+"""High-level orchestrator for creating and executing agents."""
+
+import os
+import uuid
+from typing import TypeVar
+
+from opentelemetry import trace
+from rich.prompt import Prompt
+
+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.flock_agent import FlockAgent
+from flock.core.logging.formatters.base_formatter import FormatterOptions
+from flock.core.logging.formatters.pprint_formatter import PrettyPrintFormatter
+from flock.core.logging.logging import get_logger
+from flock.core.logging.telemetry import setup_tracing
+from flock.core.registry.agent_registry import Registry
+from flock.core.util.cli_helper import display_banner
+from flock.core.util.input_resolver import top_level_to_keys
+
+T = TypeVar("T", bound=FlockAgent)
+logger = get_logger("flock")
+tracer = setup_tracing()
+tracer = trace.get_tracer(__name__)
+
+
+class Flock:
+    """High-level orchestrator for creating and executing agents.
+
+    Flock manages the registration of agents and tools, sets up the global context, and runs the agent workflows.
+    It provides an easy-to-use API for both local (debug) and production (Temporal) execution.
+    """
+
+    def __init__(
+        self,
+        model: str = "openai/gpt-4o",
+        local_debug: bool = False,
+        enable_logging: bool = False,
+        output_formatter: FormatterOptions = FormatterOptions(
+            PrettyPrintFormatter
+        ),
+    ):
+        """Initialize the Flock orchestrator.
+
+        Args:
+            model (str): The default model identifier to be used for agents. Defaults to "openai/gpt-4o".
+            local_debug (bool): If True, run the agent workflow locally for debugging purposes. Defaults to False.
+            enable_logging (bool): If True, enable verbose logging. Defaults to False.
+            output_formatter (FormatterOptions): Options for formatting output results.
+        """
+        with tracer.start_as_current_span("flock_init") as span:
+            span.set_attribute("model", model)
+            span.set_attribute("local_debug", local_debug)
+            span.set_attribute("enable_logging", enable_logging)
+            span.set_attribute(
+                "output_formatter", output_formatter.formatter.__name__
+            )
+            logger.info(
+                "Initializing Flock",
+                model=model,
+                local_debug=local_debug,
+                enable_logging=enable_logging,
+            )
+            logger.enable_logging = enable_logging
+
+            display_banner()
+
+            self.agents: dict[str, FlockAgent] = {}
+            self.registry = Registry()
+            self.context = FlockContext()
+            self.model = model
+            self.local_debug = local_debug
+            self.output_formatter = output_formatter
+
+            if local_debug:
+                os.environ["LOCAL_DEBUG"] = "1"
+                logger.debug("Set LOCAL_DEBUG environment variable")
+            elif "LOCAL_DEBUG" in os.environ:
+                del os.environ["LOCAL_DEBUG"]
+                logger.debug("Removed LOCAL_DEBUG environment variable")
+
+    def add_agent(self, agent: T) -> T:
+        """Add a new agent to the Flock system.
+
+        This method registers the agent, updates the internal registry and global context, and
+        sets default values if needed. If an agent with the same name already exists, the existing
+        agent is returned.
+
+        Args:
+            agent (FlockAgent): The agent instance to add.
+
+        Returns:
+            FlockAgent: The registered agent instance.
+        """
+        if not agent.model:
+            agent.model = self.model
+            logger.debug(
+                f"Using default model for agent {agent.name}", model=self.model
+            )
+
+        if agent.name in self.agents:
+            logger.warning(
+                f"Agent {agent.name} already exists, returning existing instance"
+            )
+            return self.agents[agent.name]
+        logger.info("Adding new agent")
+
+        self.agents[agent.name] = agent
+        self.registry.register_agent(agent)
+        self.context.add_agent_definition(
+            type(agent), agent.name, agent.to_dict()
+        )
+
+        if hasattr(agent, "tools") and agent.tools:
+            for tool in agent.tools:
+                self.registry.register_tool(tool.__name__, tool)
+                logger.debug("Registered tool", tool_name=tool.__name__)
+        logger.success("Agent added successfully")
+        return agent
+
+    def add_tool(self, tool_name: str, tool: callable):
+        """Register a tool with the Flock system.
+
+        Args:
+            tool_name (str): The name under which the tool will be registered.
+            tool (callable): The tool function to register.
+        """
+        logger.info("Registering tool", tool_name=tool_name)
+        self.registry.register_tool(tool_name, tool)
+        logger.debug("Tool registered successfully")
+
+    async def run_async(
+        self,
+        start_agent: FlockAgent | str,
+        input: dict = {},
+        context: FlockContext = None,
+        run_id: str = "",
+        box_result: bool = True,
+    ) -> dict:
+        """Entry point for running an agent system asynchronously.
+
+        This method performs the following steps:
+          1. If a string is provided for start_agent, it looks up the agent in the registry.
+          2. Optionally uses a provided global context.
+          3. Generates a unique run ID if one is not provided.
+          4. Initializes the context with standard variables (like agent name, input data, run ID, and debug flag).
+          5. Executes the agent workflow either locally (for debugging) or via Temporal (for production).
+
+        Args:
+            start_agent (FlockAgent | str): The agent instance or the name of the agent to start the workflow.
+            input (dict): A dictionary of input values required by the agent.
+            context (FlockContext, optional): A FlockContext instance to use. If not provided, a default context is used.
+            run_id (str, optional): A unique identifier for this run. If empty, one is generated automatically.
+            box_result (bool, optional): If True, wraps the output in a Box for nicer formatting. Defaults to True.
+
+        Returns:
+            dict: A dictionary containing the result of the agent workflow execution.
+
+        Raises:
+            ValueError: If the specified agent is not found in the registry.
+            Exception: For any other errors encountered during execution.
+        """
+        with tracer.start_as_current_span("run_async") as span:
+            span.set_attribute(
+                "start_agent",
+                start_agent.name
+                if hasattr(start_agent, "name")
+                else start_agent,
+            )
+
+            span.set_attribute("input", str(input))
+            span.set_attribute("context", str(context))
+            span.set_attribute("run_id", run_id)
+            span.set_attribute("box_result", box_result)
+
+            try:
+                if isinstance(start_agent, str):
+                    logger.debug(
+                        "Looking up agent by name", agent_name=start_agent
+                    )
+                    start_agent = self.registry.get_agent(start_agent)
+                    if not start_agent:
+                        logger.error("Agent not found", agent_name=start_agent)
+                        raise ValueError(
+                            f"Agent '{start_agent}' not found in registry"
+                        )
+                    start_agent.resolve_callables(context=self.context)
+                if context:
+                    logger.debug("Using provided context")
+                    self.context = context
+                if not run_id:
+                    run_id = f"{start_agent.name}_{uuid.uuid4().hex[:4]}"
+                    logger.debug("Generated run ID", run_id=run_id)
+
+                # TODO - Add a check for required input keys
+                input_keys = top_level_to_keys(start_agent.input)
+                for key in input_keys:
+                    if key.startswith("flock."):
+                        key = key[6:]  # Remove the "flock." prefix
+                    if key not in input:
+                        input[key] = Prompt.ask(
+                            f"Please enter {key} for {start_agent.name}"
+                        )
+
+                # Initialize the context with standardized variables
+                initialize_context(
+                    self.context,
+                    start_agent.name,
+                    input,
+                    run_id,
+                    self.local_debug,
+                )
+
+                logger.info(
+                    "Starting agent execution",
+                    agent=start_agent.name,
+                    local_debug=self.local_debug,
+                )
+
+                if self.local_debug:
+                    return await run_local_workflow(
+                        self.context, self.output_formatter, box_result
+                    )
+                else:
+                    return await run_temporal_workflow(
+                        self.context, self.output_formatter, box_result
+                    )
+            except Exception as e:
+                logger.exception("Execution failed", error=str(e))
+                raise